From 14936c861cf7160d6ba3a0eeb00584e3abec8d81 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 08:29:18 -0700 Subject: [PATCH 1/5] feat(env): set up the infra envs from their published images, building only what isn't published A run on local stores now registers a kind's infra images from the ones the release workflow published for the sources installed, when GHCR has every one of them, instead of building them on this host. The digest is the one the release wheel records or, with no record, the one GHCR serves under the content tag, pinned when the image is registered. A stock env built here from unchanged sources switches to the published images once they exist; an env built from someone's own Dockerfile, or in stores that aren't local, is never touched. Registering needs no Docker here. DockerImageArtifact gains an optional digest: the image is named by its content tag, which its containers start from, and pulled by the digest, then tagged. A VM that can't pull it, after waiting out a rate limit twice, builds it from the build context the artifact keeps, under the same tag, so compose and docker run find it either way. The local provider skips the pull when its Docker already holds the tag, so later runs need no registry; Modal containers pull by the digest. The dry run says which envs it would build, pull, or switch. Co-Authored-By: Claude Opus 5.5 (1M context) --- CONTRIBUTING.md | 6 +- .../artifact/artifacts/docker_image.py | 35 +++- src/agent_env/bundle/preflight.py | 19 ++- src/agent_env/bundle/run.py | 7 +- src/agent_env/cli/run.py | 4 +- src/agent_env/cli/up.py | 4 +- src/agent_env/env/bootstrap.py | 152 +++++++++++++----- src/agent_env/env/published_images.py | 24 +++ .../providers/env_state/local_postgres.py | 6 +- .../sandbox_providers/host_images.py | 8 + .../sandbox_providers/local_sandbox.py | 8 +- .../sandbox_providers/modal_sandbox.py | 9 +- .../providers/sandbox_providers/sandbox.py | 61 ++++++- .../store/image_store/registry_api.py | 13 +- tst/unit/bundle/preflight_test.py | 21 ++- tst/unit/conftest.py | 7 + tst/unit/env/bootstrap_test.py | 113 +++++++++++-- tst/unit/env/published_images_test.py | 41 ++++- .../sandbox_providers/local_sandbox_test.py | 18 +++ .../modal_context_build_test.py | 16 ++ .../tensorlake_sandbox_test.py | 4 +- .../sandbox_providers/vm_sandbox_test.py | 78 +++++++++ .../golden/artifact/docker_image.json | 1 + 23 files changed, 562 insertions(+), 93 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 870b396b..ee5848fc 100644 --- a/CONTRIBUTING.md +++ b/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-:src-` 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. diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index d61a6491..c11cc015 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/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 ``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,9 @@ 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" ) + digest: str | None = Field( + default=None, description="The digest a sandbox pulls image_name's image by, running it as image_name" + ) # No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields. @model_serializer(mode="wrap") @@ -236,6 +240,30 @@ def put_ref(cls, id: str, *, description: str, image_name: str, version = store.next_version(id) return store.put_document(cls(id=id, version=version, description=description, image_name=pinned)) + @classmethod + def put_published(cls, id: str, *, description: str, image_name: str, digest: str, context_path: str, + dockerfile_path: str) -> "DockerImageArtifact": + """Register ``image_name``, a registry's tag for the image built from ``context_path``, which it serves under + ``digest``: a sandbox pulls the image by the digest and runs it under the tag, or builds it from the uploaded + context when it can't pull it. Nothing is built, pushed or saved, so no Docker is needed here.""" + from agent_env.artifact.store import get_artifact_store + from agent_env.config import get_config + + 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, + digest=digest, + build_context_object_url=_upload_context(get_config().get_object_store_to_write(prefix, id), prefix, context), + dockerfile_path=context.dockerfile, + source_digest=context.source_digest(None), + )) + @classmethod def put_context(cls, id: str, *, description: str, context_path: str, dockerfile_path: str, platform: str | None = DEFAULT_BUILD_PLATFORM) -> "DockerImageArtifact": @@ -281,6 +309,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's ``digest`` when it records one.""" + return f"{self.image_name}@{self.digest}" if self.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 diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index eb470913..ade5a75b 100644 --- a/src/agent_env/bundle/preflight.py +++ b/src/agent_env/bundle/preflight.py @@ -42,6 +42,7 @@ InfraBuild, InfraError, default_env_id, + env_images, infra_to_build, needs_docker, put_command, @@ -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[InfraBuild, ...] # 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 @@ -168,14 +168,15 @@ def finish(self) -> Preflight: 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,8 +435,8 @@ 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.""" + """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.""" built = {build.kind for build in infra_to_build(kinds)} if stores_are_local() else set() found = [] for kind in sorted(kinds - built): @@ -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 diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index f28b08ce..ab93d9ca 100644 --- a/src/agent_env/bundle/run.py +++ b/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 @@ -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[InfraBuild, ...] = () # 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: diff --git a/src/agent_env/cli/run.py b/src/agent_env/cli/run.py index aa788518..efd16960 100644 --- a/src/agent_env/cli/run.py +++ b/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)}") diff --git a/src/agent_env/cli/up.py b/src/agent_env/cli/up.py index 6ac7dcbe..db221e01 100644 --- a/src/agent_env/cli/up.py +++ b/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: diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index e9c93c61..d3005624 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -2,9 +2,10 @@ local Postgres state runs from, and the website browser it adds for websites. Each has the bare id config names (``default_gateway_env_id``, ``default_service_db_env_id``, ``default_website_browser_env_id``). -The put commands build them one at a time; ``ensure_default_envs`` builds the ones a run needs, into local stores only. -Each image keeps its build context beside its tar.gz, so a VM or Modal running another platform than the one it was -built on builds it again from that. +The put commands build them one at a time; ``ensure_default_envs`` sets up the ones a run needs, into local stores only: +from the images agent-env publishes for the sources installed, which each deploy pulls, when all of a kind's are +published, else built on this host. Each image keeps its build context, which a sandbox builds it from when it can't +pull it, or when it runs another platform than the one it was built for. """ from __future__ import annotations @@ -20,7 +21,7 @@ from agent_env.env.envs import GatewayEnv, MCPServerEnv from agent_env.env.envs.service_db import ServiceDBEnv from agent_env.env.envs.website_browser import WEBSITE_BROWSER_ENVIRONMENT_NAME, WEBSITE_BROWSER_IMAGE_TAG -from agent_env.env.published_images import PublishedImage +from agent_env.env.published_images import PublishedImage, published_digest from agent_env.store.base import NotFoundError from agent_env.store.document_store import LocalSqliteDocumentStore from agent_env.store.ids import derive_id @@ -53,10 +54,11 @@ def _quiet(line: str) -> None: def put_gateway_env(env_id: str, *, platform: str | None, metadata: Mapping[str, str] | None = None, - say: Callable[[str], None] = _quiet) -> GatewayEnv: - """Build the gateway image for ``platform`` (this host's when None) and write it as the gateway env ``env_id``.""" - artifact = _build_and_put("gateway", derive_id(env_id, "env_image"), "Created from agent-env CLI", - GATEWAY_DOCKERFILE, GATEWAY_IMAGE_TAG, platform=platform, say=say) + say: Callable[[str], None] = _quiet, published: Mapping[Path, str] | None = None) -> GatewayEnv: + """Write the gateway env ``env_id``, its image registered from ``published``, each Dockerfile's published digest, + when given, else built for ``platform`` (this host's when None).""" + artifact = _image("gateway", derive_id(env_id, "env_image"), "Created from agent-env CLI", GATEWAY_DOCKERFILE, + GATEWAY_IMAGE_TAG, platform=platform, say=say, published=published) say("Creating GatewayEnv...") env = GatewayEnv.put(id=env_id, docker_image_artifact=artifact, metadata=_metadata(GATEWAY, metadata)) say(f"Created GatewayEnv: id={env.id} version={env.version}") @@ -64,17 +66,15 @@ def put_gateway_env(env_id: str, *, platform: str | None, metadata: Mapping[str, def put_service_db_env(env_id: str, *, platform: str | None, metadata: Mapping[str, str] | None = None, - say: Callable[[str], None] = _quiet) -> ServiceDBEnv: - """Build the service-db, db-web and db-mcp images for ``platform`` (this host's when None) and write them as the - service-db env ``env_id``.""" - db = _build_and_put("ServiceDB", derive_id(env_id, "db_image"), f"ServiceDB PostgreSQL image for {env_id}", - SERVICE_DB_DOCKERFILE, SERVICE_DB_IMAGE_NAME, platform=platform, say=say) - db_web = _build_and_put("db-web", derive_id(env_id, "db_web_image"), - "db-web lightweight web UI for database inspection", DB_WEB_DOCKERFILE, DB_WEB_IMAGE_NAME, - platform=platform, say=say) - db_mcp = _build_and_put("db-mcp", derive_id(env_id, "db_mcp_image"), - "db-mcp PostgreSQL MCP server for direct DB access", DB_MCP_DOCKERFILE, DB_MCP_IMAGE_NAME, - platform=platform, say=say) + say: Callable[[str], None] = _quiet, published: Mapping[Path, str] | None = None) -> ServiceDBEnv: + """Write the service-db env ``env_id``, its service-db, db-web and db-mcp images registered from ``published``, each + Dockerfile's published digest, when given, else built for ``platform`` (this host's when None).""" + db = _image("ServiceDB", derive_id(env_id, "db_image"), f"ServiceDB PostgreSQL image for {env_id}", + SERVICE_DB_DOCKERFILE, SERVICE_DB_IMAGE_NAME, platform=platform, say=say, published=published) + db_web = _image("db-web", derive_id(env_id, "db_web_image"), "db-web lightweight web UI for database inspection", + DB_WEB_DOCKERFILE, DB_WEB_IMAGE_NAME, platform=platform, say=say, published=published) + db_mcp = _image("db-mcp", derive_id(env_id, "db_mcp_image"), "db-mcp PostgreSQL MCP server for direct DB access", + DB_MCP_DOCKERFILE, DB_MCP_IMAGE_NAME, platform=platform, say=say, published=published) say("Creating ServiceDBEnv...") env = ServiceDBEnv.put(id=env_id, db_docker_image_artifact=db, db_web_docker_image_artifact=db_web, db_mcp_docker_image_artifact=db_mcp, metadata=_metadata(SERVICE_DB, metadata)) @@ -83,11 +83,13 @@ def put_service_db_env(env_id: str, *, platform: str | None, metadata: Mapping[s def put_website_browser_env(env_id: str, *, platform: str | None, metadata: Mapping[str, str] | None = None, - say: Callable[[str], None] = _quiet) -> MCPServerEnv: - """Build the website browser image for ``platform`` (this host's when None) and write it as the MCP server env - ``env_id``.""" - artifact = _build_and_put("website browser", derive_id(env_id, "env_image"), "Website browser MCP server", - WEBSITE_BROWSER_DOCKERFILE, WEBSITE_BROWSER_IMAGE_TAG, platform=platform, say=say) + say: Callable[[str], None] = _quiet, + published: Mapping[Path, str] | None = None) -> MCPServerEnv: + """Write the website browser as the MCP server env ``env_id``, its image registered from ``published``, each + Dockerfile's published digest, when given, else built for ``platform`` (this host's when None).""" + artifact = _image("website browser", derive_id(env_id, "env_image"), "Website browser MCP server", + WEBSITE_BROWSER_DOCKERFILE, WEBSITE_BROWSER_IMAGE_TAG, platform=platform, say=say, + published=published) say("Creating MCPServerEnv...") env = MCPServerEnv.put(id=env_id, docker_image_artifact=artifact, environment_name=WEBSITE_BROWSER_ENVIRONMENT_NAME, metadata=_metadata(WEBSITE_BROWSER, metadata)) @@ -95,6 +97,19 @@ def put_website_browser_env(env_id: str, *, platform: str | None, metadata: Mapp return env +def _image(label: str, artifact_id: str, description: str, dockerfile: Path, image_name: str, *, platform: str | None, + say: Callable[[str], None], published: Mapping[Path, str] | None) -> DockerImageArtifact: + """``dockerfile``'s image, written as ``artifact_id`` with the folder it's in as its build context: registered as + the image published under its digest in ``published`` when given, else built.""" + if published is None: + return _build_and_put(label, artifact_id, description, dockerfile, image_name, platform=platform, say=say) + ref = _PUBLISHED[dockerfile].ref() + say(f"Registering the published {label} image {ref}...") + return DockerImageArtifact.put_published(id=artifact_id, description=description, image_name=ref, + digest=published[dockerfile], context_path=str(dockerfile.parent), + dockerfile_path=dockerfile.name) + + def _build_and_put(label: str, artifact_id: str, description: str, dockerfile: Path, image_name: str, *, platform: str | None, say: Callable[[str], None]) -> DockerImageArtifact: """Build ``dockerfile``'s image for ``platform`` from the folder it's in, and write it with that folder as its build @@ -129,6 +144,25 @@ def _metadata(kind: str, extra: Mapping[str, str] | None) -> dict[str, str]: PublishedImage("agentenv-db-mcp", DB_MCP_DOCKERFILE), PublishedImage("agentenv-website-browser", WEBSITE_BROWSER_DOCKERFILE), ) +_PUBLISHED = {image.dockerfile: image for image in PUBLISHED_IMAGES} + + +def published_digests(kind: str) -> dict[Path, str] | None: + """The digest each of ``kind``'s images is published under for the sources installed, by its Dockerfile; None unless + every one is, since a kind's images are all pulled or all built.""" + digests = {} + for dockerfile in _DOCKERFILES[kind]: + if (digest := published_digest(_PUBLISHED[dockerfile])) is None: + return None + digests[dockerfile] = digest + return digests + + +def env_images(env: Env) -> list[DockerImageArtifact]: + """The images an infra env runs: a service-db's three, or a gateway's or the website browser's one.""" + if isinstance(env, ServiceDBEnv): + return [env.db_docker_image_artifact, env.db_web_docker_image_artifact, env.db_mcp_docker_image_artifact] + return [env.docker_image_artifact] def build_inputs_digest(kind: str) -> str: @@ -152,17 +186,42 @@ def __init__(self, problems: list[str]): self.problems = problems +BUILD, PULL, SWITCH = "build", "pull", "switch" # how a run sets up an infra env +_BUILT_HERE = "built on this machine" + + @dataclass(frozen=True) class InfraBuild: - """An infra env a run builds, and why.""" + """An infra env a run sets up, how, and why.""" kind: str # GATEWAY, SERVICE_DB or WEBSITE_BROWSER id: str - reason: str # "missing", or which agent-env built the one in the store + reason: str # "missing", which agent-env built the one in the store, or _BUILT_HERE + action: str = BUILD # BUILD its images on this host, or PULL them, or SWITCH an env built here to pull them + published: tuple[tuple[Path, str], ...] = () # each image's Dockerfile and published digest, when they're pulled def __str__(self) -> str: return f"{self.kind} env {self.id!r} ({self.reason})" + @property + def builds(self) -> bool: + return self.action == BUILD + + def doing(self) -> str: + """What the run does, as it says when it starts.""" + return {BUILD: "building, which can take minutes", + PULL: f"registering its published {self._images}, which each deploy pulls", + SWITCH: f"switching to its published {self._images}, which each deploy pulls"}[self.action] + + def planned(self) -> str: + """What the run would do, as a dry run says.""" + return {BUILD: "would build", PULL: f"would pull its published {self._images}", + SWITCH: f"would switch to its published {self._images}"}[self.action] + + @property + def _images(self) -> str: + return "images" if len(_DOCKERFILES[self.kind]) > 1 else "image" + def default_env_id(kind: str) -> str: config = get_config() @@ -175,9 +234,11 @@ def put_command(kind: str) -> str: def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: - """The infra envs of ``kinds`` a run builds: each one missing from the store and, in local stores, each one agent-env - built from the Dockerfile it ships whose recorded build inputs aren't this release's. Raises InfraError naming the - put command for each one missing from stores that aren't all local, which agent-env never builds into.""" + """The infra envs of ``kinds`` a run sets up, in local stores only: each one missing from the store, and each one + agent-env built from the Dockerfile it ships whose recorded build inputs aren't this release's, are pulled from the + images published for these sources or, when they aren't, built; one built on this machine from these sources + switches to them once they're published. Raises InfraError naming the put command for each one missing from stores + that aren't all local, which agent-env never sets up infra in.""" kinds = _in_order(kinds) if not kinds: return [] @@ -189,40 +250,51 @@ def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: env = Env.get(env_id) except NotFoundError: if local: - builds.append(InfraBuild(kind, env_id, "missing")) + builds.append(_pulled_or_built(kind, env_id, "missing")) else: problems.append(f"the {kind} env {env_id!r} isn't in the store, and agent-env builds infra envs only into " f"local stores; put it with `{put_command(kind)}`") continue + if not (local and _built_from_stock(env, kind)): + continue recorded = (env.metadata or {}).get(BUILD_INPUTS_KEY) - if local and _built_from_stock(env, kind) and recorded != build_inputs_digest(kind): + if recorded != build_inputs_digest(kind): reason = "its build inputs changed" if recorded else "built before agent-env recorded its build inputs" - builds.append(InfraBuild(kind, env_id, reason)) + builds.append(_pulled_or_built(kind, env_id, reason)) + elif not all(image.digest for image in env_images(env)) and (published := published_digests(kind)) is not None: + builds.append(InfraBuild(kind, env_id, _BUILT_HERE, SWITCH, tuple(published.items()))) if problems: raise InfraError(problems) return builds +def _pulled_or_built(kind: str, env_id: str, reason: str) -> InfraBuild: + if (published := published_digests(kind)) is None: + return InfraBuild(kind, env_id, reason) + return InfraBuild(kind, env_id, reason, PULL, tuple(published.items())) + + def needs_docker(builds: Iterable[InfraBuild], reason: str) -> str: """Why ``builds`` can't be built here, given why docker can't be used (``docker_unreachable``).""" return f"building the {', '.join(build.kind for build in builds)} env needs docker, and {reason}" def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _quiet) -> list[InfraBuild]: - """Build the infra envs of ``kinds`` that ``infra_to_build`` names, for this host's platform, and return them. Each - id is locked while it's checked and built, so concurrent runs build it once. Raises InfraError before building - anything when one can't be built: missing from stores that aren't local, or docker unreachable.""" + """Set up the infra envs of ``kinds`` that ``infra_to_build`` names, and return them: each registered from its + published images, or built for this host's platform. Each id is locked while it's checked and set up, so concurrent + runs set it up once. Raises InfraError before setting any up when one can't be: missing from stores that aren't + local, or to be built with docker unreachable.""" kinds = _in_order(kinds) if not kinds: return [] ids = [default_env_id(kind) for kind in kinds] - with holding_locks(ids, on_wait=lambda: say("waiting for another agent-env run to finish building the infra envs")): + with holding_locks(ids, on_wait=lambda: say("waiting for another agent-env run to finish setting up the infra envs")): builds = infra_to_build(kinds) - if builds and (reason := docker_unreachable()): - raise InfraError([needs_docker(builds, reason)]) + if (built := [build for build in builds if build.builds]) and (reason := docker_unreachable()): + raise InfraError([needs_docker(built, reason)]) for build in builds: - say(f"{build}: building, which can take minutes") - _PUTS[build.kind](build.id, platform=None) + say(f"{build}: {build.doing()}") + _PUTS[build.kind](build.id, platform=None, published=dict(build.published) or None) return builds diff --git a/src/agent_env/env/published_images.py b/src/agent_env/env/published_images.py index 2d76edad..da2a17af 100644 --- a/src/agent_env/env/published_images.py +++ b/src/agent_env/env/published_images.py @@ -5,17 +5,24 @@ from __future__ import annotations import json +import logging import re from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path +from agent_env.store.image_store.registry_api import pin_digest from agent_env.utils.docker.context import BuildContext +logger = logging.getLogger(__name__) + REGISTRY = "ghcr.io/scaleapi" # The release workflow writes this into the wheels it publishes; a checkout, or a wheel built elsewhere, has none. RELEASE_FILE = Path(__file__).with_name("published_images.json") _DIGEST = re.compile(r"^sha256:[a-f0-9]{64}$") +# Short, so a network that drops requests doesn't hold up a run's checks for long. +_LOOKUP_TIMEOUT_SECONDS = 10 +_SERVED: dict[str, str] = {} # the digest the registry served each ref under, which never changes: no tag is moved @dataclass(frozen=True) @@ -40,6 +47,23 @@ def ref(self) -> str: return f"{self.repository}:{self.tag()}" +def published_digest(image: PublishedImage) -> str | None: + """The digest ``image`` is published under for the sources installed: the one this install's release recorded, else + the one the registry serves under its tag now; None when it serves none or can't be asked.""" + ref = image.ref() + return release_digests(RELEASE_FILE).get(ref) or _registry_digest(ref) + + +def _registry_digest(ref: str) -> str | None: + if ref not in _SERVED: + try: + _SERVED[ref] = pin_digest(ref, None, timeout=_LOOKUP_TIMEOUT_SECONDS).rpartition("@")[2] + except ValueError as e: + logger.info(f"No published image for {ref}: {e}") + return None + return _SERVED[ref] + + def release_digests(path: Path = RELEASE_FILE) -> dict[str, str]: """The digest the release published under each ref (``repository:tag``), as its wheel records them; empty when the install records none. Raises ValueError when the file isn't such a record.""" diff --git a/src/agent_env/providers/env_state/local_postgres.py b/src/agent_env/providers/env_state/local_postgres.py index 2e5d5cc4..a700f0b4 100644 --- a/src/agent_env/providers/env_state/local_postgres.py +++ b/src/agent_env/providers/env_state/local_postgres.py @@ -135,9 +135,9 @@ def service_db_config(self) -> "ServiceDBConfig": return self._service_db_config def store_images_to_load(self) -> list["DockerImageArtifact"]: - """The servicedb image set (servicedb + pgweb + db-mcp) the VM preloads from stored tars — the - artifacts (with ``tar_gz_object_url``), which the image-name config can't provide. Reuses the - cached default env, so no extra ``Env.get`` beyond ``service_db_config``'s.""" + """The servicedb image set (servicedb + pgweb + db-mcp) the VM preloads, as their artifacts, which the + image-name config can't provide: each a tar.gz to load, a published image to pull, or a build context. Reuses + the cached default env, so no extra ``Env.get`` beyond ``service_db_config``'s.""" env = self._default_service_db_env() return [ env.db_docker_image_artifact, diff --git a/src/agent_env/providers/sandbox_providers/host_images.py b/src/agent_env/providers/sandbox_providers/host_images.py index d5b0dd54..6e198e7b 100644 --- a/src/agent_env/providers/sandbox_providers/host_images.py +++ b/src/agent_env/providers/sandbox_providers/host_images.py @@ -46,6 +46,14 @@ def holds(image_name: str, image_sources: str) -> bool: return _built_from(image_name) == image_sources +def present(image_name: str) -> bool: + """Whether this machine's Docker holds an image named ``image_name``.""" + try: + return subprocess.run(["docker", "image", "inspect", image_name], capture_output=True).returncode == 0 + except FileNotFoundError: # no docker on PATH, so no image here + return False + + def built_here(image_name: str) -> bool: """Whether this machine's Docker holds ``image_name`` built here from a build context, which no registry has.""" return _built_from(image_name) is not None diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 55912b85..f43bc0ed 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -38,7 +38,7 @@ from agent_env import config from agent_env.attribution import Attribution from agent_env.config.paths import state_root -from agent_env.providers.sandbox_providers.host_images import build_here, built_here +from agent_env.providers.sandbox_providers.host_images import build_here, built_here, present from agent_env.providers.sandbox_providers.sandbox import SANDBOX_LABEL, NetworkPolicy, VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_CONTAINER, @@ -228,6 +228,12 @@ async def pull_images(self, image_names: list[str]) -> None: if pulled := [name for name, here in zip(image_names, built) if not here]: await super().pull_images(pulled) + async def pull_published(self, artifacts: list) -> list: + """Pull each published image but one this machine's Docker already holds under its name, so a run after the + first needs no registry.""" + held = await asyncio.gather(*(asyncio.to_thread(present, artifact.image_name) for artifact in artifacts)) + return await super().pull_published([artifact for artifact, here in zip(artifacts, held) if not here]) + def host_port(self, port: int) -> int: """The allocated host port for a published container port (identity if unmapped).""" return self._port_map.get(port, port) diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index cb09bd87..e1dc1500 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -260,6 +260,8 @@ def __init__( # The images prepare_image builds from their build contexts, by name, so create_container builds them rather # than pull. self._context_images: dict[str, DockerImageArtifact] = {} + # What create_container pulls for each image prepare_image found recording a digest: its name pinned to it. + self._pinned: dict[str, str] = {} async def _get_client(self) -> modal.Client: if self._client is None: @@ -289,13 +291,16 @@ async def _get_app(self, app_name: str, app_tags: Optional[dict[str, str]] = Non async def prepare_image(self, image: DockerImageArtifact, *, attribution: Optional[Attribution] = None) -> None: """Build ``image`` from its build context when Modal can't run it by name: it's only a build context, its name is in a registry Modal can't reach, or it's built for another platform than Modal's; unless this process has - already built its sources. One it pulls drops an image prepared earlier under its name.""" + already built its sources. One it pulls drops an image prepared earlier under its name, and is pulled by its + digest when it records one.""" if image.buildable and (image.by_name_problem(on_this_machine=self.ON_THIS_MACHINE) or image.platform_problem(self._sandbox_cls.platform)): self._context_images[image.image_name] = image await self._context_image(image) else: self._context_images.pop(image.image_name, None) + if image.digest: + self._pinned[image.image_name] = image.pinned_name async def _context_image(self, image: DockerImageArtifact, *, stale: str | None = None) -> tuple[modal.Image, str]: @@ -351,7 +356,7 @@ async def create_container( if (prepared := self._context_images.get(image_name)) is not None: image, built_id = await self._context_image(prepared) else: - image = await self._registry_image(image_name) + image = await self._registry_image(self._pinned.get(image_name, image_name)) logger.info( f"Creating Modal sandbox from image {image_name} " diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 87f4842e..3b3d493c 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -62,6 +62,10 @@ _REGISTRY_RETRY_SECONDS = 5.0 # How Docker's legacy builder and BuildKit report a failed RUN command, the one build failure not tried again. _RUN_FAILED = ("returned a non-zero code", "did not complete successfully") +# How a registry that rate-limits a pull answers, and the seconds a VM waits before each pull it tries again: short, +# since a published image it can't pull it builds from its build context instead, which takes about a minute. +_RATE_LIMITED = ("toomanyrequests", "429 too many requests") +_RATE_LIMIT_WAITS_SECONDS = (10, 30) # A VM file is read off its host one range per exec, whose base64 every provider's exec returns in one response. A # provider serves each exec a fraction of its bandwidth, so many ranges are in flight. @@ -312,10 +316,11 @@ async def wait_for_vm(self) -> None: async def load_docker_images(self, artifacts: list) -> None: """Put each DockerImageArtifact's image on the VM: those with a tar.gz are loaded from it, in parallel, those - with only a build context are built from it, and the rest are pulled by image name; one whose tar.gz is built - for another platform than the VM's is built from the build context beside it, for the VM's. An image this VM - can't get, such as one to pull from a registry on this machine when the VM isn't on it, or can't run, being - built for another platform with no build context, is refused before anything is loaded.""" + with only a build context are built from it, those with a digest are pulled by it (``pull_published``), and + the rest are pulled by image name; one whose tar.gz is built for another platform than the VM's is built from + the build context beside it, for the VM's. An image this VM can't get, such as one to pull from a registry on + this machine when the VM isn't on it, or can't run, being built for another platform with no build context, is + refused before anything is loaded.""" if not artifacts: return if problems := self.load_problems(artifacts): @@ -325,10 +330,38 @@ async def load_docker_images(self, artifacts: list) -> None: if artifact.tar_gz_object_url and not artifact.rebuilt_for(self.platform)]: await self._load_tarballs(tarballs) if pulled := [artifact.image_name for artifact in artifacts - if not artifact.tar_gz_object_url and not artifact.context_only]: + if not artifact.tar_gz_object_url and not artifact.context_only and not artifact.digest]: await self.pull_images(pulled) if built: await self.build_images(built) + if published := [artifact for artifact in artifacts if artifact.digest]: + await self._pull_or_build(published) + + async def _pull_or_build(self, artifacts: list) -> None: + """``pull_published`` each image, and build each it couldn't pull from its build context.""" + if not (unpulled := await self.pull_published(artifacts)): + return + try: + await self.build_images([artifact for artifact, _ in unpulled]) + except RuntimeError as e: + why = "; ".join(f"{artifact.pinned_name}: {reason}" for artifact, reason in unpulled) + raise RuntimeError(f"Couldn't pull {why}. {e}") from e + + async def pull_published(self, artifacts: list) -> list[tuple[Any, str]]: + """Pull each image by its digest, concurrently, and tag it with its ``image_name``, which its containers start + from; return each the VM couldn't pull, with why, for it to build from its build context instead.""" + results = await asyncio.gather(*(_pull_published(self, artifact) for artifact in artifacts), + return_exceptions=True) + unpulled = [] + for artifact, result in zip(artifacts, results, strict=True): + if isinstance(result, RuntimeError): + reason = _last_output_line(str(result)) + logger.warning("Couldn't pull %s on %s, so it's built from its build context: %s", artifact.pinned_name, + self.sandbox_id, reason) + unpulled.append((artifact, reason)) + elif isinstance(result, BaseException): + raise result + return unpulled def load_problems(self, artifacts: list) -> list[str]: """Why this VM can't get or run each DockerImageArtifact it can't, as ``load_docker_images`` refuses them: one @@ -603,6 +636,24 @@ def _last_output_line(error: str) -> str: return next((line for line in lines if line), "") +async def _pull_published(sandbox: VmSandbox, artifact) -> None: + """``docker pull`` ``artifact``'s image by its digest and tag it with its ``image_name``. A registry that + rate-limits the VM is asked again after each of ``_RATE_LIMIT_WAITS_SECONDS``.""" + pinned = artifact.pinned_name + for wait in (*_RATE_LIMIT_WAITS_SECONDS, None): + try: + await _once_more(sandbox, f"docker pull {shlex.quote(pinned)}", f"Pulling {pinned}", + unless=("no matching manifest",)) + break + except RuntimeError as e: + if wait is None or not any(marker in str(e).lower() for marker in _RATE_LIMITED): + raise + logger.warning("%s rate-limited the pull of %s on %s; trying again in %ss", registry_host_from_ref(pinned), + pinned, sandbox.sandbox_id, wait) + await asyncio.sleep(wait) + await sandbox.exec_script(f"docker tag {shlex.quote(pinned)} {shlex.quote(artifact.image_name)}") + + async def _pull(sandbox: VmSandbox, image_name: str) -> None: """``docker pull image_name``, once more if it fails. An image built for linux/amd64 only has nothing for an arm64 host (an Apple Silicon Mac running the local provider), so that pull falls back to the amd64 image, which diff --git a/src/agent_env/store/image_store/registry_api.py b/src/agent_env/store/image_store/registry_api.py index 174180ba..403c7e5c 100644 --- a/src/agent_env/store/image_store/registry_api.py +++ b/src/agent_env/store/image_store/registry_api.py @@ -41,18 +41,19 @@ class _ImageRef: digest: str | None -def pin_digest(ref: str, auth: RegistryAuth | None) -> str: +def pin_digest(ref: str, auth: RegistryAuth | None, *, timeout: float = _TIMEOUT_SECONDS) -> str: """``ref`` pinned to the digest its registry serves for it now, as ``name:tag@sha256:...``. A tag (``latest`` when none is given) is resolved; a digest is checked to exist; a ref giving both must agree, since a pull goes by the digest and ignores the tag. ``auth`` is used where the registry asks for credentials; - without it, a registry's anonymous token is. Raises ValueError naming what the registry answered. + without it, a registry's anonymous token is. Each request waits ``timeout`` seconds at most. Raises ValueError naming + what the registry answered. """ image = _parse(ref) if image.digest is None: tag = image.tag or "latest" - return f"{image.name}:{tag}@{_check(image, tag, auth)}" - served = _check(image, image.tag or image.digest, auth) + return f"{image.name}:{tag}@{_check(image, tag, auth, timeout)}" + served = _check(image, image.tag or image.digest, auth, timeout) if image.tag and served != image.digest: raise ValueError(f"{ref}: the tag {image.tag!r} names {served} now, not {image.digest}; give the tag or the " "digest alone") @@ -86,7 +87,7 @@ def _parse(ref: str) -> _ImageRef: return _ImageRef(name=name, host=host, repository=repository, tag=tag, digest=digest or None) -def _check(image: _ImageRef, reference: str, auth: RegistryAuth | None) -> str: +def _check(image: _ImageRef, reference: str, auth: RegistryAuth | None, timeout: float) -> str: """The digest the registry serves for ``reference`` (a tag or a digest) in ``image``'s repository.""" # A registry on this machine serves plain HTTP, as Docker allows for a loopback one. scheme = "http" if is_loopback_host(image.host) else "https" @@ -94,7 +95,7 @@ def _check(image: _ImageRef, reference: str, auth: RegistryAuth | None) -> str: what = f"{image.name}{'@' if _DIGEST.match(reference) else ':'}{reference}" basic = (auth.username, auth.password) if auth is not None else None try: - with httpx.Client(timeout=_TIMEOUT_SECONDS, follow_redirects=True) as client: + with httpx.Client(timeout=timeout, follow_redirects=True) as client: headers = {"Accept": _MANIFEST_TYPES} # Basic credentials go with every request until a Bearer token replaces them. credentials = basic diff --git a/tst/unit/bundle/preflight_test.py b/tst/unit/bundle/preflight_test.py index 4595e402..2375b417 100644 --- a/tst/unit/bundle/preflight_test.py +++ b/tst/unit/bundle/preflight_test.py @@ -796,6 +796,17 @@ def test_infra_built_for_deploys_elsewhere_needs_docker_here(bundle_dir, monkeyp "building the service-db, gateway env needs docker, and docker isn't on PATH"] +def test_infra_pulled_for_deploys_elsewhere_needs_no_docker_here(bundle_dir, monkeypatch): + _env("crm", REMOTE) + _task(bundle_dir, [{"id": "env", "type": "deploy_env", "env_id": "crm"}]) + monkeypatch.setattr(preflight_module, "docker_unreachable", lambda: "docker isn't on PATH") + monkeypatch.setattr("agent_env.env.published_images._registry_digest", lambda ref: "sha256:" + "4" * 64) + + assert [f"{build}: {build.planned()}" for build in dry_run_bundle(bundle_dir, sandbox="modal_vm").infra] == [ + "service-db env 'default-db' (missing): would pull its published images", + "gateway env 'default' (missing): would pull its published image"] + + def test_a_gateway_deploy_on_another_provider_needs_infra_it_can_reach(bundle_dir): _env("crm", REMOTE) _task(bundle_dir, [{"id": "env", "type": "deploy_env", "env_id": "crm"}]) @@ -930,15 +941,19 @@ def a_local_image_lands_once_planned(*args, **kwargs): assert dry_run_bundle(bundle_dir, sandbox="modal_vm").runs -def test_the_cli_dry_run_lists_the_infra_it_would_build(bundle_dir, quiet_logs): +def test_the_cli_dry_run_lists_the_infra_it_would_set_up_and_how(bundle_dir, quiet_logs, monkeypatch): _env("crm") _task(bundle_dir, [{"id": "env", "type": "deploy_env", "env_id": "crm"}]) + gateway = next(image.ref() for image in bootstrap.PUBLISHED_IMAGES if image.dockerfile == bootstrap.GATEWAY_DOCKERFILE) + monkeypatch.setattr("agent_env.env.published_images._registry_digest", + lambda ref: "sha256:" + "4" * 64 if ref == gateway else None) result = CliRunner().invoke(cli, ["run", str(bundle_dir), "--dry-run"]) assert result.exit_code == 0, result.output - assert ("Would build first:\n service-db env 'default-db' (missing)\n gateway env 'default' (missing)\n" - "Would run:\n tasks/t.json v1\n") in result.output + assert ("Would set up first:\n service-db env 'default-db' (missing): would build\n" + " gateway env 'default' (missing): would pull its published image\nWould run:\n tasks/t.json v1\n" + ) in result.output def test_a_run_builds_its_infra_after_its_writes_and_before_its_tasks(bundle_dir, monkeypatch): diff --git a/tst/unit/conftest.py b/tst/unit/conftest.py index ee617785..2372d637 100644 --- a/tst/unit/conftest.py +++ b/tst/unit/conftest.py @@ -39,6 +39,13 @@ def _docker_answers(monkeypatch): monkeypatch.setattr("agent_env.bundle.preflight.model_endpoint_missing", lambda needed: None) +@pytest.fixture(autouse=True) +def _nothing_published(monkeypatch): + """Unit tests reach no registry, so no infra image is published for them and a run builds its infra. Tests of + pulling patch the lookup.""" + monkeypatch.setattr("agent_env.env.published_images._registry_digest", lambda ref: None) + + @pytest.fixture(autouse=True) def deprecations_fire_again(): """``warn_deprecated`` warns once per process; each test sees its own first use.""" diff --git a/tst/unit/env/bootstrap_test.py b/tst/unit/env/bootstrap_test.py index 4f557f91..04c3d527 100644 --- a/tst/unit/env/bootstrap_test.py +++ b/tst/unit/env/bootstrap_test.py @@ -1,5 +1,6 @@ -"""The infra envs a run builds: the ones missing from the store, or built by another agent-env release, into local -stores only, for this host's platform, once each however many runs ask.""" +"""The infra envs a run sets up: the ones missing from the store, or built by another agent-env release, into local +stores only, once each however many runs ask; pulled from the images published for their sources when every one of a +kind's is, else built for this host's platform. One built here switches to them once they're published.""" import importlib import re @@ -16,16 +17,21 @@ from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, WEBSITE_BROWSER, InfraError, ensure_default_envs +_DIGEST = "sha256:" + "5" * 64 + + @pytest.fixture def puts(local_stores, monkeypatch): - """Stands in for each put: records the kind, id and platform, and writes the env as this release's.""" + """Stands in for each put: records the kind, id and platform, or that its images were pulled, and writes the env as + this release's.""" calls = [] def put(kind): - def write(env_id, *, platform): - calls.append((kind, env_id, platform)) + def write(env_id, *, platform, published=None): + calls.append((kind, env_id, platform) if published is None else (kind, env_id, "pulled")) artifact = get_artifact_store().put_document(DockerImageArtifact( - id=f"{kind}-{env_id}", description=kind, image_name=f"{kind}:v1", tar_gz_s3_url="file:///x.tar.gz")) + id=f"{kind}-{env_id}", description=kind, image_name=f"{kind}:v1", + tar_gz_s3_url=None if published else "file:///x.tar.gz", digest=_DIGEST if published else None)) shipped = bootstrap._DOCKERFILES[kind][0].relative_to(bootstrap._ENV_PACKAGE.parent.parent) GatewayEnv.put(id=env_id, docker_image_artifact=artifact, metadata={ "dockerfile_path": f"/another/install/site-packages/{shipped}", @@ -62,6 +68,19 @@ def keep(id, *, description, image_name, build_context_path=None, dockerfile_pat assert [folder / name for folder, name in kept] == bootstrap._DOCKERFILES[kind] +@pytest.fixture +def published(monkeypatch): + """The kinds whose images the registry serves, all under ``_DIGEST``: those added to the set returned.""" + kinds: set[str] = set() + + def served(ref): + refs = {bootstrap._PUBLISHED[dockerfile].ref() for kind in kinds for dockerfile in bootstrap._DOCKERFILES[kind]} + return _DIGEST if ref in refs else None + + monkeypatch.setattr("agent_env.env.published_images._registry_digest", served) + return kinds + + def test_missing_infra_is_built_for_this_host_in_order_and_once(puts): said = [] @@ -108,11 +127,12 @@ def test_infra_that_records_no_shipped_dockerfile_is_left_as_it_is(puts): assert puts == [] -def test_an_env_built_from_someone_elses_dockerfile_is_theirs_to_rebuild(puts, monkeypatch): +def test_an_env_built_from_someone_elses_dockerfile_is_theirs_to_rebuild(puts, published): artifact = get_artifact_store().put_document(DockerImageArtifact( id="website-browser__env_image", description="b", image_name="b:v1", tar_gz_s3_url="file:///b.tar.gz")) GatewayEnv.put(id="website-browser", docker_image_artifact=artifact, metadata={"dockerfile_path": "/home/me/my-browser/Dockerfile", bootstrap.BUILD_INPUTS_KEY: "theirs"}) + published.add(WEBSITE_BROWSER) assert ensure_default_envs([WEBSITE_BROWSER]) == [] assert puts == [] @@ -134,20 +154,95 @@ def test_nothing_is_built_without_docker(puts, monkeypatch): assert puts == [] -def test_stores_that_arent_local_are_never_built_into(puts, monkeypatch): +def test_stores_that_arent_local_are_never_built_into(puts, published, monkeypatch): ensure_default_envs([GATEWAY]) monkeypatch.setattr(bootstrap, "stores_are_local", lambda: False) monkeypatch.setattr(bootstrap, "build_inputs_digest", lambda kind: "this release's") + published.update({GATEWAY, SERVICE_DB}) with pytest.raises(InfraError) as e: ensure_default_envs([GATEWAY, SERVICE_DB]) assert e.value.problems == ["the service-db env 'default-db' isn't in the store, and agent-env builds infra envs " "only into local stores; put it with `agent-env env service-db put --id default-db`"] - assert ensure_default_envs([GATEWAY]) == [] # built from other inputs, in a store it doesn't build into + assert ensure_default_envs([GATEWAY]) == [] # built from other inputs, in a store it doesn't set up infra in assert len(puts) == 1 +def test_missing_infra_whose_images_are_all_published_is_registered_from_them_with_no_docker(puts, published, + monkeypatch): + published.update({GATEWAY, SERVICE_DB}) + monkeypatch.setattr(bootstrap, "docker_unreachable", lambda: "docker isn't on PATH") + said = [] + + builds = ensure_default_envs([GATEWAY, SERVICE_DB], say=said.append) + + assert puts == [(SERVICE_DB, "default-db", "pulled"), (GATEWAY, "default", "pulled")] + assert [build.action for build in builds] == [bootstrap.PULL, bootstrap.PULL] + assert said == ["service-db env 'default-db' (missing): registering its published images, which each deploy pulls", + "gateway env 'default' (missing): registering its published image, which each deploy pulls"] + assert ensure_default_envs([GATEWAY, SERVICE_DB]) == [] + + +def test_a_kinds_images_are_pulled_only_when_every_one_is_published(puts, monkeypatch): + db_web = bootstrap._PUBLISHED[bootstrap.DB_WEB_DOCKERFILE].ref() + monkeypatch.setattr("agent_env.env.published_images._registry_digest", + lambda ref: _DIGEST if ref == db_web else None) + + (build,) = ensure_default_envs([SERVICE_DB]) + + assert build.builds and puts == [(SERVICE_DB, "default-db", None)] + + +def test_infra_built_here_is_kept_until_its_images_are_published_then_switched_to_them_once(puts, published): + ensure_default_envs([GATEWAY]) + assert ensure_default_envs([GATEWAY]) == [] + published.add(GATEWAY) + said = [] + + (build,) = ensure_default_envs([GATEWAY], say=said.append) + + assert (build.action, build.reason) == (bootstrap.SWITCH, "built on this machine") + assert said == ["gateway env 'default' (built on this machine): switching to its published image, which each " + "deploy pulls"] + assert ensure_default_envs([GATEWAY]) == [] + assert puts == [(GATEWAY, "default", None), (GATEWAY, "default", "pulled")] + + +def test_infra_built_from_other_inputs_is_pulled_when_its_images_are_published(puts, published, monkeypatch): + ensure_default_envs([GATEWAY]) + monkeypatch.setattr(bootstrap, "build_inputs_digest", lambda kind: "this release's") + published.add(GATEWAY) + + (build,) = ensure_default_envs([GATEWAY]) + + assert (build.action, build.reason) == (bootstrap.PULL, "its build inputs changed") + assert puts == [(GATEWAY, "default", None), (GATEWAY, "default", "pulled")] + + +def test_only_the_infra_a_run_builds_needs_docker(puts, published, monkeypatch): + published.add(GATEWAY) + monkeypatch.setattr(bootstrap, "docker_unreachable", lambda: "docker isn't on PATH") + + with pytest.raises(InfraError) as e: + ensure_default_envs([GATEWAY, SERVICE_DB]) + + assert e.value.problems == ["building the service-db env needs docker, and docker isn't on PATH"] + assert puts == [] + + +def test_a_published_image_is_registered_by_its_tag_and_digest_with_the_context_it_is_built_from(local_stores): + env = bootstrap.put_gateway_env("default", platform=None, published={bootstrap.GATEWAY_DOCKERFILE: _DIGEST}) + + image, published = env.docker_image_artifact, bootstrap._PUBLISHED[bootstrap.GATEWAY_DOCKERFILE] + assert (image.image_name, image.digest, image.tar_gz_object_url) == (published.ref(), _DIGEST, None) + assert image.pinned_name == f"{published.ref()}@{_DIGEST}" + assert (image.dockerfile_path, image.source_digest) == ("Dockerfile", f"sha256:{published.tag().removeprefix('src-')}") + assert get_artifact_store().get_object(image.build_context_object_url) + assert env.metadata[bootstrap.BUILD_INPUTS_KEY] == bootstrap.build_inputs_digest(GATEWAY) + assert bootstrap._built_from_stock(env, GATEWAY) + + def test_the_default_stores_are_local(local_stores): assert bootstrap.stores_are_local() diff --git a/tst/unit/env/published_images_test.py b/tst/unit/env/published_images_test.py index 5ba47234..c3be9925 100644 --- a/tst/unit/env/published_images_test.py +++ b/tst/unit/env/published_images_test.py @@ -7,8 +7,16 @@ import pytest -from agent_env.env import bootstrap -from agent_env.env.published_images import REGISTRY, PublishedImage, release_digests, write_release_digests +from agent_env.env import bootstrap, published_images +from agent_env.env.published_images import ( + REGISTRY, + PublishedImage, + _registry_digest, + published_digest, + release_digests, + write_release_digests, +) +from agent_env.store.image_store.registry_api import NoSuchImageError from agent_env.utils.docker.context import BuildContext _DIGEST = "sha256:" + "b" * 64 @@ -80,3 +88,32 @@ def test_a_file_that_isnt_a_record_is_refused(tmp_path, content, message): with pytest.raises(ValueError, match=message): release_digests(path) + + +def test_an_install_whose_release_recorded_an_images_digest_pulls_it_without_asking_the_registry(tmp_path, monkeypatch): + image = bootstrap.PUBLISHED_IMAGES[0] + record = tmp_path / "published_images.json" + write_release_digests({image.ref(): _DIGEST}, record) + monkeypatch.setattr(published_images, "RELEASE_FILE", record) + monkeypatch.setattr(published_images, "_registry_digest", lambda ref: pytest.fail(f"asked the registry for {ref}")) + + assert published_digest(image) == _DIGEST + + +def test_an_image_its_release_didnt_record_is_looked_up_and_only_what_was_served_is_remembered(monkeypatch): + """A published tag is never moved, so the digest it was served under holds; one not served yet may be later.""" + monkeypatch.setattr(published_images, "_registry_digest", _registry_digest) + monkeypatch.setattr(published_images, "_SERVED", {}) + served, missing = bootstrap.PUBLISHED_IMAGES[:2] + asked = [] + + def pin(ref, auth, *, timeout): + asked.append((ref, auth, timeout)) + if ref == missing.ref(): + raise NoSuchImageError(f"{ref}: ghcr.io has no such image") + return f"{ref}@{_DIGEST}" + + monkeypatch.setattr(published_images, "pin_digest", pin) + + assert [published_digest(image) for image in (served, missing, served, missing)] == [_DIGEST, None, _DIGEST, None] + assert asked == [(served.ref(), None, 10), (missing.ref(), None, 10), (missing.ref(), None, 10)] diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_test.py index 5088d276..928b39f1 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_test.py @@ -1164,3 +1164,21 @@ def fail(*args): _HOST_IPS = ls._host_ips + + +@pytest.mark.asyncio +async def test_the_local_provider_pulls_only_the_published_images_its_docker_doesnt_hold(monkeypatch, tmp_path): + """So a run after the first needs no registry.""" + images = [DockerImageArtifact(id=name, description="d", image_name=f"ghcr.io/team/{name}:src-1", + digest="sha256:" + "9" * 64) for name in ("held", "new")] + monkeypatch.setattr(ls, "present", lambda image_name: image_name.startswith("ghcr.io/team/held:")) + asked = [] + + async def pull_published(self, artifacts): + asked.extend(artifact.id for artifact in artifacts) + return [] + + monkeypatch.setattr(ls.VmSandbox, "pull_published", pull_published) + + assert await LocalSandbox(work_dir=tmp_path).pull_published(images) == [] + assert asked == ["new"] diff --git a/tst/unit/providers/sandbox_providers/modal_context_build_test.py b/tst/unit/providers/sandbox_providers/modal_context_build_test.py index 310e653b..c42cab1f 100644 --- a/tst/unit/providers/sandbox_providers/modal_context_build_test.py +++ b/tst/unit/providers/sandbox_providers/modal_context_build_test.py @@ -300,3 +300,19 @@ async def test_a_container_of_an_image_it_wasnt_told_of_is_pulled(local_stores, assert create.aio.call_args.kwargs["image"] == (pulled, "registry.example/solver:v2") assert builds.calls == [] + + +@pytest.mark.asyncio +async def test_a_published_image_is_pulled_by_its_digest_not_built(local_stores, context, builds, monkeypatch): + pulled = MagicMock() + monkeypatch.setattr(modal.Image, "from_registry", lambda name, **kwargs: (pulled, name)) + image = _context_image("gateway-image", context).model_copy( + update={"image_name": "ghcr.io/team/gateway:src-1", "digest": "sha256:" + "8" * 64}) + provider = _provider() + await provider.prepare_image(image) + + with _fake_create() as create, pytest.raises(RuntimeError, match="Modal sandbox create failed"): + await provider.create_container(image_name=image.image_name, port=8000, env={}) + + assert create.aio.call_args.kwargs["image"] == (pulled, image.pinned_name) + assert builds.calls == [] diff --git a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py index 6fd8e116..b2662866 100644 --- a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py @@ -23,8 +23,8 @@ def _artifact(url: str | None = "s3://bucket/image.tar.gz", *, context_only: bool = False) -> SimpleNamespace: return SimpleNamespace( - tar_gz_object_url=url, context_only=context_only, image_name="img:latest", load_problem=lambda **_: None, - rebuilt_for=lambda *_, **__: False, + tar_gz_object_url=url, context_only=context_only, image_name="img:latest", digest=None, + load_problem=lambda **_: None, rebuilt_for=lambda *_, **__: False, ) diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py index 7a97d6ea..2ca0a709 100644 --- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py @@ -323,6 +323,84 @@ async def test_the_first_pull_to_fail_cancels_the_rest_so_a_stalled_one_cant_hol assert sorted(sandbox.cancelled) == sorted([f"docker pull {PUBLIC}", f"docker pull {PRIVATE}"]) +_GATEWAY = "ghcr.io/team/gateway:src-" + "6" * 64 +_GATEWAY_DIGEST = "sha256:" + "7" * 64 +_PINNED = f"{_GATEWAY}@{_GATEWAY_DIGEST}" + + +def _published() -> DockerImageArtifact: + return DockerImageArtifact(id="gw", version=1, description="d", image_name=_GATEWAY, digest=_GATEWAY_DIGEST, + build_context_object_url="s3://bucket/gw/build-context.tar.gz", dockerfile_path="Dockerfile", + source_digest="sha256:" + "6" * 64) + + +class _PinnedPullFailsVm(_RecordingVmSandbox): + """Fails the first ``fails`` pulls by digest with ``error``, and every build when ``build_fails``.""" + + def __init__(self, error: str, fails: int = 1_000, build_fails: bool = False): + super().__init__() + self.error, self.fails, self.build_fails = error, fails, build_fails + + async def exec_script(self, script, *, max_retries=0): + self.scripts.append(script) + if script == f"docker pull {_PINNED}" and self.fails: + self.fails -= 1 + raise RuntimeError(f"Script failed (exit 1):\nstderr: {self.error}") + if "docker build" in script and self.build_fails: + raise RuntimeError("Script failed (exit 1):\nstderr: The command '/bin/sh -c pip install' returned a non-zero code: 1") + return "" + + +@pytest.fixture +def no_waits(monkeypatch, no_pause): + monkeypatch.setattr(sandbox_module, "_RATE_LIMIT_WAITS_SECONDS", (0, 0)) + + +@pytest.mark.asyncio +async def test_a_published_image_is_pulled_by_its_digest_and_tagged_with_the_name_its_containers_start_from(registry): + sandbox = _RecordingVmSandbox() + + await sandbox.load_docker_images([_published(), _image(None, PUBLIC)]) + + assert sandbox.scripts == [f"docker pull {PUBLIC}", f"docker pull {_PINNED}", f"docker tag {_PINNED} {_GATEWAY}"] + + +@pytest.mark.asyncio +async def test_a_published_image_the_vm_cant_pull_is_built_from_its_build_context_under_its_name(signing_store, + no_waits): + sandbox = _PinnedPullFailsVm("manifest unknown") + + await sandbox.load_docker_images([_published()]) + + assert sandbox.scripts.count(f"docker pull {_PINNED}") == 2 # tried once more, then built + assert not any(script.startswith("docker tag") for script in sandbox.scripts) + assert any(f"docker build --platform linux/amd64 -f Dockerfile -t {_GATEWAY} ." in script + for script in sandbox.scripts) + + +@pytest.mark.asyncio +async def test_a_registry_that_rate_limits_the_pull_is_asked_again_before_the_image_is_built(signing_store, no_waits): + passes = _PinnedPullFailsVm("toomanyrequests: retry-after: 30", fails=4) + await passes.load_docker_images([_published()]) + assert passes.scripts.count(f"docker pull {_PINNED}") == 5 and passes.scripts[-1] == f"docker tag {_PINNED} {_GATEWAY}" + + limited = _PinnedPullFailsVm("429 Too Many Requests") + await limited.load_docker_images([_published()]) + assert limited.scripts.count(f"docker pull {_PINNED}") == 6 + assert any("docker build" in script for script in limited.scripts) + + +@pytest.mark.asyncio +async def test_a_published_image_neither_pulled_nor_built_says_why_each_failed(signing_store, no_waits): + sandbox = _PinnedPullFailsVm("manifest unknown", build_fails=True) + + with pytest.raises(RuntimeError) as e: + await sandbox.load_docker_images([_published()]) + + assert str(e.value).startswith(f"Couldn't pull {_PINNED}: manifest unknown. Building {_GATEWAY} from its build " + "context failed: ") + + class _ScriptRecorder(VmSandbox): """Records exec_script invocations, running the real write_file_from_text.""" diff --git a/tst/unit/serialization/golden/artifact/docker_image.json b/tst/unit/serialization/golden/artifact/docker_image.json index 3b59f9bf..8465c6cd 100644 --- a/tst/unit/serialization/golden/artifact/docker_image.json +++ b/tst/unit/serialization/golden/artifact/docker_image.json @@ -2,6 +2,7 @@ "build_context_object_url": null, "build_context_s3_url": null, "description": "slack-mcp:v5 image", + "digest": null, "dockerfile_path": null, "id": "slack-mcp-image", "image_name": "slack-mcp:v5", From 7df32a096071da0ce81b54b565b247ed3368a9a2 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 09:20:52 -0700 Subject: [PATCH 2/5] fix(env): keep an infra env's own metadata when it's set up again, and drop a stale Modal pin A switch, a pull or a rebuild of an infra env in the store wrote its new version with only the metadata the put records afresh, dropping what was given with --metadata. The new version now keeps the env's own keys and records the generated ones (who built it, with which agent-env, from which Dockerfile, the git state, the build inputs) anew. Modal kept the digest it was told for an image name after the same name was prepared again without one. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/bootstrap.py | 14 +++++++++----- .../sandbox_providers/modal_sandbox.py | 2 ++ src/agent_env/utils/docker/metadata.py | 5 +++++ tst/unit/env/bootstrap_test.py | 19 ++++++++++++++++++- .../modal_context_build_test.py | 4 ++++ 5 files changed, 38 insertions(+), 6 deletions(-) diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index d3005624..11df23b6 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -12,7 +12,7 @@ import hashlib from collections.abc import Callable, Iterable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from agent_env.artifact import DockerImageArtifact @@ -30,7 +30,7 @@ from agent_env.store.object_store import LocalFilesystemObjectStore from agent_env.store.routing import configured_store from agent_env.utils.docker.context import BuildContext -from agent_env.utils.docker.metadata import detect_env_metadata +from agent_env.utils.docker.metadata import detect_env_metadata, detected from agent_env.utils.docker.build import build_image, docker_unreachable _ENV_PACKAGE = Path(__file__).parent @@ -199,6 +199,7 @@ class InfraBuild: reason: str # "missing", which agent-env built the one in the store, or _BUILT_HERE action: str = BUILD # BUILD its images on this host, or PULL them, or SWITCH an env built here to pull them published: tuple[tuple[Path, str], ...] = () # each image's Dockerfile and published digest, when they're pulled + kept: tuple[tuple[str, str], ...] = () # the metadata of the env in the store its new version keeps def __str__(self) -> str: return f"{self.kind} env {self.id!r} ({self.reason})" @@ -258,11 +259,13 @@ def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: if not (local and _built_from_stock(env, kind)): continue recorded = (env.metadata or {}).get(BUILD_INPUTS_KEY) + kept = tuple((key, value) for key, value in (env.metadata or {}).items() + if not detected(key) and key != BUILD_INPUTS_KEY) if recorded != build_inputs_digest(kind): reason = "its build inputs changed" if recorded else "built before agent-env recorded its build inputs" - builds.append(_pulled_or_built(kind, env_id, reason)) + builds.append(replace(_pulled_or_built(kind, env_id, reason), kept=kept)) elif not all(image.digest for image in env_images(env)) and (published := published_digests(kind)) is not None: - builds.append(InfraBuild(kind, env_id, _BUILT_HERE, SWITCH, tuple(published.items()))) + builds.append(InfraBuild(kind, env_id, _BUILT_HERE, SWITCH, tuple(published.items()), kept)) if problems: raise InfraError(problems) return builds @@ -294,7 +297,8 @@ def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _q raise InfraError([needs_docker(built, reason)]) for build in builds: say(f"{build}: {build.doing()}") - _PUTS[build.kind](build.id, platform=None, published=dict(build.published) or None) + _PUTS[build.kind](build.id, platform=None, metadata=dict(build.kept) or None, + published=dict(build.published) or None) return builds diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index e1dc1500..50b1df6d 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -301,6 +301,8 @@ async def prepare_image(self, image: DockerImageArtifact, *, attribution: Option self._context_images.pop(image.image_name, None) if image.digest: self._pinned[image.image_name] = image.pinned_name + else: + self._pinned.pop(image.image_name, None) async def _context_image(self, image: DockerImageArtifact, *, stale: str | None = None) -> tuple[modal.Image, str]: diff --git a/src/agent_env/utils/docker/metadata.py b/src/agent_env/utils/docker/metadata.py index 62c68758..88e1def0 100644 --- a/src/agent_env/utils/docker/metadata.py +++ b/src/agent_env/utils/docker/metadata.py @@ -18,6 +18,11 @@ def detect_base_metadata() -> dict[str, str]: return {k: v for k, v in metadata.items() if v} +def detected(key: str) -> bool: + """Whether ``key`` is one ``detect_env_metadata`` records, which describes a build rather than the env.""" + return key in ("created_by", "agent_env_version", "dockerfile_path") or key.startswith("git_") + + def detect_env_metadata(dockerfile: Path, context: Path) -> dict[str, str]: """Auto-detect metadata from a Dockerfile path and its git repo.""" metadata = detect_base_metadata() diff --git a/tst/unit/env/bootstrap_test.py b/tst/unit/env/bootstrap_test.py index 04c3d527..8e26cf6f 100644 --- a/tst/unit/env/bootstrap_test.py +++ b/tst/unit/env/bootstrap_test.py @@ -27,7 +27,7 @@ def puts(local_stores, monkeypatch): calls = [] def put(kind): - def write(env_id, *, platform, published=None): + def write(env_id, *, platform, metadata=None, published=None): calls.append((kind, env_id, platform) if published is None else (kind, env_id, "pulled")) artifact = get_artifact_store().put_document(DockerImageArtifact( id=f"{kind}-{env_id}", description=kind, image_name=f"{kind}:v1", @@ -231,6 +231,23 @@ def test_only_the_infra_a_run_builds_needs_docker(puts, published, monkeypatch): assert puts == [] +def test_an_env_set_up_again_keeps_its_own_metadata_and_records_how_it_was_set_up_afresh(local_stores, published): + artifact = get_artifact_store().put_document(DockerImageArtifact( + id="gateway-default", description="g", image_name="env-gateway:v1", tar_gz_s3_url="file:///g.tar.gz")) + GatewayEnv.put(id="default", docker_image_artifact=artifact, metadata={ + "owner": "me", "git_commit": "0ld", "agent_env_version": "0.0.1", + "dockerfile_path": str(bootstrap.GATEWAY_DOCKERFILE), + bootstrap.BUILD_INPUTS_KEY: bootstrap.build_inputs_digest(GATEWAY)}) + published.add(GATEWAY) + + (build,) = ensure_default_envs([GATEWAY]) + + metadata = GatewayEnv.get("default").metadata + assert build.action == bootstrap.SWITCH and metadata["owner"] == "me" + assert metadata.get("git_commit") != "0ld" and metadata.get("agent_env_version") != "0.0.1" + assert metadata[bootstrap.BUILD_INPUTS_KEY] == bootstrap.build_inputs_digest(GATEWAY) + + def test_a_published_image_is_registered_by_its_tag_and_digest_with_the_context_it_is_built_from(local_stores): env = bootstrap.put_gateway_env("default", platform=None, published={bootstrap.GATEWAY_DOCKERFILE: _DIGEST}) diff --git a/tst/unit/providers/sandbox_providers/modal_context_build_test.py b/tst/unit/providers/sandbox_providers/modal_context_build_test.py index c42cab1f..7cc4dfe9 100644 --- a/tst/unit/providers/sandbox_providers/modal_context_build_test.py +++ b/tst/unit/providers/sandbox_providers/modal_context_build_test.py @@ -316,3 +316,7 @@ async def test_a_published_image_is_pulled_by_its_digest_not_built(local_stores, assert create.aio.call_args.kwargs["image"] == (pulled, image.pinned_name) assert builds.calls == [] + await provider.prepare_image(image.model_copy(update={"digest": None})) + with _fake_create() as create, pytest.raises(RuntimeError, match="Modal sandbox create failed"): + await provider.create_container(image_name=image.image_name, port=8000, env={}) + assert create.aio.call_args.kwargs["image"] == (pulled, image.image_name) From 148eb5d63591c3161acff2c9b4482d903852486b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 09:30:51 -0700 Subject: [PATCH 3/5] refactor(images): name the registry's digest of a published image image_digest The artifact already has source_digest, the digest of the sources an image is built from; a bare digest beside it didn't say which one it was. image_digest is the registry's digest of the built image, as Docker and the registries call it. The field hasn't been released, so no stored document has it under the old name. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../artifact/artifacts/docker_image.py | 17 +++++++++-------- src/agent_env/env/bootstrap.py | 4 ++-- .../sandbox_providers/modal_sandbox.py | 2 +- .../providers/sandbox_providers/sandbox.py | 4 ++-- tst/unit/env/bootstrap_test.py | 4 ++-- .../sandbox_providers/local_sandbox_test.py | 2 +- .../modal_context_build_test.py | 4 ++-- .../tensorlake_sandbox_test.py | 2 +- .../sandbox_providers/vm_sandbox_test.py | 2 +- .../golden/artifact/docker_image.json | 2 +- 10 files changed, 22 insertions(+), 21 deletions(-) diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index c11cc015..210464f3 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -69,7 +69,7 @@ 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, by its ``digest`` when it records one, or builds from the build context beside it when it can't; or, + 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) @@ -88,8 +88,9 @@ 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" ) - digest: str | None = Field( - default=None, description="The digest a sandbox pulls image_name's image by, running it as image_name" + 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. @@ -241,10 +242,10 @@ def put_ref(cls, id: str, *, description: str, image_name: str, return store.put_document(cls(id=id, version=version, description=description, image_name=pinned)) @classmethod - def put_published(cls, id: str, *, description: str, image_name: str, digest: str, context_path: str, + def put_published(cls, id: str, *, description: str, image_name: str, image_digest: str, context_path: str, dockerfile_path: str) -> "DockerImageArtifact": """Register ``image_name``, a registry's tag for the image built from ``context_path``, which it serves under - ``digest``: a sandbox pulls the image by the digest and runs it under the tag, or builds it from the uploaded + ``image_digest``: a sandbox pulls the image by the digest and runs it under the tag, or builds it from the uploaded context when it can't pull it. Nothing is built, pushed or saved, so no Docker is needed here.""" from agent_env.artifact.store import get_artifact_store from agent_env.config import get_config @@ -258,7 +259,7 @@ def put_published(cls, id: str, *, description: str, image_name: str, digest: st version=version, description=description, image_name=image_name, - digest=digest, + image_digest=image_digest, build_context_object_url=_upload_context(get_config().get_object_store_to_write(prefix, id), prefix, context), dockerfile_path=context.dockerfile, source_digest=context.source_digest(None), @@ -311,8 +312,8 @@ def load(self) -> bytes: @property def pinned_name(self) -> str: - """What a sandbox pulls: ``image_name`` pinned to the image's ``digest`` when it records one.""" - return f"{self.image_name}@{self.digest}" if self.digest else self.image_name + """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: diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index 11df23b6..0079b257 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -106,7 +106,7 @@ def _image(label: str, artifact_id: str, description: str, dockerfile: Path, ima ref = _PUBLISHED[dockerfile].ref() say(f"Registering the published {label} image {ref}...") return DockerImageArtifact.put_published(id=artifact_id, description=description, image_name=ref, - digest=published[dockerfile], context_path=str(dockerfile.parent), + image_digest=published[dockerfile], context_path=str(dockerfile.parent), dockerfile_path=dockerfile.name) @@ -264,7 +264,7 @@ def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: if recorded != build_inputs_digest(kind): reason = "its build inputs changed" if recorded else "built before agent-env recorded its build inputs" builds.append(replace(_pulled_or_built(kind, env_id, reason), kept=kept)) - elif not all(image.digest for image in env_images(env)) and (published := published_digests(kind)) is not None: + elif not all(image.image_digest for image in env_images(env)) and (published := published_digests(kind)) is not None: builds.append(InfraBuild(kind, env_id, _BUILT_HERE, SWITCH, tuple(published.items()), kept)) if problems: raise InfraError(problems) diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index 50b1df6d..0d37bc3b 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -299,7 +299,7 @@ async def prepare_image(self, image: DockerImageArtifact, *, attribution: Option await self._context_image(image) else: self._context_images.pop(image.image_name, None) - if image.digest: + if image.image_digest: self._pinned[image.image_name] = image.pinned_name else: self._pinned.pop(image.image_name, None) diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 3b3d493c..3a1abe39 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -330,11 +330,11 @@ async def load_docker_images(self, artifacts: list) -> None: if artifact.tar_gz_object_url and not artifact.rebuilt_for(self.platform)]: await self._load_tarballs(tarballs) if pulled := [artifact.image_name for artifact in artifacts - if not artifact.tar_gz_object_url and not artifact.context_only and not artifact.digest]: + if not artifact.tar_gz_object_url and not artifact.context_only and not artifact.image_digest]: await self.pull_images(pulled) if built: await self.build_images(built) - if published := [artifact for artifact in artifacts if artifact.digest]: + if published := [artifact for artifact in artifacts if artifact.image_digest]: await self._pull_or_build(published) async def _pull_or_build(self, artifacts: list) -> None: diff --git a/tst/unit/env/bootstrap_test.py b/tst/unit/env/bootstrap_test.py index 8e26cf6f..60ccbd84 100644 --- a/tst/unit/env/bootstrap_test.py +++ b/tst/unit/env/bootstrap_test.py @@ -31,7 +31,7 @@ def write(env_id, *, platform, metadata=None, published=None): calls.append((kind, env_id, platform) if published is None else (kind, env_id, "pulled")) artifact = get_artifact_store().put_document(DockerImageArtifact( id=f"{kind}-{env_id}", description=kind, image_name=f"{kind}:v1", - tar_gz_s3_url=None if published else "file:///x.tar.gz", digest=_DIGEST if published else None)) + tar_gz_s3_url=None if published else "file:///x.tar.gz", image_digest=_DIGEST if published else None)) shipped = bootstrap._DOCKERFILES[kind][0].relative_to(bootstrap._ENV_PACKAGE.parent.parent) GatewayEnv.put(id=env_id, docker_image_artifact=artifact, metadata={ "dockerfile_path": f"/another/install/site-packages/{shipped}", @@ -252,7 +252,7 @@ def test_a_published_image_is_registered_by_its_tag_and_digest_with_the_context_ env = bootstrap.put_gateway_env("default", platform=None, published={bootstrap.GATEWAY_DOCKERFILE: _DIGEST}) image, published = env.docker_image_artifact, bootstrap._PUBLISHED[bootstrap.GATEWAY_DOCKERFILE] - assert (image.image_name, image.digest, image.tar_gz_object_url) == (published.ref(), _DIGEST, None) + assert (image.image_name, image.image_digest, image.tar_gz_object_url) == (published.ref(), _DIGEST, None) assert image.pinned_name == f"{published.ref()}@{_DIGEST}" assert (image.dockerfile_path, image.source_digest) == ("Dockerfile", f"sha256:{published.tag().removeprefix('src-')}") assert get_artifact_store().get_object(image.build_context_object_url) diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_test.py index 928b39f1..0d55287b 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_test.py @@ -1170,7 +1170,7 @@ def fail(*args): async def test_the_local_provider_pulls_only_the_published_images_its_docker_doesnt_hold(monkeypatch, tmp_path): """So a run after the first needs no registry.""" images = [DockerImageArtifact(id=name, description="d", image_name=f"ghcr.io/team/{name}:src-1", - digest="sha256:" + "9" * 64) for name in ("held", "new")] + image_digest="sha256:" + "9" * 64) for name in ("held", "new")] monkeypatch.setattr(ls, "present", lambda image_name: image_name.startswith("ghcr.io/team/held:")) asked = [] diff --git a/tst/unit/providers/sandbox_providers/modal_context_build_test.py b/tst/unit/providers/sandbox_providers/modal_context_build_test.py index 7cc4dfe9..2d88c462 100644 --- a/tst/unit/providers/sandbox_providers/modal_context_build_test.py +++ b/tst/unit/providers/sandbox_providers/modal_context_build_test.py @@ -307,7 +307,7 @@ async def test_a_published_image_is_pulled_by_its_digest_not_built(local_stores, pulled = MagicMock() monkeypatch.setattr(modal.Image, "from_registry", lambda name, **kwargs: (pulled, name)) image = _context_image("gateway-image", context).model_copy( - update={"image_name": "ghcr.io/team/gateway:src-1", "digest": "sha256:" + "8" * 64}) + update={"image_name": "ghcr.io/team/gateway:src-1", "image_digest": "sha256:" + "8" * 64}) provider = _provider() await provider.prepare_image(image) @@ -316,7 +316,7 @@ async def test_a_published_image_is_pulled_by_its_digest_not_built(local_stores, assert create.aio.call_args.kwargs["image"] == (pulled, image.pinned_name) assert builds.calls == [] - await provider.prepare_image(image.model_copy(update={"digest": None})) + await provider.prepare_image(image.model_copy(update={"image_digest": None})) with _fake_create() as create, pytest.raises(RuntimeError, match="Modal sandbox create failed"): await provider.create_container(image_name=image.image_name, port=8000, env={}) assert create.aio.call_args.kwargs["image"] == (pulled, image.image_name) diff --git a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py index b2662866..db1d38b4 100644 --- a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_test.py @@ -23,7 +23,7 @@ def _artifact(url: str | None = "s3://bucket/image.tar.gz", *, context_only: bool = False) -> SimpleNamespace: return SimpleNamespace( - tar_gz_object_url=url, context_only=context_only, image_name="img:latest", digest=None, + tar_gz_object_url=url, context_only=context_only, image_name="img:latest", image_digest=None, load_problem=lambda **_: None, rebuilt_for=lambda *_, **__: False, ) diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py index 2ca0a709..97e0414a 100644 --- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py @@ -329,7 +329,7 @@ async def test_the_first_pull_to_fail_cancels_the_rest_so_a_stalled_one_cant_hol def _published() -> DockerImageArtifact: - return DockerImageArtifact(id="gw", version=1, description="d", image_name=_GATEWAY, digest=_GATEWAY_DIGEST, + return DockerImageArtifact(id="gw", version=1, description="d", image_name=_GATEWAY, image_digest=_GATEWAY_DIGEST, build_context_object_url="s3://bucket/gw/build-context.tar.gz", dockerfile_path="Dockerfile", source_digest="sha256:" + "6" * 64) diff --git a/tst/unit/serialization/golden/artifact/docker_image.json b/tst/unit/serialization/golden/artifact/docker_image.json index 8465c6cd..d9e0060c 100644 --- a/tst/unit/serialization/golden/artifact/docker_image.json +++ b/tst/unit/serialization/golden/artifact/docker_image.json @@ -2,9 +2,9 @@ "build_context_object_url": null, "build_context_s3_url": null, "description": "slack-mcp:v5 image", - "digest": null, "dockerfile_path": null, "id": "slack-mcp-image", + "image_digest": null, "image_name": "slack-mcp:v5", "platform": null, "source_digest": null, From 5134ce39d52d48bf391a585da46839fcdb257333 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 09:43:27 -0700 Subject: [PATCH 4/5] refactor(images): register a published image with put_ref, given its build context and digest A published image is an image already in a registry, which put_ref registers, so it takes the two things that make one different instead of a constructor of its own: the build context, uploaded for a sandbox that can't pull the image to build it from, and the digest the tag is published under. Given them, the tag is recorded as image_name and the digest beside it, as given; without them, a ref is registered as before. The published path now also gets put_ref's checks that a ref names its registry and isn't on this machine's. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../artifact/artifacts/docker_image.py | 57 +++++++++---------- src/agent_env/env/bootstrap.py | 6 +- .../artifact/test_docker_image_put_ref.py | 30 +++++++++- 3 files changed, 59 insertions(+), 34 deletions(-) diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index 210464f3..b8a8424c 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -212,12 +212,18 @@ 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``. + 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: nothing is built, pushed or saved, 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``. + + Given the build context the image is built from (``context_path``, holding ``dockerfile_path``), which is + uploaded, and the digest the tag is published under (``image_digest``), the tag is recorded as ``image_name`` + and the digest beside it, as given: a sandbox pulls the image by the digest and runs it under the tag, or + builds it from the context when it can't pull it. ``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 @@ -233,6 +239,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: @@ -241,30 +262,6 @@ def put_ref(cls, id: str, *, description: str, image_name: str, version = store.next_version(id) return store.put_document(cls(id=id, version=version, description=description, image_name=pinned)) - @classmethod - def put_published(cls, id: str, *, description: str, image_name: str, image_digest: str, context_path: str, - dockerfile_path: str) -> "DockerImageArtifact": - """Register ``image_name``, a registry's tag for the image built from ``context_path``, which it serves under - ``image_digest``: a sandbox pulls the image by the digest and runs it under the tag, or builds it from the uploaded - context when it can't pull it. Nothing is built, pushed or saved, so no Docker is needed here.""" - from agent_env.artifact.store import get_artifact_store - from agent_env.config import get_config - - 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(get_config().get_object_store_to_write(prefix, id), prefix, context), - dockerfile_path=context.dockerfile, - source_digest=context.source_digest(None), - )) - @classmethod def put_context(cls, id: str, *, description: str, context_path: str, dockerfile_path: str, platform: str | None = DEFAULT_BUILD_PLATFORM) -> "DockerImageArtifact": diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index 0079b257..a561411c 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -105,9 +105,9 @@ def _image(label: str, artifact_id: str, description: str, dockerfile: Path, ima return _build_and_put(label, artifact_id, description, dockerfile, image_name, platform=platform, say=say) ref = _PUBLISHED[dockerfile].ref() say(f"Registering the published {label} image {ref}...") - return DockerImageArtifact.put_published(id=artifact_id, description=description, image_name=ref, - image_digest=published[dockerfile], context_path=str(dockerfile.parent), - dockerfile_path=dockerfile.name) + return DockerImageArtifact.put_ref(artifact_id, description=description, image_name=ref, + image_digest=published[dockerfile], context_path=str(dockerfile.parent), + dockerfile_path=dockerfile.name) def _build_and_put(label: str, artifact_id: str, description: str, dockerfile: Path, image_name: str, *, diff --git a/tst/unit/artifact/test_docker_image_put_ref.py b/tst/unit/artifact/test_docker_image_put_ref.py index c06a5ed2..422f1dfa 100644 --- a/tst/unit/artifact/test_docker_image_put_ref.py +++ b/tst/unit/artifact/test_docker_image_put_ref.py @@ -1,5 +1,6 @@ """`DockerImageArtifact.put_ref` registers an image already in a registry: a document with no tar.gz, its tag pinned -to a digest, refused when no sandbox elsewhere could pull it. The registry API is patched out.""" +to a digest, or, given its build context and the digest it's published under, its tag with the digest beside it and the +context uploaded; refused when no sandbox elsewhere could pull it. The registry API is patched out.""" import pytest @@ -9,6 +10,7 @@ from agent_env.store import ImageStore, OciRegistryImageStore, RegistryAuth from agent_env.store.base import NotFoundError from agent_env.store.routing import namespace_routing +from agent_env.utils.docker.context import BuildContext DIGEST = "sha256:" + "c" * 64 @@ -100,3 +102,29 @@ def unresolvable(ref, auth): with pytest.raises(ValueError, match="ghcr.io has no such image"): DockerImageArtifact.put_ref("tool-image", description="d", image_name="ghcr.io/org/tool:v9") assert get_artifact_store().next_version("tool-image") == 1 + + +def test_a_ref_given_its_build_context_runs_under_its_tag_pulled_by_the_digest_given(local_stores, resolved, tmp_path): + """So a sandbox that can't pull it builds it from the context, under the name its containers start from.""" + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + + image = DockerImageArtifact.put_ref("tool-image", description="d", image_name="ghcr.io/org/tool:src-1", + image_digest=DIGEST, context_path=str(tmp_path), dockerfile_path="Dockerfile") + + stored = DockerImageArtifact.get("tool-image", image.version) + assert (stored.image_name, stored.image_digest, stored.pinned_name) == ( + "ghcr.io/org/tool:src-1", DIGEST, f"ghcr.io/org/tool:src-1@{DIGEST}") + assert stored.buildable and not stored.context_only and stored.tar_gz_object_url is None + assert stored.source_digest == BuildContext.of(tmp_path, tmp_path / "Dockerfile").source_digest(None) + assert get_artifact_store().get_object(stored.build_context_object_url) + assert stored.load_problem(on_this_machine=False) is None + assert resolved == [] + + +def test_a_ref_given_its_build_context_must_name_its_registry_too(local_stores, resolved, tmp_path): + (tmp_path / "Dockerfile").write_text("FROM scratch\n") + + with pytest.raises(ValueError, match="'tool:src-1' doesn't name its registry"): + DockerImageArtifact.put_ref("tool-image", description="d", image_name="tool:src-1", image_digest=DIGEST, + context_path=str(tmp_path), dockerfile_path="Dockerfile") + assert get_artifact_store().next_version("tool-image") == 1 From 9318d2c3e76de67ec2b3381cb3429f69fe2c3fbe Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 10:21:42 -0700 Subject: [PATCH 5/5] refactor(env): name the infra setup for what it is, and keep each list of facts in one place - InfraBuild and infra_to_build are InfraSetup and infra_to_set_up: they pull and switch too. - How a run sets an env up is a Setup enum carrying what the run and a dry run say, in place of three strings and a reason constant; the digests it pulls and the metadata it keeps are plain mappings. - ServiceDBEnv.images lists its three images, which bootstrap and the local Postgres store read instead of each listing them; the keys detected() names are one tuple, and a test fails if detect_env_metadata records one it doesn't name. - The docstrings of put_ref, infra_to_set_up, ensure_default_envs and load_docker_images are down to their point. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../artifact/artifacts/docker_image.py | 20 ++--- src/agent_env/bundle/preflight.py | 16 ++-- src/agent_env/bundle/run.py | 4 +- src/agent_env/env/bootstrap.py | 90 +++++++++---------- src/agent_env/env/envs/service_db.py | 5 ++ .../providers/env_state/local_postgres.py | 7 +- .../sandbox_providers/modal_sandbox.py | 3 +- .../providers/sandbox_providers/sandbox.py | 9 +- src/agent_env/utils/docker/metadata.py | 6 +- tst/unit/env/bootstrap_test.py | 10 +-- tst/unit/utils/docker/metadata_test.py | 19 ++++ 11 files changed, 99 insertions(+), 90 deletions(-) create mode 100644 tst/unit/utils/docker/metadata_test.py diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index b8a8424c..c5ea0e0d 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -215,19 +215,13 @@ def put_tar( 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: nothing is built, pushed or saved, 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``. - - Given the build context the image is built from (``context_path``, holding ``dockerfile_path``), which is - uploaded, and the digest the tag is published under (``image_digest``), the tag is recorded as ``image_name`` - and the digest beside it, as given: a sandbox pulls the image by the digest and runs it under the tag, or - builds it from the context when it can't pull it. - - ``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.""" + """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 diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index ade5a75b..a09741d7 100644 --- a/src/agent_env/bundle/preflight.py +++ b/src/agent_env/bundle/preflight.py @@ -39,11 +39,11 @@ GATEWAY, SERVICE_DB, WEBSITE_BROWSER, - InfraBuild, + InfraSetup, InfraError, default_env_id, env_images, - infra_to_build, + infra_to_set_up, needs_docker, put_command, stores_are_local, @@ -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 it sets up first: missing, built from other inputs, or to switch to pulling + 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,9 +161,9 @@ 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() @@ -437,9 +437,9 @@ def _refusal(self, link: SandboxProvider, form: Accepts, image: _Image, *, vms_l 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 sets up, in local stores, isn't checked: it's written anew, pulled from its registry or built, with its build context.""" - built = {build.kind for build in infra_to_build(kinds)} if stores_are_local() else set() + 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) @@ -474,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")): diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index ab93d9ca..80b79da3 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -24,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 @@ -153,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 set up 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: diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index a561411c..2d6773ac 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -12,7 +12,8 @@ import hashlib from collections.abc import Callable, Iterable, Mapping -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace +from enum import Enum from pathlib import Path from agent_env.artifact import DockerImageArtifact @@ -160,9 +161,7 @@ def published_digests(kind: str) -> dict[Path, str] | None: def env_images(env: Env) -> list[DockerImageArtifact]: """The images an infra env runs: a service-db's three, or a gateway's or the website browser's one.""" - if isinstance(env, ServiceDBEnv): - return [env.db_docker_image_artifact, env.db_web_docker_image_artifact, env.db_mcp_docker_image_artifact] - return [env.docker_image_artifact] + return env.images if isinstance(env, ServiceDBEnv) else [env.docker_image_artifact] def build_inputs_digest(kind: str) -> str: @@ -186,38 +185,40 @@ def __init__(self, problems: list[str]): self.problems = problems -BUILD, PULL, SWITCH = "build", "pull", "switch" # how a run sets up an infra env -_BUILT_HERE = "built on this machine" +class Setup(Enum): + """How a run sets up an infra env: what it says as it starts, and what a dry run says it would do.""" + + BUILD = ("building, which can take minutes", "would build") + PULL = ("registering its published {images}, which each deploy pulls", "would pull its published {images}") + SWITCH = ("switching to its published {images}, which each deploy pulls", "would switch to its published {images}") + + def __init__(self, doing: str, planned: str): + self.doing, self.planned = doing, planned @dataclass(frozen=True) -class InfraBuild: +class InfraSetup: """An infra env a run sets up, how, and why.""" kind: str # GATEWAY, SERVICE_DB or WEBSITE_BROWSER id: str - reason: str # "missing", which agent-env built the one in the store, or _BUILT_HERE - action: str = BUILD # BUILD its images on this host, or PULL them, or SWITCH an env built here to pull them - published: tuple[tuple[Path, str], ...] = () # each image's Dockerfile and published digest, when they're pulled - kept: tuple[tuple[str, str], ...] = () # the metadata of the env in the store its new version keeps + reason: str # "missing", which agent-env built the one in the store, or "built on this machine" + action: Setup = Setup.BUILD + published: Mapping[Path, str] = field(default_factory=dict) # each image's Dockerfile and digest, when pulled + kept: Mapping[str, str] = field(default_factory=dict) # the metadata of the env in the store its new version keeps def __str__(self) -> str: return f"{self.kind} env {self.id!r} ({self.reason})" @property def builds(self) -> bool: - return self.action == BUILD + return self.action is Setup.BUILD def doing(self) -> str: - """What the run does, as it says when it starts.""" - return {BUILD: "building, which can take minutes", - PULL: f"registering its published {self._images}, which each deploy pulls", - SWITCH: f"switching to its published {self._images}, which each deploy pulls"}[self.action] + return self.action.doing.format(images=self._images) def planned(self) -> str: - """What the run would do, as a dry run says.""" - return {BUILD: "would build", PULL: f"would pull its published {self._images}", - SWITCH: f"would switch to its published {self._images}"}[self.action] + return self.action.planned.format(images=self._images) @property def _images(self) -> str: @@ -234,16 +235,14 @@ def put_command(kind: str) -> str: return f"agent-env env {kind} put --id {default_env_id(kind)}" -def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: - """The infra envs of ``kinds`` a run sets up, in local stores only: each one missing from the store, and each one - agent-env built from the Dockerfile it ships whose recorded build inputs aren't this release's, are pulled from the - images published for these sources or, when they aren't, built; one built on this machine from these sources - switches to them once they're published. Raises InfraError naming the put command for each one missing from stores - that aren't all local, which agent-env never sets up infra in.""" +def infra_to_set_up(kinds: Iterable[str]) -> list[InfraSetup]: + """The infra envs of ``kinds`` a run sets up, in local stores only: stock ones missing or built from other inputs, + pulled when their images are published for these sources and built otherwise, and stock ones built here switched + once they're published. Raises InfraError naming the put command for each one missing from stores not all local.""" kinds = _in_order(kinds) if not kinds: return [] - builds, problems = [], [] + setups, problems = [], [] local = stores_are_local() for kind in kinds: env_id = default_env_id(kind) @@ -251,7 +250,7 @@ def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: env = Env.get(env_id) except NotFoundError: if local: - builds.append(_pulled_or_built(kind, env_id, "missing")) + setups.append(_pulled_or_built(kind, env_id, "missing")) else: problems.append(f"the {kind} env {env_id!r} isn't in the store, and agent-env builds infra envs only into " f"local stores; put it with `{put_command(kind)}`") @@ -259,47 +258,44 @@ def infra_to_build(kinds: Iterable[str]) -> list[InfraBuild]: if not (local and _built_from_stock(env, kind)): continue recorded = (env.metadata or {}).get(BUILD_INPUTS_KEY) - kept = tuple((key, value) for key, value in (env.metadata or {}).items() - if not detected(key) and key != BUILD_INPUTS_KEY) + kept = {key: value for key, value in (env.metadata or {}).items() if not detected(key) and key != BUILD_INPUTS_KEY} if recorded != build_inputs_digest(kind): reason = "its build inputs changed" if recorded else "built before agent-env recorded its build inputs" - builds.append(replace(_pulled_or_built(kind, env_id, reason), kept=kept)) + setups.append(replace(_pulled_or_built(kind, env_id, reason), kept=kept)) elif not all(image.image_digest for image in env_images(env)) and (published := published_digests(kind)) is not None: - builds.append(InfraBuild(kind, env_id, _BUILT_HERE, SWITCH, tuple(published.items()), kept)) + setups.append(InfraSetup(kind, env_id, "built on this machine", Setup.SWITCH, published, kept)) if problems: raise InfraError(problems) - return builds + return setups -def _pulled_or_built(kind: str, env_id: str, reason: str) -> InfraBuild: +def _pulled_or_built(kind: str, env_id: str, reason: str) -> InfraSetup: if (published := published_digests(kind)) is None: - return InfraBuild(kind, env_id, reason) - return InfraBuild(kind, env_id, reason, PULL, tuple(published.items())) + return InfraSetup(kind, env_id, reason) + return InfraSetup(kind, env_id, reason, Setup.PULL, published) -def needs_docker(builds: Iterable[InfraBuild], reason: str) -> str: +def needs_docker(builds: Iterable[InfraSetup], reason: str) -> str: """Why ``builds`` can't be built here, given why docker can't be used (``docker_unreachable``).""" return f"building the {', '.join(build.kind for build in builds)} env needs docker, and {reason}" -def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _quiet) -> list[InfraBuild]: - """Set up the infra envs of ``kinds`` that ``infra_to_build`` names, and return them: each registered from its - published images, or built for this host's platform. Each id is locked while it's checked and set up, so concurrent - runs set it up once. Raises InfraError before setting any up when one can't be: missing from stores that aren't +def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _quiet) -> list[InfraSetup]: + """Set up the infra envs of ``kinds`` that ``infra_to_set_up`` names, each id locked so concurrent runs set it up + once, and return them. Raises InfraError before setting any up when one can't be: missing from stores that aren't local, or to be built with docker unreachable.""" kinds = _in_order(kinds) if not kinds: return [] ids = [default_env_id(kind) for kind in kinds] with holding_locks(ids, on_wait=lambda: say("waiting for another agent-env run to finish setting up the infra envs")): - builds = infra_to_build(kinds) - if (built := [build for build in builds if build.builds]) and (reason := docker_unreachable()): + setups = infra_to_set_up(kinds) + if (built := [setup for setup in setups if setup.builds]) and (reason := docker_unreachable()): raise InfraError([needs_docker(built, reason)]) - for build in builds: - say(f"{build}: {build.doing()}") - _PUTS[build.kind](build.id, platform=None, metadata=dict(build.kept) or None, - published=dict(build.published) or None) - return builds + for setup in setups: + say(f"{setup}: {setup.doing()}") + _PUTS[setup.kind](setup.id, platform=None, metadata=setup.kept or None, published=setup.published or None) + return setups def stores_are_local() -> bool: diff --git a/src/agent_env/env/envs/service_db.py b/src/agent_env/env/envs/service_db.py index 9a912d71..c976e8bf 100644 --- a/src/agent_env/env/envs/service_db.py +++ b/src/agent_env/env/envs/service_db.py @@ -48,6 +48,11 @@ def __init__( self.db_web_docker_image_artifact = db_web_docker_image_artifact self.db_mcp_docker_image_artifact = db_mcp_docker_image_artifact + @property + def images(self) -> list[DockerImageArtifact]: + """Its three images: the database, db-web and db-mcp.""" + return [self.db_docker_image_artifact, self.db_web_docker_image_artifact, self.db_mcp_docker_image_artifact] + def to_config(self) -> ServiceDBConfig: return ServiceDBConfig( db_image=self.db_docker_image_artifact.image_name, diff --git a/src/agent_env/providers/env_state/local_postgres.py b/src/agent_env/providers/env_state/local_postgres.py index a700f0b4..45c348d2 100644 --- a/src/agent_env/providers/env_state/local_postgres.py +++ b/src/agent_env/providers/env_state/local_postgres.py @@ -138,12 +138,7 @@ def store_images_to_load(self) -> list["DockerImageArtifact"]: """The servicedb image set (servicedb + pgweb + db-mcp) the VM preloads, as their artifacts, which the image-name config can't provide: each a tar.gz to load, a published image to pull, or a build context. Reuses the cached default env, so no extra ``Env.get`` beyond ``service_db_config``'s.""" - env = self._default_service_db_env() - return [ - env.db_docker_image_artifact, - env.db_web_docker_image_artifact, - env.db_mcp_docker_image_artifact, - ] + return self._default_service_db_env().images @classmethod def default_instance(cls) -> EnvStateInstance: diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index 0d37bc3b..2b91ca42 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -291,8 +291,7 @@ async def _get_app(self, app_name: str, app_tags: Optional[dict[str, str]] = Non async def prepare_image(self, image: DockerImageArtifact, *, attribution: Optional[Attribution] = None) -> None: """Build ``image`` from its build context when Modal can't run it by name: it's only a build context, its name is in a registry Modal can't reach, or it's built for another platform than Modal's; unless this process has - already built its sources. One it pulls drops an image prepared earlier under its name, and is pulled by its - digest when it records one.""" + already built its sources. One it pulls, by its digest when it has one, drops an image prepared under its name.""" if image.buildable and (image.by_name_problem(on_this_machine=self.ON_THIS_MACHINE) or image.platform_problem(self._sandbox_cls.platform)): self._context_images[image.image_name] = image diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 3a1abe39..9ec67073 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -315,12 +315,9 @@ async def wait_for_vm(self) -> None: ) async def load_docker_images(self, artifacts: list) -> None: - """Put each DockerImageArtifact's image on the VM: those with a tar.gz are loaded from it, in parallel, those - with only a build context are built from it, those with a digest are pulled by it (``pull_published``), and - the rest are pulled by image name; one whose tar.gz is built for another platform than the VM's is built from - the build context beside it, for the VM's. An image this VM can't get, such as one to pull from a registry on - this machine when the VM isn't on it, or can't run, being built for another platform with no build context, is - refused before anything is loaded.""" + """Put each DockerImageArtifact's image on the VM: a tar.gz is loaded, a build context built, an image with a + digest pulled by it (``pull_published``), and the rest pulled by name; a tar.gz for another platform is built + from its context instead. One the VM can't get or run (``load_problems``) is refused before anything loads.""" if not artifacts: return if problems := self.load_problems(artifacts): diff --git a/src/agent_env/utils/docker/metadata.py b/src/agent_env/utils/docker/metadata.py index 88e1def0..ae9b283c 100644 --- a/src/agent_env/utils/docker/metadata.py +++ b/src/agent_env/utils/docker/metadata.py @@ -18,9 +18,13 @@ def detect_base_metadata() -> dict[str, str]: return {k: v for k, v in metadata.items() if v} +# The keys detect_env_metadata records, besides its git_* ones. +_DETECTED = ("created_by", "agent_env_version", "dockerfile_path") + + def detected(key: str) -> bool: """Whether ``key`` is one ``detect_env_metadata`` records, which describes a build rather than the env.""" - return key in ("created_by", "agent_env_version", "dockerfile_path") or key.startswith("git_") + return key in _DETECTED or key.startswith("git_") def detect_env_metadata(dockerfile: Path, context: Path) -> dict[str, str]: diff --git a/tst/unit/env/bootstrap_test.py b/tst/unit/env/bootstrap_test.py index 60ccbd84..3bb93bcd 100644 --- a/tst/unit/env/bootstrap_test.py +++ b/tst/unit/env/bootstrap_test.py @@ -141,7 +141,7 @@ def test_an_env_built_from_someone_elses_dockerfile_is_theirs_to_rebuild(puts, p def test_a_run_that_needs_no_infra_reads_no_store(monkeypatch): monkeypatch.setattr(bootstrap, "stores_are_local", lambda: pytest.fail("read the configured stores")) - assert bootstrap.infra_to_build([]) == [] + assert bootstrap.infra_to_set_up([]) == [] def test_nothing_is_built_without_docker(puts, monkeypatch): @@ -178,7 +178,7 @@ def test_missing_infra_whose_images_are_all_published_is_registered_from_them_wi builds = ensure_default_envs([GATEWAY, SERVICE_DB], say=said.append) assert puts == [(SERVICE_DB, "default-db", "pulled"), (GATEWAY, "default", "pulled")] - assert [build.action for build in builds] == [bootstrap.PULL, bootstrap.PULL] + assert [build.action for build in builds] == [bootstrap.Setup.PULL, bootstrap.Setup.PULL] assert said == ["service-db env 'default-db' (missing): registering its published images, which each deploy pulls", "gateway env 'default' (missing): registering its published image, which each deploy pulls"] assert ensure_default_envs([GATEWAY, SERVICE_DB]) == [] @@ -202,7 +202,7 @@ def test_infra_built_here_is_kept_until_its_images_are_published_then_switched_t (build,) = ensure_default_envs([GATEWAY], say=said.append) - assert (build.action, build.reason) == (bootstrap.SWITCH, "built on this machine") + assert (build.action, build.reason) == (bootstrap.Setup.SWITCH, "built on this machine") assert said == ["gateway env 'default' (built on this machine): switching to its published image, which each " "deploy pulls"] assert ensure_default_envs([GATEWAY]) == [] @@ -216,7 +216,7 @@ def test_infra_built_from_other_inputs_is_pulled_when_its_images_are_published(p (build,) = ensure_default_envs([GATEWAY]) - assert (build.action, build.reason) == (bootstrap.PULL, "its build inputs changed") + assert (build.action, build.reason) == (bootstrap.Setup.PULL, "its build inputs changed") assert puts == [(GATEWAY, "default", None), (GATEWAY, "default", "pulled")] @@ -243,7 +243,7 @@ def test_an_env_set_up_again_keeps_its_own_metadata_and_records_how_it_was_set_u (build,) = ensure_default_envs([GATEWAY]) metadata = GatewayEnv.get("default").metadata - assert build.action == bootstrap.SWITCH and metadata["owner"] == "me" + assert build.action is bootstrap.Setup.SWITCH and metadata["owner"] == "me" assert metadata.get("git_commit") != "0ld" and metadata.get("agent_env_version") != "0.0.1" assert metadata[bootstrap.BUILD_INPUTS_KEY] == bootstrap.build_inputs_digest(GATEWAY) diff --git a/tst/unit/utils/docker/metadata_test.py b/tst/unit/utils/docker/metadata_test.py new file mode 100644 index 00000000..6181ad87 --- /dev/null +++ b/tst/unit/utils/docker/metadata_test.py @@ -0,0 +1,19 @@ +"""What a put records about how an image was built, which a new version of an infra env records afresh rather than +keeps (``detected``).""" + +from pathlib import Path + +from agent_env.utils.docker import metadata + +REPO = Path(__file__).resolve().parents[4] + + +def test_every_key_a_put_records_about_the_build_is_one_it_records_afresh(monkeypatch): + """So a key added to the detection is never carried over from an env's earlier version.""" + monkeypatch.setenv("USER", "me") + + recorded = metadata.detect_env_metadata(REPO / "src/agent_env/env/gateway/Dockerfile", REPO) + + assert {"created_by", "agent_env_version", "dockerfile_path"} <= set(recorded) + assert [key for key in recorded if not metadata.detected(key)] == [] + assert not metadata.detected("owner")