#000061: stream packs to disk, not in-memory BytesIO

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-<hash>
                 (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.
This commit is contained in:
russell@unturf.com 2026-05-25 21:07:06 -04:00
parent edb4e2d29e
commit 57f89894e2
No known key found for this signature in database
2 changed files with 217 additions and 88 deletions

View file

@ -43,8 +43,11 @@ from __future__ import annotations
import abc import abc
import io import io
import json import json
import os
import tarfile import tarfile
import tempfile
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path
from typing import Iterator from typing import Iterator
import zstandard import zstandard
@ -105,6 +108,17 @@ class ObjectStoreBackend(abc.ABC):
def put(self, key: str, body: bytes, *, content_type: str = "application/octet-stream") -> None: 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 @abc.abstractmethod
def get(self, key: str) -> bytes: def get(self, key: str) -> bytes:
... ...
@ -130,6 +144,15 @@ class ObjectStoreBackend(abc.ABC):
def put_pack(self, pack_hash: str, body: bytes) -> None: def put_pack(self, pack_hash: str, body: bytes) -> None:
self.put(pack_key(pack_hash), body, content_type="application/zstd") 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: def get_pack(self, pack_hash: str) -> bytes:
return self.get(pack_key(pack_hash)) return self.get(pack_key(pack_hash))
@ -228,6 +251,19 @@ class S3CompatibleBackend(ObjectStoreBackend):
Config=self._transfer_config, 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: def get(self, key: str) -> bytes:
resp = self._client.get_object(Bucket=self._bucket, Key=key) resp = self._client.get_object(Bucket=self._bucket, Key=key)
return resp["Body"].read() 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: def put(self, key: str, body: bytes, *, content_type: str = "application/octet-stream") -> None:
self._store[key] = bytes(body) 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: def get(self, key: str) -> bytes:
if key not in self._store: if key not in self._store:
raise KeyError(key) raise KeyError(key)
@ -337,12 +377,34 @@ class PackEntry:
@dataclass(frozen=True) @dataclass(frozen=True)
class Pack: 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 pack_hash: str
manifest_bytes: bytes manifest_bytes: bytes
body_bytes: bytes body_bytes: bytes
entries: tuple[PackEntry, ...] 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: def build_pack(chunks: list[tuple[str, bytes]]) -> Pack:
"""Bundle a list of (leaf_hash, raw_utf8_body) pairs into one 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, max_compressed_bytes: int,
level: int = 3, level: int = 3,
) -> Iterator[Pack]: work_dir: Path | str | None = None,
"""Stream (leaf_hash, raw_utf8_body) chunks into packs that each fit in ) -> Iterator[FilePack]:
`max_compressed_bytes` *compressed* bytes. """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, Memory bound: tens of MB per worker (zstd internal buffers + tar
which means a 4.4 GB DVD-R cap produces ~1.5 GB compressed packs — header staging) regardless of pack size. The full pack lives on
every disc burns at ~35 % of capacity. Streaming lets us peek the disk — not in RAM — so 4-way parallel × 4.4 GB packs takes ~tens of
compressed-output size as chunks accumulate and cut at the cap, so MB of RAM total, not 17 GB. boto3 `upload_file` reads parts from
every disc fills. disk during multipart upload so the upload path is also memory-
bounded.
How: a `ZstdCompressor.stream_writer` writes into an in-memory buffer; Disk requirement: `work_dir` (default `tempfile.gettempdir()`) needs
a streaming tar writes into the compressor. After each chunk we call enough free space to hold one in-progress pack per parallel worker.
`flush(FLUSH_BLOCK)` to force pending compressed bytes out of zstd's For 4-way at 4.4 GB cap = ~18 GB. The temp file is the caller's to
internal buffer (preserving the dictionary for the next block — small manage — `push_pack` moves it to `local_dir` or deletes after upload.
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.
Manifest goes last because the chunk set isn't known until the cap is How: a `ZstdCompressor.stream_writer` writes directly into an open
hit — `open_pack` iterates by member name and ignores order so this is file handle; a streaming tar writes into the compressor. After each
transparent to consumers. 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 Manifest goes last because the chunk set isn't known until the cap
(~10 B) AFTER our last size check, so the actual compressed size can is hit — `open_pack` iterates by member name and ignores order so
overshoot the cap by ~1 KB. Trivial for a 4.4 GB cap; documented so this is transparent to consumers.
future tightening is intentional.
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 Single chunks larger than the cap still get their own pack (we don't
drop data — the alternative is silently losing it). drop data — the alternative is silently losing it).
""" """
if max_compressed_bytes <= 0: if max_compressed_bytes <= 0:
raise ValueError("max_compressed_bytes must be positive") 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(): def _start_pack() -> dict:
compressed_buf = io.BytesIO() # 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) cctx = zstandard.ZstdCompressor(level=level)
# closefd=False: closing the stream_writer must not close our writer = cctx.stream_writer(file_obj, closefd=False)
# underlying BytesIO (we still need to read getvalue() from it).
writer = cctx.stream_writer(compressed_buf, closefd=False)
tar = tarfile.open(fileobj=writer, mode="w|") tar = tarfile.open(fileobj=writer, mode="w|")
return { return {
"buf": compressed_buf, "path": Path(path_str),
"file_obj": file_obj,
"writer": writer, "writer": writer,
"tar": tar, "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()) entries_sorted = sorted(state["entries_by_hash"].keys())
manifest_lines = [] manifest_lines = []
entries = [] entries = []
@ -470,37 +547,70 @@ def stream_packs(
state["tar"].addfile(info, io.BytesIO(manifest_bytes)) state["tar"].addfile(info, io.BytesIO(manifest_bytes))
state["tar"].close() state["tar"].close()
state["writer"].close() state["writer"].close()
body_bytes = state["buf"].getvalue() state["file_obj"].close()
pack_hash = hash_leaf(manifest_bytes).hex() pack_hash = hash_leaf(manifest_bytes).hex()
return Pack( body_size = state["path"].stat().st_size
return FilePack(
pack_hash=pack_hash, pack_hash=pack_hash,
manifest_bytes=manifest_bytes, manifest_bytes=manifest_bytes,
body_bytes=body_bytes, body_path=state["path"],
body_size=body_size,
entries=tuple(entries), entries=tuple(entries),
) )
state = _start_pack() state = _start_pack()
for leaf_hash, body in chunks: try:
if len(leaf_hash) != 64: for leaf_hash, body in chunks:
raise ValueError(f"bad leaf_hash length: {leaf_hash}") if len(leaf_hash) != 64:
# Dedupe within a pack — same hash → same body by construction; raise ValueError(f"bad leaf_hash length: {leaf_hash}")
# the second copy is wasted bytes. # Dedupe within a pack — same hash → same body by construction.
if leaf_hash in state["entries_by_hash"]: if leaf_hash in state["entries_by_hash"]:
continue continue
info = tarfile.TarInfo(name=f"blobs/{leaf_hash[:2]}/{leaf_hash[2:]}") info = tarfile.TarInfo(name=f"blobs/{leaf_hash[:2]}/{leaf_hash[2:]}")
info.size = len(body) info.size = len(body)
state["tar"].addfile(info, io.BytesIO(body)) state["tar"].addfile(info, io.BytesIO(body))
state["entries_by_hash"][leaf_hash] = len(body) state["entries_by_hash"][leaf_hash] = len(body)
# Force the compressor to emit pending bytes so .tell() reflects # FLUSH_BLOCK pushes pending bytes out of zstd's internal buffer
# reality. FLUSH_BLOCK preserves dictionary state (vs FLUSH_FRAME # to the file handle, preserving dictionary state. Then flush
# which would close the frame and lose compression continuity). # the Python file buffer so the OS sees them and tell() returns
state["writer"].flush(zstandard.FLUSH_BLOCK) # the real on-disk size.
if state["buf"].tell() >= max_compressed_bytes: state["writer"].flush(zstandard.FLUSH_BLOCK)
yield _finalize_pack(state) state["file_obj"].flush()
state = _start_pack() if state["file_obj"].tell() >= max_compressed_bytes:
yield _finalize_pack(state)
state = _start_pack()
if state["entries_by_hash"]: if state["entries_by_hash"]:
yield _finalize_pack(state) 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]]: def open_pack(body_bytes: bytes) -> Iterator[tuple[str, bytes]]:

View file

@ -394,44 +394,63 @@ def push_pack(
# easily and we want one number per pack at finalize time. # easily and we want one number per pack at finalize time.
nonlocal_skipped = [0] 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): for pack in stream_packs(_chunk_source(), max_compressed_bytes=max_pack_bytes):
pack_uncompressed = uncompressed_running[0] - last_finalized_uncompressed[0] pack_uncompressed = uncompressed_running[0] - last_finalized_uncompressed[0]
last_finalized_uncompressed[0] = uncompressed_running[0] last_finalized_uncompressed[0] = uncompressed_running[0]
if push_to_bucket: try:
backend.put_pack(pack.pack_hash, pack.body_bytes) if push_to_bucket:
backend.put_pack_manifest(pack.pack_hash, pack.manifest_bytes) backend.put_pack_file(pack.pack_hash, pack.body_path)
if out_dir is not None: backend.put_pack_manifest(pack.pack_hash, pack.manifest_bytes)
short = pack.pack_hash[:16] if out_dir is not None:
(out_dir / f"arborist-pack-{short}.tar.zst").write_bytes(pack.body_bytes) short = pack.pack_hash[:16]
(out_dir / f"arborist-pack-{short}.manifest.ndjson").write_bytes( final_path = out_dir / f"arborist-pack-{short}.tar.zst"
pack.manifest_bytes # Move (cheap if on same filesystem) instead of read+write.
) shutil.move(str(pack.body_path), str(final_path))
with transaction(conn): (out_dir / f"arborist-pack-{short}.manifest.ndjson").write_bytes(
event_hash = append_audit( pack.manifest_bytes
conn, )
event_type="cold_pack_pushed", with transaction(conn):
subject_root=document_root, event_hash = append_audit(
body={ conn,
"pack_hash": pack.pack_hash, event_type="cold_pack_pushed",
"chunk_count": len(pack.entries), subject_root=document_root,
"uncompressed_bytes": pack_uncompressed, body={
"compressed_bytes": len(pack.body_bytes), "pack_hash": pack.pack_hash,
"snapshot_root": snapshot_root, "chunk_count": len(pack.entries),
"snapshot_doc_count": doc_count_at_pack, "uncompressed_bytes": pack_uncompressed,
"snapshot_ts": pack_ts, "compressed_bytes": pack.body_size,
"pushed_to_bucket": push_to_bucket, "snapshot_root": snapshot_root,
"local_dir": str(out_dir) if out_dir else None, "snapshot_doc_count": doc_count_at_pack,
"backend": backend_id if push_to_bucket else None, "snapshot_ts": pack_ts,
}, "pushed_to_bucket": push_to_bucket,
) "local_dir": str(out_dir) if out_dir else None,
pack_results.append({ "backend": backend_id if push_to_bucket else None,
"pack_hash": pack.pack_hash, },
"chunk_count": len(pack.entries), )
"uncompressed_bytes": pack_uncompressed, pack_results.append({
"compressed_bytes": len(pack.body_bytes), "pack_hash": pack.pack_hash,
"snapshot_root": snapshot_root, "chunk_count": len(pack.entries),
"audit_event_hash": event_hash, "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] skipped_hash_mismatch = nonlocal_skipped[0]