feat(env): set up the infra envs from their published images, building only what isn't published

CONTRIBUTING.md

@@ -91,8 +91,10 @@ the previous release, so a clear pull request title is also the release note. A
 A release also publishes the infra images agent-env ships, once the `PUBLISH_INFRA_IMAGES` repository variable is
 `true`: one per Dockerfile under `src/agent_env/env`, at `ghcr.io/scaleapi/agentenv-<name>:src-<hex>` for linux/amd64
 and linux/arm64, where the tag is the build context's source digest, so an unchanged Dockerfile is not built again.
-They're published so installs can pull them instead of building the infra images, and the wheel and the release
-notes record each image's digest. A failed image job keeps the release off PyPI. Maintainers can publish the images
+Installs pull them instead of building the infra images: a run on local stores registers a kind's images when GHCR has
+every one for the sources installed, by the digest the wheel records or, with no record, the one GHCR serves, and
+builds them otherwise; a sandbox that can't pull one builds it from its build context. The wheel and the release notes
+record each image's digest. A failed image job keeps the release off PyPI. Maintainers can publish the images
 alone by running the workflow by hand on `main`. GHCR keeps a new package private until a package admin makes it
 public, and the workflow warns about each one until then.
 

src/agent_env/artifact/artifacts/docker_image.py

@@ -69,7 +69,8 @@ def _git_clone_commands(owner: str, repo: str, ref: str | None, token: str | Non
 
 class DockerImageArtifact(Artifact):
     """A Docker image artifact: a tar.gz of the image in the object store; without one, a registry reference a
-    sandbox pulls; or, with neither, only its build context."""
+    sandbox pulls, by its ``image_digest`` when it records one, or builds from the build context beside it when it can't; or,
+    with neither, only its build context."""
 
     model_config = ConfigDict(populate_by_name=True)
 
@@ -87,6 +88,10 @@ class DockerImageArtifact(Artifact):
     source_digest: str | None = Field(
         default=None, description="sha256 over the build context's files and modes, the Dockerfile's path and the platform"
     )
+    image_digest: str | None = Field(
+        default=None, description="The registry's digest of the built image, not of its sources (source_digest): a "
+                                  "sandbox pulls the image by it and runs it as image_name"
+    )
 
     # No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields.
     @model_serializer(mode="wrap")
@@ -207,16 +212,16 @@ def put_tar(
         return store.put_document(instance)
 
     @classmethod
-    def put_ref(cls, id: str, *, description: str, image_name: str,
-                earlier_pin: str | None = None) -> "DockerImageArtifact":
-        """Register ``image_name``, an image already in a registry, writing this document only: nothing is pushed,
-        saved or uploaded, and a sandbox pulls the image. A tag is pinned to the digest the registry serves for it
-        now, recorded as ``name:tag@sha256:...``, so every run gets the same image; a digest is checked to exist.
-        ``earlier_pin``, the pin an earlier version recorded, is kept instead when it pins ``image_name``.
-
-        ``image_name`` must name its registry. The registry is read with the configured image store's credentials
-        where it holds that registry, anonymously otherwise. One on this machine is refused unless ``id``'s image
-        store is the local registry, since nothing else could pull from it."""
+    def put_ref(cls, id: str, *, description: str, image_name: str, earlier_pin: str | None = None,
+                image_digest: str | None = None, context_path: str | None = None,
+                dockerfile_path: str | None = None) -> "DockerImageArtifact":
+        """Register ``image_name``, an image already in a registry, which a sandbox pulls; nothing is built, pushed or
+        saved. A tag is pinned to the digest the registry serves now, as ``name:tag@sha256:...`` (``earlier_pin`` kept
+        when it pins the same ref). Given the image's build context and ``image_digest``, the tag stays the name, the
+        digest is recorded beside it, and the context is uploaded for a sandbox that can't pull the image to build it.
+
+        ``image_name`` must name its registry, read with the configured image store's credentials where it holds them;
+        one on this machine is refused unless ``id``'s image store is the local registry."""
         from agent_env.artifact.store import get_artifact_store
         from agent_env.config import get_config
 
@@ -228,6 +233,21 @@ def put_ref(cls, id: str, *, description: str, image_name: str,
                 and not isinstance(config.get_image_store_for(id), LocalRegistryImageStore)):
             raise ValueError(f"{image_name} is in a registry on this machine, which only an image in the local "
                              "registry's store, such as an @local one, can name")
+        if context_path is not None:
+            context = _context_of(Path(context_path), dockerfile_path)
+            store = get_artifact_store()
+            version = store.next_version(id)
+            prefix = store.attempt_prefix("docker_image", id)
+            return store.put_document(cls(
+                id=id,
+                version=version,
+                description=description,
+                image_name=image_name,
+                image_digest=image_digest,
+                build_context_object_url=_upload_context(config.get_object_store_to_write(prefix, id), prefix, context),
+                dockerfile_path=context.dockerfile,
+                source_digest=context.source_digest(None),
+            ))
         if earlier_pin is not None and pinned_from(earlier_pin, image_name):
             pinned = earlier_pin
         else:
@@ -281,6 +301,11 @@ def load(self) -> bytes:
             raise ValueError(f"{self.id!r} v{self.version} has no tar.gz; its image is {how}")
         return get_artifact_store().get_object(self.tar_gz_object_url)
 
+    @property
+    def pinned_name(self) -> str:
+        """What a sandbox pulls: ``image_name`` pinned to the ``image_digest`` when it records one."""
+        return f"{self.image_name}@{self.image_digest}" if self.image_digest else self.image_name
+
     @property
     def context_only(self) -> bool:
         """Whether only the image's build context is stored: no tar.gz, and an image name that names no registry, the

src/agent_env/bundle/preflight.py

@@ -39,10 +39,11 @@
     GATEWAY,
     SERVICE_DB,
     WEBSITE_BROWSER,
-    InfraBuild,
+    InfraSetup,
     InfraError,
     default_env_id,
-    infra_to_build,
+    env_images,
+    infra_to_set_up,
     needs_docker,
     put_command,
     stores_are_local,
@@ -51,7 +52,6 @@
 from agent_env.env.envs._deployment import provider_or_class
 from agent_env.env.envs.mcp_server import MCPServerEnv
 from agent_env.env.envs.multi_env import MultiEnv
-from agent_env.env.envs.service_db import ServiceDBEnv
 from agent_env.env.envs.website import WebsiteEnv
 from agent_env.env.registry import get_env_registry
 from agent_env.providers.env_providers.env_gateway_provider import EnvironmentGatewayProvider, _gateway_topology
@@ -89,7 +89,7 @@ class Preflight:
     """What a run needs before its tasks start, beyond its writes: the infra envs it builds."""
 
     infra_kinds: frozenset[str]  # the infra its deploys on the local provider run on
-    infra: tuple[InfraBuild, ...]  # those the store doesn't hold yet, or holds built from other inputs
+    infra: tuple[InfraSetup, ...]  # those it sets up first: missing, built from other inputs, or to switch to pulling
     built_here: frozenset[str] = frozenset()  # the ids of the images the bundle builds that a deploy here runs
 
 
@@ -161,21 +161,22 @@ def task(self, task: ResolvedEntry) -> None:
                 self._problem(where, str(e))
 
     def finish(self) -> Preflight:
-        infra: list[InfraBuild] = []
+        infra: list[InfraSetup] = []
         try:
-            infra = infra_to_build(self.infra)
+            infra = infra_to_set_up(self.infra)
         except InfraError as e:
             self.run_problems.extend(e.problems)
         built_here = self._built_here()
         self._default_agent()
-        reason = docker_unreachable() if self.docker_users or infra else None
+        built = [build for build in infra if build.builds]
+        reason = docker_unreachable() if self.docker_users or built else None
         if reason and self.docker_users:
             self.run_problems.append(f"the local sandbox provider runs containers for {_listed(self.docker_users)}, "
                                      f"and {reason}")
         elif reason:  # infra the run builds for deploys elsewhere
-            self.run_problems.append(needs_docker(infra, reason))
+            self.run_problems.append(needs_docker(built, reason))
         else:
-            self._docker_plugins(built_here, infra)
+            self._docker_plugins(built_here, built)
         needed = frozenset().union(*(names for _, names in self.model_users))
         if needed and (missing := model_endpoint_missing(needed)):
             name, why = missing
@@ -434,11 +435,11 @@ def _refusal(self, link: SandboxProvider, form: Accepts, image: _Image, *, vms_l
                 f"on a provider {builds}")
 
     def _remote_infra(self, link: SandboxProvider, kinds: set[str]) -> list[str]:
-        """Why ``link``, a provider not on this machine, can't run the infra envs of ``kinds``. One the run builds, in
-        local stores missing or built from other inputs, isn't checked: it's built anew, with its build context."""
-        built = {build.kind for build in infra_to_build(kinds)} if stores_are_local() else set()
+        """Why ``link``, a provider not on this machine, can't run the infra envs of ``kinds``. One the run sets up, in
+        local stores, isn't checked: it's written anew, pulled from its registry or built, with its build context."""
+        set_up = {setup.kind for setup in infra_to_set_up(kinds)} if stores_are_local() else set()
         found = []
-        for kind in sorted(kinds - built):
+        for kind in sorted(kinds - set_up):
             env_id = default_env_id(kind)
             try:
                 env = Env.get(env_id)
@@ -447,10 +448,8 @@ def _remote_infra(self, link: SandboxProvider, kinds: set[str]) -> list[str]:
                              f"and the store doesn't hold it; agent-env builds infra envs only into local stores, so put "
                              f"it with `{put_command(kind)}`")
                 continue
-            artifacts = ([env.db_docker_image_artifact, env.db_web_docker_image_artifact,
-                          env.db_mcp_docker_image_artifact] if isinstance(env, ServiceDBEnv) else [env.docker_image_artifact])
             accepts = EnvironmentGatewayProvider.accepts(link)
-            found += [refusal for image in artifacts
+            found += [refusal for image in env_images(env)
                       if (refusal := self._refusal(link, accepts, _Image(f"the {kind} env {env_id!r}'s image {image.id!r}",
                                                                          image), vms_load=True))]
         return found
@@ -475,7 +474,7 @@ def _built_here(self) -> frozenset[str]:
                                              "modal")
         return frozenset(here)
 
-    def _docker_plugins(self, built_here: frozenset[str], infra: list[InfraBuild]) -> None:
+    def _docker_plugins(self, built_here: frozenset[str], infra: list[InfraSetup]) -> None:
         """The Docker plugins the run uses on this machine: Compose for each gateway stack it runs, Buildx for each
         image it builds, as ``docker build`` does with BuildKit."""
         if self.compose_users and (missing := docker_plugin_missing("compose")):

src/agent_env/bundle/run.py

@@ -6,8 +6,9 @@
 SIGTERM cancels the runs: each one that started is marked cancelled and torn down, and a second one stops the
 teardown. A bundle's evals run only the bundle's own tasks for now, so one naming a store task is refused
 before anything is written, and so is a deploy its sandbox provider can't serve (``preflight``). The infra envs a
-gateway deploy on the local provider needs are built once the writes are done, before any task starts. A dry run
-makes every check the run makes before its first task, and writes, builds and runs nothing.
+gateway deploy on the local provider needs are set up once the writes are done, before any task starts: pulled from
+their published images, or built. A dry run makes every check the run makes before its first task, and writes, builds
+and runs nothing.
 """
 
 from __future__ import annotations
@@ -23,7 +24,7 @@
 from pathlib import Path
 
 from agent_env.artifact.artifacts.docker_image import DockerImageArtifact
-from agent_env.env.bootstrap import InfraBuild, ensure_default_envs
+from agent_env.env.bootstrap import InfraSetup, ensure_default_envs
 from agent_env.providers import build_sandbox_provider
 from agent_env.store.base import NotFoundError
 from agent_env.store.routing import namespace_routing, run_scope
@@ -152,7 +153,7 @@ class DryRun:
     materialization: Materialization  # each write at the version it would leave, and why
     runs: tuple[BundleEntry, ...]  # the tasks that would run, each once, in the plan's order
     skipped: tuple[BundleEntry, ...]  # the tasks no eval names, which running every eval leaves out
-    infra: tuple[InfraBuild, ...] = ()  # the infra envs the run would build first
+    infra: tuple[InfraSetup, ...] = ()  # the infra envs the run would set up first
     built_here: frozenset[str] = frozenset()  # the images the run would build on this machine, for its deploys here
 
     def path(self, entry: BundleEntry) -> str:

src/agent_env/cli/run.py

@@ -211,9 +211,9 @@ def _report_dry_run(dry: DryRun) -> None:
     for line in images:
         click.echo(f"  {line}")
     if dry.infra:
-        click.echo("Would build first:")
+        click.echo("Would set up first:")
     for build in dry.infra:
-        click.echo(f"  {build}")
+        click.echo(f"  {build}: {build.planned()}")
     click.echo("Would run:")
     for entry in dry.runs:
         click.echo(f"  {dry.path(entry)} v{materialization.version_of('task', entry.id)}")

src/agent_env/cli/up.py

@@ -27,8 +27,8 @@ def _require_explorer_deps() -> None:
 
 
 def _bootstrap_envs() -> None:
-    """Build the two envs ``deploy_env`` resolves by id (service-db, gateway) when they're missing or were built from
-    other inputs than this release's, so `up` is cheap after the first run, which builds their images."""
+    """Set up the two envs ``deploy_env`` resolves by id (service-db, gateway) when they're missing or were built from
+    other inputs than this release's, so `up` is cheap after the first run, which pulls or builds their images."""
     try:
         builds = ensure_default_envs((SERVICE_DB, GATEWAY), say=lambda line: click.echo(f"  {line}"))
     except InfraError as e: