tensorblob
tensorblob
A lightweight, dynamic-sized, memory-mapped tensor storage with file-like APIs, while also supporting integer indexing and slicing, built with MemoryMappedTensor from tensordict.
Features
- 🔗 Memory-mapped storage: Efficient storage of large collections of same-shaped tensors
- 💾 File-like APIs: Read, write, and seek like a file, while also supporting integer indexing and slicing
- âš¡ Dynamic-sized: No need to specify the total number of tensors upfront
- 🔄 Extend and truncate: Extend the blob with another blob or truncate the blob to a specific position
- 🚀 LRU cache: Automatic management of memory-mapped blocks for scalability with large blobs
- 🧩 Multi-field databases:
TensorDBmanages several row-aligned blobs for heterogeneous data, e.g., multivariate time series or event streams
Installation
From PyPI:
pip install tensorblob
If you are interested in the experimental (i.e., unstable and undertested) version, you can install it from GitHub:
pip install git+https://github.com/Guest400123064/tensorblob.git
Core Use Cases
Quick Start
The example below shows how to create a new storage for a collection of randomly generated fake embeddings, and how to access them by index. Since the storage is memory-mapped, no need to read all tensors into memory; just access them by index.
import torch
from tensorblob import TensorBlob
# Create a new storage for a collection of randomly generated fake embeddings;
# need to specify the data type and shape of each tensor for creation
with TensorBlob.open("embeddings.blob", "w", dtype="float32", shape=768) as blob:
blob.write(torch.randn(100_000, 768))
print(f"Wrote {len(blob)} embeddings")
# No need to specify the configurations again after creation
with TensorBlob.open("embeddings.blob", "r") as blob:
e1 = blob[42]
e2 = blob[-1:16384:-12345]
print(f"Similarity: {torch.cosine_similarity(e1, e2)}")
Processing Large Datasets
Store and preprocess datasets larger than RAM using memory mapping can be useful to accelerate the training process by reducing the time spent on data loading and transformation.
with TensorBlob.open("data/images.blob", "w", dtype="float32", shape=(3, 224, 224)) as blob:
for image_batch in data_loader:
blob.write(preprocess(image_batch))
with TensorBlob.open("data/images.blob", "r") as blob:
for image in blob:
result = model(image)
Incremental Data Collection
Append new data to existing blobs can be useful with streaming data collection.
with TensorBlob.open("positions.blob", "w", dtype="float32", shape=3) as blob:
blob.write(initial_position)
# Later: append more data by opening the blob in append mode
with TensorBlob.open("positions.blob", "a") as blob:
for pos in trajectory_queue.get():
blob.write(pos)
print(f"Total trajectory recorded: {len(blob)}")
Random Access and Updates with File-Like APIs
Read and modify specific tensors starting from a specific position.
import io
with TensorBlob.open("data/features.blob", "r+") as blob:
blob.seek(1000)
print(f"Current position: {blob.tell()}")
batch = blob.read(size=100)
print(f"Read {batch.shape} tensors")
# Update specific positions, whence is also supported
blob.seek(-500, whence=io.SEEK_END)
blob.write(updated_features)
# Append new data
blob.seek(len(blob))
blob.write(additional_features)
Extend and Truncate
Extend the blob with another blob or truncate the blob to a specific position. Extension could be useful if we want to merge two blobs into one, e.g., results from two different processes. Note that extension operation does not delete the original data.
with TensorBlob.open("data/features.blob", "a") as blob:
blob.extend(other_blob)
# Extension without maintaining the order is faster
with TensorBlob.open("data/features.blob", "r+") as blob:
blob.extend(other_blob, maintain_order=False)
with TensorBlob.open("data/features.blob", "r+") as blob:
blob.truncate(1000)
print(f"Truncated to {len(blob)} tensors")
Heterogeneous Data with TensorDB
For multi-modal or multi-field data (e.g., multivariate time series, event streams), TensorDB manages several TensorBlobs under the hood — one per field — with row orders always aligned. Each field has its own dtype and shape, and each field's storage gets its own independent LRU cache and block files.
from tensorblob import TensorDB
# Create a database with a fixed schema mapping field names to (dtype, shape)
with TensorDB.open("events.db", "w",
schema={"price": ("float32", 1),
"embed": ("float16", 768)}) as db:
# Rows are dense: every write must supply every field with the same row count
db.write({"price": torch.randn(100_000, 1),
"embed": torch.randn(100_000, 768).half()})
print(f"Wrote {len(db)} rows")
# No need to specify the schema again after creation
with TensorDB.open("events.db", "r") as db:
row = db[42] # {"price": tensor of shape (1,), "embed": (768,)}
batch = db[10:100] # {"price": (90, 1), "embed": (90, 768)}
print(f"Fields: {list(batch)}, price range: {batch['price'].min()}..{batch['price'].max()}")
TensorDB supports the same file-like APIs as TensorBlob, applied row-wise across all fields:
with TensorDB.open("events.db", "r+") as db:
db.seek(1000)
batch = db.read(size=100) # dict of (100, ...) tensors
db.seek(-500, whence=io.SEEK_END)
db.write({"price": new_prices, "embed": new_embeds}) # overwrite in place
db.truncate(10_000) # truncate all fields at once
db.extend(other_db, maintain_order=False) # merge another db with the same schema
# Cleanup removes the whole database directory
TensorDB.unlink("events.db")
Consistency guarantee: a write commits the row count only after all fields are written. If a crash interrupts a write mid-way, the next open reports the last committed (fully written) row count, and writable opens automatically truncate the stray partial rows, so row alignment is always preserved.
Performance and Scalability
Memory Management
TensorBlob uses an LRU (Least Recently Used) cache to manage memory-mapped blocks efficiently. This allows you to work with blobs containing millions of tensors without loading everything into memory.
Default behavior:
- Automatically caches up to ~4,000 blocks (1/16 of system's VMA limit)
- Blocks loaded on-demand when accessed
- Least recently used blocks automatically evicted when cache is full
For large-scale workloads:
# Increase cache for better random access performance
with TensorBlob.open("large.blob", "r", max_cached_blocks=10_000) as blob:
for idx in random_indices:
tensor = blob[idx] # Cached blocks reused efficiently
# Decrease cache for memory-constrained environments
with TensorBlob.open("data.blob", "r", max_cached_blocks=100) as blob:
for tensor in blob: # Sequential access works fine with small cache
process(tensor)
Performance tips:
- Sequential access patterns work well with any cache size
- Random access benefits from larger cache sizes
- Each cached block consumes ~200 bytes of kernel memory (VMA overhead)
- System limit: typically ~65,000 memory-mapped regions per process
- To avoid frequent cache evictions, one can also increase the block size to reduce the total number of blocks
Contributing
Contributions welcome! Please submit a Pull Request.
License
Apache License 2.0 - see LICENSE file for details.
class
TensorBlob(configmixin._core.ConfigMixin):
43class TensorBlob(ConfigMixin): 44 _m_rd = False 45 _m_wr = False 46 _m_ap = False 47 48 status_name = ".stat" 49 config_name = ".conf" 50 ignore_for_config: ClassVar[list[str]] = ["filename", "mode", "max_cached_blocks"] 51 52 @classmethod 53 def open( 54 cls, 55 filename, 56 mode="r", 57 *, 58 dtype=None, 59 shape=None, 60 block_size=8192, 61 max_cached_blocks=None, 62 ): 63 r"""Open a TensorBlob with file-like interface for tensor storage. 64 65 TensorBlob provides persistent, memory-mapped storage for large collections 66 of same-shaped tensors. It uses a block-based architecture where tensors are 67 organized into fixed-size blocks for efficient I/O and memory management. 68 69 The blob is stored as a directory containing: 70 - ``.conf``: Configuration file (dtype, shape, block_size) 71 - ``.stat``: State file (length, block list) 72 - Block files: UUID-named memory-mapped tensor files 73 74 Parameters 75 ---------- 76 filename : str or Path 77 Directory path for blob storage. Supports tilde expansion (~) and 78 relative paths. 79 mode : str, default="r" 80 File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). See below for details. 81 dtype : str or torch.dtype, optional 82 Data type for tensors. Required for new blobs (modes 'w', 'w+'). 83 shape : tuple of int or int, optional 84 Shape of individual tensors. Required for new blobs (modes 'w', 'w+'). 85 block_size : int, default=8192 86 Number of tensors per memory-mapped block file. 87 max_cached_blocks : int, optional 88 Maximum number of memory-mapped blocks to keep cached. When exceeded, 89 least recently used blocks are unmapped. If None (default), uses 1/16 90 of system's max_map_count limit (typically ~4000). This limits kernel 91 VMA overhead for blobs with many blocks. 92 93 Returns 94 ------- 95 TensorBlob 96 Opened blob object. Use with context manager for automatic cleanup. 97 98 Raises 99 ------ 100 FileNotFoundError 101 If mode is 'r', 'r+', 'a', or 'a+' and blob doesn't exist. 102 ValueError 103 If creating new blob without dtype or shape, or if mode is invalid. 104 TypeError 105 If dtype is neither string nor torch.dtype. 106 107 Examples 108 -------- 109 Creating a new blob and writing data: 110 111 >>> import torch 112 >>> from tensorblob import TensorBlob 113 >>> 114 >>> with TensorBlob.open("data/embeddings", "w", 115 ... dtype="float32", shape=(768,)) as blob: 116 ... embeddings = torch.randn(1000, 768) 117 ... blob.write(embeddings) 118 ... print(f"Wrote {len(blob)} tensors") 119 Wrote 1000 tensors 120 121 Reading from existing blob: 122 123 >>> with TensorBlob.open("data/embeddings", "r") as blob: 124 ... all_data = blob.read() 125 ... print(all_data.shape) 126 torch.Size([1000, 768]) 127 128 Appending to existing blob: 129 130 >>> with TensorBlob.open("data/embeddings", "a") as blob: 131 ... new_data = torch.randn(100, 768) 132 ... blob.write(new_data) 133 ... print(f"Total: {len(blob)}") 134 Total: 1100 135 136 Read and update with r+ mode: 137 138 >>> with TensorBlob.open("data/embeddings", "r+") as blob: 139 ... first_10 = blob.read(size=10) 140 ... blob.seek(5) 141 ... blob.write(torch.ones(3, 768)) # Overwrite at position 5 142 143 Custom block size for large tensors: 144 145 >>> with TensorBlob.open("data/images", "w", 146 ... dtype=torch.float32, 147 ... shape=(3, 1024, 1024), 148 ... block_size=256) as blob: 149 ... images = torch.randn(1000, 3, 1024, 1024) 150 ... blob.write(images) 151 152 Custom cache size for large-scale random access: 153 154 >>> # Increase cache for better random access performance 155 >>> with TensorBlob.open("data/embeddings", "r", 156 ... max_cached_blocks=10000) as blob: 157 ... for idx in random_indices: 158 ... embedding = blob[idx] # Frequently accessed blocks stay cached 159 160 >>> # Decrease cache for memory-constrained environments 161 >>> with TensorBlob.open("data/features", "r", 162 ... max_cached_blocks=100) as blob: 163 ... for feature in blob: # Sequential access works fine 164 ... process(feature) 165 166 File Access Modes 167 ----------------- 168 Similar to Python's built-in open(), supports the following modes: 169 170 Basic modes: 171 - 'r' : Read-only. Blob must exist. Position starts at beginning. 172 - 'w' : Write-only. Creates new or truncates existing. Position at start. **If the blob already exists, 173 truncation will ignore any other parameters supplied and rely on existing configuration.** 174 - 'a' : Append-only. Blob must exist. Position starts at end. 175 All writes go to end regardless of seek position. 176 177 Update modes (with '+'): 178 - 'r+' : Read and write. Blob must exist. Position at start. 179 Can overwrite existing data or extend at end. 180 - 'w+' : Read and write. Creates new or truncates existing. Position at start. 181 - 'a+' : Read and append. Blob must exist. Position at end. 182 Reads allowed anywhere, writes always append to end. 183 184 Data Type and Shape 185 ------------------- 186 All tensors in a blob must have the same dtype and shape. These are 187 specified when creating a new blob (modes 'w', 'w+') and stored in 188 the configuration file. When opening existing blobs, dtype and shape 189 are loaded automatically. 190 191 Supported dtypes: "float32", "float64", "int32", "int64", "bool", etc. 192 Can also use torch.dtype objects like torch.float32. 193 194 Shape can be: 195 - Single integer: shape=10 creates 1D tensors of shape (10,) 196 - Tuple: shape=(3, 224, 224) creates 3D tensors 197 """ 198 modes = set(mode) 199 if modes - set("raw+") or len(mode) > len(modes): 200 raise ValueError(f"Invalid mode: {mode}") 201 if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1: 202 raise ValueError( 203 f"Must have exactly one of read/write/append mode and at most one plus: {mode}" 204 ) 205 206 filename = Path(filename).expanduser().resolve() 207 if not filename.exists(): 208 if "r" in modes or "a" in modes: 209 raise FileNotFoundError(f"Blob not found: {filename!r}") 210 if dtype is None or shape is None: 211 raise ValueError( 212 f"Arguments ``dtype`` and ``shape`` are required for new blob; got: {dtype!r} and {shape!r}" 213 ) 214 if isinstance(dtype, torch.dtype): 215 dtype = str(dtype).split(".").pop() 216 elif not isinstance(dtype, str): 217 raise TypeError( 218 f"dtype must be str or torch.dtype, got {type(dtype).__name__!r}" 219 ) 220 shape = (shape,) if isinstance(shape, int) else tuple(shape) 221 return cls( 222 os.fspath(filename), dtype, shape, block_size, mode, max_cached_blocks 223 ) 224 225 return cls.from_config( 226 save_directory=filename, 227 runtime_kwargs={ 228 "mode": mode, 229 "filename": os.fspath(filename), 230 "max_cached_blocks": max_cached_blocks, 231 }, 232 ) 233 234 @classmethod 235 def unlink(cls, filename): 236 filename = Path(filename).expanduser().resolve() 237 if filename.exists(): 238 try: 239 with cls.open(filename, "w") as _: 240 pass 241 os.unlink(filename / cls.config_name) 242 os.unlink(filename / cls.status_name) 243 os.rmdir(os.fspath(filename)) 244 except (OSError, ValueError) as exc: 245 warnings.warn(f"Failed to unlink blob at {filename!r}: {exc}") 246 return False 247 return True 248 249 @classmethod 250 def apply_param_hooks(cls, jdict): 251 jdict["shape"] = tuple(jdict["shape"]) 252 return jdict 253 254 @classmethod 255 def _getsyscachesize(cls) -> int: 256 # Get default cache size for memory-mapped blocks. Returns 1/16 of system's 257 # max_map_count to be conservative, typically ~4000, leaving room for other 258 # VMAs in the process. 259 maxsize = 65536 260 try: 261 with open("/proc/sys/vm/max_map_count", "r") as f: 262 maxsize = int(f.read().strip()) 263 except (FileNotFoundError, ValueError, PermissionError): 264 pass 265 return max(maxsize // 16, 128) 266 267 @register_to_config 268 def __init__( 269 self, 270 filename: str, 271 dtype: str, 272 shape: tuple[int, ...], 273 block_size: int, 274 mode: str, 275 max_cached_blocks: int | None = None, 276 ) -> None: 277 self.filename = filename 278 self.dtype = dtype 279 self.shape = shape 280 self.block_size = block_size 281 self.mode = mode 282 self.max_cached_blocks = max_cached_blocks or self._getsyscachesize() 283 284 self._pos = 0 285 self._closed = False 286 287 if "+" in mode: 288 self._m_rd = True 289 self._m_wr = True 290 match mode.replace("+", ""): 291 case "r": 292 self._m_rd = True 293 case "w": 294 self._m_wr = True 295 self._trunc() 296 case "a": 297 self._m_wr = True 298 self._m_ap = True 299 self._create() 300 301 self._loadstatus() 302 303 @property 304 def configpath(self) -> str: 305 return os.path.join(self.filename, self.config_name) 306 307 @property 308 def statuspath(self) -> str: 309 return os.path.join(self.filename, self.status_name) 310 311 @property 312 def closed(self) -> bool: 313 return self._closed 314 315 def __enter__(self) -> Self: 316 return self 317 318 def __exit__(self, *_) -> None: 319 self.close() 320 321 def __len__(self) -> int: 322 return self._status.len 323 324 def __getitem__(self, idx: int | slice) -> torch.Tensor: 325 if not isinstance(idx, (int, slice)): 326 raise TypeError(f"Index must be int or slice, got {type(idx).__name__!r}!") 327 if isinstance(idx, int): 328 if idx >= len(self) or idx < -len(self): 329 raise IndexError( 330 f"Index out of bounds: {idx!r} (length: {len(self)})" 331 ) 332 i, o = divmod(idx + len(self) if idx < 0 else idx, self.block_size) 333 return self._getblock(i)[o].clone() 334 335 # Although the current implementation may not be efficient, it is very easy to 336 # understand and debug. More efficient implementation requires much more complex 337 # edge case handling and is error prone. Also, I think the primary cost here is 338 # still the I/O operations, not the Python code. 339 ret = [ 340 self._getblock(bd)[[i % self.block_size for i in _is]] 341 for bd, _is in groupby( 342 range(*idx.indices(len(self))), key=lambda i: i // self.block_size 343 ) 344 ] 345 if not ret: 346 return torch.empty(0, *self.shape, dtype=getattr(torch, self.dtype)) 347 return torch.cat(ret, dim=0) 348 349 def __iter__(self) -> Iterator[torch.Tensor]: 350 for i in range(self._pos, len(self)): 351 self._pos += 1 352 yield self[i] 353 354 def _trunc(self) -> None: 355 if os.path.exists(self.filename): 356 try: 357 st = TensorBlobStatus.load(self.statuspath) 358 except FileNotFoundError as exc: 359 raise FileNotFoundError( 360 f"Status file missing for blob at {self.statuspath!r}; file corrupted!" 361 ) from exc 362 for bd in st.bds: 363 os.remove(os.path.join(self.filename, bd)) 364 self.save_config(save_directory=self.filename, overwrite=True) 365 TensorBlobStatus().dump(self.statuspath) 366 367 def _create(self) -> None: 368 if not os.path.exists(self.filename): 369 self.save_config(save_directory=self.filename) 370 TensorBlobStatus().dump(self.statuspath) 371 372 def _getblock(self, bd: str | int = -1) -> MemoryMappedTensor: 373 if not self._status.bds: 374 self._addblock() 375 if isinstance(bd, int): 376 bd = self._status.bds[bd] 377 if bd in self._memmap: 378 return self._memmap[bd] 379 380 # If cache no hit, a block is lazy-loaded into the cache. We need to 381 # avoid the __getitem__ call during return here to not increase the 382 # cache hit count a second time. 383 block = self._memmap[bd] = MemoryMappedTensor.from_filename( 384 os.path.join(self.filename, bd), 385 dtype=getattr(torch, self.dtype), 386 shape=(self.block_size, *self.shape), 387 ) 388 return block 389 390 def _isfull(self) -> bool: 391 return (not len(self) % self.block_size) and bool(len(self)) 392 393 def _addblock(self) -> MemoryMappedTensor: 394 if self._status.bds and not self._isfull(): 395 raise RuntimeError( 396 "Attempt to create a new block when working block " 397 f"is not full: length <{len(self) % self.block_size}> " 398 f"< capacity <{self.block_size}>." 399 ) 400 name = str(uuid.uuid4()) 401 mmap = MemoryMappedTensor.empty( 402 self.block_size, 403 *self.shape, 404 dtype=getattr(torch, self.dtype), 405 filename=os.path.join(self.filename, name), 406 ) 407 self._status.bds.append(name) 408 self._memmap[name] = mmap 409 return mmap 410 411 def _loadstatus(self) -> None: 412 try: 413 self._status = TensorBlobStatus.load(self.statuspath) 414 self._memmap = LRUCache(maxsize=self.max_cached_blocks) 415 if self._m_ap: 416 self._pos = len(self) 417 except FileNotFoundError as exc: 418 raise FileNotFoundError( 419 f"status file missing for blob at {self.statuspath!r}; file corrupted!" 420 ) from exc 421 422 def _checkclosed(self) -> None: 423 if self._closed: 424 raise OSError("I/O operation on closed blob.") 425 426 def _checkwritable(self) -> None: 427 if not self._m_wr: 428 raise OSError(f"Blob is not open for writing (mode='{self.mode}')") 429 self._checkclosed() 430 431 def _checkreadable(self) -> None: 432 if not self._m_rd: 433 raise OSError(f"Blob is not open for reading (mode='{self.mode}')") 434 self._checkclosed() 435 436 def tell(self) -> int: 437 self._checkclosed() 438 return self._pos 439 440 def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int: 441 self._checkclosed() 442 match whence: 443 case io.SEEK_SET: 444 _pos = pos 445 case io.SEEK_CUR: 446 _pos = self._pos + pos 447 case io.SEEK_END: 448 _pos = len(self) + pos 449 case _: 450 raise ValueError(f"Invalid whence: {whence!r}") 451 self._pos = max(min(_pos, len(self)), 0) 452 return self.tell() 453 454 def close(self) -> None: 455 if not self._closed and self._m_wr: 456 self.flush() 457 self._closed = True 458 459 def flush(self) -> None: 460 self._checkwritable() 461 self._status.dump(self.statuspath) 462 463 def read(self, size: int | None = None) -> torch.Tensor: 464 self._checkreadable() 465 end = min(self._pos + (size if size is not None else len(self)), len(self)) 466 ret = self[self._pos : end] 467 self.seek(end) 468 return ret 469 470 def write(self, ts: torch.Tensor) -> int: 471 self._checkwritable() 472 if self._m_ap: 473 self.seek(whence=io.SEEK_END) 474 ts = ts.view(-1, *self.shape) 475 nt = ts.size(0) 476 477 cnt = 0 478 while cnt < nt: 479 if self._isfull() and self._pos >= len(self): 480 self._addblock() 481 i, o = divmod(self._pos, self.block_size) 482 incr = min(self.block_size - o, nt - cnt) 483 self._getblock(i)[o : o + incr] = ts[cnt : cnt + incr] 484 485 # Update status length for new tensors exceeding the original range only, because 486 # the cursor may not always be at the EOF and the number of tensors written could 487 # be smaller than change in length 488 self._pos += incr 489 self._status.len += max(0, self._pos - len(self)) 490 491 cnt += incr 492 493 assert cnt == nt, f"Write incomplete: wrote {cnt} of {nt} tensors!" 494 return cnt 495 496 def truncate(self, pos: int | None = None) -> int: 497 self._checkwritable() 498 self.seek(pos if pos is not None else self.tell()) 499 brk = ceil(self.tell() / self.block_size) 500 for bd in self._status.bds[brk:]: 501 if bd in self._memmap: 502 del self._memmap[bd] 503 os.remove(os.path.join(self.filename, bd)) 504 self._status.bds = self._status.bds[:brk] 505 self._status.len = self.tell() 506 self.flush() 507 return self.tell() 508 509 def extend(self, other: TensorBlob, maintain_order: bool = False) -> None: 510 if self.dtype != other.dtype or self.shape != other.shape: 511 raise ValueError("Blob data types and shapes must match to extend blobs!") 512 513 self._checkwritable() 514 self.seek(whence=io.SEEK_END) 515 516 # TODO: Honestly this is a bit inefficient but I think this is rarely used. 517 if maintain_order: 518 for i in range(len(other)): 519 self.write(other[i]) 520 return 521 522 # If order is not important, we can simply copy over the complete blocks from 523 # the other blob and merge incomplete blocks. 524 if self.block_size != other.block_size: 525 raise ValueError( 526 "Block sizes must match to extend blobs in non-order-preserving mode!" 527 ) 528 529 comb = [] 530 sbrk = len(self) // self.block_size * self.block_size 531 if sbrk < len(self): 532 comb.append(self[sbrk:]) 533 obrk = len(other) // other.block_size * other.block_size 534 if obrk < len(other): 535 comb.append(other[obrk:]) 536 537 # TODO: We are directly accessing internal data structures of the other blob here. 538 self.truncate(sbrk) 539 for obd in other._status.bds[: len(other) // other.block_size]: 540 sbd = str(uuid.uuid4()) 541 shutil.copy( 542 os.path.join(other.filename, obd), os.path.join(self.filename, sbd) 543 ) 544 self._status.bds.append(sbd) 545 self._status.len += self.block_size 546 self._memmap[sbd] = MemoryMappedTensor.from_filename( 547 os.path.join(self.filename, sbd), 548 dtype=getattr(torch, self.dtype), 549 shape=(self.block_size, *self.shape), 550 ) 551 552 self.seek(whence=io.SEEK_END) 553 if comb: 554 self.write(torch.cat(comb, dim=0)) 555 self.flush()
Mixin class for automated configuration registration and IO.
Attributes
- config_name (str, default=None):
Class attribute that specifies the filename under which the config should be stored when calling
save_config. Should be overridden by the subclass. - ignore_for_config (list[str], default=[]): Class attribute that specifies a list of attributes that should not be saved in the config. Should be overridden by the subclass.
Examples
In this example, we have a model with 3 arguments:
hidden_size: The hidden size of the model._num_layers: The number of layers in the model.dropout: The dropout rate of the model.
Among the three arguments, the number of layers is implicitly ignored by the decorator because of the leading
underscore; the dropout argument is explicitly based on the specification in ignore_for_config class
variable. The hidden_size argument is registered to the config.
>>> class MyModel(ConfigMixin):
... config_name = "my_model_config.json"
... ignore_for_config = ["dropout"]
...
... @register_to_config
... def __init__(self, hidden_size: int = 768, _num_layers: int = 12, dropout: float = 0.1):
... self.hidden_size = hidden_size
... self.num_layers = _num_layers
... self.dropout = dropout # This will be ignored because of the specification in `ignore_for_config`
...
>>> model = MyModel(hidden_size=1024, _num_layers=20, dropout=0.2)
>>> model.config
mappingproxy({'__notes__': {'class_name': '__main__.MyModel', 'using_default_values': [], 'args': (), 'kwargs': {}}, 'hidden_size': 1024})
>>> model.num_layers
20
>>> model.dropout
0.2
@register_to_config
TensorBlob( filename: str, dtype: str, shape: tuple[int, ...], block_size: int, mode: str, max_cached_blocks: int | None = None)
267 @register_to_config 268 def __init__( 269 self, 270 filename: str, 271 dtype: str, 272 shape: tuple[int, ...], 273 block_size: int, 274 mode: str, 275 max_cached_blocks: int | None = None, 276 ) -> None: 277 self.filename = filename 278 self.dtype = dtype 279 self.shape = shape 280 self.block_size = block_size 281 self.mode = mode 282 self.max_cached_blocks = max_cached_blocks or self._getsyscachesize() 283 284 self._pos = 0 285 self._closed = False 286 287 if "+" in mode: 288 self._m_rd = True 289 self._m_wr = True 290 match mode.replace("+", ""): 291 case "r": 292 self._m_rd = True 293 case "w": 294 self._m_wr = True 295 self._trunc() 296 case "a": 297 self._m_wr = True 298 self._m_ap = True 299 self._create() 300 301 self._loadstatus()
@classmethod
def
open( cls, filename, mode='r', *, dtype=None, shape=None, block_size=8192, max_cached_blocks=None):
52 @classmethod 53 def open( 54 cls, 55 filename, 56 mode="r", 57 *, 58 dtype=None, 59 shape=None, 60 block_size=8192, 61 max_cached_blocks=None, 62 ): 63 r"""Open a TensorBlob with file-like interface for tensor storage. 64 65 TensorBlob provides persistent, memory-mapped storage for large collections 66 of same-shaped tensors. It uses a block-based architecture where tensors are 67 organized into fixed-size blocks for efficient I/O and memory management. 68 69 The blob is stored as a directory containing: 70 - ``.conf``: Configuration file (dtype, shape, block_size) 71 - ``.stat``: State file (length, block list) 72 - Block files: UUID-named memory-mapped tensor files 73 74 Parameters 75 ---------- 76 filename : str or Path 77 Directory path for blob storage. Supports tilde expansion (~) and 78 relative paths. 79 mode : str, default="r" 80 File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). See below for details. 81 dtype : str or torch.dtype, optional 82 Data type for tensors. Required for new blobs (modes 'w', 'w+'). 83 shape : tuple of int or int, optional 84 Shape of individual tensors. Required for new blobs (modes 'w', 'w+'). 85 block_size : int, default=8192 86 Number of tensors per memory-mapped block file. 87 max_cached_blocks : int, optional 88 Maximum number of memory-mapped blocks to keep cached. When exceeded, 89 least recently used blocks are unmapped. If None (default), uses 1/16 90 of system's max_map_count limit (typically ~4000). This limits kernel 91 VMA overhead for blobs with many blocks. 92 93 Returns 94 ------- 95 TensorBlob 96 Opened blob object. Use with context manager for automatic cleanup. 97 98 Raises 99 ------ 100 FileNotFoundError 101 If mode is 'r', 'r+', 'a', or 'a+' and blob doesn't exist. 102 ValueError 103 If creating new blob without dtype or shape, or if mode is invalid. 104 TypeError 105 If dtype is neither string nor torch.dtype. 106 107 Examples 108 -------- 109 Creating a new blob and writing data: 110 111 >>> import torch 112 >>> from tensorblob import TensorBlob 113 >>> 114 >>> with TensorBlob.open("data/embeddings", "w", 115 ... dtype="float32", shape=(768,)) as blob: 116 ... embeddings = torch.randn(1000, 768) 117 ... blob.write(embeddings) 118 ... print(f"Wrote {len(blob)} tensors") 119 Wrote 1000 tensors 120 121 Reading from existing blob: 122 123 >>> with TensorBlob.open("data/embeddings", "r") as blob: 124 ... all_data = blob.read() 125 ... print(all_data.shape) 126 torch.Size([1000, 768]) 127 128 Appending to existing blob: 129 130 >>> with TensorBlob.open("data/embeddings", "a") as blob: 131 ... new_data = torch.randn(100, 768) 132 ... blob.write(new_data) 133 ... print(f"Total: {len(blob)}") 134 Total: 1100 135 136 Read and update with r+ mode: 137 138 >>> with TensorBlob.open("data/embeddings", "r+") as blob: 139 ... first_10 = blob.read(size=10) 140 ... blob.seek(5) 141 ... blob.write(torch.ones(3, 768)) # Overwrite at position 5 142 143 Custom block size for large tensors: 144 145 >>> with TensorBlob.open("data/images", "w", 146 ... dtype=torch.float32, 147 ... shape=(3, 1024, 1024), 148 ... block_size=256) as blob: 149 ... images = torch.randn(1000, 3, 1024, 1024) 150 ... blob.write(images) 151 152 Custom cache size for large-scale random access: 153 154 >>> # Increase cache for better random access performance 155 >>> with TensorBlob.open("data/embeddings", "r", 156 ... max_cached_blocks=10000) as blob: 157 ... for idx in random_indices: 158 ... embedding = blob[idx] # Frequently accessed blocks stay cached 159 160 >>> # Decrease cache for memory-constrained environments 161 >>> with TensorBlob.open("data/features", "r", 162 ... max_cached_blocks=100) as blob: 163 ... for feature in blob: # Sequential access works fine 164 ... process(feature) 165 166 File Access Modes 167 ----------------- 168 Similar to Python's built-in open(), supports the following modes: 169 170 Basic modes: 171 - 'r' : Read-only. Blob must exist. Position starts at beginning. 172 - 'w' : Write-only. Creates new or truncates existing. Position at start. **If the blob already exists, 173 truncation will ignore any other parameters supplied and rely on existing configuration.** 174 - 'a' : Append-only. Blob must exist. Position starts at end. 175 All writes go to end regardless of seek position. 176 177 Update modes (with '+'): 178 - 'r+' : Read and write. Blob must exist. Position at start. 179 Can overwrite existing data or extend at end. 180 - 'w+' : Read and write. Creates new or truncates existing. Position at start. 181 - 'a+' : Read and append. Blob must exist. Position at end. 182 Reads allowed anywhere, writes always append to end. 183 184 Data Type and Shape 185 ------------------- 186 All tensors in a blob must have the same dtype and shape. These are 187 specified when creating a new blob (modes 'w', 'w+') and stored in 188 the configuration file. When opening existing blobs, dtype and shape 189 are loaded automatically. 190 191 Supported dtypes: "float32", "float64", "int32", "int64", "bool", etc. 192 Can also use torch.dtype objects like torch.float32. 193 194 Shape can be: 195 - Single integer: shape=10 creates 1D tensors of shape (10,) 196 - Tuple: shape=(3, 224, 224) creates 3D tensors 197 """ 198 modes = set(mode) 199 if modes - set("raw+") or len(mode) > len(modes): 200 raise ValueError(f"Invalid mode: {mode}") 201 if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1: 202 raise ValueError( 203 f"Must have exactly one of read/write/append mode and at most one plus: {mode}" 204 ) 205 206 filename = Path(filename).expanduser().resolve() 207 if not filename.exists(): 208 if "r" in modes or "a" in modes: 209 raise FileNotFoundError(f"Blob not found: {filename!r}") 210 if dtype is None or shape is None: 211 raise ValueError( 212 f"Arguments ``dtype`` and ``shape`` are required for new blob; got: {dtype!r} and {shape!r}" 213 ) 214 if isinstance(dtype, torch.dtype): 215 dtype = str(dtype).split(".").pop() 216 elif not isinstance(dtype, str): 217 raise TypeError( 218 f"dtype must be str or torch.dtype, got {type(dtype).__name__!r}" 219 ) 220 shape = (shape,) if isinstance(shape, int) else tuple(shape) 221 return cls( 222 os.fspath(filename), dtype, shape, block_size, mode, max_cached_blocks 223 ) 224 225 return cls.from_config( 226 save_directory=filename, 227 runtime_kwargs={ 228 "mode": mode, 229 "filename": os.fspath(filename), 230 "max_cached_blocks": max_cached_blocks, 231 }, 232 )
Open a TensorBlob with file-like interface for tensor storage.
TensorBlob provides persistent, memory-mapped storage for large collections of same-shaped tensors. It uses a block-based architecture where tensors are organized into fixed-size blocks for efficient I/O and memory management.
The blob is stored as a directory containing:
.conf: Configuration file (dtype, shape, block_size).stat: State file (length, block list)- Block files: UUID-named memory-mapped tensor files
Parameters
- filename (str or Path): Directory path for blob storage. Supports tilde expansion (~) and relative paths.
- mode (str, default="r"): File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). See below for details.
- dtype (str or torch.dtype, optional): Data type for tensors. Required for new blobs (modes 'w', 'w+').
- shape (tuple of int or int, optional): Shape of individual tensors. Required for new blobs (modes 'w', 'w+').
- block_size (int, default=8192): Number of tensors per memory-mapped block file.
- max_cached_blocks (int, optional): Maximum number of memory-mapped blocks to keep cached. When exceeded, least recently used blocks are unmapped. If None (default), uses 1/16 of system's max_map_count limit (typically ~4000). This limits kernel VMA overhead for blobs with many blocks.
Returns
- TensorBlob: Opened blob object. Use with context manager for automatic cleanup.
Raises
- FileNotFoundError: If mode is 'r', 'r+', 'a', or 'a+' and blob doesn't exist.
- ValueError: If creating new blob without dtype or shape, or if mode is invalid.
- TypeError: If dtype is neither string nor torch.dtype.
Examples
Creating a new blob and writing data:
>>> import torch
>>> from tensorblob import TensorBlob
>>>
>>> with TensorBlob.open("data/embeddings", "w",
... dtype="float32", shape=(768,)) as blob:
... embeddings = torch.randn(1000, 768)
... blob.write(embeddings)
... print(f"Wrote {len(blob)} tensors")
Wrote 1000 tensors
Reading from existing blob:
>>> with TensorBlob.open("data/embeddings", "r") as blob:
... all_data = blob.read()
... print(all_data.shape)
torch.Size([1000, 768])
Appending to existing blob:
>>> with TensorBlob.open("data/embeddings", "a") as blob:
... new_data = torch.randn(100, 768)
... blob.write(new_data)
... print(f"Total: {len(blob)}")
Total: 1100
Read and update with r+ mode:
>>> with TensorBlob.open("data/embeddings", "r+") as blob:
... first_10 = blob.read(size=10)
... blob.seek(5)
... blob.write(torch.ones(3, 768)) # Overwrite at position 5
Custom block size for large tensors:
>>> with TensorBlob.open("data/images", "w",
... dtype=torch.float32,
... shape=(3, 1024, 1024),
... block_size=256) as blob:
... images = torch.randn(1000, 3, 1024, 1024)
... blob.write(images)
Custom cache size for large-scale random access:
>>> # Increase cache for better random access performance
>>> with TensorBlob.open("data/embeddings", "r",
... max_cached_blocks=10000) as blob:
... for idx in random_indices:
... embedding = blob[idx] # Frequently accessed blocks stay cached
>>> # Decrease cache for memory-constrained environments
>>> with TensorBlob.open("data/features", "r",
... max_cached_blocks=100) as blob:
... for feature in blob: # Sequential access works fine
... process(feature)
File Access Modes
Similar to Python's built-in open(), supports the following modes:
Basic modes:
- 'r' : Read-only. Blob must exist. Position starts at beginning.
- 'w' : Write-only. Creates new or truncates existing. Position at start. If the blob already exists, truncation will ignore any other parameters supplied and rely on existing configuration.
- 'a' : Append-only. Blob must exist. Position starts at end. All writes go to end regardless of seek position.
Update modes (with '+'):
- 'r+' : Read and write. Blob must exist. Position at start. Can overwrite existing data or extend at end.
- 'w+' : Read and write. Creates new or truncates existing. Position at start.
- 'a+' : Read and append. Blob must exist. Position at end. Reads allowed anywhere, writes always append to end.
Data Type and Shape
All tensors in a blob must have the same dtype and shape. These are specified when creating a new blob (modes 'w', 'w+') and stored in the configuration file. When opening existing blobs, dtype and shape are loaded automatically.
Supported dtypes: "float32", "float64", "int32", "int64", "bool", etc. Can also use torch.dtype objects like torch.float32.
Shape can be:
- Single integer: shape=10 creates 1D tensors of shape (10,)
- Tuple: shape=(3, 224, 224) creates 3D tensors
@classmethod
def
unlink(cls, filename):
234 @classmethod 235 def unlink(cls, filename): 236 filename = Path(filename).expanduser().resolve() 237 if filename.exists(): 238 try: 239 with cls.open(filename, "w") as _: 240 pass 241 os.unlink(filename / cls.config_name) 242 os.unlink(filename / cls.status_name) 243 os.rmdir(os.fspath(filename)) 244 except (OSError, ValueError) as exc: 245 warnings.warn(f"Failed to unlink blob at {filename!r}: {exc}") 246 return False 247 return True
@classmethod
def
apply_param_hooks(cls, jdict):
249 @classmethod 250 def apply_param_hooks(cls, jdict): 251 jdict["shape"] = tuple(jdict["shape"]) 252 return jdict
Apply post-processing hooks to the JSON dictionary.
orjson.loads only decode configs to primitive types, which may not be directly
consumable by the class initializer. For instance, a dataclass object will be
loaded as a dictionary. Therefore, this method is intended to be overridden by the
subclass to perform additional post-processing on the loaded config dictionary.
Note that, it is highly discouraged to abuse this method to deserialize complex objects
and one should consider using runtime_kwargs argument of from_config instead,
to explicitly pass the complex objects to the class initializer.
By default, this method returns the input dictionary unchanged.
Parameters
- jdict (dict[str, Any]): The config dictionary after deserialization.
Returns
- dict[str, Any]: The config dictionary after post-processing.
def
seek(self, pos: int = 0, whence: int = 0) -> int:
440 def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int: 441 self._checkclosed() 442 match whence: 443 case io.SEEK_SET: 444 _pos = pos 445 case io.SEEK_CUR: 446 _pos = self._pos + pos 447 case io.SEEK_END: 448 _pos = len(self) + pos 449 case _: 450 raise ValueError(f"Invalid whence: {whence!r}") 451 self._pos = max(min(_pos, len(self)), 0) 452 return self.tell()
def
write(self, ts: torch.Tensor) -> int:
470 def write(self, ts: torch.Tensor) -> int: 471 self._checkwritable() 472 if self._m_ap: 473 self.seek(whence=io.SEEK_END) 474 ts = ts.view(-1, *self.shape) 475 nt = ts.size(0) 476 477 cnt = 0 478 while cnt < nt: 479 if self._isfull() and self._pos >= len(self): 480 self._addblock() 481 i, o = divmod(self._pos, self.block_size) 482 incr = min(self.block_size - o, nt - cnt) 483 self._getblock(i)[o : o + incr] = ts[cnt : cnt + incr] 484 485 # Update status length for new tensors exceeding the original range only, because 486 # the cursor may not always be at the EOF and the number of tensors written could 487 # be smaller than change in length 488 self._pos += incr 489 self._status.len += max(0, self._pos - len(self)) 490 491 cnt += incr 492 493 assert cnt == nt, f"Write incomplete: wrote {cnt} of {nt} tensors!" 494 return cnt
def
truncate(self, pos: int | None = None) -> int:
496 def truncate(self, pos: int | None = None) -> int: 497 self._checkwritable() 498 self.seek(pos if pos is not None else self.tell()) 499 brk = ceil(self.tell() / self.block_size) 500 for bd in self._status.bds[brk:]: 501 if bd in self._memmap: 502 del self._memmap[bd] 503 os.remove(os.path.join(self.filename, bd)) 504 self._status.bds = self._status.bds[:brk] 505 self._status.len = self.tell() 506 self.flush() 507 return self.tell()
509 def extend(self, other: TensorBlob, maintain_order: bool = False) -> None: 510 if self.dtype != other.dtype or self.shape != other.shape: 511 raise ValueError("Blob data types and shapes must match to extend blobs!") 512 513 self._checkwritable() 514 self.seek(whence=io.SEEK_END) 515 516 # TODO: Honestly this is a bit inefficient but I think this is rarely used. 517 if maintain_order: 518 for i in range(len(other)): 519 self.write(other[i]) 520 return 521 522 # If order is not important, we can simply copy over the complete blocks from 523 # the other blob and merge incomplete blocks. 524 if self.block_size != other.block_size: 525 raise ValueError( 526 "Block sizes must match to extend blobs in non-order-preserving mode!" 527 ) 528 529 comb = [] 530 sbrk = len(self) // self.block_size * self.block_size 531 if sbrk < len(self): 532 comb.append(self[sbrk:]) 533 obrk = len(other) // other.block_size * other.block_size 534 if obrk < len(other): 535 comb.append(other[obrk:]) 536 537 # TODO: We are directly accessing internal data structures of the other blob here. 538 self.truncate(sbrk) 539 for obd in other._status.bds[: len(other) // other.block_size]: 540 sbd = str(uuid.uuid4()) 541 shutil.copy( 542 os.path.join(other.filename, obd), os.path.join(self.filename, sbd) 543 ) 544 self._status.bds.append(sbd) 545 self._status.len += self.block_size 546 self._memmap[sbd] = MemoryMappedTensor.from_filename( 547 os.path.join(self.filename, sbd), 548 dtype=getattr(torch, self.dtype), 549 shape=(self.block_size, *self.shape), 550 ) 551 552 self.seek(whence=io.SEEK_END) 553 if comb: 554 self.write(torch.cat(comb, dim=0)) 555 self.flush()
class
TensorDB(configmixin._core.ConfigMixin):
38class TensorDB(ConfigMixin): 39 _m_rd = False 40 _m_wr = False 41 _m_ap = False 42 43 status_name = ".stat" 44 config_name = ".conf" 45 ignore_for_config: ClassVar[list[str]] = ["filename", "mode", "max_cached_blocks"] 46 47 @classmethod 48 def open( 49 cls, 50 filename, 51 mode="r", 52 *, 53 schema=None, 54 block_size=8192, 55 max_cached_blocks=None, 56 ): 57 r"""Open a TensorDB with file-like interface for multi-field tensor storage. 58 59 TensorDB provides persistent, row-aligned storage for heterogeneous 60 (multi-field) tensor collections. Each field is stored as a plain 61 :class:`TensorBlob` in a subdirectory, and TensorDB keeps the row 62 orders of all fields aligned. Rows are dense: every write must supply 63 every field with the same row count. 64 65 The database is stored as a directory containing: 66 - ``.conf``: Schema file (field names, dtypes, shapes) 67 - ``.stat``: State file (committed row count) 68 - Field-named subdirectories: One TensorBlob per field 69 70 Parameters 71 ---------- 72 filename : str or Path 73 Directory path for database storage. Supports tilde expansion (~) 74 and relative paths. 75 mode : str, default="r" 76 File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). Behaves like 77 :meth:`TensorBlob.open`; the mode is applied to every field. 78 schema : dict, optional 79 Mapping of field names to ``(dtype, shape)`` pairs, e.g., 80 ``{"price": ("float32", 1), "embed": (torch.float16, (768,))}``. 81 Required for new databases (modes 'w', 'w+'). Fixed at creation; 82 loaded automatically when opening existing databases. 83 block_size : int, default=8192 84 Number of rows per memory-mapped block file, applied to all fields. 85 max_cached_blocks : int, optional 86 Maximum number of memory-mapped blocks to keep cached per field. 87 If None (default), uses 1/16 of system's max_map_count limit. 88 89 Returns 90 ------- 91 TensorDB 92 Opened database object. Use with context manager for automatic 93 cleanup. 94 95 Raises 96 ------ 97 FileNotFoundError 98 If mode is 'r', 'r+', 'a', or 'a+' and database doesn't exist. 99 ValueError 100 If creating new database without schema, if the schema is 101 malformed, or if mode is invalid. 102 TypeError 103 If a dtype is neither string nor torch.dtype. 104 105 Examples 106 -------- 107 Creating a new database and writing dense rows: 108 109 >>> import torch 110 >>> from tensorblob import TensorDB 111 >>> 112 >>> with TensorDB.open("events.db", "w", 113 ... schema={"price": ("float32", 1), 114 ... "embed": ("float32", 768)}) as db: 115 ... db.write({"price": torch.randn(1000, 1), 116 ... "embed": torch.randn(1000, 768)}) 117 ... print(f"Wrote {len(db)} rows") 118 Wrote 1000 rows 119 120 Reading rows back, aligned across fields: 121 122 >>> with TensorDB.open("events.db", "r") as db: 123 ... row = db[42] # {"price": (1,), "embed": (768,)} 124 ... batch = db[10:100] # {"price": (90, 1), "embed": (90, 768)} 125 126 Notes 127 ----- 128 Writes commit the row count only after all fields are written. If a 129 crash leaves fields longer than the committed count, the next open 130 reports the committed (minimum) length and, for writable modes, 131 truncates the stray rows back to restore alignment. 132 """ 133 modes = set(mode) 134 if modes - set("raw+") or len(mode) > len(modes): 135 raise ValueError(f"Invalid mode: {mode}") 136 if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1: 137 raise ValueError( 138 f"Must have exactly one of read/write/append mode and at most one plus: {mode}" 139 ) 140 141 filename = Path(filename).expanduser().resolve() 142 if not filename.exists(): 143 if "r" in modes or "a" in modes: 144 raise FileNotFoundError(f"Database not found: {filename!r}") 145 if schema is None: 146 raise ValueError("Argument ``schema`` is required for new database!") 147 schema = cls._normalize_schema(schema) 148 return cls(os.fspath(filename), schema, block_size, mode, max_cached_blocks) 149 150 return cls.from_config( 151 save_directory=filename, 152 runtime_kwargs={ 153 "mode": mode, 154 "filename": os.fspath(filename), 155 "max_cached_blocks": max_cached_blocks, 156 }, 157 ) 158 159 @classmethod 160 def unlink(cls, filename): 161 filename = Path(filename).expanduser().resolve() 162 if filename.exists(): 163 try: 164 shutil.rmtree(filename) 165 except OSError as exc: 166 warnings.warn(f"Failed to unlink database at {filename!r}: {exc}") 167 return False 168 return True 169 170 @classmethod 171 def _normalize_schema(cls, schema): 172 if not isinstance(schema, dict) or not schema: 173 raise ValueError( 174 "Schema must be a non-empty dict mapping field names to (dtype, shape)!" 175 ) 176 norm = {} 177 for name, spec in schema.items(): 178 if ( 179 not isinstance(name, str) 180 or not name 181 or name.startswith(".") 182 or "/" in name 183 or os.sep in name 184 or (os.altsep and os.altsep in name) 185 ): 186 raise ValueError(f"Invalid field name: {name!r}") 187 dtype, shape = spec 188 if isinstance(dtype, torch.dtype): 189 dtype = str(dtype).split(".").pop() 190 elif not isinstance(dtype, str): 191 raise TypeError( 192 f"dtype must be str or torch.dtype, got {type(dtype).__name__!r}" 193 ) 194 shape = (shape,) if isinstance(shape, int) else tuple(shape) 195 norm[name] = (dtype, shape) 196 return norm 197 198 @classmethod 199 def apply_param_hooks(cls, jdict): 200 jdict["schema"] = { 201 name: (dtype, tuple(shape)) 202 for name, (dtype, shape) in jdict["schema"].items() 203 } 204 return jdict 205 206 @register_to_config 207 def __init__( 208 self, 209 filename: str, 210 schema: dict[str, tuple[str, tuple[int, ...]]], 211 block_size: int, 212 mode: str, 213 max_cached_blocks: int | None = None, 214 ) -> None: 215 self.filename = filename 216 self.schema = schema 217 self.block_size = block_size 218 self.mode = mode 219 self.max_cached_blocks = max_cached_blocks 220 221 self._closed = False 222 223 if "+" in mode: 224 self._m_rd = True 225 self._m_wr = True 226 match mode.replace("+", ""): 227 case "r": 228 self._m_rd = True 229 case "w": 230 self._m_wr = True 231 case "a": 232 self._m_wr = True 233 self._m_ap = True 234 235 isnew = not os.path.exists(self.filename) 236 if isnew: 237 os.makedirs(self.filename) 238 self.save_config(save_directory=self.filename) 239 self._cols = { 240 name: TensorBlob.open( 241 os.path.join(self.filename, name), 242 mode, 243 dtype=dtype, 244 shape=shape, 245 block_size=block_size, 246 max_cached_blocks=max_cached_blocks, 247 ) 248 for name, (dtype, shape) in self.schema.items() 249 } 250 251 # For new or truncated databases the committed count starts at zero; no 252 # repair is needed since the columns were just (re)initialized above. 253 if isnew or "w" in mode: 254 self._status = TensorDBStatus() 255 self._status.dump(self.statuspath) 256 else: 257 self._loadstatus() 258 259 @property 260 def configpath(self) -> str: 261 return os.path.join(self.filename, self.config_name) 262 263 @property 264 def statuspath(self) -> str: 265 return os.path.join(self.filename, self.status_name) 266 267 @property 268 def closed(self) -> bool: 269 return self._closed 270 271 def __enter__(self) -> Self: 272 return self 273 274 def __exit__(self, *_) -> None: 275 self.close() 276 277 def __len__(self) -> int: 278 return self._status.len 279 280 def __getitem__(self, idx: int | slice) -> dict[str, torch.Tensor]: 281 if not isinstance(idx, (int, slice)): 282 raise TypeError(f"Index must be int or slice, got {type(idx).__name__!r}!") 283 return {name: col[idx] for name, col in self._cols.items()} 284 285 def __iter__(self) -> Iterator[dict[str, torch.Tensor]]: 286 for i in range(self.tell(), len(self)): 287 self.seek(i + 1) 288 yield self[i] 289 290 def _loadstatus(self) -> None: 291 try: 292 self._status = TensorDBStatus.load(self.statuspath) 293 except FileNotFoundError as exc: 294 raise FileNotFoundError( 295 f"Status file missing for database at {self.statuspath!r}; file corrupted!" 296 ) from exc 297 298 # The committed row count is the source of truth. A crash mid-write can 299 # leave some columns longer than the committed count; report the minimum 300 # and, if writable, truncate stray rows back to restore alignment. 301 target = min([self._status.len] + [len(col) for col in self._cols.values()]) 302 if target != self._status.len or any( 303 len(col) != target for col in self._cols.values() 304 ): 305 warnings.warn( 306 f"Inconsistent column lengths detected for database at {self.filename!r}; " 307 f"reporting {target} committed rows." 308 ) 309 self._status.len = target 310 if self._m_wr: 311 for col in self._cols.values(): 312 if len(col) != target: 313 col.truncate(target) 314 self._status.dump(self.statuspath) 315 316 def _checkclosed(self) -> None: 317 if self._closed: 318 raise OSError("I/O operation on closed database.") 319 320 def _checkwritable(self) -> None: 321 if not self._m_wr: 322 raise OSError(f"Database is not open for writing (mode='{self.mode}')") 323 self._checkclosed() 324 325 def _checkreadable(self) -> None: 326 if not self._m_rd: 327 raise OSError(f"Database is not open for reading (mode='{self.mode}')") 328 self._checkclosed() 329 330 def tell(self) -> int: 331 self._checkclosed() 332 return next(iter(self._cols.values())).tell() 333 334 def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int: 335 self._checkclosed() 336 for col in self._cols.values(): 337 col.seek(pos, whence) 338 return self.tell() 339 340 def close(self) -> None: 341 if self._closed: 342 return 343 for col in self._cols.values(): 344 col.close() 345 if self._m_wr: 346 self._status.dump(self.statuspath) 347 self._closed = True 348 349 def flush(self) -> None: 350 self._checkwritable() 351 for col in self._cols.values(): 352 col.flush() 353 self._status.dump(self.statuspath) 354 355 def read(self, size: int | None = None) -> dict[str, torch.Tensor]: 356 self._checkreadable() 357 # Clamp at the committed row count so uncommitted trailing rows left by 358 # an interrupted write are never visible. 359 remaining = len(self) - self.tell() 360 size = remaining if size is None else min(size, remaining) 361 if size <= 0: 362 return { 363 name: torch.empty(0, *shape, dtype=getattr(torch, dtype)) 364 for name, (dtype, shape) in self.schema.items() 365 } 366 return {name: col.read(size) for name, col in self._cols.items()} 367 368 def write(self, rows: dict[str, torch.Tensor]) -> int: 369 self._checkwritable() 370 if not isinstance(rows, dict): 371 raise TypeError( 372 f"Rows must be a dict mapping field names to tensors, got {type(rows).__name__!r}!" 373 ) 374 missing = sorted(self.schema.keys() - rows.keys()) 375 extra = sorted(rows.keys() - self.schema.keys()) 376 if missing or extra: 377 raise ValueError( 378 "Dense writes require exactly the schema fields; " 379 f"missing: {missing!r}, unexpected: {extra!r}" 380 ) 381 382 nts = { 383 name: ts.view(-1, *self.schema[name][1]).size(0) 384 for name, ts in rows.items() 385 } 386 if len(set(nts.values())) != 1: 387 raise ValueError(f"All fields must have the same row count; got: {nts!r}") 388 nt = next(iter(nts.values())) 389 390 # Columns are written first and the committed row count is bumped only 391 # afterwards, so an interrupted write is rolled back on the next open. 392 for name, col in self._cols.items(): 393 col.write(rows[name]) 394 self._status.len = len(next(iter(self._cols.values()))) 395 return nt 396 397 def truncate(self, pos: int | None = None) -> int: 398 self._checkwritable() 399 for col in self._cols.values(): 400 col.truncate(pos) 401 self._status.len = self.tell() 402 self._status.dump(self.statuspath) 403 return self.tell() 404 405 def extend(self, other: TensorDB, maintain_order: bool = False) -> None: 406 if set(self.schema) != set(other.schema): 407 raise ValueError("Schema fields must match to extend databases!") 408 self._checkwritable() 409 for name, col in self._cols.items(): 410 col.extend(other._cols[name], maintain_order=maintain_order) 411 self._status.len = len(next(iter(self._cols.values()))) 412 self._status.dump(self.statuspath)
Mixin class for automated configuration registration and IO.
Attributes
- config_name (str, default=None):
Class attribute that specifies the filename under which the config should be stored when calling
save_config. Should be overridden by the subclass. - ignore_for_config (list[str], default=[]): Class attribute that specifies a list of attributes that should not be saved in the config. Should be overridden by the subclass.
Examples
In this example, we have a model with 3 arguments:
hidden_size: The hidden size of the model._num_layers: The number of layers in the model.dropout: The dropout rate of the model.
Among the three arguments, the number of layers is implicitly ignored by the decorator because of the leading
underscore; the dropout argument is explicitly based on the specification in ignore_for_config class
variable. The hidden_size argument is registered to the config.
>>> class MyModel(ConfigMixin):
... config_name = "my_model_config.json"
... ignore_for_config = ["dropout"]
...
... @register_to_config
... def __init__(self, hidden_size: int = 768, _num_layers: int = 12, dropout: float = 0.1):
... self.hidden_size = hidden_size
... self.num_layers = _num_layers
... self.dropout = dropout # This will be ignored because of the specification in `ignore_for_config`
...
>>> model = MyModel(hidden_size=1024, _num_layers=20, dropout=0.2)
>>> model.config
mappingproxy({'__notes__': {'class_name': '__main__.MyModel', 'using_default_values': [], 'args': (), 'kwargs': {}}, 'hidden_size': 1024})
>>> model.num_layers
20
>>> model.dropout
0.2
@register_to_config
TensorDB( filename: str, schema: dict[str, tuple[str, tuple[int, ...]]], block_size: int, mode: str, max_cached_blocks: int | None = None)
206 @register_to_config 207 def __init__( 208 self, 209 filename: str, 210 schema: dict[str, tuple[str, tuple[int, ...]]], 211 block_size: int, 212 mode: str, 213 max_cached_blocks: int | None = None, 214 ) -> None: 215 self.filename = filename 216 self.schema = schema 217 self.block_size = block_size 218 self.mode = mode 219 self.max_cached_blocks = max_cached_blocks 220 221 self._closed = False 222 223 if "+" in mode: 224 self._m_rd = True 225 self._m_wr = True 226 match mode.replace("+", ""): 227 case "r": 228 self._m_rd = True 229 case "w": 230 self._m_wr = True 231 case "a": 232 self._m_wr = True 233 self._m_ap = True 234 235 isnew = not os.path.exists(self.filename) 236 if isnew: 237 os.makedirs(self.filename) 238 self.save_config(save_directory=self.filename) 239 self._cols = { 240 name: TensorBlob.open( 241 os.path.join(self.filename, name), 242 mode, 243 dtype=dtype, 244 shape=shape, 245 block_size=block_size, 246 max_cached_blocks=max_cached_blocks, 247 ) 248 for name, (dtype, shape) in self.schema.items() 249 } 250 251 # For new or truncated databases the committed count starts at zero; no 252 # repair is needed since the columns were just (re)initialized above. 253 if isnew or "w" in mode: 254 self._status = TensorDBStatus() 255 self._status.dump(self.statuspath) 256 else: 257 self._loadstatus()
@classmethod
def
open( cls, filename, mode='r', *, schema=None, block_size=8192, max_cached_blocks=None):
47 @classmethod 48 def open( 49 cls, 50 filename, 51 mode="r", 52 *, 53 schema=None, 54 block_size=8192, 55 max_cached_blocks=None, 56 ): 57 r"""Open a TensorDB with file-like interface for multi-field tensor storage. 58 59 TensorDB provides persistent, row-aligned storage for heterogeneous 60 (multi-field) tensor collections. Each field is stored as a plain 61 :class:`TensorBlob` in a subdirectory, and TensorDB keeps the row 62 orders of all fields aligned. Rows are dense: every write must supply 63 every field with the same row count. 64 65 The database is stored as a directory containing: 66 - ``.conf``: Schema file (field names, dtypes, shapes) 67 - ``.stat``: State file (committed row count) 68 - Field-named subdirectories: One TensorBlob per field 69 70 Parameters 71 ---------- 72 filename : str or Path 73 Directory path for database storage. Supports tilde expansion (~) 74 and relative paths. 75 mode : str, default="r" 76 File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). Behaves like 77 :meth:`TensorBlob.open`; the mode is applied to every field. 78 schema : dict, optional 79 Mapping of field names to ``(dtype, shape)`` pairs, e.g., 80 ``{"price": ("float32", 1), "embed": (torch.float16, (768,))}``. 81 Required for new databases (modes 'w', 'w+'). Fixed at creation; 82 loaded automatically when opening existing databases. 83 block_size : int, default=8192 84 Number of rows per memory-mapped block file, applied to all fields. 85 max_cached_blocks : int, optional 86 Maximum number of memory-mapped blocks to keep cached per field. 87 If None (default), uses 1/16 of system's max_map_count limit. 88 89 Returns 90 ------- 91 TensorDB 92 Opened database object. Use with context manager for automatic 93 cleanup. 94 95 Raises 96 ------ 97 FileNotFoundError 98 If mode is 'r', 'r+', 'a', or 'a+' and database doesn't exist. 99 ValueError 100 If creating new database without schema, if the schema is 101 malformed, or if mode is invalid. 102 TypeError 103 If a dtype is neither string nor torch.dtype. 104 105 Examples 106 -------- 107 Creating a new database and writing dense rows: 108 109 >>> import torch 110 >>> from tensorblob import TensorDB 111 >>> 112 >>> with TensorDB.open("events.db", "w", 113 ... schema={"price": ("float32", 1), 114 ... "embed": ("float32", 768)}) as db: 115 ... db.write({"price": torch.randn(1000, 1), 116 ... "embed": torch.randn(1000, 768)}) 117 ... print(f"Wrote {len(db)} rows") 118 Wrote 1000 rows 119 120 Reading rows back, aligned across fields: 121 122 >>> with TensorDB.open("events.db", "r") as db: 123 ... row = db[42] # {"price": (1,), "embed": (768,)} 124 ... batch = db[10:100] # {"price": (90, 1), "embed": (90, 768)} 125 126 Notes 127 ----- 128 Writes commit the row count only after all fields are written. If a 129 crash leaves fields longer than the committed count, the next open 130 reports the committed (minimum) length and, for writable modes, 131 truncates the stray rows back to restore alignment. 132 """ 133 modes = set(mode) 134 if modes - set("raw+") or len(mode) > len(modes): 135 raise ValueError(f"Invalid mode: {mode}") 136 if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1: 137 raise ValueError( 138 f"Must have exactly one of read/write/append mode and at most one plus: {mode}" 139 ) 140 141 filename = Path(filename).expanduser().resolve() 142 if not filename.exists(): 143 if "r" in modes or "a" in modes: 144 raise FileNotFoundError(f"Database not found: {filename!r}") 145 if schema is None: 146 raise ValueError("Argument ``schema`` is required for new database!") 147 schema = cls._normalize_schema(schema) 148 return cls(os.fspath(filename), schema, block_size, mode, max_cached_blocks) 149 150 return cls.from_config( 151 save_directory=filename, 152 runtime_kwargs={ 153 "mode": mode, 154 "filename": os.fspath(filename), 155 "max_cached_blocks": max_cached_blocks, 156 }, 157 )
Open a TensorDB with file-like interface for multi-field tensor storage.
TensorDB provides persistent, row-aligned storage for heterogeneous
(multi-field) tensor collections. Each field is stored as a plain
TensorBlob in a subdirectory, and TensorDB keeps the row
orders of all fields aligned. Rows are dense: every write must supply
every field with the same row count.
The database is stored as a directory containing:
.conf: Schema file (field names, dtypes, shapes).stat: State file (committed row count)- Field-named subdirectories: One TensorBlob per field
Parameters
- filename (str or Path): Directory path for database storage. Supports tilde expansion (~) and relative paths.
- mode (str, default="r"):
File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). Behaves like
TensorBlob.open(); the mode is applied to every field. - schema (dict, optional):
Mapping of field names to
(dtype, shape)pairs, e.g.,{"price": ("float32", 1), "embed": (torch.float16, (768,))}. Required for new databases (modes 'w', 'w+'). Fixed at creation; loaded automatically when opening existing databases. - block_size (int, default=8192): Number of rows per memory-mapped block file, applied to all fields.
- max_cached_blocks (int, optional): Maximum number of memory-mapped blocks to keep cached per field. If None (default), uses 1/16 of system's max_map_count limit.
Returns
- TensorDB: Opened database object. Use with context manager for automatic cleanup.
Raises
- FileNotFoundError: If mode is 'r', 'r+', 'a', or 'a+' and database doesn't exist.
- ValueError: If creating new database without schema, if the schema is malformed, or if mode is invalid.
- TypeError: If a dtype is neither string nor torch.dtype.
Examples
Creating a new database and writing dense rows:
>>> import torch
>>> from tensorblob import TensorDB
>>>
>>> with TensorDB.open("events.db", "w",
... schema={"price": ("float32", 1),
... "embed": ("float32", 768)}) as db:
... db.write({"price": torch.randn(1000, 1),
... "embed": torch.randn(1000, 768)})
... print(f"Wrote {len(db)} rows")
Wrote 1000 rows
Reading rows back, aligned across fields:
>>> with TensorDB.open("events.db", "r") as db:
... row = db[42] # {"price": (1,), "embed": (768,)}
... batch = db[10:100] # {"price": (90, 1), "embed": (90, 768)}
Notes
Writes commit the row count only after all fields are written. If a crash leaves fields longer than the committed count, the next open reports the committed (minimum) length and, for writable modes, truncates the stray rows back to restore alignment.
@classmethod
def
apply_param_hooks(cls, jdict):
198 @classmethod 199 def apply_param_hooks(cls, jdict): 200 jdict["schema"] = { 201 name: (dtype, tuple(shape)) 202 for name, (dtype, shape) in jdict["schema"].items() 203 } 204 return jdict
Apply post-processing hooks to the JSON dictionary.
orjson.loads only decode configs to primitive types, which may not be directly
consumable by the class initializer. For instance, a dataclass object will be
loaded as a dictionary. Therefore, this method is intended to be overridden by the
subclass to perform additional post-processing on the loaded config dictionary.
Note that, it is highly discouraged to abuse this method to deserialize complex objects
and one should consider using runtime_kwargs argument of from_config instead,
to explicitly pass the complex objects to the class initializer.
By default, this method returns the input dictionary unchanged.
Parameters
- jdict (dict[str, Any]): The config dictionary after deserialization.
Returns
- dict[str, Any]: The config dictionary after post-processing.
def
read(self, size: int | None = None) -> dict[str, torch.Tensor]:
355 def read(self, size: int | None = None) -> dict[str, torch.Tensor]: 356 self._checkreadable() 357 # Clamp at the committed row count so uncommitted trailing rows left by 358 # an interrupted write are never visible. 359 remaining = len(self) - self.tell() 360 size = remaining if size is None else min(size, remaining) 361 if size <= 0: 362 return { 363 name: torch.empty(0, *shape, dtype=getattr(torch, dtype)) 364 for name, (dtype, shape) in self.schema.items() 365 } 366 return {name: col.read(size) for name, col in self._cols.items()}
def
write(self, rows: dict[str, torch.Tensor]) -> int:
368 def write(self, rows: dict[str, torch.Tensor]) -> int: 369 self._checkwritable() 370 if not isinstance(rows, dict): 371 raise TypeError( 372 f"Rows must be a dict mapping field names to tensors, got {type(rows).__name__!r}!" 373 ) 374 missing = sorted(self.schema.keys() - rows.keys()) 375 extra = sorted(rows.keys() - self.schema.keys()) 376 if missing or extra: 377 raise ValueError( 378 "Dense writes require exactly the schema fields; " 379 f"missing: {missing!r}, unexpected: {extra!r}" 380 ) 381 382 nts = { 383 name: ts.view(-1, *self.schema[name][1]).size(0) 384 for name, ts in rows.items() 385 } 386 if len(set(nts.values())) != 1: 387 raise ValueError(f"All fields must have the same row count; got: {nts!r}") 388 nt = next(iter(nts.values())) 389 390 # Columns are written first and the committed row count is bumped only 391 # afterwards, so an interrupted write is rolled back on the next open. 392 for name, col in self._cols.items(): 393 col.write(rows[name]) 394 self._status.len = len(next(iter(self._cols.values()))) 395 return nt
405 def extend(self, other: TensorDB, maintain_order: bool = False) -> None: 406 if set(self.schema) != set(other.schema): 407 raise ValueError("Schema fields must match to extend databases!") 408 self._checkwritable() 409 for name, col in self._cols.items(): 410 col.extend(other._cols[name], maintain_order=maintain_order) 411 self._status.len = len(next(iter(self._cols.values()))) 412 self._status.dump(self.statuspath)