Coverage for src/loman/serialization/blobs.py: 95%
196 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-07 00:32 +0000
1"""Blob storage: where a value's bytes live when they are too big to inline.
3A saved computation has one logical layout::
5 manifest.json
6 blobs/0000.npy
7 blobs/0001.npy
9A ``.loman`` file is that tree inside a zip. A directory container is that tree
10on disk. One specification, two ways of writing it down --- which is asserted by
11a test comparing the manifests byte for byte, rather than left to discipline.
13Blob keys are never derived from node keys. Zero-padded integer ids sidestep
14``/`` in hierarchical keys, Windows reserved names, case-insensitive collisions
15and unicode normalisation in one move; the blob table records which node a blob
16came from so a human can still find their way around.
18Two roles are kept apart here, which is what makes storage pluggable:
20:class:`BlobStore` is *where bytes go*. It has two methods, and a user
21implementing one for S3 or a database inherits compression, deduplication,
22checksums and blob-table bookkeeping without writing any of it.
24:class:`BlobWriter` and :class:`BlobReader` are *the bookkeeping*. One writer per
25save, holding the blob table and routing each blob to whichever store its node
26asked for. A single save can therefore span several stores --- most values in the
27container, one node's frames in S3 --- and the manifest records which is which.
29Bytes reaching a store are already compressed, if compressing them paid. The
30container never compresses again: that would double-compress an already-compact
31payload, and would put every member behind a decompression step, foreclosing a
32future zero-copy read.
33"""
35from __future__ import annotations
37import hashlib
38import io
39import json
40import posixpath
41import shutil
42import zipfile
43from abc import ABC, abstractmethod
44from collections.abc import Callable
45from pathlib import Path
46from typing import Any, BinaryIO, cast
48from loman.exception import SerializationError
50from .compression import compress_blob, decompress_blob
52MANIFEST_NAME = "manifest.json"
53BLOB_DIR = "blobs"
55#: Name of the store that writes into the container itself. A blob entry without
56#: a ``store`` field means this one, which keeps the common manifest short.
57CONTAINER_STORE = "container"
59# Key marking a blob reference inside an encoded value. Reserved, so a user dict
60# containing it is escaped rather than mistaken for one.
61BLOB_REF_KEY = "$blob"
63# Zip timestamp written for every member. Zip stores no timezone and defaults to
64# "now", which would make two saves of the same computation differ. A fixed
65# stamp makes saves byte-reproducible, which is what lets them be compared,
66# cached by content, or diffed.
67_FIXED_ZIP_TIMESTAMP = (1980, 1, 1, 0, 0, 0)
69BlobPayload = bytes | bytearray | memoryview | Callable[[BinaryIO], None]
72def blob_ref(blob_id: int) -> dict[str, int]:
73 """Return the encoded reference to blob *blob_id*."""
74 return {BLOB_REF_KEY: blob_id}
77def _payload_to_bytes(payload: BlobPayload) -> bytes:
78 """Materialise *payload* as bytes, running a writer callable if given."""
79 if callable(payload):
80 writer = cast("Callable[[BinaryIO], None]", payload)
81 buf = io.BytesIO()
82 writer(buf)
83 return buf.getvalue()
84 return bytes(payload)
87class BlobStore(ABC):
88 """Where a blob's bytes are kept. Implement two methods.
90 Everything else --- ids, compression, deduplication, checksums, the blob
91 table --- belongs to :class:`BlobWriter` and is inherited, so a store for S3
92 or a database is genuinely just these two::
94 class S3Store(BlobStore):
95 def __init__(self, bucket, prefix, client):
96 self.bucket, self.prefix, self.client = bucket, prefix, client
98 def write_blob(self, key, data):
99 self.client.put_object(Bucket=self.bucket, Key=f"{self.prefix}/{key}", Body=data)
101 def read_blob(self, key):
102 return self.client.get_object(Bucket=self.bucket, Key=f"{self.prefix}/{key}")["Body"].read()
104 Pass it by name to both ends::
106 comp.save('run.loman', stores={'s3': S3Store(...)})
107 Computation.load('run.loman', stores={'s3': S3Store(...)})
109 The key is a short relative path such as ``blobs/0000.npy``. A store is free
110 to place it under a prefix of its own, as above; it just has to hand the same
111 bytes back for the same key.
113 The manifest records the store's *name*, never its configuration, so a saved
114 file never contains a bucket, a connection string or a credential. The
115 consequence is that a file cannot resolve its own external blobs: whoever
116 loads it supplies the matching store.
117 """
119 @abstractmethod
120 def write_blob(self, key: str, data: bytes) -> None:
121 """Store *data* under *key*."""
123 @abstractmethod
124 def read_blob(self, key: str) -> bytes:
125 """Return the bytes stored under *key*."""
127 def key_for(self, blob_id: int, codec: str, node: str | None) -> str:
128 """Return the key to store blob *blob_id* under.
130 Overriding this lets a store lay its keys out differently --- by node
131 name, or partitioned by date. The key is recorded in the manifest, so
132 whatever is returned here is what :meth:`read_blob` is later asked for.
133 """
134 return f"{BLOB_DIR}/{blob_id:04d}.{codec}"
137class ZipBlobStore(BlobStore):
138 """Blobs as stored (uncompressed) members of a zip archive."""
140 def __init__(self, zf: zipfile.ZipFile) -> None:
141 """Read and write blobs in the already-open archive *zf*."""
142 self._zf = zf
144 def write_blob(self, key: str, data: bytes) -> None:
145 """Write *data* as a stored zip member."""
146 info = zipfile.ZipInfo(key, date_time=_FIXED_ZIP_TIMESTAMP)
147 info.compress_type = zipfile.ZIP_STORED
148 self._zf.writestr(info, data)
150 def read_blob(self, key: str) -> bytes:
151 """Return the raw bytes of a zip member."""
152 return self._zf.read(_validate_member_path(key))
155class DirBlobStore(BlobStore):
156 """Blobs as files under a directory."""
158 def __init__(self, root: Path) -> None:
159 """Read and write blobs under *root*."""
160 self._root = root
162 def write_blob(self, key: str, data: bytes) -> None:
163 """Write *data* to a file under the container root."""
164 target = self._root / key
165 target.parent.mkdir(parents=True, exist_ok=True)
166 target.write_bytes(data)
168 def read_blob(self, key: str) -> bytes:
169 """Return the raw bytes of a file under the container root."""
170 return (self._root / _validate_member_path(key)).read_bytes()
173class MemoryBlobStore(BlobStore):
174 """Blobs in a dict. Useful for tests, and as a worked example of the interface."""
176 def __init__(self, data: dict[str, bytes] | None = None) -> None:
177 """Keep blobs in *data*, or in a fresh dict."""
178 self.data: dict[str, bytes] = {} if data is None else data
180 def write_blob(self, key: str, data: bytes) -> None:
181 """Store *data* under *key*."""
182 self.data[key] = data
184 def read_blob(self, key: str) -> bytes:
185 """Return the bytes stored under *key*."""
186 try:
187 return self.data[key]
188 except KeyError:
189 msg = f"No blob stored under {key!r}"
190 raise SerializationError(msg) from None
193class BlobWriter:
194 """Blob bookkeeping for one save: ids, compression, dedup, and the table.
196 Routes each blob to a named :class:`BlobStore`. A save can use several ---
197 the container for most values, an external store for the nodes that asked
198 for one --- and the manifest records which store holds each blob.
199 """
201 def __init__(
202 self,
203 stores: dict[str, BlobStore],
204 *,
205 compression: str = "none",
206 dedupe: str = "identity",
207 checksums: bool = False,
208 ) -> None:
209 """Write blobs into *stores*, keyed by name."""
210 self._stores = stores
211 self._entries: list[dict[str, Any]] = []
212 self._compression = compression
213 self._dedupe = dedupe
214 self._checksums = checksums
215 self._by_identity: dict[int, int] = {}
216 self._by_content: dict[str, int] = {}
217 # Keeps deduplicated objects alive so their ids stay unique; see put().
218 self._dedupe_refs: list[Any] = []
220 def can_store(self, store: str | None) -> bool:
221 """Return whether blobs can be written to the named store.
223 ``None`` means the container's own store. A container that holds no
224 blobs --- the single JSON document --- has none, but an external store
225 named on a node still works there, which is how a readable manifest can
226 sit alongside data in S3.
227 """
228 return (store or CONTAINER_STORE) in self._stores
230 @property
231 def accepts_blobs(self) -> bool:
232 """Whether any store at all is available."""
233 return bool(self._stores)
235 def table(self) -> list[dict[str, Any]]:
236 """Return the manifest's blob table."""
237 return self._entries
239 def put(
240 self,
241 payload: BlobPayload,
242 *,
243 codec: str,
244 node: str | None,
245 store: str | None = None,
246 compressible: bool = True,
247 dedupe_on: Any = None,
248 ) -> int:
249 """Store *payload*, compressing and deduplicating as configured."""
250 store_name = store or CONTAINER_STORE
251 target = self._stores.get(store_name)
252 if target is None:
253 known = sorted(self._stores) or ["(none)"]
254 msg = (
255 f"Node {node!r} asks for blob store {store_name!r}, which was not supplied. "
256 f"Pass it as save(..., stores={{{store_name!r}: ...}}). Available: {known}"
257 )
258 raise SerializationError(msg)
260 if dedupe_on is not None and self._dedupe == "identity":
261 existing = self._by_identity.get(id(dedupe_on))
262 if existing is not None:
263 return existing
265 raw = _payload_to_bytes(payload)
266 stored, compression = compress_blob(raw, self._compression, compressible=compressible)
268 digest = None
269 if self._checksums or self._dedupe == "content":
270 digest = hashlib.sha256(stored).hexdigest()
271 if self._dedupe == "content":
272 existing = self._by_content.get(digest)
273 if existing is not None:
274 return existing
276 blob_id = len(self._entries)
277 key = target.key_for(blob_id, codec, node)
278 target.write_blob(key, stored)
280 entry: dict[str, Any] = {
281 "id": blob_id,
282 "path": key,
283 "codec": codec,
284 "compression": compression,
285 "size": len(raw),
286 }
287 if store_name != CONTAINER_STORE:
288 entry["store"] = store_name
289 if compression != "none":
290 entry["stored_size"] = len(stored)
291 if node is not None:
292 entry["node"] = node
293 if self._checksums and digest is not None:
294 entry["sha256"] = digest
295 self._entries.append(entry)
297 if dedupe_on is not None:
298 # The object is kept alive alongside its id. Without that reference a
299 # temporary --- a column's `.to_numpy()`, say --- would be collected
300 # as soon as this returns, and CPython would hand the same id to the
301 # next temporary, deduplicating two unrelated values onto one blob.
302 self._by_identity[id(dedupe_on)] = blob_id
303 self._dedupe_refs.append(dedupe_on)
304 if digest is not None:
305 self._by_content.setdefault(digest, blob_id)
306 return blob_id
309class BlobReader:
310 """Resolves blob references from an already-written container."""
312 def __init__(self, entries: list[dict[str, Any]], stores: dict[str, BlobStore]) -> None:
313 """Resolve blobs listed in *entries* against *stores*."""
314 self._entries = {int(entry["id"]): entry for entry in entries}
315 self._stores = stores
317 def get(self, blob_id: int) -> bytes:
318 """Return the bytes of blob *blob_id*, decompressed."""
319 entry = self._entries.get(int(blob_id))
320 if entry is None:
321 msg = f"Blob reference {blob_id} does not resolve: the manifest lists no such blob"
322 raise SerializationError(msg)
324 store_name = entry.get("store", CONTAINER_STORE)
325 store = self._stores.get(store_name)
326 if store is None:
327 known = sorted(self._stores) or ["(none)"]
328 node = entry.get("node")
329 where = f" (node {node!r})" if node else ""
330 msg = (
331 f"Blob {blob_id}{where} is held in store {store_name!r}, which was not supplied. "
332 f"Pass it as load(..., stores={{{store_name!r}: ...}}). Available: {known}. "
333 "A saved file records a store's name but never its configuration, so it cannot "
334 "resolve external blobs on its own."
335 )
336 raise SerializationError(msg)
338 return decompress_blob(store.read_blob(entry["path"]), entry.get("compression", "none"))
341def _validate_member_path(path: str) -> str:
342 """Return *path* if it is safe to resolve inside a container.
344 A manifest is data, and its blob paths are chosen by whoever wrote the file.
345 An absolute path or one climbing out with ``..`` would read from anywhere on
346 the machine, so both are refused before the path is used.
347 """
348 if not path:
349 msg = "Blob path is empty"
350 raise SerializationError(msg)
351 normalised = posixpath.normpath(path)
352 if posixpath.isabs(path) or normalised.startswith("..") or Path(path).is_absolute():
353 msg = f"Refusing blob path outside the container: {path!r}"
354 raise SerializationError(msg)
355 return normalised
358def write_zip_container(
359 path: Path,
360 build: Callable[[dict[str, BlobStore]], dict[str, Any]],
361) -> None:
362 """Write a ``.loman`` archive at *path*.
364 *build* is handed the container's store and returns the manifest. Blobs are
365 written first, as the manifest is only complete once every blob has an id ---
366 so the manifest member is added last, even though it is read first.
367 """
368 _require_parent(path)
369 tmp = path.with_name(path.name + ".tmp")
370 try:
371 with zipfile.ZipFile(tmp, "w") as zf:
372 manifest = build({CONTAINER_STORE: ZipBlobStore(zf)})
373 info = zipfile.ZipInfo(MANIFEST_NAME, date_time=_FIXED_ZIP_TIMESTAMP)
374 info.compress_type = zipfile.ZIP_DEFLATED
375 zf.writestr(info, dump_manifest(manifest))
376 tmp.replace(path)
377 finally:
378 if tmp.exists(): # pragma: no cover - only on a failed write
379 tmp.unlink()
382def write_dir_container(
383 path: Path,
384 build: Callable[[dict[str, BlobStore]], dict[str, Any]],
385) -> None:
386 """Write a directory container at *path*, replacing any existing one.
388 Built alongside the target and swapped in at the end, so a save that fails
389 part way leaves whatever was there before intact. Clearing the old blobs
390 first would be simpler, but then an unserializable value in the middle of a
391 graph would destroy the last good checkpoint --- losing data on the strength
392 of an operation that did not even succeed.
393 """
394 if path.exists() and not path.is_dir():
395 msg = f"Cannot write a directory container over the existing file {str(path)!r}"
396 raise SerializationError(msg)
397 _require_parent(path)
399 staging = path.with_name(path.name + ".tmp")
400 shutil.rmtree(staging, ignore_errors=True)
401 staging.mkdir(parents=True)
402 try:
403 manifest = build({CONTAINER_STORE: DirBlobStore(staging)})
404 (staging / MANIFEST_NAME).write_text(dump_manifest(manifest), encoding="utf-8")
406 # Swap: move the old aside, put the new in place, then discard the old.
407 # Two renames rather than one because a directory cannot be replaced
408 # atomically the way a file can.
409 previous = path.with_name(path.name + ".previous")
410 shutil.rmtree(previous, ignore_errors=True)
411 if path.exists():
412 path.rename(previous)
413 try:
414 staging.rename(path)
415 except OSError: # pragma: no cover - put the old one back and re-raise
416 if previous.exists():
417 previous.rename(path)
418 raise
419 shutil.rmtree(previous, ignore_errors=True)
420 finally:
421 shutil.rmtree(staging, ignore_errors=True)
424def _require_parent(path: Path) -> None:
425 """Fail early, and about the right path, when the destination has no parent.
427 Both containers write to a sibling temporary first, so without this the
428 error names a ``.tmp`` file the caller never asked for.
429 """
430 parent = path.parent
431 if not parent.exists():
432 msg = f"Cannot save to {str(path)!r}: the directory {str(parent)!r} does not exist"
433 raise SerializationError(msg)
436def dump_manifest(manifest: dict[str, Any]) -> str:
437 """Return the manifest as JSON text.
439 ``allow_nan`` is off so a non-finite float can never reach the file as a
440 bare ``NaN`` token, which Python reads back and no other JSON parser will.
441 """
442 return json.dumps(manifest, allow_nan=False)
445def read_zip_manifest(zf: zipfile.ZipFile) -> dict[str, Any]:
446 """Return the manifest from an open archive."""
447 try:
448 raw = zf.read(MANIFEST_NAME)
449 except KeyError as exc:
450 msg = f"Not a loman container: no {MANIFEST_NAME} inside the archive"
451 raise SerializationError(msg) from exc
452 return json.loads(raw.decode("utf-8"))
455def read_dir_manifest(root: Path) -> dict[str, Any]:
456 """Return the manifest from a directory container."""
457 manifest_path = root / MANIFEST_NAME
458 if not manifest_path.is_file():
459 msg = f"Not a loman container: {str(root)!r} has no {MANIFEST_NAME}"
460 raise SerializationError(msg)
461 return json.loads(manifest_path.read_text(encoding="utf-8"))