Skip to content

Commit 3d83f45

Browse files
feat(env): set up the infra envs from their published images, building only what isn't published (#277)
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1 parent 95a3a16 commit 3d83f45

27 files changed

Lines changed: 667 additions & 131 deletions

File tree

‎CONTRIBUTING.md‎

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

‎src/agent_env/artifact/artifacts/docker_image.py‎

Lines changed: 36 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ def _git_clone_commands(owner: str, repo: str, ref: str | None, token: str | Non
6969

7070
class DockerImageArtifact(Artifact):
7171
"""A Docker image artifact: a tar.gz of the image in the object store; without one, a registry reference a
72-
sandbox pulls; or, with neither, only its build context."""
72+
sandbox pulls, by its ``image_digest`` when it records one, or builds from the build context beside it when it can't; or,
73+
with neither, only its build context."""
7374

7475
model_config = ConfigDict(populate_by_name=True)
7576

@@ -87,6 +88,10 @@ class DockerImageArtifact(Artifact):
8788
source_digest: str | None = Field(
8889
default=None, description="sha256 over the build context's files and modes, the Dockerfile's path and the platform"
8990
)
91+
image_digest: str | None = Field(
92+
default=None, description="The registry's digest of the built image, not of its sources (source_digest): a "
93+
"sandbox pulls the image by it and runs it as image_name"
94+
)
9095

9196
# No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields.
9297
@model_serializer(mode="wrap")
@@ -207,16 +212,16 @@ def put_tar(
207212
return store.put_document(instance)
208213

209214
@classmethod
210-
def put_ref(cls, id: str, *, description: str, image_name: str,
211-
earlier_pin: str | None = None) -> "DockerImageArtifact":
212-
"""Register ``image_name``, an image already in a registry, writing this document only: nothing is pushed,
213-
saved or uploaded, and a sandbox pulls the image. A tag is pinned to the digest the registry serves for it
214-
now, recorded as ``name:tag@sha256:...``, so every run gets the same image; a digest is checked to exist.
215-
``earlier_pin``, the pin an earlier version recorded, is kept instead when it pins ``image_name``.
216-
217-
``image_name`` must name its registry. The registry is read with the configured image store's credentials
218-
where it holds that registry, anonymously otherwise. One on this machine is refused unless ``id``'s image
219-
store is the local registry, since nothing else could pull from it."""
215+
def put_ref(cls, id: str, *, description: str, image_name: str, earlier_pin: str | None = None,
216+
image_digest: str | None = None, context_path: str | None = None,
217+
dockerfile_path: str | None = None) -> "DockerImageArtifact":
218+
"""Register ``image_name``, an image already in a registry, which a sandbox pulls; nothing is built, pushed or
219+
saved. A tag is pinned to the digest the registry serves now, as ``name:tag@sha256:...`` (``earlier_pin`` kept
220+
when it pins the same ref). Given the image's build context and ``image_digest``, the tag stays the name, the
221+
digest is recorded beside it, and the context is uploaded for a sandbox that can't pull the image to build it.
222+
223+
``image_name`` must name its registry, read with the configured image store's credentials where it holds them;
224+
one on this machine is refused unless ``id``'s image store is the local registry."""
220225
from agent_env.artifact.store import get_artifact_store
221226
from agent_env.config import get_config
222227

@@ -228,6 +233,21 @@ def put_ref(cls, id: str, *, description: str, image_name: str,
228233
and not isinstance(config.get_image_store_for(id), LocalRegistryImageStore)):
229234
raise ValueError(f"{image_name} is in a registry on this machine, which only an image in the local "
230235
"registry's store, such as an @local one, can name")
236+
if context_path is not None:
237+
context = _context_of(Path(context_path), dockerfile_path)
238+
store = get_artifact_store()
239+
version = store.next_version(id)
240+
prefix = store.attempt_prefix("docker_image", id)
241+
return store.put_document(cls(
242+
id=id,
243+
version=version,
244+
description=description,
245+
image_name=image_name,
246+
image_digest=image_digest,
247+
build_context_object_url=_upload_context(config.get_object_store_to_write(prefix, id), prefix, context),
248+
dockerfile_path=context.dockerfile,
249+
source_digest=context.source_digest(None),
250+
))
231251
if earlier_pin is not None and pinned_from(earlier_pin, image_name):
232252
pinned = earlier_pin
233253
else:
@@ -281,6 +301,11 @@ def load(self) -> bytes:
281301
raise ValueError(f"{self.id!r} v{self.version} has no tar.gz; its image is {how}")
282302
return get_artifact_store().get_object(self.tar_gz_object_url)
283303

304+
@property
305+
def pinned_name(self) -> str:
306+
"""What a sandbox pulls: ``image_name`` pinned to the ``image_digest`` when it records one."""
307+
return f"{self.image_name}@{self.image_digest}" if self.image_digest else self.image_name
308+
284309
@property
285310
def context_only(self) -> bool:
286311
"""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‎

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,11 @@
3939
GATEWAY,
4040
SERVICE_DB,
4141
WEBSITE_BROWSER,
42-
InfraBuild,
42+
InfraSetup,
4343
InfraError,
4444
default_env_id,
45-
infra_to_build,
45+
env_images,
46+
infra_to_set_up,
4647
needs_docker,
4748
put_command,
4849
stores_are_local,
@@ -51,7 +52,6 @@
5152
from agent_env.env.envs._deployment import provider_or_class
5253
from agent_env.env.envs.mcp_server import MCPServerEnv
5354
from agent_env.env.envs.multi_env import MultiEnv
54-
from agent_env.env.envs.service_db import ServiceDBEnv
5555
from agent_env.env.envs.website import WebsiteEnv
5656
from agent_env.env.registry import get_env_registry
5757
from agent_env.providers.env_providers.env_gateway_provider import EnvironmentGatewayProvider, _gateway_topology
@@ -89,7 +89,7 @@ class Preflight:
8989
"""What a run needs before its tasks start, beyond its writes: the infra envs it builds."""
9090

9191
infra_kinds: frozenset[str] # the infra its deploys on the local provider run on
92-
infra: tuple[InfraBuild, ...] # those the store doesn't hold yet, or holds built from other inputs
92+
infra: tuple[InfraSetup, ...] # those it sets up first: missing, built from other inputs, or to switch to pulling
9393
built_here: frozenset[str] = frozenset() # the ids of the images the bundle builds that a deploy here runs
9494

9595

@@ -161,21 +161,22 @@ def task(self, task: ResolvedEntry) -> None:
161161
self._problem(where, str(e))
162162

163163
def finish(self) -> Preflight:
164-
infra: list[InfraBuild] = []
164+
infra: list[InfraSetup] = []
165165
try:
166-
infra = infra_to_build(self.infra)
166+
infra = infra_to_set_up(self.infra)
167167
except InfraError as e:
168168
self.run_problems.extend(e.problems)
169169
built_here = self._built_here()
170170
self._default_agent()
171-
reason = docker_unreachable() if self.docker_users or infra else None
171+
built = [build for build in infra if build.builds]
172+
reason = docker_unreachable() if self.docker_users or built else None
172173
if reason and self.docker_users:
173174
self.run_problems.append(f"the local sandbox provider runs containers for {_listed(self.docker_users)}, "
174175
f"and {reason}")
175176
elif reason: # infra the run builds for deploys elsewhere
176-
self.run_problems.append(needs_docker(infra, reason))
177+
self.run_problems.append(needs_docker(built, reason))
177178
else:
178-
self._docker_plugins(built_here, infra)
179+
self._docker_plugins(built_here, built)
179180
needed = frozenset().union(*(names for _, names in self.model_users))
180181
if needed and (missing := model_endpoint_missing(needed)):
181182
name, why = missing
@@ -434,11 +435,11 @@ def _refusal(self, link: SandboxProvider, form: Accepts, image: _Image, *, vms_l
434435
f"on a provider {builds}")
435436

436437
def _remote_infra(self, link: SandboxProvider, kinds: set[str]) -> list[str]:
437-
"""Why ``link``, a provider not on this machine, can't run the infra envs of ``kinds``. One the run builds, in
438-
local stores missing or built from other inputs, isn't checked: it's built anew, with its build context."""
439-
built = {build.kind for build in infra_to_build(kinds)} if stores_are_local() else set()
438+
"""Why ``link``, a provider not on this machine, can't run the infra envs of ``kinds``. One the run sets up, in
439+
local stores, isn't checked: it's written anew, pulled from its registry or built, with its build context."""
440+
set_up = {setup.kind for setup in infra_to_set_up(kinds)} if stores_are_local() else set()
440441
found = []
441-
for kind in sorted(kinds - built):
442+
for kind in sorted(kinds - set_up):
442443
env_id = default_env_id(kind)
443444
try:
444445
env = Env.get(env_id)
@@ -447,10 +448,8 @@ def _remote_infra(self, link: SandboxProvider, kinds: set[str]) -> list[str]:
447448
f"and the store doesn't hold it; agent-env builds infra envs only into local stores, so put "
448449
f"it with `{put_command(kind)}`")
449450
continue
450-
artifacts = ([env.db_docker_image_artifact, env.db_web_docker_image_artifact,
451-
env.db_mcp_docker_image_artifact] if isinstance(env, ServiceDBEnv) else [env.docker_image_artifact])
452451
accepts = EnvironmentGatewayProvider.accepts(link)
453-
found += [refusal for image in artifacts
452+
found += [refusal for image in env_images(env)
454453
if (refusal := self._refusal(link, accepts, _Image(f"the {kind} env {env_id!r}'s image {image.id!r}",
455454
image), vms_load=True))]
456455
return found
@@ -475,7 +474,7 @@ def _built_here(self) -> frozenset[str]:
475474
"modal")
476475
return frozenset(here)
477476

478-
def _docker_plugins(self, built_here: frozenset[str], infra: list[InfraBuild]) -> None:
477+
def _docker_plugins(self, built_here: frozenset[str], infra: list[InfraSetup]) -> None:
479478
"""The Docker plugins the run uses on this machine: Compose for each gateway stack it runs, Buildx for each
480479
image it builds, as ``docker build`` does with BuildKit."""
481480
if self.compose_users and (missing := docker_plugin_missing("compose")):

‎src/agent_env/bundle/run.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@
66
SIGTERM cancels the runs: each one that started is marked cancelled and torn down, and a second one stops the
77
teardown. A bundle's evals run only the bundle's own tasks for now, so one naming a store task is refused
88
before anything is written, and so is a deploy its sandbox provider can't serve (``preflight``). The infra envs a
9-
gateway deploy on the local provider needs are built once the writes are done, before any task starts. A dry run
10-
makes every check the run makes before its first task, and writes, builds and runs nothing.
9+
gateway deploy on the local provider needs are set up once the writes are done, before any task starts: pulled from
10+
their published images, or built. A dry run makes every check the run makes before its first task, and writes, builds
11+
and runs nothing.
1112
"""
1213

1314
from __future__ import annotations
@@ -23,7 +24,7 @@
2324
from pathlib import Path
2425

2526
from agent_env.artifact.artifacts.docker_image import DockerImageArtifact
26-
from agent_env.env.bootstrap import InfraBuild, ensure_default_envs
27+
from agent_env.env.bootstrap import InfraSetup, ensure_default_envs
2728
from agent_env.providers import build_sandbox_provider
2829
from agent_env.store.base import NotFoundError
2930
from agent_env.store.routing import namespace_routing, run_scope
@@ -152,7 +153,7 @@ class DryRun:
152153
materialization: Materialization # each write at the version it would leave, and why
153154
runs: tuple[BundleEntry, ...] # the tasks that would run, each once, in the plan's order
154155
skipped: tuple[BundleEntry, ...] # the tasks no eval names, which running every eval leaves out
155-
infra: tuple[InfraBuild, ...] = () # the infra envs the run would build first
156+
infra: tuple[InfraSetup, ...] = () # the infra envs the run would set up first
156157
built_here: frozenset[str] = frozenset() # the images the run would build on this machine, for its deploys here
157158

158159
def path(self, entry: BundleEntry) -> str:

‎src/agent_env/cli/run.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -211,9 +211,9 @@ def _report_dry_run(dry: DryRun) -> None:
211211
for line in images:
212212
click.echo(f" {line}")
213213
if dry.infra:
214-
click.echo("Would build first:")
214+
click.echo("Would set up first:")
215215
for build in dry.infra:
216-
click.echo(f" {build}")
216+
click.echo(f" {build}: {build.planned()}")
217217
click.echo("Would run:")
218218
for entry in dry.runs:
219219
click.echo(f" {dry.path(entry)} v{materialization.version_of('task', entry.id)}")

‎src/agent_env/cli/up.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ def _require_explorer_deps() -> None:
2727

2828

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

0 commit comments

Comments
 (0)