feat(computer)!: compose on the put's sandbox provider, in a builder VM per arch

src/agent_env/bundle/authoring.py

@@ -7,7 +7,7 @@
 import stat
 import tempfile
 import unicodedata
-from collections.abc import Callable, Iterator, Mapping, Sequence
+from collections.abc import Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
@@ -18,6 +18,7 @@
 from agent_env.entity.refs import EntityKind, parse_toml_ref
 from agent_env.env.env import Env
 from agent_env.env.envs._deployment import provider_refusal
+from agent_env.env.envs.computer.utils.composer import Composing
 from agent_env.providers.env_providers.constants import GATEWAY_SERVICE_NAMES
 from agent_env.providers.env_providers.env_gateway_provider import EnvironmentGatewayProvider, environment_name_problem
 from agent_env.providers.env_providers.env_provider import _env_provider_class
@@ -48,8 +49,8 @@ class AuthoringContext:
 
     bundle: Bundle
     entry: BundleEntry
-    #: Called just before writing the entry builds an image with this machine's Docker, which a run says it does.
-    on_build: Callable[[], None] | None = None
+    #: How writing a computer env composes it.
+    composing: Composing = Composing()
 
     @property
     def id(self) -> str:

src/agent_env/bundle/materialize.py

@@ -19,14 +19,18 @@
 import copy
 from collections.abc import Callable, Iterator, Mapping
 from contextlib import contextmanager, nullcontext
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any
 
 from agent_env.a2a_agent import A2AAgent
 from agent_env.artifact import Artifact
 from agent_env.artifact.artifacts.docker_image import DockerImageArtifact, push_context_image
 from agent_env.artifact.registry import canonical_type, get_artifact_registry
 from agent_env.entity.refs import EntityKind, EntityRef, RefRole, ref_sites
+from agent_env.env import Env
+from agent_env.env.envs.computer.env import ComputerEnv
+from agent_env.env.envs.computer.utils import composer
+from agent_env.env.envs.computer.utils.composer import Composing
 from agent_env.env.registry import get_env_registry
 from agent_env.eval import Eval, EvalTask
 from agent_env.providers.sandbox_providers.host_images import build, build_staged
@@ -64,6 +68,8 @@ class Materialization:
     writes: tuple[Materialized, ...]  # in the plan's order
     # A dry run's steps with a preflight that read what it would write first, so the store can't check them yet.
     not_preflighted: tuple[tuple[Write, TaskStep], ...] = ()
+    # The computer envs reused with pairs left to compose, which a run composes and a dry run says it would.
+    composing: tuple[str, ...] = ()
 
     def version_of(self, store: str, id: str) -> int:
         for done in self.writes:
@@ -81,6 +87,8 @@ def materialize(
     on_write: Callable[[Materialized], None] | None = None,
     build_here: frozenset[str] = frozenset(),
     ledger: Ledger | None = None,
+    composing: Composing = Composing(),
+    on_progress: Callable[[Write, str], None] | None = None,
 ) -> Materialization:
     """Write ``plan``'s entities, tasks and evals. The first write that fails stops it, and every earlier
     one stays: they're in the ledger, so the next run reuses them. ``on_wait`` is called when another run
@@ -94,6 +102,9 @@ def materialize(
     Each write gets the version it would reuse or the store's next, which another run can take first. A step
     reading what the run would write first isn't preflighted, since the store doesn't hold it yet.
 
+    Each computer env written is composed as ``composing`` says, and so is what one reused has left, as ``agent-env
+    env build`` does; ``on_progress``, when given, is told how each goes, a line at a time with its write.
+
     ``ledger`` keeps the rows, this run's in the ``@local`` namespace's store unless given; it also says whether each
     task and eval names the versions it reads (``Ledger.pins_references``) and whether each image the bundle builds is
     built here and pushed (``Ledger.pushes_images``), as a registration's does."""
@@ -106,6 +117,7 @@ def materialize(
     tasks = [write for write in plan.writes if write.kind is BundleKind.TASK]
     evals = [write for write in plan.writes if write.kind is BundleKind.EVAL]
     done: dict[tuple[str, str], Materialized] = {}
+    left: list[str] = []
 
     def through_ledger(write: Write, write_fn: Callable[[], int]) -> None:
         with _noted(plan, write, "checking" if dry_run else "writing"):
@@ -128,9 +140,13 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None:
             if ledger.pushes_images and isinstance(write.source, BuiltImage) and write.id not in plan.published:
                 through_ledger(write, lambda: _write_pushed_image(plan, write, on_build))
             else:
-                through_ledger(write, lambda: _write_entity(plan, write, here, on_build))
+                through_ledger(write, lambda: _write_entity(plan, write, here, on_build,
+                                                            _composing(composing, on_progress, write)))
             if here and done[_key(write)].reused:
                 _build_reused(plan, write, done[_key(write)].version, on_build)
+            if write.kind is BundleKind.ENV and done[_key(write)].reused and _compose_reused(
+                    write, done[_key(write)].version, _composing(composing, on_progress, write), dry_run):
+                left.append(write.id)
         unwritten = {_key(item.write) for item in done.values() if not item.reused} if dry_run else set()
         written = done if ledger.pins_references else None
         built, problems, unchecked = {}, [], []
@@ -147,18 +163,24 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None:
                                                    requires=write.requires).version)
         for write in evals:
             through_ledger(write, lambda: _write_eval(plan, write, written))
-    return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked))
+    return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked), tuple(left))
+
+
+def _composing(composing: Composing, on_progress: Callable[[Write, str], None] | None, write: Write) -> Composing:
+    """``composing`` for ``write``, telling ``on_progress`` its lines with it when given."""
+    return composing if on_progress is None else replace(composing, on_progress=lambda line: on_progress(write, line))
 
 
-def _write_entity(plan: Plan, write: Write, here: bool, on_build: Callable[[Write], None] | None) -> int:
+def _write_entity(plan: Plan, write: Write, here: bool, on_build: Callable[[Write], None] | None,
+                  composing: Composing) -> int:
     if isinstance(write.source, ReferencedImage):
         return _write_referenced_image(plan, write)
     if isinstance(write.source, BuiltImage):
         if write.id in plan.published:
             return _write_published_image(plan, write)
         return _write_build_context(plan, write, here, on_build)
     if write.kind is BundleKind.ENV:
-        return _write_env(plan, write, _before(on_build, write))
+        return _write_env(plan, write, composing)
     return _WRITERS[write.kind](plan, write)
 
 
@@ -292,11 +314,29 @@ def _write_agent(plan: Plan, write: Write) -> int:
                               AuthoringContext(plan.bundle_of(entry), entry)).version
 
 
-def _write_env(plan: Plan, write: Write, on_build: Callable[[], None] | None = None) -> int:
+def _compose_reused(write: Write, version: int, composing: Composing, dry_run: bool) -> bool:
+    """Whether the computer env ``write`` reuses has pairs left to compose on ``composing``'s provider (``pending``:
+    failed, stopped, or built but unprobed though it runs their arch natively) or its computer to write again;
+    composed unless ``dry_run``, which says what it would compose."""
+    env = Env.get(write.id, version)
+    if not isinstance(env, ComputerEnv) or env.computer is None:
+        return False
+    pairs = composer.pending(env, composing.sandbox_provider)
+    if not pairs and composer.current(env):
+        return False
+    if not dry_run:
+        composer.compose(env, pairs, composing)
+    elif composing.on_progress is not None:
+        composing.on_progress(
+            f"would compose {', '.join(f'{form}/{arch}' for form, arch in pairs) or 'its computer again'}")
+    return True
+
+
+def _write_env(plan: Plan, write: Write, composing: Composing = Composing()) -> int:
     entry = write.source.entry
     cls = get_env_registry()[entry.type]
     return cls.from_toml(_pinned(plan, write, cls.toml_refs),
-                         AuthoringContext(plan.bundle_of(entry), entry, on_build=on_build)).version
+                         AuthoringContext(plan.bundle_of(entry), entry, composing=composing)).version
 
 
 def _pinned(plan: Plan, write: Write, refs: tuple[EntityRef, ...]) -> Any:

src/agent_env/bundle/register.py

@@ -31,6 +31,8 @@
 from agent_env.config import get_config
 from agent_env.entity.refs import RefRole, ref_sites
 from agent_env.plugins._bundles import AGENT_ALIASES
+from agent_env.env.envs.computer.utils.composer import Composing
+from agent_env.providers import build_sandbox_provider
 from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore
 from agent_env.store.ids import LOCAL_PREFIX, is_local_id
 from agent_env.store.image_store.oci_registry_credentials import is_loopback_host
@@ -77,16 +79,20 @@ def register_bundle(
     dry_run: bool = False,
     declaring: Declaring | None = None,
     on_progress: Callable[[str], None] | None = None,
+    sandbox: str | None = None,
+    emulate: bool = True,
 ) -> Registration:
     """Register the bundle at ``root``: ``tasks`` and ``evals`` by name, or with neither every task and eval it has,
     or every entity when it has none, with what they reference. ``prefix`` opens each entry's registry id, the folder's
     name and a dash unless given. An id the store holds already takes a new version, whoever wrote it. ``on_progress``
     is given a line for each write, as it's checked or written, and before each image is built. A folder with no
     bundle.toml that uses plugin packages is given one, or told what to declare, as ``declaring`` says: by default as a
-    run of the folder would (``declaring_for``).
+    run of the folder would (``declaring_for``). ``sandbox`` is the provider its computer envs compose on, by default
+    the env's, emulating an arch no link runs natively unless ``emulate`` is False, which leaves it skipped.
 
     It turns namespace routing on while it runs. Raises BundleError listing every problem found before anything is
     written; a write that fails stops it, and every earlier one stays, in the ledger, for the next run to reuse."""
+    composing = Composing(build_sandbox_provider(sandbox) if sandbox else None, emulate)
     with namespace_routing():
         bundle = parse_bundle(Path(root))
         resolved = resolve_bundle(bundle)
@@ -108,7 +114,8 @@ def register_bundle(
             registered, dry_run=dry_run, ledger=ledger,
             on_wait=lambda: say("waiting for another agent-env run to finish writing this bundle's ids"),
             on_build=lambda write: say(f"  {_label(registered, write)}: building with docker and pushing"),
-            on_write=lambda done: say(f"  {registered_line(registered, done)}"))
+            on_write=lambda done: say(f"  {registered_line(registered, done)}"), composing=composing,
+            on_progress=lambda write, line: say(f"  {_label(registered, write)}: {line}"))
     return Registration(registered, materialization, ids, tuple(_floating(registered)), identity, shared, requires)
 
 

src/agent_env/bundle/run.py

@@ -25,6 +25,7 @@
 
 from agent_env.artifact.artifacts.docker_image import DockerImageArtifact
 from agent_env.env.bootstrap import InfraSetup, ensure_default_envs
+from agent_env.env.envs.computer.utils.composer import Composing
 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
@@ -214,12 +215,14 @@ def run_bundle(
     on_progress: Callable[[str], None] | None = None,
     id_root: str | None = None,
     keep: bool = False,
+    emulate: bool = True,
 ) -> BundleRun:
     """Write what the tasks and evals selected in ``root`` need, as ``plan_bundle`` selects them, then run
     each of their tasks once. ``model``, ``user_model`` and ``judge_model`` replace the models the solver, the
     user simulator and the judges run on, and ``agent`` the agent the solver runs as (``Plan.agent``), leaving the
-    tasks as written; ``sandbox`` is the provider that envs,
-    agents, sandboxes and the judge deploy on. ``on_progress`` is given a line of text
+    tasks as written; ``sandbox`` is the provider that computer envs compose on and that envs,
+    agents, sandboxes and the judge deploy on; a computer env composing there emulates an arch no link runs natively
+    unless ``emulate`` is False, which leaves it skipped. ``on_progress`` is given a line of text
     at each write and step. ``id_root`` roots the ids of a bundle an installed package provides
     (``InstalledBundle.id_root``); a folder's ids otherwise come from its path.
 
@@ -239,12 +242,11 @@ def run_bundle(
     written, before any task is. Building an infra env the run needs raises what the build raises, before any task
     runs."""
     _refuse_a_running_loop()
-    if sandbox:
-        build_sandbox_provider(sandbox)
+    composing = _composing(sandbox, emulate)
     say = _progress(on_progress)
     with namespace_routing():
         prepared = _prepare(root, tasks, evals, sandbox, id_root, agent, False, say,
-                            declaring_for(Path(root), installed=id_root is not None, dry_run=False))
+                            declaring_for(Path(root), installed=id_root is not None, dry_run=False), composing)
         plan, entries = prepared.plan, [entry for entry, _ in prepared.tasks]
         with Interrupts() as interrupts:
             overrides = _overrides(sandbox, user_model, judge_model)
@@ -280,9 +282,9 @@ def prepare_bundle(
     It doesn't turn namespace routing on: the caller keeps it on for as long as it runs the tasks
     (``agent_env.store.routing.namespace_routing``), which suits a service that runs nothing else. Raises what
     ``run_bundle`` raises before its first task runs, except for a running event loop, which it doesn't need."""
-    if sandbox:
-        build_sandbox_provider(sandbox)
-    return _prepare(root, tasks, evals, sandbox, id_root, None, played, _progress(on_progress), _hint(id_root))
+    composing = _composing(sandbox)
+    return _prepare(root, tasks, evals, sandbox, id_root, None, played, _progress(on_progress), _hint(id_root),
+                    composing)
 
 
 def _write_bundle(
@@ -297,14 +299,18 @@ def _write_bundle(
         return _written_with_infra(plan, preflight, say)
 
 
-def _written_with_infra(plan: Plan, preflight: Preflight, say: Callable[[str], None]) -> Materialization:
-    """Write ``plan``, then set up the infra envs its runs deploy on."""
+def _written_with_infra(plan: Plan, preflight: Preflight, say: Callable[[str], None],
+                        composing: Composing = Composing()) -> Materialization:
+    """Write ``plan``, its computer envs composed as ``composing`` says, then set up the infra envs its runs deploy
+    on."""
     materialization = materialize(
         plan,
         on_wait=lambda: say("waiting for another agent-env run to finish writing this bundle's ids"),
         on_build=lambda write: say(f"{_label(plan, write)}: building with docker, which can take minutes"),
         on_write=lambda done: say(_written(plan, done)),
         build_here=preflight.built_here,
+        composing=composing,
+        on_progress=lambda write, line: say(f"{_label(plan, write)}: {line}"),
     )
     if preflight.infra:
         ensure_default_envs(preflight.infra_kinds, say=say)
@@ -320,22 +326,25 @@ def dry_run_bundle(
     agent: str | None = None,
     on_progress: Callable[[str], None] | None = None,
     id_root: str | None = None,
+    emulate: bool = True,
 ) -> DryRun:
     """Check what ``run_bundle`` with these arguments checks before its first task runs, and say what it would
     write and run, writing and running nothing: no entity, ledger row, lock or instance. Each write's version is
     the one it would reuse or the store's next, which another run writing the id first can take. ``on_progress``
-    is given the line ``run_bundle`` gives for each write, as it is checked.
+    is given the line ``run_bundle`` gives for each write, as it is checked, and for each computer env it reuses
+    with pairs left, what it would compose on ``sandbox`` (``emulate`` as ``run_bundle`` takes it).
 
     It raises what ``run_bundle`` raises before its first task runs, a failed preflight included. Like
     ``run_bundle``, it turns namespace routing on for the whole process while it runs, so it suits the CLI and
     scripts, not a service; it runs no event loop, so a running one isn't refused."""
-    if sandbox:
-        build_sandbox_provider(sandbox)
+    composing = _composing(sandbox, emulate)
     say = _progress(on_progress)
     with namespace_routing():
         plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent)
         _say_plan(plan, say, declaring_for(Path(root), installed=id_root is not None, dry_run=True))
-        materialization = materialize(plan, dry_run=True, on_write=lambda done: say(_written(plan, done)))
+        materialization = materialize(plan, dry_run=True, on_write=lambda done: say(_written(plan, done)),
+                                      composing=composing,
+                                      on_progress=lambda write, line: say(f"{_label(plan, write)}: {line}"))
         runs = _to_run(plan)
         return DryRun(materialization, runs, _skipped(plan, runs, every=not tasks and not evals), preflight.infra,
                       preflight.built_here)
@@ -363,10 +372,11 @@ def _bundle_run(plan: Plan, materialization: Materialization, runs: tuple[TaskRu
 
 
 def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandbox: str | None, id_root: str | None,
-             agent: str | None, played: bool, say: Callable[[str], None], declaring: Declaring) -> PreparedBundle:
+             agent: str | None, played: bool, say: Callable[[str], None], declaring: Declaring,
+             composing: Composing = Composing()) -> PreparedBundle:
     plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played)
     _say_plan(plan, say, declaring)
-    materialization = _written_with_infra(plan, preflight, say)
+    materialization = _written_with_infra(plan, preflight, say, composing)
     to_run = tuple((entry, Task.get(entry.id, materialization.version_of("task", entry.id))) for entry in _to_run(plan))
     return PreparedBundle(plan, materialization, to_run)
 
@@ -383,6 +393,12 @@ def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_ro
                                played=played)
 
 
+def _composing(sandbox: str | None, emulate: bool = True) -> Composing:
+    """How a run composes its computer envs: on the provider ``sandbox`` names, and emulating unless ``emulate`` is
+    False. Raises ValueError when ``sandbox`` names no provider."""
+    return Composing(build_sandbox_provider(sandbox) if sandbox else None, emulate)
+
+
 def _to_run(plan: Plan) -> tuple[BundleEntry, ...]:
     """The tasks selected on their own and every task a selected eval names, each once, in the plan's order."""
     wanted = {entry.entry.id for entry in plan.tasks}