experiment: slow tier two tests at a time per runner (do not merge)

.github/scripts/prune_computer_images.py

@@ -0,0 +1,112 @@
+"""Delete the CI computer image mirror's images no run has used for ``--days`` days.
+
+``share_computer_images.py`` tags each image a passing main run used ``<tag>-seen-YYYYMMDD``. An image is kept while its
+newest such tag, or its upload when it has none, is within ``--days`` of now, with every platform manifest a kept index
+lists, since GHCR lists those as versions of their own. Every other version is deleted through GitHub's packages API,
+read and written with ``GITHUB_TOKEN``, as GHCR is. A version GitHub won't delete warns. ``--dry-run`` prints what
+would go and deletes nothing.
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import datetime
+import json
+import os
+import re
+import sys
+import urllib.error
+import urllib.parse
+import urllib.request
+
+SEEN = re.compile(r"^[0-9a-f]{16}-[0-9a-f]{64}-seen-(\d{8})$")
+_INDEXES = {"application/vnd.oci.image.index.v1+json", "application/vnd.docker.distribution.manifest.list.v2+json"}
+_ACCEPT = ", ".join([*_INDEXES, "application/vnd.oci.image.manifest.v1+json",
+                     "application/vnd.docker.distribution.manifest.v2+json"])
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument("org")
+    parser.add_argument("package", help="the container package, as agentenv-ci-computers")
+    parser.add_argument("--days", type=int, default=7)
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+
+    api = f"https://api.github.com/orgs/{args.org}/packages/container/{args.package}/versions"
+    try:
+        versions = _versions(api)
+    except urllib.error.HTTPError as e:
+        if e.code != 404:
+            raise
+        print(f"::warning::prune computer images: no package {args.org}/{args.package}, or this repository has no "
+              "access to it, so nothing is pruned")
+        return 0
+    cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=args.days)
+    stale = stale_versions(versions, cutoff, lambda digest: _children(args.org, args.package, digest))
+    failed = 0
+    for version in stale:
+        print(f"{'would delete' if args.dry_run else 'deleting'} {version['name']} "
+              f"{version['metadata']['container']['tags'] or '(untagged)'}")
+        if not args.dry_run:
+            try:
+                _github("DELETE", f"{api}/{version['id']}")
+            except urllib.error.HTTPError as e:
+                failed += 1
+                print(f"::warning::couldn't delete {version['name']}: HTTP {e.code} {e.read().decode()[:200]}")
+    print(f"prune computer images: {len(stale) - failed} of {len(versions)} versions "
+          f"{'would be ' if args.dry_run else ''}deleted, none used since {cutoff:%Y-%m-%d}")
+    return 0
+
+
+def stale_versions(versions: list[dict], cutoff: datetime.datetime, children) -> list[dict]:
+    """The ``versions`` to delete: neither used nor uploaded since ``cutoff``, nor a platform manifest of one kept.
+    ``children(digest)`` lists the digests of the manifests the index of that digest lists, or none for an image."""
+    kept = {version["name"] for version in versions if _last_used(version) >= cutoff}
+    for digest in list(kept):
+        kept.update(children(digest))
+    return [version for version in versions if version["name"] not in kept]
+
+
+def _last_used(version: dict) -> datetime.datetime:
+    created = datetime.datetime.fromisoformat(version["created_at"].replace("Z", "+00:00"))
+    seen = [datetime.datetime.strptime(m.group(1), "%Y%m%d").replace(tzinfo=datetime.timezone.utc)
+            for tag in version["metadata"]["container"]["tags"] if (m := SEEN.match(tag))]
+    return max([created, *(day + datetime.timedelta(days=1) for day in seen)])
+
+
+def _versions(api: str) -> list[dict]:
+    versions, page = [], 1
+    while batch := _github("GET", f"{api}?per_page=100&page={page}"):
+        versions += batch
+        page += 1
+    return versions
+
+
+def _children(org: str, package: str, digest: str) -> list[str]:
+    """The platform manifests the index ``digest`` lists, read from GHCR with ``GITHUB_TOKEN``; none for an image."""
+    repository = f"{org}/{package}"
+    signed_in = base64.b64encode(f"x:{os.environ['GITHUB_TOKEN']}".encode()).decode()
+    with urllib.request.urlopen(urllib.request.Request(f"https://ghcr.io/token?scope=repository:{repository}:pull",
+                                                       headers={"Authorization": f"Basic {signed_in}"})) as response:
+        token = json.load(response)["token"]
+    request = urllib.request.Request(f"https://ghcr.io/v2/{repository}/manifests/{digest}",
+                                     headers={"Authorization": f"Bearer {token}", "Accept": _ACCEPT})
+    with urllib.request.urlopen(request) as response:
+        manifest = json.load(response)
+        media = response.headers.get("Content-Type", "").split(";")[0].strip()
+    return [child["digest"] for child in manifest.get("manifests", [])] if media in _INDEXES else []
+
+
+def _github(method: str, url: str):
+    request = urllib.request.Request(url, method=method, headers={
+        "Authorization": f"Bearer {os.environ['GITHUB_TOKEN']}", "Accept": "application/vnd.github+json",
+        "X-GitHub-Api-Version": "2022-11-28"})
+    with urllib.request.urlopen(request) as response:
+        body = response.read()
+    return json.loads(body) if body else None
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/scripts/share_computer_images.py

@@ -0,0 +1,53 @@
+"""Copy every composed computer image a slow-tier runner's registry holds to the CI computer image mirror, the one
+``AGENT_ENV_COMPUTER_IMAGE_MIRROR`` names, so later runs copy them back rather than build them.
+
+Each image goes under its mirror tag (``composed.mirror_tag``: the composer's digest, then the key), and again with
+today's date, ``<tag>-seen-YYYYMMDD``, which is how ``prune_computer_images.py`` tells an image a recent run used from one
+none has. An image already in the mirror is only tagged again: no blob it holds is sent twice. The mirror is written as
+``GITHUB_ACTOR`` with ``GITHUB_TOKEN``; one the composer can't read anonymously, as GHCR makes a new package, warns.
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import json
+import os
+import re
+import sys
+import urllib.request
+
+from agent_env.env.envs.computer.utils.composed import REPOSITORY, mirror_tag
+from agent_env.store.image_store.oci_registry_credentials import RegistryAuth
+from agent_env.store.image_store.registry_api import copy_image, pin_digest, pin_if_served
+
+KEY = re.compile(r"^[0-9a-f]{64}$")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument("mirror", help="the mirror repository, as ghcr.io/org/agentenv-ci-computers")
+    parser.add_argument("--registry", default="localhost:5000", help="the runner's registry")
+    args = parser.parse_args()
+
+    with urllib.request.urlopen(f"http://{args.registry}/v2/{REPOSITORY}/tags/list") as response:
+        keys = sorted(tag for tag in json.load(response).get("tags") or [] if KEY.match(tag))
+    auth = RegistryAuth(args.mirror.partition("/")[0], os.environ["GITHUB_ACTOR"], os.environ["GITHUB_TOKEN"])
+    seen = f"seen-{datetime.datetime.now(datetime.timezone.utc):%Y%m%d}"
+    for key in keys:
+        source = pin_digest(f"{args.registry}/{REPOSITORY}:{key}", None)
+        for tag in (mirror_tag(key), f"{mirror_tag(key)}-{seen}"):
+            copy_image(source, f"{args.mirror}:{tag}", None, auth)
+        print(f"shared {key}")
+    print(f"share computer images: {len(keys)} images in {args.mirror}, each tagged {seen}")
+    if keys:
+        try:
+            pin_if_served(f"{args.mirror}:{mirror_tag(keys[0])}", None)
+        except ValueError as e:
+            print(f"::warning::the composer reads {args.mirror} anonymously and can't, so every run builds its images; "
+                  f"make the package public: {e}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/computer-image-mirror.yml

@@ -0,0 +1,31 @@
+name: Computer image mirror
+
+# Deletes the slow tier's shared computer images (local-backends.yml) that no passing main run has used for a week.
+on:
+  schedule:
+    - cron: "17 6 * * *"
+  workflow_dispatch:
+    inputs:
+      dry_run:
+        description: List what would be deleted, and delete nothing
+        type: boolean
+        default: false
+
+permissions:
+  contents: read
+  packages: write
+
+jobs:
+  prune:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+        with:
+          sparse-checkout: .github/scripts
+      - name: Delete the images unused for a week
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        run: |
+          python3 .github/scripts/prune_computer_images.py scaleapi agentenv-ci-computers --days 7 \
+            ${{ inputs.dry_run && '--dry-run' || '' }}

.github/workflows/local-backends.yml

@@ -22,6 +22,8 @@ concurrency:
 
 env:
   AGENT_ENV_ENVIRONMENT: dev
+  # The slow tier's composed computer images, shared between its runners (computer-image-mirror.yml prunes it).
+  COMPUTER_IMAGE_MIRROR: ghcr.io/scaleapi/agentenv-ci-computers
   PIP_DISABLE_PIP_VERSION_CHECK: "1"
 
 jobs:
@@ -154,8 +156,14 @@ jobs:
         shard: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
     runs-on: ubuntu-latest
     timeout-minutes: 40
+    permissions:
+      contents: read
+      packages: write
     steps:
       - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+        # The tests run with packages: write; only the share step, on main, gets the token.
+        with:
+          persist-credentials: false
       - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: "3.12"
@@ -168,63 +176,40 @@ jobs:
       # starts afterwards.
       - 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.
-      - name: Week of the computer images' cache
-        id: week
-        run: echo "week=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT"
-      - name: Restore the computer images
-        id: computers
-        uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
-        with:
-          path: ${{ runner.temp }}/registry
-          key: computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}
-          restore-keys: computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}-
       - name: Start the local registry
         run: |
-          mkdir -p "$RUNNER_TEMP/registry"
-          docker run --detach --name registry --publish 5000:5000 --volume "$RUNNER_TEMP/registry:/var/lib/registry" \
-            mirror.gcr.io/library/registry:2
+          docker run --detach --name registry --publish 5000:5000 mirror.gcr.io/library/registry:2
           timeout 30 sh -c 'until curl -sf http://localhost:5000/v2/ >/dev/null; do sleep 1; done'
       # x64 hosted runners have /dev/kvm; the disk form's tests boot under it, and skip without it, which fails below.
       - name: Let the local provider boot disks with KVM
         run: sudo chmod 666 /dev/kvm
+      # A computer image or disk is tagged with a key of everything it is built from. A compose that finds no image of
+      # its key in the local registry copies the one the mirror holds, which main's runs share, rather than build it.
+      # Experiment: sample the runner's used memory every 5s while the tests run two at a time.
+      - name: Sample memory
+        run: nohup sh -c 'while true; do free -m | awk "NR==2{print \$3}" >> "$RUNNER_TEMP/mem.log"; sleep 5; done' >/dev/null 2>&1 &
       - name: Slow integration tier on the local defaults
         id: tests
         env:
           AGENTENV_TEST_SHARD: ${{ matrix.shard }}/12
+          AGENT_ENV_COMPUTER_IMAGE_MIRROR: ${{ env.COMPUTER_IMAGE_MIRROR }}
         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
-        id: kept
+          .venv/bin/python -m pytest -m int_test_slow -n 2 --dist loadfile -rA --timeout=1200 \
+            --junitxml=junit-integration-slow.xml tst/integration --ignore=tst/integration/env/gateway/gateway_test.py
+      # Once main's tests pass, every image this runner built or copied goes to the mirror under its key and today's
+      # date, which the prune of computer-image-mirror.yml keeps an image a run used within the week by.
+      - name: Share the computer images
+        continue-on-error: true
+        if: >-
+          always() && steps.tests.outcome == 'success' && github.event_name == 'push'
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        run: .venv/bin/python .github/scripts/share_computer_images.py "$COMPUTER_IMAGE_MIRROR"
+      - name: Peak memory
         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 stop registry
-          sudo find "$repositories" -mindepth 1 -maxdepth 1 ! -name agentenv-computers -exec rm -rf {} +
-          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"
-          echo "computers=true" >> "$GITHUB_OUTPUT"
-      - name: Save the computer images
-        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
-          key: ${{ steps.kept.outputs.key }}
+          echo "peak used MB: $(sort -n "$RUNNER_TEMP/mem.log" | tail -1) of $(free -m | awk 'NR==2{print $2}')"
+          if sudo dmesg | grep -q -i -E "out of memory|oom-kill"; then echo "OOM KILL SEEN"; else echo "no OOM kills"; fi
       - name: Skips must be declared capability gaps
         if: always()
         run: |

src/agent_env/config/describe.py

@@ -673,6 +673,7 @@ def env_sources() -> tuple[EnvSource, ...]:
             EnvSource("AGENT_ENV_MODAL_REGION", None, "no file equivalent"),
             EnvSource("AGENT_ENV_MODAL_APP_NAME", None, "no file equivalent"),
             EnvSource("AGENT_ENV_FIXTURE_PREFIX", None, "no file equivalent"),
+            EnvSource("AGENT_ENV_COMPUTER_IMAGE_MIRROR", None, "no file equivalent"),
             EnvSource("MODAL_TOKEN_ID", None, "credential, no file equivalent"),
             EnvSource("MODAL_TOKEN_SECRET", None, "credential, no file equivalent"),
     ]

src/agent_env/config/runtime.py

@@ -24,6 +24,11 @@
 objects, image builds, env and agent snapshots, changelogs, default agent and judge
 trajectories, verifier outputs and the A2A validator's fixtures, so a fresh control-plane
 backend can write without colliding with existing objects in a shared store. Empty in prod.
+
+AGENT_ENV_COMPUTER_IMAGE_MIRROR (default ""): a registry repository, read anonymously, holding composed computer images
+under their keys and the composer's own digest, as ``ghcr.io/org/computers``. A compose that finds no image of a key in the
+env's image store copies the mirror's into it instead of building it. CI sets it so its runners share the images main
+built.
 """
 
 from __future__ import annotations
@@ -226,6 +231,9 @@ class Config:
     fixture_prefix: str = field(
         default_factory=lambda: os.getenv("AGENT_ENV_FIXTURE_PREFIX", "")
     )
+    computer_image_mirror: str = field(
+        default_factory=lambda: os.getenv("AGENT_ENV_COMPUTER_IMAGE_MIRROR", "")
+    )
     agent_sandbox_mode: str = field(
         default_factory=lambda: os.getenv("AGENT_SANDBOX_MODE") or _default_sandbox_mode()
     )

src/agent_env/env/envs/computer/utils/composed.py

@@ -117,6 +117,12 @@ def composition() -> str:
     return "sha256:" + digest.hexdigest()
 
 
+def mirror_tag(key: str) -> str:
+    """The tag of the image of ``key`` in a computer image mirror: the key, under the start of ``composition()``, so a
+    change to the code that composes an image builds it again rather than reuse one the old code built."""
+    return f"{composition().removeprefix('sha256:')[:16]}-{key.removeprefix('sha256:')}"
+
+
 def _rank(result: dict) -> int:
     """How far a pair got: ready, then built and either unprobed or failed at its probe, then anything short of an
     image, and below that one another composer recorded. A result takes a stored one's place only when it got at least

src/agent_env/env/envs/computer/utils/composer.py

@@ -49,15 +49,15 @@
     thin_disk,
     windows_disk,
 )
-from agent_env.env.envs.computer.utils.composed import BUILT, REPOSITORY, Record, failed, now, stale
+from agent_env.env.envs.computer.utils.composed import BUILT, REPOSITORY, Record, failed, mirror_tag, now, stale
 from agent_env.env.envs.computer.utils.install import Install, for_os, platform_of
 from agent_env.env.published_images import DISK, DISK_TOOLS_IMAGE, published_disk, published_pinned
 from agent_env.providers.sandbox_providers.host_images import platform_tag
 from agent_env.store.ids import derive_id
 from agent_env.store.image_store.image_contents import image_config, labels
 from agent_env.store.image_store.oci_registry_credentials import names_registry
 from agent_env.store.image_store.image_store import ImageStore
-from agent_env.store.image_store.registry_api import pin_digest, pin_if_served, write_index
+from agent_env.store.image_store.registry_api import copy_image, pin_digest, pin_if_served, write_index
 from agent_env.store.local_state import holding_locks
 from agent_env.utils.docker.build import DockerBuildError
 
@@ -691,9 +691,20 @@ def _key(**inputs: Any) -> str:
 
 def _stored(env: ComputerEnv, key: str) -> tuple[ImageStore, str, str | None]:
     """The env's image store, the ref there of the image of ``key``, and that ref pinned when the store serves one
-    already, for any env that built it."""
+    already, for any env that built it, or once it's copied there from the configured computer image mirror."""
     store = _image_store(env)
-    ref = store.image_ref(REPOSITORY, key.removeprefix("sha256:"))
+    tag = key.removeprefix("sha256:")
+    ref = store.image_ref(REPOSITORY, tag)
+    if (pinned := pin_if_served(ref, store.auth(ref))) is not None or not (mirror := get_config().computer_image_mirror):
+        return store, ref, pinned
+    try:
+        if (mirrored := pin_if_served(f"{mirror}:{mirror_tag(tag)}", None)) is None:
+            return store, ref, None
+        copy_image(mirrored, ref, None, store.auth(ref))
+    except (ValueError, KeyError, OSError) as e:
+        logger.warning(f"the computer image mirror {mirror} couldn't give {tag}, so it's built: {e}")
+        return store, ref, None
+    logger.info(f"copied {tag} from the computer image mirror {mirror}")
     return store, ref, pin_if_served(ref, store.auth(ref))