diff --git a/.github/scripts/prune_registry.py b/.github/scripts/prune_registry.py new file mode 100644 index 000000000..0b2ace7f8 --- /dev/null +++ b/.github/scripts/prune_registry.py @@ -0,0 +1,98 @@ +"""Keep only the images of a local registry's repository that a run used, before the registry's storage is cached. + +``prune`` drops each tag the registry's log has no manifest request for, then each manifest no remaining tag reaches, +following an image index to the manifests it lists. The registry's ``garbage-collect`` then frees the layers only those +manifests used; with ``--delete-untagged`` it would also delete a kept index's untagged platform manifests. ``links``, +run after it, drops the repository's links to layers it freed. +""" + +from __future__ import annotations + +import argparse +import json +import re +import shutil +import sys +from pathlib import Path + +_V2 = Path("docker") / "registry" / "v2" + + +def main(argv: list[str]) -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + pruning = commands.add_parser("prune", help="drop the tags and manifests the run didn't use") + linking = commands.add_parser("links", help="drop links to layers garbage-collect freed") + for command in (pruning, linking): + command.add_argument("storage", type=Path, help="the registry's storage root") + command.add_argument("repository") + pruning.add_argument("log", type=Path, help="the registry's log of the run") + args = parser.parse_args(argv) + if args.command == "prune": + kept, dropped = prune(args.storage, args.repository, used_tags(args.log.read_text(errors="replace"), + args.repository)) + print(f"kept {kept} tags of {args.repository}, dropped {dropped}") + else: + print(f"dropped {drop_freed_links(args.storage, args.repository)} links to freed layers") + return 0 + + +def used_tags(log: str, repository: str) -> set[str]: + """The tags the log names a request for a manifest of ``repository`` by.""" + return set(re.findall(rf"/v2/{re.escape(repository)}/manifests/([A-Za-z0-9_][A-Za-z0-9_.-]{{0,127}})(?=[\s\"?])", + log)) + + +def prune(storage: Path, repository: str, used: set[str]) -> tuple[int, int]: + """Drop the tags of ``repository`` not in ``used``, then the manifests no remaining tag reaches; how many tags it + kept and dropped.""" + manifests = storage / _V2 / "repositories" / repository / "_manifests" + tags = _children(manifests / "tags") + for tag in tags: + if tag.name not in used: + shutil.rmtree(tag) + kept = [tag for tag in tags if tag.name in used] + reached = _reached(storage, [(tag / "current" / "link").read_text().strip() for tag in kept]) + for revision in _children(manifests / "revisions" / "sha256"): + if f"sha256:{revision.name}" not in reached: + shutil.rmtree(revision) + return len(kept), len(tags) - len(kept) + + +def drop_freed_links(storage: Path, repository: str) -> int: + """Drop each link of ``repository`` to a layer that's no longer stored; how many.""" + dropped = 0 + for layer in _children(storage / _V2 / "repositories" / repository / "_layers" / "sha256"): + if not _blob(storage, f"sha256:{layer.name}").is_file(): + shutil.rmtree(layer) + dropped += 1 + return dropped + + +def _reached(storage: Path, digests: list[str]) -> set[str]: + """``digests`` and every manifest an image index among them lists, recursively.""" + reached, pending = set(), list(digests) + while pending: + digest = pending.pop() + if digest in reached: + continue + reached.add(digest) + try: + manifest = json.loads(_blob(storage, digest).read_bytes()) + except (FileNotFoundError, ValueError): + continue + pending += [listed["digest"] for listed in manifest.get("manifests", []) if "digest" in listed] + return reached + + +def _children(folder: Path) -> list[Path]: + return [child for child in folder.iterdir() if child.is_dir()] if folder.is_dir() else [] + + +def _blob(storage: Path, digest: str) -> Path: + hex_digest = digest.removeprefix("sha256:") + return storage / _V2 / "blobs" / "sha256" / hex_digest[:2] / hex_digest / "data" + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/.github/workflows/local-backends.yml b/.github/workflows/local-backends.yml index ed7fe161d..71e12e08e 100644 --- a/.github/workflows/local-backends.yml +++ b/.github/workflows/local-backends.yml @@ -169,17 +169,21 @@ jobs: - name: Pull the base images from mirror.gcr.io run: python3 .github/scripts/mirror_base_images.py src tst # A computer image or disk is tagged with a key of everything it is built from, and the composer reuses one the - # image store already holds for its key. The registry keeps its storage on the runner, restored from the computer - # images main saved latest this week, so a run reuses them rather than building them again. A shard that runs no - # computer tests misses. + # image store already holds for its key. That key includes the lock of the packages the image installs, which the + # composer resolves from Ubuntu's archive, PyPI and download URLs unless the store holds it. The registry keeps + # its storage on the runner, and the computer tests share their locks and files in the packages folder + # (tst/util/package_cache.py); both are restored from what main saved latest this week, so a run neither resolves + # nor builds again what hasn't changed. A shard that runs no computer tests misses. - name: Week of the computer images' cache id: week run: echo "week=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT" - - name: Restore the computer images + - name: Restore the computer images and packages id: computers uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - path: ${{ runner.temp }}/registry + path: | + ${{ runner.temp }}/registry + ${{ runner.temp }}/packages key: computers-${{ matrix.shard }}-${{ steps.week.outputs.week }} restore-keys: computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}- - name: Start the local registry @@ -195,35 +199,49 @@ jobs: id: tests env: AGENTENV_TEST_SHARD: ${{ matrix.shard }}/12 + AGENTENV_TEST_PACKAGES: ${{ runner.temp }}/packages run: | .venv/bin/python -m pytest -m int_test_slow -rA --timeout=1200 --junitxml=junit-integration-slow.xml \ tst/integration --ignore=tst/integration/env/gateway/gateway_test.py - # Only the computer images are kept, and only main saves them, once its tests pass, so pull requests reuse main's - # without filling the repository's cache. The entry is named by the image keys the registry holds, so a run that - # built any new image saves a new one. The garbage collection frees what only the deleted repositories - # referenced; with --delete-untagged it would also delete the untagged platform manifests of an image index, and - # their layers. - - name: Keep only the computer images + # Only the computer images this run used are kept, with the packages, and only main saves them, once its tests + # pass, so pull requests reuse main's without filling the repository's cache. An image is used when the registry's + # log has a request for it by its tag (prune_registry.py), so what a week's earlier runs built and this one no + # longer asks for is dropped. The garbage collection frees what only the dropped tags and repositories referenced; + # with --delete-untagged it would also delete the untagged platform manifests of an image index, and their layers. + # The entry is named by the image keys the registry holds and the locks the packages hold, so a run that built a + # new image or resolved a new lock saves a new one. + - name: Keep only the computer images this run used id: kept if: always() run: | repositories="$RUNNER_TEMP/registry/docker/registry/v2/repositories" if [ ! -d "$repositories/agentenv-computers" ]; then echo "computers=false" >> "$GITHUB_OUTPUT"; exit 0; fi + docker logs registry > "$RUNNER_TEMP/registry.log" 2>&1 docker stop registry sudo find "$repositories" -mindepth 1 -maxdepth 1 ! -name agentenv-computers -exec rm -rf {} + + sudo python3 .github/scripts/prune_registry.py prune "$RUNNER_TEMP/registry" agentenv-computers \ + "$RUNNER_TEMP/registry.log" docker run --rm --volume "$RUNNER_TEMP/registry:/var/lib/registry" mirror.gcr.io/library/registry:2 \ garbage-collect /etc/docker/registry/config.yml > /dev/null - du -sh "$RUNNER_TEMP/registry" - tags=$(sudo ls "$repositories/agentenv-computers/_manifests/tags" | sort | sha256sum | cut -c1-16) - echo "key=computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}-$tags" >> "$GITHUB_OUTPUT" + sudo python3 .github/scripts/prune_registry.py links "$RUNNER_TEMP/registry" agentenv-computers + mkdir -p "$RUNNER_TEMP/packages" + du -sh "$RUNNER_TEMP/registry" "$RUNNER_TEMP/packages" + if [ -z "$(sudo ls "$repositories/agentenv-computers/_manifests/tags")" ]; then + echo "computers=false" >> "$GITHUB_OUTPUT"; exit 0 + fi + key=$({ sudo ls "$repositories/agentenv-computers/_manifests/tags" + ls "$RUNNER_TEMP/packages/agentenv-packages/locks" 2>/dev/null; } | sort | sha256sum | cut -c1-16) + echo "key=computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}-$key" >> "$GITHUB_OUTPUT" echo "computers=true" >> "$GITHUB_OUTPUT" - - name: Save the computer images + - name: Save the computer images and packages if: >- always() && github.event_name == 'push' && steps.tests.outcome == 'success' && steps.kept.outputs.computers == 'true' && steps.kept.outputs.key != steps.computers.outputs.cache-matched-key uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - path: ${{ runner.temp }}/registry + path: | + ${{ runner.temp }}/registry + ${{ runner.temp }}/packages key: ${{ steps.kept.outputs.key }} - name: Skips must be declared capability gaps if: always() diff --git a/AGENTS.md b/AGENTS.md index 817dac87f..780c14bd8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -389,6 +389,11 @@ orchestrators call it the same way. What callers rely on: `tst/integration/slow_durations.json`. The unit and installer tiers are split the same way. When the shards' times drift apart, rewrite the table from all of a `main` run's `junit-*` artifacts with `.github/scripts/slow_test_durations.py`. +- The computer tests compose in stores of their own, but share the package locks and files a compose + resolves (`tst/util/package_cache.py`) through `AGENTENV_TEST_PACKAGES`, a folder for the run unless + it's set. The slow tier caches that folder with the local registry's computer images, so a run + resolves from Ubuntu's archive, PyPI or a download URL only what changed since `main` last saved + them. - `tst/manual/` and `tst/benchmarks/` run only by hand, as each module's docstring says; no CI job collects them. - `tst/formal/` holds TLA+ models, run only by hand with each model's `check.sh`; the `AGENTS.md` diff --git a/tst/conftest.py b/tst/conftest.py index e33f3a589..9dbe4d0aa 100644 --- a/tst/conftest.py +++ b/tst/conftest.py @@ -12,7 +12,7 @@ from agent_env.store.document_store import document_store as document_store_module from agent_env.store.object_store.local import grant_server as grant_server_module from agent_env.store.routing import disable_namespace_routing -from tst.util import shards +from tst.util import package_cache, shards def _reset_store_singletons() -> None: @@ -27,6 +27,7 @@ def _reset_store_singletons() -> None: _STATE_DIR = pytest.StashKey[str]() +_PACKAGES_DIR = pytest.StashKey[str]() _state_home = pytest.MonkeyPatch() @@ -36,11 +37,17 @@ def pytest_configure(config): (the capability checks behind skip marks).""" config.stash[_STATE_DIR] = tempfile.mkdtemp(prefix="agent-env-state-") _state_home.setenv("XDG_STATE_HOME", config.stash[_STATE_DIR]) + # The computer tests' resolved packages, shared by their stores for the run unless CI names a folder it keeps. + if not os.environ.get(package_cache.ENV_VAR): + config.stash[_PACKAGES_DIR] = tempfile.mkdtemp(prefix="agent-env-packages-") + _state_home.setenv(package_cache.ENV_VAR, config.stash[_PACKAGES_DIR]) def pytest_unconfigure(config): _state_home.undo() shutil.rmtree(config.stash[_STATE_DIR], ignore_errors=True) + if packages := config.stash.get(_PACKAGES_DIR, None): + shutil.rmtree(packages, ignore_errors=True) @pytest.hookimpl(trylast=True) diff --git a/tst/integration/cli/run_bundle_browser_use_test.py b/tst/integration/cli/run_bundle_browser_use_test.py index e28c56c1d..e932c9fc5 100644 --- a/tst/integration/cli/run_bundle_browser_use_test.py +++ b/tst/integration/cli/run_bundle_browser_use_test.py @@ -27,6 +27,7 @@ from agent_env.task_step.task_steps.browser.connect_cdp import BrowserConnectCdpTaskStep from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep +from tst.util import package_cache from tst.util.fake_openai import FakeOpenAI, Text, ToolCall pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -69,11 +70,13 @@ def state(monkeypatch, tmp_path, endpoint): monkeypatch.delenv(name, raising=False) configure() reset_artifact_store() + package_cache.seed() logging.disable(logging.CRITICAL) # pytest's live logging would take CliRunner's stdout try: yield tmp_path assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a run left a sandbox work folder" finally: + package_cache.keep() logging.disable(logging.NOTSET) grant_server(None, None).close() reset_artifact_store() diff --git a/tst/integration/computer/conftest.py b/tst/integration/computer/conftest.py index 543d63d12..a95538ded 100644 --- a/tst/integration/computer/conftest.py +++ b/tst/integration/computer/conftest.py @@ -2,7 +2,8 @@ shipped desktop example's computer env, the cua-driver test bundle's, the chromium test bundle's, the website-browser test bundle's, the TLS clients test bundle's and agentenv's chromium deployed alone, each composed once for a module's tests; -``state``, local stores for a test that runs a bundle; ``docker``, which runs the docker CLI for them; and +``state``, local stores for a test that runs a bundle; each of these stores starts with the packages the run has resolved +(``tst.util.package_cache``); ``docker``, which runs the docker CLI for them; and ``cdp_clients``, Playwright and Puppeteer installed to connect to a CDP address.""" import logging @@ -27,6 +28,7 @@ from agent_env.store.object_store.local.grant_server import grant_server from agent_env.store.routing import namespace_routing from agent_env.task.store import get_task_instance_store +from tst.util import package_cache from tst.util.published_image import pull_published_else_build BASE = Path(agent_env.__file__).parent / "computers" / "artifacts" / "ubuntu-24.04" @@ -118,9 +120,13 @@ def _composed(root: Path, bundle, env_id: str | None = None): patch.setenv("AGENT_ENV_OBJECT_STORE", "local") configure() reset_artifact_store() + package_cache.seed() with namespace_routing(): plan = plan_bundle(resolve_bundle(bundle), installed=installed_bundles) - materialize(plan, build_here=frozenset(w.id for w in plan.writes if isinstance(w.source, BuiltImage))) + try: + materialize(plan, build_here=frozenset(w.id for w in plan.writes if isinstance(w.source, BuiltImage))) + finally: + package_cache.keep() yield Env.get(env_id or f"{bundle.id_root}/desktop") reset_artifact_store() reset_config() @@ -142,11 +148,13 @@ def state(monkeypatch, tmp_path): monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) configure() reset_artifact_store() + package_cache.seed() logging.disable(logging.CRITICAL) try: yield tmp_path assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a run left a sandbox work folder" finally: + package_cache.keep() logging.disable(logging.NOTSET) grant_server(None, None).close() reset_artifact_store() diff --git a/tst/unit/test_package_cache.py b/tst/unit/test_package_cache.py new file mode 100644 index 000000000..7b7f8c6dd --- /dev/null +++ b/tst/unit/test_package_cache.py @@ -0,0 +1,119 @@ +"""``tst.util.package_cache``: the computer tests' stores share their package locks and files through one folder.""" + +from __future__ import annotations + +import errno +import hashlib + +import pytest + +from agent_env.config import configure, get_config, reset_config +from agent_env.env.envs.computer.utils.packages import PREFIX, Lock, LockedDeb, LockedDownload, LockedFile +from tst.util import package_cache + +_INPUTS = "a" * 64 + + +@pytest.fixture +def stores(monkeypatch, tmp_path): + """A local object store under this test's folder and the shared folder; the store's root.""" + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv(package_cache.ENV_VAR, str(tmp_path / "shared")) + configure() + yield get_config().get_object_store() + reset_config() + + +def test_a_lock_a_compose_stored_is_shared_with_its_files_and_seeds_the_next_store(stores, tmp_path, monkeypatch): + files = _store_lock(stores, {"apt.deb": b"deb", "wheel.whl": b"wheel", "chrome.zip": b"zip"}) + + package_cache.keep() + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "next")) + reset_config() + configure() + package_cache.seed() + + seeded = get_config().get_object_store() + assert seeded.root != stores.root + assert Lock.model_validate_json(seeded.read(_lock_key())).inputs == _INPUTS + assert seeded.get_object_metadata(_lock_key()).content_type == "application/json" + assert {sha256: seeded.read(f"{PREFIX}/files/{sha256}") for sha256 in files} == files + + +def test_where_folders_cant_share_a_hard_link_a_copied_lock_keeps_its_content_type(stores, tmp_path, monkeypatch): + """The store names an object's content type by the object's size and modification time, so a copy keeps both.""" + files = _store_lock(stores, {"apt.deb": b"deb"}) + + def cross_device(*_args): + raise OSError(errno.EXDEV, "Invalid cross-device link") + + with monkeypatch.context() as across: + across.setattr(package_cache.os, "link", cross_device) + package_cache.keep() + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "next")) + reset_config() + configure() + package_cache.seed() + + seeded = get_config().get_object_store() + assert seeded.get_object_metadata(_lock_key()).content_type == "application/json" + assert {sha256: seeded.read(f"{PREFIX}/files/{sha256}") for sha256 in files} == files + assert not (tmp_path / "shared" / PREFIX / "files" / next(iter(files))).samefile( + stores.root / PREFIX / "files" / next(iter(files))) + + +def test_a_lock_whose_files_arent_all_shared_isnt_seeded(stores, tmp_path, monkeypatch): + files = _store_lock(stores, {"apt.deb": b"deb", "wheel.whl": b"wheel"}) + package_cache.keep() + (tmp_path / "shared" / PREFIX / "files" / next(iter(files))).unlink() + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "next")) + reset_config() + configure() + + package_cache.seed() + + assert not get_config().get_object_store().exists(_lock_key()) + + +def test_keeping_never_replaces_what_is_shared(stores, tmp_path): + files = _store_lock(stores, {"apt.deb": b"deb"}) + (sha256,) = files + shared = tmp_path / "shared" / PREFIX / "files" / sha256 + shared.parent.mkdir(parents=True) + shared.write_bytes(b"already shared") + + package_cache.keep() + + assert shared.read_bytes() == b"already shared" + + +def test_without_a_shared_folder_nothing_is_shared(stores, tmp_path, monkeypatch): + monkeypatch.delenv(package_cache.ENV_VAR) + _store_lock(stores, {"apt.deb": b"deb"}) + + package_cache.keep() + package_cache.seed() + + assert not (tmp_path / "shared").exists() + + +def _store_lock(store, named: dict[str, bytes]) -> dict[str, bytes]: + """A lock of ``named``'s files, stored with them as a compose stores them; the files by their sha256.""" + files = {hashlib.sha256(data).hexdigest(): data for data in named.values()} + sha = {name: hashlib.sha256(data).hexdigest() for name, data in named.items()} + for sha256, data in files.items(): + store.put(f"{PREFIX}/files/{sha256}", data) + debs = [LockedDeb(file=name, sha256=sha[name], size=len(named[name]), package="p", version="1", arch="amd64") + for name in named if name.endswith(".deb")] + wheels = [LockedFile(file=name, sha256=sha[name], size=len(named[name])) for name in named if name.endswith(".whl")] + downloads = [LockedDownload(file=name, sha256=sha[name], size=len(named[name]), url=f"https://example.com/{name}", + to=name) for name in named if name.endswith(".zip")] + lock = Lock(inputs=_INPUTS, arch="amd64", apt=debs, wheels={"driver": wheels} if wheels else {}, + downloads={"app": downloads} if downloads else {}) + store.put(_lock_key(), lock.model_dump_json().encode(), content_type="application/json") + return files + + +def _lock_key() -> str: + return f"{PREFIX}/locks/{_INPUTS}.json" diff --git a/tst/unit/test_prune_registry_script.py b/tst/unit/test_prune_registry_script.py new file mode 100644 index 000000000..2b842562c --- /dev/null +++ b/tst/unit/test_prune_registry_script.py @@ -0,0 +1,97 @@ +"""``.github/scripts/prune_registry.py``, which keeps only the computer images a slow-tier run used before the local +registry's storage is cached: their tags, the manifests those reach, and the layer links garbage collection leaves.""" + +from __future__ import annotations + +import ast +import hashlib +import json +import sys +from pathlib import Path + +import pytest + +from tst.util.scripts import SCRIPTS, load_script + +SCRIPT = SCRIPTS / "prune_registry.py" +REPOSITORY = "agentenv-computers" +_USED, _STALE = "a" * 64, "b" * 64 + + +@pytest.fixture +def script(): + return load_script(SCRIPT.name) + + +def test_it_imports_only_the_standard_library(): + """It runs under the runner's python3 with sudo, outside the project's environment.""" + imported = set() + for node in ast.walk(ast.parse(SCRIPT.read_text())): + if isinstance(node, ast.Import): + imported |= {alias.name.split(".")[0] for alias in node.names} + elif isinstance(node, ast.ImportFrom) and node.level == 0: + imported.add(node.module.split(".")[0]) + assert imported <= set(sys.stdlib_module_names) | {"__future__"} + + +def test_a_tag_is_used_when_the_log_has_a_request_for_its_manifest_by_it(script): + log = "\n".join([ + f'time="t" level=info msg="response completed" http.request.method=HEAD ' + f'http.request.uri="/v2/{REPOSITORY}/manifests/{_USED}" http.response.status=200', + f'172.17.0.1 - - [10/Nov/2026:05:00:00 +0000] "PUT /v2/{REPOSITORY}/manifests/{"c" * 64} HTTP/1.1" 201 0', + f'172.17.0.1 - - [10/Nov/2026:05:00:01 +0000] "GET /v2/{REPOSITORY}/manifests/sha256:{"d" * 64} HTTP/1.1" 200', + f'172.17.0.1 - - [10/Nov/2026:05:00:02 +0000] "HEAD /v2/other/manifests/{_STALE} HTTP/1.1" 200 0', + ]) + + assert script.used_tags(log, REPOSITORY) == {_USED, "c" * 64} + + +def test_pruning_keeps_the_used_tags_and_every_manifest_their_index_lists(script, tmp_path): + platforms = [_blob(tmp_path, {"layers": [], "n": n}) for n in range(2)] + index = _blob(tmp_path, {"manifests": [{"digest": digest} for digest in platforms]}) + stale = _blob(tmp_path, {"layers": [], "n": "stale"}) + _tag(tmp_path, _USED, index) + _tag(tmp_path, _STALE, stale) + for digest in (index, *platforms, stale): + _revision(tmp_path, digest) + + kept, dropped = script.prune(tmp_path, REPOSITORY, {_USED}) + + manifests = tmp_path / "docker/registry/v2/repositories" / REPOSITORY / "_manifests" + assert (kept, dropped) == (1, 1) + assert sorted(p.name for p in (manifests / "tags").iterdir()) == [_USED] + assert {p.name for p in (manifests / "revisions/sha256").iterdir()} == {d.removeprefix("sha256:") + for d in (index, *platforms)} + + +def test_links_to_layers_garbage_collection_freed_are_dropped(script, tmp_path): + kept = _blob(tmp_path, {"layer": "kept"}) + layers = tmp_path / "docker/registry/v2/repositories" / REPOSITORY / "_layers/sha256" + for digest in (kept, "sha256:" + "e" * 64): + (layers / digest.removeprefix("sha256:")).mkdir(parents=True) + (layers / digest.removeprefix("sha256:") / "link").write_text(digest) + + assert script.drop_freed_links(tmp_path, REPOSITORY) == 1 + assert [p.name for p in layers.iterdir()] == [kept.removeprefix("sha256:")] + + +def _blob(storage: Path, content: dict) -> str: + data = json.dumps(content).encode() + hex_digest = hashlib.sha256(data).hexdigest() + path = storage / "docker/registry/v2/blobs/sha256" / hex_digest[:2] / hex_digest / "data" + path.parent.mkdir(parents=True) + path.write_bytes(data) + return f"sha256:{hex_digest}" + + +def _tag(storage: Path, tag: str, digest: str) -> None: + current = storage / "docker/registry/v2/repositories" / REPOSITORY / "_manifests/tags" / tag / "current" + current.mkdir(parents=True) + (current / "link").write_text(digest) + + +def _revision(storage: Path, digest: str) -> None: + revision = (storage / "docker/registry/v2/repositories" / REPOSITORY / "_manifests/revisions/sha256" + / digest.removeprefix("sha256:")) + revision.mkdir(parents=True) + (revision / "link").write_text(digest) diff --git a/tst/util/package_cache.py b/tst/util/package_cache.py new file mode 100644 index 000000000..3c7042f3f --- /dev/null +++ b/tst/util/package_cache.py @@ -0,0 +1,104 @@ +"""The package locks a computer's compose resolves, and the files they name, shared between the fresh local stores the +tests compose in, so a lock resolved once in a run, or restored from an earlier run's cache, isn't resolved from the +network again. + +Everything under the packages prefix is content-addressed, a lock by a digest of its inputs and a file by its sha256, so +any store may take any of it. The rest of an object store isn't: its keys follow entity ids and versions. The shared +folder is ``AGENTENV_TEST_PACKAGES``, which tst/conftest.py sets for the run when it's unset; it mirrors the store's +layout under the prefix, content types included.""" + +from __future__ import annotations + +import errno +import os +import shutil +from pathlib import Path + +from agent_env.config import get_config +from agent_env.env.envs.computer.utils.packages import PREFIX, Lock +from agent_env.store.object_store import LocalFilesystemObjectStore + +ENV_VAR = "AGENTENV_TEST_PACKAGES" + + +def seed() -> None: + """Give the configured local object store every shared lock whose files are all shared: its files, then the lock.""" + if (shared := _shared()) is None: + return + store = _store() + for lock in sorted((shared.packages / "locks").glob("*.json")): + files = _files(lock) + if files is None or not all((shared.packages / "files" / sha256).is_file() for sha256 in files): + continue + for sha256 in files: + _link(shared.packages / "files" / sha256, store.packages / "files" / sha256) + _link_lock(lock.name, shared, store) + + +def keep() -> None: + """Share what the configured local object store holds under the prefix: its files first, then its locks, so the + shared folder never has a lock without its files.""" + if (shared := _shared()) is None: + return + store = _store() + for path in sorted((store.packages / "files").glob("*")): + _link(path, shared.packages / "files" / path.name) + for lock in sorted((store.packages / "locks").glob("*.json")): + _link_lock(lock.name, store, shared) + + +class _Tree: + """Where a folder keeps the objects under the prefix, and their content types (the local store's + ``.agentenv-meta``).""" + + def __init__(self, root: Path, prefix: str): + self.packages, self.meta = root / prefix, root / ".agentenv-meta" / prefix + + +def _shared() -> _Tree | None: + folder = os.environ.get(ENV_VAR) + return _Tree(Path(folder), PREFIX) if folder else None + + +def _store() -> _Tree: + config = get_config() + store = config.get_object_store() + if not isinstance(store, LocalFilesystemObjectStore): + raise TypeError(f"packages are shared with a local object store, not {type(store).__name__}") + return _Tree(store.root, f"{config.get_artifact_key_prefix()}{PREFIX}") + + +def _files(lock: Path) -> list[str] | None: + """The sha256 of every file ``lock`` names; None for one in a shape this version doesn't read.""" + try: + parsed = Lock.model_validate_json(lock.read_bytes()) + except ValueError: + return None + named = [*parsed.apt, *(f for files in parsed.wheels.values() for f in files), + *(f for files in parsed.downloads.values() for f in files)] + return [f.sha256 for f in named] + + +def _link_lock(name: str, source: _Tree, target: _Tree) -> None: + if (meta := source.meta / "locks" / name).is_file(): + _link(meta, target.meta / "locks" / name) + _link(source.packages / "locks" / name, target.packages / "locks" / name) + + +def _link(source: Path, target: Path) -> None: + """``source`` at ``target`` unless something is there: a hard link, safe as the store never writes an object in + place, or a copy where the two folders can't share one, keeping its modification time, which the store's record of + an object's content type names it by.""" + if target.exists(): + return + target.parent.mkdir(parents=True, exist_ok=True) + try: + os.link(source, target) + except FileExistsError: + pass + except OSError as e: + if e.errno not in (errno.EXDEV, errno.EPERM, errno.EMLINK, errno.ENOTSUP, errno.EOPNOTSUPP): + raise + staged = target.with_name(f".{target.name}.{os.getpid()}") + shutil.copy2(source, staged) + os.replace(staged, target)