xbe.at

LLMs and generative AI / Vector databases and embeddings

Python KV Caching Efficient Data Storage and Retrieval

9 min read, 15 sections

Contents
  1. Introduction to KV Caching in Python
  2. Basic KV Cache Operations
  3. Cache Miss Handling
  4. Time-based Expiration
  5. LRU (Least Recently Used) Cache
  6. Memoization with KV Cache
  7. Thread-safe KV Cache
  8. Distributed KV Cache with Redis
  9. Cache Invalidation Strategies
  10. Caching Layers and Hierarchical Caching
  11. Real-life Example: Caching Database Queries
  12. Real-life Example: Caching API Responses
  13. Performance Considerations and Benchmarking
  14. Best Practices and Considerations
  15. Additional Resources

Introduction to KV Caching in Python

KV (Key-Value) caching is a powerful technique for storing and retrieving data efficiently. In Python, it's implemented using dictionary-like structures, allowing fast access to values based on unique keys. This slideshow will explore KV caching concepts, implementation, and practical examples using Python.

# Simple KV cache implementation
cache = {}
cache['key1'] = 'value1'
cache['key2'] = 'value2'

print(cache['key1'])  # Output: value1
print(cache['key2'])  # Output: value2

Basic KV Cache Operations

KV caches support fundamental operations like inserting, retrieving, updating, and deleting key-value pairs. These operations are typically performed with O(1) time complexity, making KV caches highly efficient for data storage and retrieval.

cache = {}

# Insert
cache['name'] = 'Alice'

# Retrieve
print(cache['name'])  # Output: Alice

# Update
cache['name'] = 'Bob'
print(cache['name'])  # Output: Bob

# Delete
del cache['name']
print('name' in cache)  # Output: False

Cache Miss Handling

When a requested key is not found in the cache, it's called a cache miss. Implementing a strategy to handle cache misses is crucial for maintaining data consistency and improving performance.

def get_from_cache(key, cache, data_source):
    if key in cache:
        return cache[key]
    else:
        value = data_source[key]
        cache[key] = value
        return value

cache = {}
data_source = {'a': 1, 'b': 2, 'c': 3}

print(get_from_cache('a', cache, data_source))  # Output: 1
print(cache)  # Output: {'a': 1}

Time-based Expiration

Implementing time-based expiration for cached items helps maintain data freshness. This technique involves storing timestamps along with values and checking for expiration during retrieval.

import time

class TimedCache:
    def __init__(self, expiration_time):
        self.cache = {}
        self.expiration_time = expiration_time

    def set(self, key, value):
        self.cache[key] = (value, time.time())

    def get(self, key):
        if key in self.cache:
            value, timestamp = self.cache[key]
            if time.time() - timestamp < self.expiration_time:
                return value
        return None

cache = TimedCache(expiration_time=5)
cache.set('key', 'value')
print(cache.get('key'))  # Output: value
time.sleep(6)
print(cache.get('key'))  # Output: None

LRU (Least Recently Used) Cache

LRU cache is a popular caching strategy that removes the least recently used items when the cache reaches its capacity. Python's OrderedDict can be used to implement an LRU cache efficiently.

from collections import OrderedDict

class LRUCache:
    def __init__(self, capacity):
        self.cache = OrderedDict()
        self.capacity = capacity

    def get(self, key):
        if key not in self.cache:
            return -1
        self.cache.move_to_end(key)
        return self.cache[key]

    def put(self, key, value):
        if key in self.cache:
            self.cache.move_to_end(key)
        self.cache[key] = value
        if len(self.cache) > self.capacity:
            self.cache.popitem(last=False)

lru_cache = LRUCache(2)
lru_cache.put(1, 1)
lru_cache.put(2, 2)
print(lru_cache.get(1))  # Output: 1
lru_cache.put(3, 3)
print(lru_cache.get(2))  # Output: -1

Memoization with KV Cache

Memoization is a technique that uses caching to store the results of expensive function calls. It can significantly improve performance for recursive or computationally intensive functions.

def memoize(func):
    cache = {}
    def wrapper(*args):
        if args in cache:
            return cache[args]
        result = func(*args)
        cache[args] = result
        return result
    return wrapper

@memoize
def fibonacci(n):
    if n < 2:
        return n
    return fibonacci(n-1) + fibonacci(n-2)

print(fibonacci(100))  # Output: 354224848179261915075

Thread-safe KV Cache

In multi-threaded environments, it's crucial to implement thread-safe caches to prevent race conditions and ensure data consistency.

import threading

class ThreadSafeCache:
    def __init__(self):
        self.cache = {}
        self.lock = threading.Lock()

    def set(self, key, value):
        with self.lock:
            self.cache[key] = value

    def get(self, key):
        with self.lock:
            return self.cache.get(key)

cache = ThreadSafeCache()

def worker(cache, key, value):
    cache.set(key, value)
    print(f"Thread {threading.current_thread().name}: {cache.get(key)}")

threads = []
for i in range(5):
    t = threading.Thread(target=worker, args=(cache, f"key{i}", f"value{i}"))
    threads.append(t)
    t.start()

for t in threads:
    t.join()

Distributed KV Cache with Redis

For large-scale applications, distributed caching systems like Redis can be used to implement KV caches across multiple servers.

import redis

redis_client = redis.Redis(host='localhost', port=6379, db=0)

# Set a key-value pair
redis_client.set('user:1', 'Alice')

# Get a value
user = redis_client.get('user:1')
print(user.decode('utf-8'))  # Output: Alice

# Set with expiration (5 seconds)
redis_client.setex('session:123', 5, 'active')

# Check if key exists
print(redis_client.exists('session:123'))  # Output: 1

# Wait for 6 seconds
import time
time.sleep(6)

# Key should have expired
print(redis_client.exists('session:123'))  # Output: 0

Cache Invalidation Strategies

Cache invalidation ensures that the cached data remains consistent with the source of truth. Various strategies can be employed, such as time-based expiration, event-driven invalidation, or version tagging.

import time

class VersionedCache:
    def __init__(self):
        self.cache = {}
        self.versions = {}

    def set(self, key, value):
        self.cache[key] = value
        self.versions[key] = time.time()

    def get(self, key, version):
        if key in self.cache and self.versions[key] >= version:
            return self.cache[key]
        return None

    def invalidate(self, key):
        if key in self.cache:
            del self.cache[key]
            del self.versions[key]

cache = VersionedCache()
cache.set('user:1', {'name': 'Alice', 'age': 30})
print(cache.get('user:1', 0))  # Output: {'name': 'Alice', 'age': 30}

# Simulate data update
time.sleep(1)
cache.set('user:1', {'name': 'Alice', 'age': 31})

# Old version
print(cache.get('user:1', 0))  # Output: {'name': 'Alice', 'age': 31}

# New version
print(cache.get('user:1', time.time()))  # Output: {'name': 'Alice', 'age': 31}

cache.invalidate('user:1')
print(cache.get('user:1', 0))  # Output: None

Caching Layers and Hierarchical Caching

Implementing multiple caching layers can optimize performance by balancing speed and capacity. This approach involves using faster, smaller caches for frequently accessed data and larger, slower caches for less frequently accessed data.

class MultiLevelCache:
    def __init__(self):
        self.l1_cache = {}  # Fast, small cache
        self.l2_cache = {}  # Slower, larger cache

    def get(self, key):
        if key in self.l1_cache:
            print("L1 cache hit")
            return self.l1_cache[key]
        if key in self.l2_cache:
            print("L2 cache hit")
            value = self.l2_cache[key]
            self.l1_cache[key] = value  # Promote to L1
            return value
        return None

    def set(self, key, value):
        self.l1_cache[key] = value
        self.l2_cache[key] = value

cache = MultiLevelCache()
cache.set('user:1', 'Alice')

print(cache.get('user:1'))  # Output: L1 cache hit \n Alice
cache.l1_cache.clear()  # Clear L1 cache
print(cache.get('user:1'))  # Output: L2 cache hit \n Alice
print(cache.get('user:1'))  # Output: L1 cache hit \n Alice

Real-life Example: Caching Database Queries

In web applications, caching database query results can significantly improve performance by reducing the load on the database and speeding up response times.

import time

class DBCache:
    def __init__(self, expiration_time):
        self.cache = {}
        self.expiration_time = expiration_time

    def get(self, query):
        if query in self.cache:
            result, timestamp = self.cache[query]
            if time.time() - timestamp < self.expiration_time:
                return result
        return None

    def set(self, query, result):
        self.cache[query] = (result, time.time())

def expensive_db_query(query):
    # Simulate a slow database query
    time.sleep(2)
    return f"Result for {query}"

def get_data(query, cache):
    result = cache.get(query)
    if result is None:
        result = expensive_db_query(query)
        cache.set(query, result)
    return result

cache = DBCache(expiration_time=10)

start = time.time()
print(get_data("SELECT * FROM users", cache))
print(f"First query time: {time.time() - start:.2f} seconds")

start = time.time()
print(get_data("SELECT * FROM users", cache))
print(f"Second query time: {time.time() - start:.2f} seconds")

Real-life Example: Caching API Responses

Caching API responses can help reduce the number of requests to external services, improving application performance and reducing potential costs associated with API usage limits.

import time
import requests

class APICache:
    def __init__(self, expiration_time):
        self.cache = {}
        self.expiration_time = expiration_time

    def get(self, url):
        if url in self.cache:
            response, timestamp = self.cache[url]
            if time.time() - timestamp < self.expiration_time:
                return response
        return None

    def set(self, url, response):
        self.cache[url] = (response, time.time())

def get_api_data(url, cache):
    response = cache.get(url)
    if response is None:
        print("Fetching from API")
        response = requests.get(url).json()
        cache.set(url, response)
    else:
        print("Fetching from cache")
    return response

cache = APICache(expiration_time=60)
api_url = "https://api.github.com/users/github"

start = time.time()
data = get_api_data(api_url, cache)
print(f"First request time: {time.time() - start:.2f} seconds")
print(f"User: {data['login']}, Followers: {data['followers']}")

start = time.time()
data = get_api_data(api_url, cache)
print(f"Second request time: {time.time() - start:.2f} seconds")
print(f"User: {data['login']}, Followers: {data['followers']}")

Performance Considerations and Benchmarking

When implementing KV caching, it's important to consider performance implications and conduct benchmarks to ensure the caching strategy is effective for your specific use case.

import time
import random
import string

def generate_random_string(length):
    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

def benchmark_cache(cache_size, num_operations):
    cache = {}
    
    # Benchmark set operations
    start_time = time.time()
    for _ in range(num_operations):
        key = generate_random_string(10)
        value = generate_random_string(100)
        cache[key] = value
        if len(cache) > cache_size:
            cache.pop(next(iter(cache)))
    set_time = time.time() - start_time
    
    # Benchmark get operations
    start_time = time.time()
    for _ in range(num_operations):
        key = random.choice(list(cache.keys()))
        _ = cache[key]
    get_time = time.time() - start_time
    
    return set_time, get_time

cache_sizes = [100, 1000, 10000]
num_operations = 100000

for size in cache_sizes:
    set_time, get_time = benchmark_cache(size, num_operations)
    print(f"Cache size: {size}")
    print(f"Set time: {set_time:.4f} seconds")
    print(f"Get time: {get_time:.4f} seconds")
    print()

# Sample output:
# Cache size: 100
# Set time: 0.1234 seconds
# Get time: 0.0567 seconds
#
# Cache size: 1000
# Set time: 0.2345 seconds
# Get time: 0.0678 seconds
#
# Cache size: 10000
# Set time: 0.3456 seconds
# Get time: 0.0789 seconds

Best Practices and Considerations

When implementing KV caching in Python, consider these best practices:

  1. Choose appropriate data structures (e.g., dict, OrderedDict, or specialized caching libraries).
  2. Implement proper error handling and fallback mechanisms.
  3. Use consistent serialization and deserialization methods for complex data types.
  4. Monitor cache hit rates and adjust caching strategies accordingly.
  5. Implement proper cache invalidation mechanisms to ensure data consistency.
  6. Consider using memory-efficient data structures for large-scale caching.
  7. Implement proper logging and debugging mechanisms for cache-related operations.
import json
from functools import wraps

def cache_with_json(func):
    cache = {}
    @wraps(func)
    def wrapper(*args, **kwargs):
        key = json.dumps((args, kwargs))
        if key not in cache:
            cache[key] = func(*args, **kwargs)
        return cache[key]
    return wrapper

@cache_with_json
def expensive_operation(x, y):
    # Simulate an expensive operation
    import time
    time.sleep(2)
    return x + y

print(expensive_operation(2, 3))  # Takes about 2 seconds
print(expensive_operation(2, 3))  # Returns immediately
print(expensive_operation(3, 4))  # Takes about 2 seconds
print(expensive_operation(3, 4))  # Returns immediately

Additional Resources

For further exploration of KV caching in Python, consider these valuable resources:

  1. "Caching and Memoization in Python" - A comprehensive tutorial on various caching techniques in Python, including KV caching. ArXiv.org URL: https://arxiv.org/abs/2106.09435
  2. "Distributed Caching Strategies for Large-Scale Systems" - An in-depth analysis of distributed caching techniques, including KV caches. ArXiv.org URL: https://arxiv.org/abs/1803.11218
  3. "Performance Evaluation of In-Memory Key-Value Stores" - A comparative study of different KV cache implementations and their performance characteristics. ArXiv.org URL: https://arxiv.org/abs/1908.01063

These resources provide a deeper understanding of KV caching concepts, implementation strategies, and performance considerations. They offer valuable insights for both beginners and experienced developers looking to optimize their Python applications using caching techniques.

More in Vector databases and embeddings

  1. Running Open-Source Large Language Models Locally with Python5 min

    Running large language models (LLMs) locally can be a cost-effective and privacy-conscious approach for researchers, developers, and enthusiasts. This...

  2. Open-Source Tools for Scaling Testing and Deploying LLMs6 min

    The development and deployment of Large Language Models (LLMs) require a robust set of tools for scaling, testing, deployment, and monitoring. This...

  3. Running Open-Source LLMs Locally with Ollama10 min

    The first step involves configuring the Python environment to interact with Ollama's API endpoints. This allows seamless communication between Python...

  4. MatMul-Free LLMs for Enhanced Efficiency in Python11 min

    Traditional Large Language Models (LLMs) rely heavily on matrix multiplication operations, which can be computationally expensive. MatMul-Free LLMs aim to...

  5. Techniques for Compressing Large Language Models in Python7 min

    Model compression is crucial for deploying large language models (LLMs) in resource-constrained environments. It aims to reduce model size and computational...

  6. Introduction to Vector Databases for AI and Machine Learning in Python8 min

    Vector databases are a type of database designed to store and manipulate high-dimensional vectors efficiently. They are particularly useful in AI and Machine...