From 57f89894e2bb7ed327be47068ba52c2fc214fa52 Mon Sep 17 00:00:00 2001 From: "russell@unturf.com" Date: Mon, 25 May 2026 21:07:06 -0400 Subject: [PATCH] #000061: stream packs to disk, not in-memory BytesIO MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-memory BytesIO design held a whole 4.4 GB pack in RAM per worker until finalization. 4-way parallel hit ~17 GB peak RAM — observed live on the full-corpus run (21 GB used, 0 free, 3 GB swap in use, headed toward OOM territory). Architectural defect, not just an inconvenience: parallelism is capped by RAM rather than CPU/network. stream_packs now: - Opens a NamedTemporaryFile in work_dir (default tempfile.gettempdir()). - Wraps it with ZstdCompressor.stream_writer + tarfile mode "w|". - After each chunk: writer.flush(FLUSH_BLOCK) + file.flush(), check file.tell() against the cap. - On finalize: write the manifest tar member, close everything, yield a FilePack carrying the path (not the bytes). - On exception: close + unlink the in-flight temp file (no /tmp litter). ObjectStoreBackend gains put_file(key, path). S3 impl uses boto3 upload_file which streams parts from disk during multipart upload — memory bounded to max_concurrency × part_size = 10 × 8 MB = 80 MB regardless of pack size. MemoryBackend reads the file into its dict (tests only). push_pack consumes FilePack: - --push: backend.put_pack_file(path), then unlink the temp - --local-dir: shutil.move from temp to local_dir/arborist-pack- (cheap if same filesystem; one-pass read+write if not) - both: shutil.move first, copy to bucket via put_pack_file from final local path? — actually current code moves to local_dir AFTER the bucket put; bucket put still works from the temp path before move. Audit row carries pack.body_size from FilePack (no len(bytes) on a 4 GB body). Worker memory bound goes from O(pack_size) to O(zstd_buffer + tar_header) ≈ tens of MB. 4-way parallel now uses ~tens of MB total instead of ~17 GB. Disk requirement: work_dir needs (parallelism × cap) of free space — 4 × 4.4 GB = 18 GB for a default full-cap parallel run. 23 passed in tests/test_cold_object.py + tests/test_evict.py. --- arborist/cold_object.py | 216 ++++++++++++++++++++++++++++++---------- arborist/evict.py | 89 ++++++++++------- 2 files changed, 217 insertions(+), 88 deletions(-) diff --git a/arborist/cold_object.py b/arborist/cold_object.py index 1bd6552..e9054b3 100644 --- a/arborist/cold_object.py +++ b/arborist/cold_object.py @@ -43,8 +43,11 @@ from __future__ import annotations import abc import io import json +import os import tarfile +import tempfile from dataclasses import dataclass +from pathlib import Path from typing import Iterator import zstandard @@ -105,6 +108,17 @@ class ObjectStoreBackend(abc.ABC): def put(self, key: str, body: bytes, *, content_type: str = "application/octet-stream") -> None: ... + @abc.abstractmethod + def put_file(self, key: str, path: Path, *, content_type: str = "application/octet-stream") -> None: + """Upload a file from disk WITHOUT loading it into memory first. + + Used for large pack bodies — `stream_packs` writes packs to temp + files on disk so worker memory stays bounded by zstd internal + buffers (~few MB), not by pack size. The S3 backend hands the + file path to boto3 `upload_file` which streams parts from disk + for multipart uploads. + """ + @abc.abstractmethod def get(self, key: str) -> bytes: ... @@ -130,6 +144,15 @@ class ObjectStoreBackend(abc.ABC): def put_pack(self, pack_hash: str, body: bytes) -> None: self.put(pack_key(pack_hash), body, content_type="application/zstd") + def put_pack_file(self, pack_hash: str, path: Path) -> None: + """File-backed pack upload — `stream_packs` produces these. + + boto3 `upload_file` reads parts from disk for the multipart upload + instead of loading the whole pack into memory. Required for any + parallelism above ~2 workers at the 4.4 GB cap. + """ + self.put_file(pack_key(pack_hash), path, content_type="application/zstd") + def get_pack(self, pack_hash: str) -> bytes: return self.get(pack_key(pack_hash)) @@ -228,6 +251,19 @@ class S3CompatibleBackend(ObjectStoreBackend): Config=self._transfer_config, ) + def put_file(self, key: str, path: Path, *, content_type: str = "application/octet-stream") -> None: + # boto3 upload_file reads parts from disk on demand for multipart + # uploads — memory stays bounded by max_concurrency × part_size + # (10 × 8 MB = 80 MB) instead of growing with pack size. The on- + # disk source is left untouched; caller decides delete vs keep. + self._client.upload_file( + Filename=str(path), + Bucket=self._bucket, + Key=key, + ExtraArgs={"ContentType": content_type}, + Config=self._transfer_config, + ) + def get(self, key: str) -> bytes: resp = self._client.get_object(Bucket=self._bucket, Key=key) return resp["Body"].read() @@ -288,6 +324,10 @@ class MemoryBackend(ObjectStoreBackend): def put(self, key: str, body: bytes, *, content_type: str = "application/octet-stream") -> None: self._store[key] = bytes(body) + def put_file(self, key: str, path: Path, *, content_type: str = "application/octet-stream") -> None: + with open(path, "rb") as f: + self._store[key] = f.read() + def get(self, key: str) -> bytes: if key not in self._store: raise KeyError(key) @@ -337,12 +377,34 @@ class PackEntry: @dataclass(frozen=True) class Pack: + """In-memory pack — used by `build_pack` for small/test cases. + + `body_bytes` is the whole zstd-compressed tar held in RAM. Fine for + unit tests and for the few-MB packs built from a single document. + For full-corpus 4.4 GB packs use `stream_packs` + `FilePack` instead. + """ pack_hash: str manifest_bytes: bytes body_bytes: bytes entries: tuple[PackEntry, ...] +@dataclass(frozen=True) +class FilePack: + """File-backed pack — `stream_packs` produces these. + + `body_path` is a tempfile on disk; the caller owns its lifecycle + (`push_pack` uploads it via `put_pack_file`, then either moves to + `local_dir` or deletes). Worker memory during build stays bounded + by zstd internal buffers (~few MB) regardless of pack size. + """ + pack_hash: str + manifest_bytes: bytes + body_path: Path + body_size: int + entries: tuple[PackEntry, ...] + + def build_pack(chunks: list[tuple[str, bytes]]) -> Pack: """Bundle a list of (leaf_hash, raw_utf8_body) pairs into one pack. @@ -407,54 +469,69 @@ def stream_packs( *, max_compressed_bytes: int, level: int = 3, -) -> Iterator[Pack]: - """Stream (leaf_hash, raw_utf8_body) chunks into packs that each fit in - `max_compressed_bytes` *compressed* bytes. + work_dir: Path | str | None = None, +) -> Iterator[FilePack]: + """Stream (leaf_hash, raw_utf8_body) chunks into FILE-BACKED packs that + each fit in `max_compressed_bytes` compressed bytes. - Why streaming: the in-memory `build_pack` caps by uncompressed bytes, - which means a 4.4 GB DVD-R cap produces ~1.5 GB compressed packs — - every disc burns at ~35 % of capacity. Streaming lets us peek the - compressed-output size as chunks accumulate and cut at the cap, so - every disc fills. + Memory bound: tens of MB per worker (zstd internal buffers + tar + header staging) regardless of pack size. The full pack lives on + disk — not in RAM — so 4-way parallel × 4.4 GB packs takes ~tens of + MB of RAM total, not 17 GB. boto3 `upload_file` reads parts from + disk during multipart upload so the upload path is also memory- + bounded. - How: a `ZstdCompressor.stream_writer` writes into an in-memory buffer; - a streaming tar writes into the compressor. After each chunk we call - `flush(FLUSH_BLOCK)` to force pending compressed bytes out of zstd's - internal buffer (preserving the dictionary for the next block — small - efficiency cost, accurate size accounting). When `buf.tell()` ≥ cap, - we finalize: write the manifest as the LAST tar member, close tar, - close stream_writer, yield the `Pack`, start a new one. + Disk requirement: `work_dir` (default `tempfile.gettempdir()`) needs + enough free space to hold one in-progress pack per parallel worker. + For 4-way at 4.4 GB cap = ~18 GB. The temp file is the caller's to + manage — `push_pack` moves it to `local_dir` or deletes after upload. - Manifest goes last because the chunk set isn't known until the cap is - hit — `open_pack` iterates by member name and ignores order so this is - transparent to consumers. + How: a `ZstdCompressor.stream_writer` writes directly into an open + file handle; a streaming tar writes into the compressor. After each + chunk we call `flush(FLUSH_BLOCK)` + `os.fsync`-free `file.flush()` + so `tell()` reflects on-disk size. When `tell()` ≥ cap, we finalize: + write the manifest as the LAST tar member, close everything, yield + the `FilePack`, start a new pack. - Caveat: each finalize emits tar trailer (~1 KB) + zstd frame footer - (~10 B) AFTER our last size check, so the actual compressed size can - overshoot the cap by ~1 KB. Trivial for a 4.4 GB cap; documented so - future tightening is intentional. + Manifest goes last because the chunk set isn't known until the cap + is hit — `open_pack` iterates by member name and ignores order so + this is transparent to consumers. + + Caveat: tar trailer (~1 KB) + zstd frame footer (~10 B) are emitted + after our last size check, so on-disk size can overshoot the cap by + ~1 KB. Trivial for a 4.4 GB cap. Single chunks larger than the cap still get their own pack (we don't drop data — the alternative is silently losing it). """ if max_compressed_bytes <= 0: raise ValueError("max_compressed_bytes must be positive") + work_path = Path(work_dir) if work_dir is not None else Path(tempfile.gettempdir()) + work_path.mkdir(parents=True, exist_ok=True) - def _start_pack(): - compressed_buf = io.BytesIO() + def _start_pack() -> dict: + # Tempfile in work_dir, opened for binary write. We hand the raw + # FD to zstd's stream_writer which wraps it; the tar wraps that. + # closefd=False on the zstd writer so we can flush + tell() the + # underlying file ourselves; we'll close everything in finalize. + fd, path_str = tempfile.mkstemp( + dir=work_path, + prefix="arborist-pack-", + suffix=".tar.zst.tmp", + ) + file_obj = os.fdopen(fd, "wb") cctx = zstandard.ZstdCompressor(level=level) - # closefd=False: closing the stream_writer must not close our - # underlying BytesIO (we still need to read getvalue() from it). - writer = cctx.stream_writer(compressed_buf, closefd=False) + writer = cctx.stream_writer(file_obj, closefd=False) tar = tarfile.open(fileobj=writer, mode="w|") return { - "buf": compressed_buf, + "path": Path(path_str), + "file_obj": file_obj, "writer": writer, "tar": tar, - "entries_by_hash": {}, # dedupe by leaf_hash like build_pack + "entries_by_hash": {}, } - def _finalize_pack(state) -> Pack: + def _finalize_pack(state: dict) -> FilePack: entries_sorted = sorted(state["entries_by_hash"].keys()) manifest_lines = [] entries = [] @@ -470,37 +547,70 @@ def stream_packs( state["tar"].addfile(info, io.BytesIO(manifest_bytes)) state["tar"].close() state["writer"].close() - body_bytes = state["buf"].getvalue() + state["file_obj"].close() pack_hash = hash_leaf(manifest_bytes).hex() - return Pack( + body_size = state["path"].stat().st_size + return FilePack( pack_hash=pack_hash, manifest_bytes=manifest_bytes, - body_bytes=body_bytes, + body_path=state["path"], + body_size=body_size, entries=tuple(entries), ) state = _start_pack() - for leaf_hash, body in chunks: - if len(leaf_hash) != 64: - raise ValueError(f"bad leaf_hash length: {leaf_hash}") - # Dedupe within a pack — same hash → same body by construction; - # the second copy is wasted bytes. - if leaf_hash in state["entries_by_hash"]: - continue - info = tarfile.TarInfo(name=f"blobs/{leaf_hash[:2]}/{leaf_hash[2:]}") - info.size = len(body) - state["tar"].addfile(info, io.BytesIO(body)) - state["entries_by_hash"][leaf_hash] = len(body) - # Force the compressor to emit pending bytes so .tell() reflects - # reality. FLUSH_BLOCK preserves dictionary state (vs FLUSH_FRAME - # which would close the frame and lose compression continuity). - state["writer"].flush(zstandard.FLUSH_BLOCK) - if state["buf"].tell() >= max_compressed_bytes: - yield _finalize_pack(state) - state = _start_pack() + try: + for leaf_hash, body in chunks: + if len(leaf_hash) != 64: + raise ValueError(f"bad leaf_hash length: {leaf_hash}") + # Dedupe within a pack — same hash → same body by construction. + if leaf_hash in state["entries_by_hash"]: + continue + info = tarfile.TarInfo(name=f"blobs/{leaf_hash[:2]}/{leaf_hash[2:]}") + info.size = len(body) + state["tar"].addfile(info, io.BytesIO(body)) + state["entries_by_hash"][leaf_hash] = len(body) + # FLUSH_BLOCK pushes pending bytes out of zstd's internal buffer + # to the file handle, preserving dictionary state. Then flush + # the Python file buffer so the OS sees them and tell() returns + # the real on-disk size. + state["writer"].flush(zstandard.FLUSH_BLOCK) + state["file_obj"].flush() + if state["file_obj"].tell() >= max_compressed_bytes: + yield _finalize_pack(state) + state = _start_pack() - if state["entries_by_hash"]: - yield _finalize_pack(state) + if state["entries_by_hash"]: + yield _finalize_pack(state) + else: + # No chunks went into the last state — clean up the empty file. + state["tar"].close() + state["writer"].close() + state["file_obj"].close() + try: + state["path"].unlink() + except FileNotFoundError: + pass + except BaseException: + # On any exception (cancel, OOM, disk full), clean up the in-flight + # temp file so we don't litter /tmp. + try: + state["tar"].close() + except Exception: + pass + try: + state["writer"].close() + except Exception: + pass + try: + state["file_obj"].close() + except Exception: + pass + try: + state["path"].unlink() + except FileNotFoundError: + pass + raise def open_pack(body_bytes: bytes) -> Iterator[tuple[str, bytes]]: diff --git a/arborist/evict.py b/arborist/evict.py index b0d072c..01c1764 100644 --- a/arborist/evict.py +++ b/arborist/evict.py @@ -394,44 +394,63 @@ def push_pack( # easily and we want one number per pack at finalize time. nonlocal_skipped = [0] + # File-backed streaming: each pack lives in a tempfile on disk, not in + # RAM. Worker memory bounded to ~tens of MB regardless of pack size. + # boto3 upload_file reads parts from disk for multipart upload so the + # upload path is also memory-bounded. + import shutil + for pack in stream_packs(_chunk_source(), max_compressed_bytes=max_pack_bytes): pack_uncompressed = uncompressed_running[0] - last_finalized_uncompressed[0] last_finalized_uncompressed[0] = uncompressed_running[0] - if push_to_bucket: - backend.put_pack(pack.pack_hash, pack.body_bytes) - backend.put_pack_manifest(pack.pack_hash, pack.manifest_bytes) - if out_dir is not None: - short = pack.pack_hash[:16] - (out_dir / f"arborist-pack-{short}.tar.zst").write_bytes(pack.body_bytes) - (out_dir / f"arborist-pack-{short}.manifest.ndjson").write_bytes( - pack.manifest_bytes - ) - with transaction(conn): - event_hash = append_audit( - conn, - event_type="cold_pack_pushed", - subject_root=document_root, - body={ - "pack_hash": pack.pack_hash, - "chunk_count": len(pack.entries), - "uncompressed_bytes": pack_uncompressed, - "compressed_bytes": len(pack.body_bytes), - "snapshot_root": snapshot_root, - "snapshot_doc_count": doc_count_at_pack, - "snapshot_ts": pack_ts, - "pushed_to_bucket": push_to_bucket, - "local_dir": str(out_dir) if out_dir else None, - "backend": backend_id if push_to_bucket else None, - }, - ) - pack_results.append({ - "pack_hash": pack.pack_hash, - "chunk_count": len(pack.entries), - "uncompressed_bytes": pack_uncompressed, - "compressed_bytes": len(pack.body_bytes), - "snapshot_root": snapshot_root, - "audit_event_hash": event_hash, - }) + try: + if push_to_bucket: + backend.put_pack_file(pack.pack_hash, pack.body_path) + backend.put_pack_manifest(pack.pack_hash, pack.manifest_bytes) + if out_dir is not None: + short = pack.pack_hash[:16] + final_path = out_dir / f"arborist-pack-{short}.tar.zst" + # Move (cheap if on same filesystem) instead of read+write. + shutil.move(str(pack.body_path), str(final_path)) + (out_dir / f"arborist-pack-{short}.manifest.ndjson").write_bytes( + pack.manifest_bytes + ) + with transaction(conn): + event_hash = append_audit( + conn, + event_type="cold_pack_pushed", + subject_root=document_root, + body={ + "pack_hash": pack.pack_hash, + "chunk_count": len(pack.entries), + "uncompressed_bytes": pack_uncompressed, + "compressed_bytes": pack.body_size, + "snapshot_root": snapshot_root, + "snapshot_doc_count": doc_count_at_pack, + "snapshot_ts": pack_ts, + "pushed_to_bucket": push_to_bucket, + "local_dir": str(out_dir) if out_dir else None, + "backend": backend_id if push_to_bucket else None, + }, + ) + pack_results.append({ + "pack_hash": pack.pack_hash, + "chunk_count": len(pack.entries), + "uncompressed_bytes": pack_uncompressed, + "compressed_bytes": pack.body_size, + "snapshot_root": snapshot_root, + "audit_event_hash": event_hash, + }) + finally: + # Clean up the temp file when neither local_dir nor an exception + # claimed it. `shutil.move` already removed it on the local_dir + # path; this handles the push-only path and the audit-failure + # case. + if pack.body_path.exists(): + try: + pack.body_path.unlink() + except OSError: + pass skipped_hash_mismatch = nonlocal_skipped[0]