feat(computer): the disk form on Modal VM sandboxes

.github/scripts/infra_build.py

@@ -122,7 +122,8 @@
                             "command -v disk image-env mkfs.ext4 mkfs.vfat mcopy sfdisk qemu-img"],
     "agentenv-qemu-launcher": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c",
                                "qemu-system-x86_64 --version && qemu-system-aarch64 --version"
-                               " && test -f /usr/share/OVMF/OVMF_CODE_4M.fd && test -f /usr/share/AAVMF/AAVMF_CODE.fd"],
+                               " && test -f /usr/share/OVMF/OVMF_CODE_4M.fd && test -f /usr/share/AAVMF/AAVMF_CODE.fd"
+                               " && command -v socat zstd"],
 }
 
 

src/agent_env/bundle/materialize.py

@@ -89,6 +89,7 @@ def materialize(
     ledger: Ledger | None = None,
     composing: Composing = Composing(),
     on_progress: Callable[[Write, str], None] | None = None,
+    sandbox: str | 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
@@ -99,6 +100,8 @@ def materialize(
     sources (``Plan.published``) is never built here: a deploy pulls it.
 
     ``dry_run`` checks and preflights what a run would, and writes nothing: no entity, ledger row or lock.
+    ``sandbox``, the provider ``agent-env run --sandbox`` names, is where each step's preflight checks its deploys,
+    as the run deploys them there.
     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.
 
@@ -153,7 +156,7 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None:
         for write in tasks:
             with _noted(plan, write, "preflighting"):
                 built[write.id] = _build(write, plan, written)
-                found, skipped = _preflight(plan, write, built[write.id], unwritten)
+                found, skipped = _preflight(plan, write, built[write.id], unwritten, sandbox)
             problems.extend(f"{_path(plan, write)}: {problem}" for problem in found)
             unchecked.extend((write, step) for step in skipped)
         if problems:
@@ -434,18 +437,26 @@ def pinned_steps(plan: Plan, write: Write, written: Mapping[tuple[str, str], Mat
     return steps
 
 
-def _preflight(plan: Plan, write: Write, task: Task,
-               unwritten: set[tuple[str, str]]) -> tuple[list[str], list[TaskStep]]:
+#: The fields a step names the provider it deploys on in, which a run's ``--sandbox`` overrides (``sandbox_overrides``).
+_SANDBOX_FIELDS = ("sandbox_type", "judge_sandbox_type")
+
+
+def _preflight(plan: Plan, write: Write, task: Task, unwritten: set[tuple[str, str]],
+               sandbox: str | None = None) -> tuple[list[str], list[TaskStep]]:
     """``task``'s preflight problems, and the steps with a preflight left unchecked because they read a
     (store, id) in ``unwritten``, which a dry run would write first; an artifact and an agent can share an id.
     A step reading one of the task's own outputs is skipped too, and not listed: the output only exists once the
-    task runs. A step the solver deploys or installs at is checked for the agent the run gives it (``Plan.agent``)."""
+    task runs. A step the solver deploys or installs at is checked for the agent the run gives it (``Plan.agent``), and
+    a step that deploys on a provider for the run's ``sandbox``, which overrides it."""
     outputs = {output.id for output in write.source.outputs}
     solver = plan.agent.steps.get(write.id, ()) if plan.agent is not None else ()
     problems, unchecked = [], []
     for config, step in zip(write.source.config, task.steps):
+        overrides = {key: sandbox for key in _SANDBOX_FIELDS if sandbox is not None and hasattr(step, key)}
         if step.id in solver:
-            config = {**config, "a2a_agent_id": plan.agent.agent_id, "a2a_agent_version": plan.agent.version}
+            overrides |= {"a2a_agent_id": plan.agent.agent_id, "a2a_agent_version": plan.agent.version}
+        if overrides:
+            config = {**config, **overrides}
             step = build_step(config)
         problems.extend(step.config_problems())
         reads = _reads(config)

src/agent_env/bundle/run.py

@@ -224,9 +224,10 @@ def run_bundle(
     """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 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
+    tasks as written; ``sandbox`` is the provider that envs, agents, sandboxes and the judge deploy on, whose deploys
+    pick the forms a computer env composes, on the env's default provider, this machine unless the config names
+    another; that compose emulates an arch no link of it 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``) or one downloaded from the Hub (``HubBundle.id_root``); a folder's ids otherwise come
     from its path. ``source`` says where a Hub bundle came from (``HubBundle.source``), which each run's context
@@ -310,9 +311,9 @@ def _write_bundle(
 
 
 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."""
+                        composing: Composing = Composing(), sandbox: str | None = None) -> Materialization:
+    """Write ``plan``, its computer envs composed as ``composing`` says and its steps preflighted for ``sandbox``,
+    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"),
@@ -321,6 +322,7 @@ def _written_with_infra(plan: Plan, preflight: Preflight, say: Callable[[str], N
         build_here=preflight.built_here,
         composing=composing,
         on_progress=lambda write, line: say(f"{_label(plan, write)}: {line}"),
+        sandbox=sandbox,
     )
     if preflight.infra:
         ensure_default_envs(preflight.infra_kinds, say=say)
@@ -353,7 +355,7 @@ def dry_run_bundle(
         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)),
-                                      composing=composing,
+                                      composing=composing, sandbox=sandbox,
                                       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,
@@ -387,7 +389,7 @@ def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandb
     plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played,
                                solver_endpoint=solver_endpoint)
     _say_plan(plan, say, declaring)
-    materialization = _written_with_infra(plan, preflight, say, composing)
+    materialization = _written_with_infra(plan, preflight, say, composing, sandbox)
     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)
 
@@ -406,8 +408,9 @@ def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_ro
 
 
 def _composing(sandbox: str | None, emulate: bool = True) -> Composing:
-    """How a run composes its computer envs: on the provider ``sandbox`` names, emulating unless ``emulate`` is False,
-    only the forms its deploys there boot. Raises ValueError when ``sandbox`` names no provider."""
+    """How a run composes its computer envs: on the env's default provider, emulating unless ``emulate`` is False, only
+    the forms its deploys on the provider ``sandbox`` names boot. Raises ValueError when ``sandbox`` names no
+    provider."""
     return Composing.for_run(build_sandbox_provider(sandbox) if sandbox else None, emulate)
 
 

src/agent_env/cli/run.py

@@ -44,15 +44,16 @@
               help="The agent the solver runs as, whatever the tasks name: an alias such as pi, an agent of this "
                    "bundle or an installed one, or the store's.")
 @click.option("--sandbox", default=None,
-              help="The sandbox provider computer envs compose on and envs, agents, sandboxes and the judge deploy "
-                   "on; comma-separated for a fallback chain.")
+              help="The sandbox provider envs, agents, sandboxes and the judge deploy on, which picks the forms "
+                   "computer envs compose on this machine; comma-separated for a fallback chain.")
 @click.option("--keep", is_flag=True,
-              help="Keep the sandboxes up after the runs, print their endpoints, and tear them down on Ctrl-C.")
+              help="Keep the sandboxes up after the runs, print their endpoints, and tear them down on Ctrl-C; a "
+                   "provider's idle timeout still ends one nothing is connected to.")
 @click.option("--dry-run", is_flag=True,
               help="Check the bundle as the run does and show what it would write and run; write and run nothing.")
 @click.option("--no-emulate", is_flag=True,
-              help="Skip, rather than emulate, a computer env's arch no link of the sandbox provider runs natively; "
-                   "a later run or env build builds it.")
+              help="Skip, rather than emulate, a computer env's arch the provider it composes on doesn't run "
+                   "natively; a later run or env build builds it.")
 @click.pass_context
 def run(ctx: click.Context, bundle: str | None, tasks: tuple[str, ...], evals: tuple[str, ...], model: str | None,
         user_model: str | None, judge_model: str | None, agent: str | None, sandbox: str | None, keep: bool,

src/agent_env/computers/README.md

@@ -18,13 +18,24 @@ Composing a computer env makes each arch's disk from its composed container, the
 `mkfs.ext4 -d` by the disk tools (`disk/linux/tools/`) on the builder's own platform and the base disk's EFI system
 partition kept, so nothing is emulated or privileged. A disk boots where `/dev/kvm` is, in the QEMU launcher
 (`disk/linux/launcher/`): gate 3 probes it there, and elsewhere it stays `unverified` until a deploy's first boot with
-KVM probes it. A deploy boots it with `env_params = { form = "disk" }`.
+KVM probes it. A deploy boots it with `env_params = { form = "disk" }`. The launcher serves each of the computer's ports
+through socat, whose listen backlog takes a burst of connections, in front of QEMU's own forward, whose backlog is 1, and
+a reset boots the disk afresh on a new overlay in the same launcher, at the same URLs. It boots on local Docker where
+Docker can pass it `/dev/kvm`, and on `modal_vm`, amd64 only, whose sandbox runs the launcher as its image under nested
+virtualization: the provider exports the disk from the image store with this machine's Docker, compressed, once for each
+composed disk, into the per-user cache, and mounts it on the sandbox at create, which Modal uploads once for each
+content and the launcher expands onto the sandbox's disk. Such a sandbox is ended once it has been idle, with no
+connection open, for the provider's `idle_timeout`, 30 minutes unless its config sets another, as a running QEMU isn't
+activity. A gate 3 probe that ends with the connection to the driver dropped, refused or timed out runs again on a fresh
+boot, up to three in all, and the record keeps each attempt; one with a failed check is recorded as it is.
 
 A put, `env build` and the release compose every form; a run composes only the forms its sandbox's deploys boot: on each
 link of a chain, the first form it boots natively, as a deploy picks, so local Docker composes containers, and a sandbox
-that boots only disks composes disks. A compose of a disk also composes the container it's made from, unless that's
-built. A deploy that boots a form no compose has started, a disk named by `env_params` included, composes it first,
-reusing what's built, as every build is reused by its key; deploys of one env at once compose it once, the others
+that boots only disks composes disks. A run composes them on the env's default provider, this machine unless the config
+names another, never on the sandbox it deploys on. A compose of a disk also composes the container it's made from,
+unless that's built. A deploy that boots a form no compose has started, a disk named by `env_params` included, composes
+it first on that same provider, saying which, reusing what's built, as every build is reused by its key; deploys of one
+env at once compose it once, the others
 waiting for it, a deploy that finds it building included. In `agent-env run` it says why and each stage as it goes, and
 a deploy waiting for another's compose says what it waits for and since when. A local run doesn't compose a registered
 env's form at deploy where a configured store isn't the per-user one, since the run's fence refuses its record: the

src/agent_env/computers/disk/linux/launcher/Dockerfile

@@ -1,9 +1,10 @@
 # The QEMU launcher: boots a Linux computer's disk form, mounted or copied at /disk/disk.img, with its settings in a
-# NoCloud seed (see launch). It boots either arch's disk, with KVM where /dev/kvm is passed in and the arch is its own.
+# NoCloud seed (see launch). It boots either arch's disk, with KVM where /dev/kvm is passed in and the arch is its own,
+# behind socat on each forwarded port, and expands a zstd-compressed disk.
 FROM public.ecr.aws/docker/library/ubuntu:24.04@sha256:24737cef17f4e22b1977df7d9a05e77c4e28b7da5b8911f8d8ccda136bf3573c
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && apt-get install -y --no-install-recommends \
-        qemu-system-x86 qemu-system-arm qemu-utils ovmf qemu-efi-aarch64 dosfstools mtools \
+        qemu-system-x86 qemu-system-arm qemu-utils ovmf qemu-efi-aarch64 dosfstools mtools socat zstd \
     && rm -rf /var/lib/apt/lists/*
 COPY launch seed /usr/local/bin/
 ENTRYPOINT ["/usr/local/bin/launch"]