Repository navigation
Archive the trace store weekly and pull it to the GPU box #220
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| """Weekly off-machine archive of the trace store. | ||
|
|
||
| Traces are kept forever; a single bot host must not be their only copy. Once | ||
| a week the scheduler packs everything not yet archived in | ||
| ``<state_dir>/traces`` into ``<state_dir>/archive/traces-<seq>-<time>.tar.gz`` | ||
| with a ``.sha256`` manifest (the archive's hash, then one line per member), | ||
| written last to mark it complete. What was archived is tracked explicitly in | ||
| ``archive/archived.json`` (path -> content hash), never by timestamps: a | ||
| blob renamed into place while an archive was being made is simply picked up by | ||
| the next one, and a record file that grew since is archived again. Names carry | ||
| a monotonic sequence FIRST, so name order is creation order whatever the | ||
| wall clock does, and are never reused or overwritten. ``index.db`` is | ||
| rebuilt from the records and never archived. | ||
|
|
||
| The publisher on the GPU box, which already reaches the bot host over SSH, | ||
| pulls every archive it does not have and keeps it only if the hash verifies. | ||
| Restoring is extracting the archives in name order into an empty trace root | ||
| (later copies of a record file replace earlier ones) and running | ||
| ``TraceStore.rebuild_index``. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import hashlib | ||
| import io | ||
| import json | ||
| import re | ||
| import tarfile | ||
| import time | ||
| from pathlib import Path | ||
|
|
||
| SEQUENCE = "trace_archive_seq" | ||
| NAME = re.compile(r"traces-\d{6}-\d{8}-\d{6}\.tar\.gz") # sequence first: name order = restore order | ||
| INDEX = "archived.json" | ||
|
|
||
|
|
||
| def _sha(data: bytes) -> str: | ||
| return hashlib.sha256(data).hexdigest() | ||
|
|
||
|
|
||
| def make_archive(rt) -> Path | None: | ||
| """Pack everything not yet archived; None when nothing is new.""" | ||
| root = rt.state_dir / "traces" | ||
| out_dir = rt.state_dir / "archive" | ||
| index_path = out_dir / INDEX | ||
| archived: dict[str, str] = json.loads(index_path.read_text(encoding="utf-8")) if index_path.exists() else {} | ||
| members: list[tuple[str, bytes]] = [] | ||
| for path in sorted(root.rglob("*")) if root.is_dir() else []: | ||
| if not path.is_file() or path.name == "index.db" or path.name.endswith(".tmp"): | ||
| continue | ||
| rel = path.relative_to(root).as_posix() | ||
| if rel.startswith("blobs/") and rel in archived: | ||
| continue # content-addressed: never changes once written | ||
| data = path.read_bytes() | ||
| digest = _sha(data) | ||
| if archived.get(rel) != digest: | ||
| members.append((rel, data)) | ||
| archived[rel] = digest | ||
| if not members: | ||
| return None | ||
| sequence = int(rt.ledger.get_cursor("*", SEQUENCE) or 0) + 1 | ||
| rt.ledger.set_cursor("*", SEQUENCE, str(sequence)) # consumed even if packing fails: never reused | ||
| name = f"traces-{sequence:06d}-" + time.strftime("%Y%m%d-%H%M%S", time.gmtime(rt.clock())) + ".tar.gz" | ||
| out_dir.mkdir(parents=True, exist_ok=True) | ||
| if (out_dir / name).exists(): | ||
| raise FileExistsError(f"refusing to overwrite archive {name}") | ||
| buffer = io.BytesIO() | ||
| with tarfile.open(fileobj=buffer, mode="w:gz") as tar: | ||
| for rel, data in members: | ||
| info = tarfile.TarInfo(rel) | ||
| info.size, info.mode, info.mtime = len(data), 0o644, int(rt.clock()) | ||
| tar.addfile(info, io.BytesIO(data)) | ||
| blob = buffer.getvalue() | ||
| manifest = f"{_sha(blob)} {name}\n" + "".join(f"{_sha(data)} {rel}\n" for rel, data in members) | ||
| tmp = out_dir / f".{name}.tmp" | ||
| tmp.write_bytes(blob) | ||
| tmp.replace(out_dir / name) | ||
| (out_dir / f"{name}.sha256").write_text(manifest, encoding="utf-8") # written last: marks it complete | ||
| tmp_index = out_dir / f".{INDEX}.tmp" | ||
| tmp_index.write_text(json.dumps(archived, sort_keys=True), encoding="utf-8") | ||
| tmp_index.replace(index_path) # only after the archive is complete | ||
| return out_dir / name | ||
|
|
||
|
|
||
| def verify_archive(data: bytes, manifest: str, name: str) -> bool: | ||
| first = (manifest.splitlines() or [""])[0].split() | ||
| return len(first) == 2 and first[1] == name and _sha(data) == first[0] | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -44,6 +44,7 @@ | |
| from .outbox import CONTROL_MAX_AGE, OutboxError, OutboxItem, check_item | ||
|
|
||
| ITEM_ID = re.compile(r"[0-9]+-[0-9a-f]{12}") | ||
| ARCHIVE_NAME = re.compile(r"traces-\d{6}-\d{8}-\d{6}\.tar\.gz") | ||
| GOVERNED_PATH = re.compile(r"knowledge/(?:repos/[A-Za-z0-9._-]+|general)/(?:[A-Za-z0-9._-]+/)*[A-Za-z0-9._-]+\.(?:md|yaml)") | ||
| BRANCH = re.compile(r"kb/[A-Za-z0-9._-]+/[A-Za-z0-9._-]+") | ||
| HOLD_LABEL = "kb:hold" | ||
|
|
@@ -63,6 +64,7 @@ class Transport(Protocol): | |
| def list_items(self) -> list[str]: ... | ||
| def read(self, rel: str) -> bytes: ... | ||
| def write_ack(self, item_id: str, data: bytes) -> None: ... | ||
| def list_archives(self) -> list[str]: ... | ||
|
|
||
|
|
||
| @dataclass | ||
|
|
@@ -79,6 +81,12 @@ def list_items(self) -> list[str]: | |
| def read(self, rel: str) -> bytes: | ||
| return (self.root / rel).read_bytes() | ||
|
|
||
| def list_archives(self) -> list[str]: | ||
| folder = self.root / "archive" | ||
| return sorted(p.name for p in folder.glob("traces-*.tar.gz") | ||
| if ARCHIVE_NAME.fullmatch(p.name) and (folder / f"{p.name}.sha256").exists()) \ | ||
| if folder.is_dir() else [] | ||
|
|
||
| def write_ack(self, item_id: str, data: bytes) -> None: | ||
| folder = self.root / "inbox" / "acks" | ||
| folder.mkdir(parents=True, exist_ok=True) | ||
|
|
@@ -116,6 +124,11 @@ def list_items(self) -> list[str]: | |
| def read(self, rel: str) -> bytes: | ||
| return self._ssh(f"cat {shlex.quote(self.root + '/' + rel)}") | ||
|
|
||
| def list_archives(self) -> list[str]: | ||
| out = self._ssh(f"ls -1 {shlex.quote(self.root + '/archive')} 2>/dev/null || true").decode().split() | ||
| names = set(out) | ||
| return sorted(n for n in names if ARCHIVE_NAME.fullmatch(n) and f"{n}.sha256" in names) | ||
|
|
||
| def write_ack(self, item_id: str, data: bytes) -> None: | ||
| folder = shlex.quote(self.root + "/inbox/acks") | ||
| tmp, final = shlex.quote(f"{self.root}/inbox/acks/.{item_id}.tmp"), \ | ||
|
|
@@ -187,6 +200,39 @@ def _send_ack(self, ack: dict) -> bool: | |
| return False | ||
|
|
||
| # one round ----------------------------------------------------------------- | ||
| def sync_archives(self) -> dict: | ||
| """Pull every weekly trace archive not yet here, keeping only those whose | ||
| hash verifies (the GPU box is the off-machine copy).""" | ||
| from .archive import verify_archive | ||
|
|
||
| dest = self.state_dir / "archive" | ||
| counts = {"pulled": 0, "rejected": 0} | ||
| try: | ||
| names = self.transport.list_archives() | ||
| except (OutboxError, OSError) as exc: | ||
| self._trace(event="archive_unreachable", error=str(exc)) | ||
| return counts | ||
| for name in names: | ||
| if (dest / f"{name}.sha256").exists(): | ||
| continue | ||
|
Comment on lines
+215
to
+217
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
If a previously pulled tarball is deleted or corrupted while its Useful? React with 👍 / 👎. |
||
| try: | ||
| data = self.transport.read(f"archive/{name}") | ||
| manifest = self.transport.read(f"archive/{name}.sha256").decode("utf-8", "replace") | ||
| except (OutboxError, OSError) as exc: | ||
| self._trace(event="archive_unreachable", archive=name, error=str(exc)) | ||
| continue | ||
| if not verify_archive(data, manifest, name): | ||
| counts["rejected"] += 1 | ||
| self._trace(event="archive_rejected", archive=name) | ||
| continue | ||
| dest.mkdir(parents=True, exist_ok=True) | ||
| (dest / f".{name}.tmp").write_bytes(data) | ||
| (dest / f".{name}.tmp").replace(dest / name) | ||
| (dest / f"{name}.sha256").write_text(manifest, encoding="utf-8") # last: marks it complete | ||
|
Comment on lines
+228
to
+231
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a remote archive exists but Useful? React with 👍 / 👎. |
||
| counts["pulled"] += 1 | ||
| self._trace(event="archive_pulled", archive=name, bytes=len(data)) | ||
| return counts | ||
|
|
||
| def run_once(self) -> dict: | ||
| summary = {"performed": 0, "failed": 0, "resent": 0, "undelivered": 0, "dry_run": 0, "skipped": 0} | ||
| try: | ||
|
|
@@ -444,5 +490,5 @@ def serve(self, *, interval: float, stop=None) -> None: | |
|
|
||
| stop = stop or threading.Event() | ||
| while not stop.is_set(): | ||
| self._trace(event="round", **self.run_once()) | ||
| self._trace(event="round", **self.run_once(), **self.sync_archives()) | ||
| stop.wait(interval) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,105 @@ | ||
| """Weekly trace archives: made on the bot host, pulled and verified by the publisher.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import io | ||
| import tarfile | ||
| import time | ||
|
|
||
| from infermatrix_copilot.kb_service.archive import make_archive, verify_archive | ||
| from infermatrix_copilot.trace_store import TraceStore | ||
| from test_kb_flow import _flow_runtime | ||
|
|
||
|
|
||
| def _runtime(tmp_path): | ||
| rt, lifecycle = _flow_runtime(tmp_path) | ||
| rt.clock = lambda: time.time() + 5 # files written "now" are before the archive time | ||
| rt.traces = TraceStore(rt.state_dir / "traces", environ={}) | ||
| return rt | ||
|
|
||
|
|
||
| def test_an_archive_holds_new_records_and_blobs_and_restores_the_store(tmp_path): | ||
| rt = _runtime(tmp_path) | ||
| record = rt.traces.append("decision", inputs={"prompt": "why"}, result={"status": "pass"}) | ||
| path = make_archive(rt) | ||
| assert path is not None and (path.parent / f"{path.name}.sha256").exists() | ||
| manifest = (path.parent / f"{path.name}.sha256").read_text() | ||
| assert verify_archive(path.read_bytes(), manifest, path.name) | ||
| with tarfile.open(fileobj=io.BytesIO(path.read_bytes())) as tar: | ||
| names = tar.getnames() | ||
| assert "index.db" not in names and any(n.startswith("records/") for n in names) | ||
| restored = tmp_path / "restored" | ||
| tar.extractall(restored, filter="data") | ||
| store = TraceStore(restored) | ||
| assert store.rebuild_index() == 1 and store.get(record["id"]) == record | ||
| assert store.blob(record["inputs"]["prompt"]) == "why" | ||
|
|
||
|
|
||
| def test_archives_are_incremental_and_never_reuse_a_name(tmp_path): | ||
| rt = _runtime(tmp_path) | ||
| rt.clock = lambda: 1790000000.0 # every archive in the same second | ||
| rt.traces.append("decision", inputs={"prompt": "first"}, result={"status": "pass"}) | ||
| first = make_archive(rt) | ||
| assert make_archive(rt) is None # nothing new since | ||
| rt.traces.append("decision", inputs={"prompt": "second"}, result={"status": "fail"}) | ||
| second = make_archive(rt) | ||
| assert second is not None and second.name != first.name and first.exists() | ||
| with tarfile.open(fileobj=io.BytesIO(second.read_bytes())) as tar: | ||
| names = tar.getnames() | ||
| assert len([n for n in names if n.startswith("blobs/")]) == 1 # only the new blob | ||
| assert any(n.startswith("records/") for n in names) # the grown record file again | ||
|
|
||
|
|
||
| def test_a_blob_that_lands_during_archiving_is_in_the_next_archive(tmp_path): | ||
| import os | ||
|
|
||
| rt = _runtime(tmp_path) | ||
| rt.traces.append("decision", inputs={"prompt": "early"}, result={"status": "pass"}) | ||
| blobs = rt.state_dir / "traces" / "blobs" / "ab" | ||
| blobs.mkdir(parents=True, exist_ok=True) | ||
| pending = blobs / ".late.tmp" # put_blob's temp file, mid-write | ||
| pending.write_bytes(b"late") | ||
| os.utime(pending, (1, 1)) # an old mtime survives the rename | ||
| first = make_archive(rt) | ||
| pending.replace(blobs / ("ab" + "0" * 62 + ".gz")) # renamed after the scan | ||
| second = make_archive(rt) | ||
| assert second is not None | ||
| with tarfile.open(fileobj=io.BytesIO(second.read_bytes())) as tar: | ||
| assert tar.getnames() == ["blobs/ab/ab" + "0" * 62 + ".gz"] | ||
| assert first.name != second.name | ||
|
|
||
|
|
||
| def test_the_publisher_pulls_verified_archives_once_and_rejects_tampered_ones(tmp_path): | ||
| from test_kb_publisher import _setup | ||
|
|
||
| rt, lifecycle, _changeset_id, pub, _gh = _setup(tmp_path) | ||
| rt.clock = lambda: time.time() + 5 | ||
| rt.traces = TraceStore(rt.state_dir / "traces", environ={}) | ||
| rt.traces.append("decision", result={"status": "pass"}) | ||
| path = make_archive(rt) | ||
| assert pub.sync_archives() == {"pulled": 1, "rejected": 0} | ||
| assert (tmp_path / "publisher" / "archive" / path.name).read_bytes() == path.read_bytes() | ||
| assert pub.sync_archives() == {"pulled": 0, "rejected": 0} # once | ||
| other = path.parent / "traces-999999-20000101-000000.tar.gz" | ||
| other.write_bytes(b"not what the manifest says") | ||
| (path.parent / f"{other.name}.sha256").write_text(f"{'0' * 64} {other.name}\n") | ||
| assert pub.sync_archives() == {"pulled": 0, "rejected": 1} | ||
| assert not (tmp_path / "publisher" / "archive" / other.name).exists() | ||
|
|
||
|
|
||
| def test_restoring_in_name_order_survives_a_clock_moving_backwards(tmp_path): | ||
| rt = _runtime(tmp_path) | ||
| rt.clock = lambda: 1790000000.0 | ||
| first_record = rt.traces.append("decision", result={"status": "pass"}) | ||
| first = make_archive(rt) | ||
| rt.clock = lambda: 1780000000.0 # the clock jumped back | ||
| second_record = rt.traces.append("decision", result={"status": "fail"}) | ||
| second = make_archive(rt) | ||
| assert sorted([first.name, second.name]) == [first.name, second.name] | ||
| restored = tmp_path / "restored" | ||
| for archive in sorted((rt.state_dir / "archive").glob("traces-*.tar.gz")): | ||
| with tarfile.open(archive) as tar: | ||
| tar.extractall(restored, filter="data") | ||
| store = TraceStore(restored) | ||
| store.rebuild_index() | ||
| assert store.get(first_record["id"]) and store.get(second_record["id"]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
On the initial archive or after a busy week, every unarchived prompt/output blob is read into and retained by
members, after which the complete compressed tar is also built in aBytesIO. Memory therefore scales to roughly the raw weekly trace volume plus the archive, so a mature trace store can exhaust memory or make every archive attempt fail; write members directly to a temporary tar file and compute hashes while streaming.Useful? React with 👍 / 👎.