tensordict-cache

A persistent, memory-mapped cache for tensor data built on TensorDict.

Store and retrieve TensorDict objects on disk using memory-mapped files. Cached entries survive process restarts and are loaded lazily without copying data into RAM.

Installation

pip install tensordict-cache

Quick start

import torch
from tensordict import TensorDict
from tensordict_cache import TensorCache

# Create a cache (directory is created if it doesn't exist)
cache = TensorCache("./my_cache")

# Store embeddings keyed by prompt text
cache["hello world"] = TensorDict(
    {"embedding": torch.randn(768)},
    batch_size=[],
)

# Retrieve them
embedding = cache["hello world"]["embedding"]

Usage

Creating a cache

from tensordict_cache import TensorCache

# Open or create a cache directory
cache = TensorCache("/path/to/cache")

# Open without loading existing entries
cache = TensorCache("/path/to/cache", load_existing=False)

Storing entries

Keys can be strings or integers. Values must be TensorDict instances.

import torch
from tensordict import TensorDict

td = TensorDict({
    "hidden_state": torch.randn(512),
    "logits": torch.randn(10),
}, batch_size=[])

cache["my_prompt"] = td
cache[42] = td  # integer keys work too

Retrieving entries

# Dict-style access (raises KeyError if missing)
result = cache["my_prompt"]

# Safe access with default
result = cache.get("my_prompt")          # returns None if missing
result = cache.get("my_prompt", default) # returns default if missing

Checking membership and length

if "my_prompt" in cache:
    print("Hit!")

print(f"Cache has {len(cache)} entries")

Listing keys and inspecting the cache

# Keys are SHA-256 hashes of the original key
print(cache.keys())

# Human-readable representation
print(cache)
# TensorCache(prefix=/path/to/cache, n_cache=3)

Deleting entries

# Remove a single entry from memory and disk
del cache["my_prompt"]

# Raises KeyError if the key doesn't exist

Size-limited cache

# Cap total disk usage at 1 GB — oldest entries are evicted first
cache = TensorCache("./my_cache", max_size_bytes=1_000_000_000)

cache["a"] = TensorDict({"x": torch.randn(768)}, batch_size=[])
cache["b"] = TensorDict({"x": torch.randn(768)}, batch_size=[])
# When a new insert exceeds the limit, the oldest entry is removed automatically

Clearing the cache

# Remove all entries from memory and disk
cache.clear()

Persistence across sessions

Data is written to disk as memory-mapped files. Reopening the same directory automatically loads all previously stored entries:

# Session 1
cache = TensorCache("./my_cache")
cache["prompt_a"] = TensorDict({"v": torch.tensor(1.0)}, batch_size=[])

# Session 2 (new process)
cache = TensorCache("./my_cache")
assert "prompt_a" in cache  # still there

How it works

Each entry is stored in a subdirectory under the cache prefix:

my_cache/
  a1b2c3d4.../   # SHA-256 hash of the key
    meta.json
    *.memmap
  e5f6g7h8.../
    meta.json
    *.memmap

Under the hood, TensorCache calls TensorDict.memmap() to write each entry and TensorDict.load_memmap() to read it back. Memory-mapped storage means tensors are not loaded into RAM until accessed, keeping memory usage low even for large caches.

API reference

Requirements

• Python >= 3.10
• tensordict >= 0.11.0

License

MIT