diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index 0b7eefca..49a05400 100644 --- a/src/agent_env/bundle/ledger.py +++ b/src/agent_env/bundle/ledger.py @@ -10,6 +10,7 @@ import hashlib import json +import stat from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass @@ -32,7 +33,7 @@ from .authoring import build_context_files, entry_files from .parse import Bundle, BundleKind -from .plan import Plan, Write, folder_walk, keeps_base_from_toml +from .plan import Plan, Write, folder_walk, keeps_base_from_toml, unpinned_store_refs from .resolve import BuiltImage LEDGER_COLLECTION = "bundle_ledger" @@ -61,6 +62,7 @@ class Check: reasons: tuple[str, ...] stored: int | None # the store's latest version of the id when checked needs: Mapping[tuple[str, str], int] # the version hashed for each earlier write it needs, by (store, id) + adopted: bool = False # ``version`` was written by an interrupted run of this bundle, which didn't record it @property def unchanged(self) -> bool: @@ -93,6 +95,10 @@ def check(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Check: if digest is None: untracked = ("its inputs aren't tracked yet, so it is written every run",) return Check(write, None, None, untracked, stored, needs) + orphan = self._orphan(write, stored, digest) + if orphan is not None: + return Check(write, digest, orphan, ("written by an interrupted run that didn't record it",), stored, needs, + adopted=True) row = self._latest(write.kind.store, write.id) recorded = row["version"] if row else None reasons = [] @@ -115,24 +121,25 @@ def digest(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Digest | inputs = {"type": _type(write), "config": _sha256(_canonical(config)), "files": {}, "needs": {}, "store_refs": {}} if write.kind is BundleKind.ARTIFACT: - files = build_context_files if isinstance(write.source, BuiltImage) else entry_files - for key, path in files(self._plan.bundle.bundle, write.source.entry).items(): - inputs["files"][key] = _file_sha256(path) + built = isinstance(write.source, BuiltImage) + for key, path in (build_context_files if built else entry_files)(self._plan.bundle.bundle, + write.source.entry).items(): + inputs["files"][key] = _file_sha256(path) + (_mode(path) if built else "") elif write.kind in (BundleKind.ENV, BundleKind.AGENT): - # An env's or agent's document records the versions of what it references, so one written anew, or - # a store entity it names without a version getting a new one, means it must be written again. A - # task or eval names its references without a version. + # An env's or agent's document records the versions of what it references, so one written anew means + # it must be written again. A task or eval names its references without a version. for store, id in write.needs: inputs["needs"][f"{store} {id}"] = str(needs[store, id]) - for ref in write.source.references: - if ref.local is None and ref.version is None: - inputs["store_refs"][f"{ref.kind} {ref.id}"] = str(self._plan.store_latest[ref.kind, ref.id]) + if write.kind not in (BundleKind.TASK, BundleKind.EVAL): + for (kind, id), version in unpinned_store_refs(self._plan, write).items(): + inputs["store_refs"][f"{kind} {id}"] = str(version) value = _sha256(_canonical({"scheme": SCHEME, "store": write.kind.store, "id": write.id, "inputs": inputs})) return Digest(value, inputs) def record(self, check: Check, write: Callable[[], int]) -> int: """Write ``check``'s entity with ``write``, which returns the version it wrote, and record it. - A pending row marks the attempt until it's done; one a crash of this bundle left is dropped here.""" + A pending row marks the attempt until it's done, with its version once ``write`` returns. One an + interrupted run of this bundle left is kept by ``check`` when its version can be, and dropped here.""" key = {"store": check.write.kind.store, "id": check.write.id, "bundle": self._bundle} self._store.ensure_index(LEDGER_COLLECTION, ["store", "id", "status"]) while self._store.delete(LEDGER_COLLECTION, Filter.of(**key, status="pending")): @@ -142,6 +149,9 @@ def record(self, check: Check, write: Callable[[], int]) -> int: "inputs": check.digest.inputs if check.digest else None, "at": _now()} self._store.insert(LEDGER_COLLECTION, pending) version = write() + # Stamped before the check below, which can take seconds, so a run interrupted during it leaves the version + # this attempt wrote, and the next run keeps it rather than writing it again. + self._store.replace(LEDGER_COLLECTION, Filter.of(**key, status="pending"), {**pending, "version": version}) done = {**pending, "status": "done", "version": version, "at": _now()} if check.digest is not None and self.digest(check.write, check.needs) != check.digest: # A file changed during the write, so what the version holds is unknown: the next run writes it again. @@ -149,6 +159,22 @@ def record(self, check: Check, write: Callable[[], int]) -> int: self._store.replace(LEDGER_COLLECTION, Filter.of(**key, status="pending"), done, upsert=True) return version + def adopt(self, check: Check) -> None: + """Record the version an interrupted run wrote, which ``check`` keeps, as that run would have.""" + key = {"store": check.write.kind.store, "id": check.write.id, "bundle": self._bundle} + done = {**key, "status": "done", "scheme": SCHEME, "digest": check.digest.value, "inputs": check.digest.inputs, + "version": check.version, "at": _now()} + self._store.replace(LEDGER_COLLECTION, Filter.of(**key, status="pending"), done, upsert=True) + + def _orphan(self, write: Write, stored: int | None, digest: Digest) -> int | None: + """The version an interrupted run of this bundle wrote and didn't record, when it's still the store's + latest and was made from what ``write`` is now made from.""" + pending = self._find(LEDGER_COLLECTION, Filter.of(store=write.kind.store, id=write.id, bundle=self._bundle, + status="pending")) + if pending is None or pending.get("version") is None or pending["version"] != stored: + return None + return stored if (pending["scheme"], pending["digest"]) == (SCHEME, digest.value) else None + def _latest(self, store: str, id: str) -> dict | None: return self._find(LEDGER_COLLECTION, Filter.of(store=store, id=id, status="done")) @@ -245,6 +271,16 @@ def _sha256(data: bytes) -> str: return "sha256:" + hashlib.sha256(data).hexdigest() +# A build copies each file's permission bits into the image, so they're part of what it's made from; a file with the +# usual rw-r--r-- is hashed as its content alone. +_USUAL_MODE = 0o644 + + +def _mode(path: Path) -> str: + mode = stat.S_IMODE(path.stat().st_mode) + return "" if mode == _USUAL_MODE else f"+{mode:o}" + + def _file_sha256(path: Path) -> str: digest = hashlib.sha256() with open(path, "rb") as file: diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index c4fefd60..84bc69c3 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -38,7 +38,7 @@ from .authoring import AuthoringContext, build_context_files from .ledger import Ledger, materializing from .parse import BundleError, BundleKind -from .plan import Plan, Write, folder_walk +from .plan import Plan, Write, folder_walk, unpinned_store_refs from .resolve import BuiltImage, build_step @dataclass(frozen=True) @@ -99,6 +99,8 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None: check = ledger.check(write, {need: done[need].version for need in write.needs}) if check.unchanged: version = check.version + if check.adopted and not dry_run: + ledger.adopt(check) elif dry_run: version = check.next_version else: @@ -184,10 +186,11 @@ def _write_agent(plan: Plan, write: Write) -> int: def _pinned(plan: Plan, write: Write, refs: tuple[EntityRef, ...]) -> Any: """A copy of ``write``'s resolved toml with each store ref that names no version pinned to the version - the plan checked, which the ledger hashed, so one the store gains before the write isn't written.""" + the plan read, which the ledger hashed (``unpinned_store_refs``).""" config = copy.deepcopy(write.source.config) + pins = unpinned_store_refs(plan, write) for site in ref_sites(refs, config, inline_pins=True): - planned = plan.store_latest.get((site.ref.kind, site.value)) if site.version is None else None + planned = pins.get((site.ref.kind, site.value)) if site.version is None else None if planned is None: continue if site.version_key is None: diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index b9ba3358..4d3e48fb 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -341,6 +341,16 @@ def _path(self, entry: BundleEntry) -> str: return relative(self.resolved.bundle.root, entry.path) +def unpinned_store_refs(plan: Plan, write: Write) -> dict[tuple[EntityKind, str], int]: + """The version the plan read for each store entity ``write`` names without a version, by (kind, id). Its writer + pins each to that version and the ledger hashes them, so a version the store gains in between isn't written, + and one it gains later rewrites ``write``.""" + if isinstance(write.source, BuiltImage): + return {} + return {(ref.kind, ref.id): plan.store_latest[ref.kind, ref.id] + for ref in write.source.references if ref.local is None and ref.version is None} + + def folder_walk(cls: type | None) -> Any: """The listing a type's write runs: a subclass keeping its base's ``from_toml`` lists the same way, one that overrides it may read its folder differently, so it gets none.""" diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index c4bc7e8b..93f00fa8 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -422,7 +422,9 @@ def _task_run_command(ref: Reference) -> str: def _written(plan: Plan, done: Materialized) -> str: what = f"{_label(plan, done.write)}: v{done.version}" - return f"{what}, unchanged" if done.reused else f"{what} ({'; '.join(done.reasons)})" + if done.reused: + return f"{what}, unchanged" + (f" ({'; '.join(done.reasons)})" if done.reasons else "") + return f"{what} ({'; '.join(done.reasons)})" def _label(plan: Plan, write: Write) -> str: diff --git a/tst/unit/bundle/ledger_test.py b/tst/unit/bundle/ledger_test.py index fa647abe..c66e9820 100644 --- a/tst/unit/bundle/ledger_test.py +++ b/tst/unit/bundle/ledger_test.py @@ -1,6 +1,7 @@ """The bundle ledger: a re-run reuses every version whose inputs haven't changed, and says why the rest are written anew.""" +import hashlib import json import shutil import subprocess @@ -11,13 +12,15 @@ import pytest +from agent_env.a2a_agent import A2AAgent from agent_env.artifact.artifacts.file import FileArtifact from agent_env.bundle import parse_bundle from agent_env.bundle import plan as plan_module from agent_env.bundle import resolve as resolve_module from agent_env.bundle.ledger import LEDGER_COLLECTION, Ledger, materializing +from agent_env.bundle.materialize import _pinned from agent_env.config import configure, get_config -from agent_env.entity_refs import EntityRef +from agent_env.entity_refs import EntityKind, EntityRef from agent_env.env.env import Env from agent_env.store import Filter, Sort, UpdateSpec from tst.unit.bundle._support import RefusingStore, local_store, plan_of @@ -123,6 +126,8 @@ def _record_changed(ledger, writes): checks, versions = {}, {} for write in writes: check = ledger.check(write, {need: versions[need] for need in write.needs}) + if check.adopted: + ledger.adopt(check) version = check.version if check.unchanged else ledger.record(check, lambda: _written(write)) versions[write.kind.store, write.id] = version checks[write.id] = check @@ -379,6 +384,21 @@ def test_an_agent_is_rewritten_when_its_toml_changes_or_its_unpinned_store_image assert checks[f"{ROOT}/t"].unchanged +def test_a_store_ref_named_without_a_version_is_pinned_at_the_version_the_ledger_hashes(bundle_dir): + get_config().get_document_store().insert("artifacts", _image_document(3)) + layout = {"agents/solver/agent.toml": 'image = "claude-image"\n', + "tasks/t.json": _steps({"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "solver"})} + for rel, text in layout.items(): + (bundle_dir / rel).parent.mkdir(parents=True, exist_ok=True) + (bundle_dir / rel).write_text(text) + plan = plan_of(bundle_dir) + agent = next(write for write in plan.writes if write.id == f"{ROOT}/solver") + + assert plan_module.unpinned_store_refs(plan, agent) == {(EntityKind.ARTIFACT, "claude-image"): 3} + assert _pinned(plan, agent, A2AAgent.toml_refs)["image"] == {"artifact": "claude-image", "version": 3} + assert Ledger.for_plan(plan).digest(agent, {}).inputs["store_refs"] == {"artifact claude-image": "3"} + + def _image_document(version): return {"id": "claude-image", "version": version, "type": "docker_image", "description": "claude", "image_name": f"claude:v{version}", "tar_gz_s3_url": f"file:///claude-v{version}.tar.gz"} @@ -443,6 +463,84 @@ def test_a_built_image_is_made_from_every_file_of_its_folder_the_toml_too(bundle assert not rebuilt[agent].unchanged +def test_a_built_images_permission_bits_are_inputs_and_a_file_artifacts_arent(bundle_dir): + (bundle_dir / "agents/solver").mkdir(parents=True) + (bundle_dir / "agents/solver/Dockerfile").write_text("FROM scratch\nCOPY run.sh /\n") + run_sh = bundle_dir / "agents/solver/run.sh" + run_sh.write_text("echo hi\n") + run_sh.chmod(0o644) + (bundle_dir / "tasks/t.json").write_text( + _steps({"id": "agent", "type": "deploy_agent", "env_ids": ["tickets"], "a2a_agent_id": "solver"})) + image = f"{ROOT}/solver__agent_image" + plan = plan_of(bundle_dir) + built = next(write for write in plan.writes if write.id == image) + # A file with the usual rw-r--r-- is hashed as its content alone, as before, so earlier rows stay reused. + assert Ledger.for_plan(plan).digest(built, {}).inputs["files"]["run.sh"] == _file_digest(run_sh) + _run(bundle_dir) + + rebuilt = [] + for mode in (0o755, 0o700): + run_sh.chmod(mode) + (bundle_dir / "artifacts/greeting/hello.txt").chmod(mode) + rebuilt.append(_run(bundle_dir)) + + assert [checks[image].reasons for checks in rebuilt] == [("files changed: run.sh",)] * 2 + assert all(checks[GREETING].unchanged for checks in rebuilt) + assert _run(bundle_dir)[image].unchanged + + +def _file_digest(path): + return "sha256:" + hashlib.sha256(path.read_bytes()).hexdigest() + + +def _interrupted_after_its_write(ledger, write, monkeypatch): + """Record ``write`` as a run does, interrupted once its write has returned, while the ledger checks what + the write held.""" + def interrupted(*_): + raise KeyboardInterrupt + + check = ledger.check(write, {}) + monkeypatch.setattr(ledger, "digest", interrupted) + with pytest.raises(KeyboardInterrupt): + ledger.record(check, lambda: _written(write)) + + +def test_a_version_an_interrupted_run_wrote_is_kept_and_recorded_not_written_again(bundle_dir, monkeypatch): + plan = plan_of(bundle_dir) + write = next(w for w in plan.writes if w.id == GREETING) + _interrupted_after_its_write(Ledger.for_plan(plan), write, monkeypatch) + + predicted = _checked(Ledger.for_plan(plan), plan.writes)[GREETING] + kept = _record_changed(Ledger.for_plan(plan), plan.writes)[GREETING] + + for check in (predicted, kept): + assert (check.version, check.adopted) == (1, True) + assert check.reasons == ("written by an interrupted run that didn't record it",) + assert _latest_version("artifacts", GREETING) == 1 + after = _record_changed(Ledger.for_plan(plan), plan.writes)[GREETING] + assert (after.version, after.adopted, after.reasons) == (1, False, ()) + assert [row["status"] for row in local_store().query(LEDGER_COLLECTION, Filter.of(id=GREETING))] == ["done"] + + +@pytest.mark.parametrize("since", ["another-write", "an-edit"]) +def test_an_interrupted_runs_version_is_written_over_once_the_store_or_the_files_have_moved_on( + bundle_dir, monkeypatch, since, +): + plan = plan_of(bundle_dir) + write = next(w for w in plan.writes if w.id == GREETING) + _interrupted_after_its_write(Ledger.for_plan(plan), write, monkeypatch) + if since == "another-write": + _written(write) + else: + (bundle_dir / "artifacts/greeting/hello.txt").write_text("edited\n") + + check = Ledger.for_plan(plan).check(write, {}) + + stored = 2 if since == "another-write" else 1 + assert not check.unchanged and not check.adopted + assert check.reasons == (f"the store's latest, v{stored}, wasn't recorded by this bundle",) + + @pytest.mark.parametrize("other", [ "the-same-folder", "another-folder-with-its-id-root", "another-bundle-declaring-its-id", ]) diff --git a/tst/unit/bundle/materialize_test.py b/tst/unit/bundle/materialize_test.py index 6a082343..c912d39b 100644 --- a/tst/unit/bundle/materialize_test.py +++ b/tst/unit/bundle/materialize_test.py @@ -21,6 +21,8 @@ from agent_env.bundle import materialize as materialize_module from agent_env.bundle.ledger import LEDGER_COLLECTION, Ledger from agent_env.bundle.materialize import materialize +from agent_env.bundle.parse import BundleKind +from agent_env.bundle.run import _written from agent_env.config import configure from agent_env.config.paths import state_root from agent_env.config.runtime import Config @@ -712,6 +714,39 @@ def agree(): assert agree()[f"{ROOT}/greeting"] == (3, False, ("the store's latest, v2, wasn't recorded by this bundle",)) +def test_a_version_a_run_wrote_before_it_was_interrupted_is_kept_by_the_next_run_and_its_dry_run( + bundle_dir, monkeypatch, +): + interrupted, write_artifact, digest = [], materialize_module._WRITERS[BundleKind.ARTIFACT], Ledger.digest + + def written(plan, write): # the greeting's write lands, and then the run is interrupted + version = write_artifact(plan, write) + interrupted.append(write.id == f"{ROOT}/greeting") + return version + + def digest_until_interrupted(self, *args): + if any(interrupted): + raise KeyboardInterrupt + return digest(self, *args) + + monkeypatch.setitem(materialize_module._WRITERS, BundleKind.ARTIFACT, written) + monkeypatch.setattr(Ledger, "digest", digest_until_interrupted) + with pytest.raises(KeyboardInterrupt): + _run(bundle_dir) + monkeypatch.setitem(materialize_module._WRITERS, BundleKind.ARTIFACT, write_artifact) + monkeypatch.setattr(Ledger, "digest", digest) + + kept = (1, True, ("written by an interrupted run that didn't record it",)) + assert _summary(_run(bundle_dir, dry_run=True))[f"{ROOT}/greeting"] == kept + done = _run(bundle_dir) + assert _summary(done)[f"{ROOT}/greeting"] == kept + greeting = next(item for item in done.writes if item.write.id == f"{ROOT}/greeting") + assert _written(done.plan, greeting) == ( + "artifacts/greeting: v1, unchanged (written by an interrupted run that didn't record it)") + assert _summary(_run(bundle_dir))[f"{ROOT}/greeting"] == (1, True, ()) + assert [d["version"] for d in local_store().query("artifacts", Filter.of(id=f"{ROOT}/greeting"))] == [1] + + def test_a_dry_run_calls_no_writer_and_takes_no_lock(bundle_dir, monkeypatch): layout(bundle_dir, {"evals/regression.toml": 'tasks = ["t"]\n'})