engineering/arborist
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

216
arborist/cold_object.py
View file

@ -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]]:

89
arborist/evict.py
View file

@ -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]