diff --git a/.github/scripts/infra_build.py b/.github/scripts/infra_build.py index 15d6e69a1..79b3428a4 100644 --- a/.github/scripts/infra_build.py +++ b/.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"], } diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index d0cd024de..e20fffda8 100644 --- a/src/agent_env/bundle/materialize.py +++ b/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) diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index 1994d4601..3d5d94f25 100644 --- a/src/agent_env/bundle/run.py +++ b/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) diff --git a/src/agent_env/cli/run.py b/src/agent_env/cli/run.py index c66cd289c..85f22f82f 100644 --- a/src/agent_env/cli/run.py +++ b/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, diff --git a/src/agent_env/computers/README.md b/src/agent_env/computers/README.md index 8ebce59df..b4189e5c8 100644 --- a/src/agent_env/computers/README.md +++ b/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 diff --git a/src/agent_env/computers/disk/linux/launcher/Dockerfile b/src/agent_env/computers/disk/linux/launcher/Dockerfile index 9d4f5a1f5..8a3d108a9 100644 --- a/src/agent_env/computers/disk/linux/launcher/Dockerfile +++ b/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"] diff --git a/src/agent_env/computers/disk/linux/launcher/launch b/src/agent_env/computers/disk/linux/launcher/launch index e28505419..48e84eab1 100755 --- a/src/agent_env/computers/disk/linux/launcher/launch +++ b/src/agent_env/computers/disk/linux/launcher/launch @@ -1,46 +1,128 @@ -#!/bin/sh -# launch: boot /disk/disk.img, a Linux computer's disk form, under QEMU until it powers off, on a copy-on-write overlay -# so the disk stays as it is: with KVM when /dev/kvm is here and the disk is this machine's arch, else with TCG, from -# UEFI firmware. The guest's console is this container's output. +#!/bin/bash +# launch: boot a Linux computer's disk form under QEMU until it powers off, on a copy-on-write overlay so the disk stays +# as it is: with KVM when /dev/kvm is here and the disk is this machine's arch, else with TCG, from UEFI firmware. The +# guest's console is this process's output. +# VM_DISK the disk, /disk/disk.img by default; one ending .zst is expanded onto this machine's disk first # VM_ENV_ the computer's variable NAME, its value base64'd, for its NoCloud seed (CIDATA; see seed) -# VM_PORTS guest TCP ports, each forwarded from the same port here +# VM_PORTS guest TCP ports, each served on the same port here by socat, whose listen backlog takes a burst of +# connections, in front of QEMU's forward on 127.0.0.1 at the port plus 10000, whose backlog is 1; +# socat retries a connection that forward drops every half second, not on the kernel's SYN backoff # VM_ARCH the disk's arch, amd64 or arm64; this machine's by default # VM_CPUS, VM_MEMORY_MB its size: 2 and 4096 by default -set -eu +# VM_REQUIRE_KVM when set, fail rather than emulate where KVM can't run the disk +# VM_STATE where its overlay, its seed and an expanded disk go: /var/lib/agentenv-vm by default +# launch reset: boot the running computer afresh on a new overlay, its environment the VM_ENV_= lines on +# stdin, returning once its new QEMU has started; its ports stay up throughout. +set -euo pipefail +state=${VM_STATE:-/var/lib/agentenv-vm} +# A reset and the boot loop meet here: the reset writes next.env, creates reset and stops the QEMU qemu.pid names; the +# loop, finding reset, takes next.env's environment, removes reset and boots afresh; the reset returns once qemu.pid +# names the new QEMU. +control=$state/control + +if [ "${1:-}" = reset ]; then + [ -s "$control/qemu.pid" ] || { echo "launch reset: no computer is running here" >&2; exit 1; } + old=$(cat "$control/qemu.pid") + grep -E '^VM_ENV_[A-Za-z_][A-Za-z0-9_]*=[A-Za-z0-9+/=]*$' > "$control/next.env" || true + touch "$control/reset" + kill "$old" + for _ in $(seq 600); do + if [ ! -e "$control/reset" ] && [ "$(cat "$control/qemu.pid")" != "$old" ]; then + exit 0 + fi + sleep 0.1 + done + echo "launch reset: the computer didn't start again within 60 s" >&2 + exit 1 +fi + here=$(dpkg --print-architecture) arch=${VM_ARCH:-$here} -work=/run/vm -rm -rf "$work" -mkdir -p "$work" -qemu-img create -q -f qcow2 -b /disk/disk.img -F qcow2 "$work/disk.qcow2" -seed "$work" -mkfs.vfat -C -n CIDATA "$work/seed.img" 2048 >/dev/null -mcopy -i "$work/seed.img" "$work/user-data" "$work/meta-data" ::/ -forwards="" -for port in ${VM_PORTS:-}; do - forwards="$forwards,hostfwd=tcp::$port-:$port" -done +disk=${VM_DISK:-/disk/disk.img} +work=$state/boot +mkdir -p "$control" +rm -f "$control/qemu.pid" "$control/reset" "$control/next.env" accel=tcg cpu=max if [ -w /dev/kvm ] && [ "$arch" = "$here" ]; then accel=kvm cpu=host +elif [ -n "${VM_REQUIRE_KVM:-}" ]; then + echo "launch: KVM can't run this $arch disk on this $here machine, and VM_REQUIRE_KVM refuses to emulate it" \ + "($([ -e /dev/kvm ] && echo "/dev/kvm isn't writable" || echo "there's no /dev/kvm"))" >&2 + exit 3 fi case "$arch" in - amd64) - cp /usr/share/OVMF/OVMF_VARS_4M.fd "$work/vars.fd" - set -- qemu-system-x86_64 -machine "q35,accel=$accel" \ - -drive if=pflash,format=raw,readonly=on,file=/usr/share/OVMF/OVMF_CODE_4M.fd ;; - arm64) - cp /usr/share/AAVMF/AAVMF_CODE.fd "$work/code.fd" - truncate -s 64M "$work/code.fd" "$work/vars.fd" - set -- qemu-system-aarch64 -machine "virt,accel=$accel" \ - -drive if=pflash,format=raw,readonly=on,file="$work/code.fd" ;; + amd64) machine=(qemu-system-x86_64 -machine "q35,accel=$accel" + -drive if=pflash,format=raw,readonly=on,file=/usr/share/OVMF/OVMF_CODE_4M.fd) ;; + arm64) machine=(qemu-system-aarch64 -machine "virt,accel=$accel" + -drive if=pflash,format=raw,readonly=on,file="$work/code.fd") ;; *) echo "launch: no architecture $arch" >&2 exit 2 ;; esac -exec "$@" -cpu "$cpu" -smp "${VM_CPUS:-2}" -m "${VM_MEMORY_MB:-4096}" \ - -drive if=pflash,format=raw,file="$work/vars.fd" \ - -drive file="$work/disk.qcow2",format=qcow2,if=virtio \ - -drive file="$work/seed.img",format=raw,if=virtio,readonly=on \ - -netdev "user,id=net$forwards" -device virtio-net-pci,netdev=net -device virtio-rng-pci \ - -display none -serial stdio -no-reboot +forwards="" fronts=() +for port in ${VM_PORTS:-}; do + behind=$((port + 10000)) + if [ "$behind" -gt 65535 ] || [[ " $VM_PORTS " == *" $behind "* ]]; then + echo "launch: port $port can't be forwarded: QEMU's forward behind its proxy would take port $behind, which" \ + "is $([ "$behind" -gt 65535 ] && echo "above 65535" || echo "another of VM_PORTS")" >&2 + exit 2 + fi + forwards="$forwards,hostfwd=tcp:127.0.0.1:$behind-:$port" + socat "TCP-LISTEN:$port,fork,reuseaddr,backlog=1024" "TCP:127.0.0.1:$behind,connect-timeout=0.5,retry=600,interval=0.05" & + fronts+=("$!") +done +trap 'kill "${fronts[@]}" 2>/dev/null || true' EXIT +trap 'kill "$(cat "$control/qemu.pid" 2>/dev/null)" 2>/dev/null; exit 143' TERM INT +if [[ "$disk" == *.zst ]]; then + zstd -q -d -f "$disk" -o "$state/disk.img" + disk=$state/disk.img +fi + +boot() { + rm -rf "$work" + mkdir -p "$work" + qemu-img create -q -f qcow2 -b "$disk" -F qcow2 "$work/disk.qcow2" + seed "$work" + mkfs.vfat -C -n CIDATA "$work/seed.img" 2048 >/dev/null + mcopy -i "$work/seed.img" "$work/user-data" "$work/meta-data" ::/ + if [ "$arch" = amd64 ]; then + cp /usr/share/OVMF/OVMF_VARS_4M.fd "$work/vars.fd" + else + cp /usr/share/AAVMF/AAVMF_CODE.fd "$work/code.fd" + truncate -s 64M "$work/code.fd" "$work/vars.fd" + fi + "${machine[@]}" -cpu "$cpu" -smp "${VM_CPUS:-2}" -m "${VM_MEMORY_MB:-4096}" \ + -drive if=pflash,format=raw,file="$work/vars.fd" \ + -drive file="$work/disk.qcow2",format=qcow2,if=virtio \ + -drive file="$work/seed.img",format=raw,if=virtio,readonly=on \ + -netdev "user,id=net$forwards" -device virtio-net-pci,netdev=net -device virtio-rng-pci \ + -display none -serial stdio -no-reboot & + echo "$!" > "$control/qemu.pid" +} + +while true; do + boot + qemu=$(cat "$control/qemu.pid") + while kill -0 "$qemu" 2>/dev/null; do + for front in "${fronts[@]}"; do + if ! kill -0 "$front" 2>/dev/null; then + status=0 + wait "$front" || status=$? + echo "launch: the proxy in front of a forwarded port exited ($status), so the computer is stopped" >&2 + kill "$qemu" 2>/dev/null || true + exit 1 + fi + done + sleep 0.2 + done + status=0 + wait "$qemu" || status=$? + [ -e "$control/reset" ] || exit "$status" + for name in $(compgen -e); do + case "$name" in VM_ENV_*) unset "$name" ;; esac + done + while IFS= read -r line; do + export "$line" + done < "$control/next.env" + rm -f "$control/reset" +done diff --git a/src/agent_env/env/envs/computer/env.py b/src/agent_env/env/envs/computer/env.py index 5e36d9a8f..9847303c2 100644 --- a/src/agent_env/env/envs/computer/env.py +++ b/src/agent_env/env/envs/computer/env.py @@ -184,6 +184,7 @@ def computer_artifact(self) -> ComputerArtifact: async def compose_unstarted(self, sandbox_provider: SandboxProvider, form: Optional[str] = None) -> None: """Compose first what a deploy on ``sandbox_provider`` boots that no compose started, once for deploys at once, + on the env's default sandbox provider, this machine unless the config names another, never on the deploy's, unless the run's fence would refuse the record (rules in ``computers/README.md``). A cancelled deploy holds the lock until its compose's thread, which no cancel stops, ends, so the next deploy waits for what it builds.""" if get_config().local_run_refuses_write(self.stored_id): @@ -204,10 +205,11 @@ async def waiting(seconds: float) -> None: self.metadata = (await asyncio.to_thread(Env.get, self.stored_id, self.version)).metadata if pairs := composer.unstarted(self, sandbox_provider, form): them = "it" if len(pairs) == 1 else "them" - say(f"composing {_pairs(pairs)} before this deploy (no compose built {them} for " + on = Composing(on_progress=progress) + say(f"composing {_pairs(pairs)} before this deploy on " + f"{' or '.join(type(link).__name__ for link in on.provider().links)} (no compose built {them} for " f"{' or '.join(type(link).__name__ for link in sandbox_provider.links)} yet); later deploys reuse {them}") - composing = asyncio.ensure_future(asyncio.to_thread( - composer.compose, self, pairs, Composing(sandbox_provider, on_progress=progress))) + composing = asyncio.ensure_future(asyncio.to_thread(composer.compose, self, pairs, on)) try: await asyncio.shield(composing) except asyncio.CancelledError: @@ -262,9 +264,13 @@ def deploy_refusal(self, *, sandbox_type: str | None = None, env_params: Optiona sandbox_provider = build_sandbox_provider(sandbox_type) if sandbox_type else get_env_sandbox_provider() links, arches = sandbox_provider.links, self.composed_arches() forms = [params.form] if params.form else gate.forms(self.computer) - if not any(boot.native_links(sandbox_provider, arch, form) for form in forms for arch in arches): + if not any(boot.native_links(sandbox_provider, arch, form, self.stored_id) for form in forms + for arch in arches): + unpulled = [type(link).__name__ for link in links if not boot.pulls_containers(link, self.stored_id)] problems.append(f"{', '.join(type(link).__name__ for link in links)} runs none of its arches " - f"({', '.join(arches)}) natively as {' or '.join(forms)}") + f"({', '.join(arches)}) natively as {' or '.join(forms)}" + ( + f", and {', '.join(unpulled)} can't pull its containers from the image store's registry " + "on this machine" if "container" in forms and unpulled else "")) policy = NetworkPolicy.from_dict(self.network) if self.network else None if policy is not None and not any(link.supports_network_policy(policy) for link in links): problems.append(f"{', '.join(type(link).__name__ for link in links)} can't enforce its network") diff --git a/src/agent_env/env/envs/computer/utils/boot.py b/src/agent_env/env/envs/computer/utils/boot.py index 634638db6..bc6a48b2f 100644 --- a/src/agent_env/env/envs/computer/utils/boot.py +++ b/src/agent_env/env/envs/computer/utils/boot.py @@ -18,8 +18,10 @@ from agentenv_protocol.computer.apps import ANSWER_TIMEOUT_S from agent_env.config import get_config +from agent_env.env.envs.computer.utils.composed import REPOSITORY from agent_env.providers.env_providers.env_provider import _size_kwargs from agent_env.providers.sandbox_providers.local_docker.provider import remove_local_work_dir +from agent_env.store.image_store.oci_registry_credentials import is_loopback_host, registry_host_from_ref if TYPE_CHECKING: from agent_env.artifact.artifacts.computer import ComputerArtifact @@ -50,17 +52,27 @@ class Booted: card: dict -def native_links(sandbox_provider: SandboxProvider, arch: str, form: str = "container") -> list[SandboxProvider]: - """The providers a boot of ``form`` for ``arch`` may run on, in the order a chain tries them: those that run its - container, or boot its disk, natively.""" - return [link for link in sandbox_provider.links if f"linux/{arch}" in ( - link.disk_image_platforms() if form == "disk" else link.native_platforms())] +def native_links(sandbox_provider: SandboxProvider, arch: str, form: str, store_id: str) -> list[SandboxProvider]: + """The providers a boot of ``form`` for ``arch`` of the env stored as ``store_id`` may run on, in the order a chain + tries them: those that boot its disk natively, which they carry themselves, or run its container natively and pull + it from the env's image store.""" + if form == "disk": + return [link for link in sandbox_provider.links if f"linux/{arch}" in link.disk_image_platforms()] + return [link for link in sandbox_provider.links if f"linux/{arch}" in link.native_platforms() + and pulls_containers(link, store_id)] -def first_native(link: SandboxProvider, pairs: Iterable[tuple[str, str]]) -> tuple[str, str] | None: +def first_native(link: SandboxProvider, pairs: Iterable[tuple[str, str]], store_id: str) -> tuple[str, str] | None: """The first of ``pairs``, forms and arches in the base's order, that ``link`` boots natively: the one a deploy on it boots, and so the one a run composes for it.""" - return next(((form, arch) for form, arch in pairs if native_links(link, arch, form)), None) + return next(((form, arch) for form, arch in pairs if native_links(link, arch, form, store_id)), None) + + +def pulls_containers(link: SandboxProvider, store_id: str) -> bool: + """Whether ``link`` pulls the containers the env stored as ``store_id`` composes: from a registry on this machine's + loopback, only a link on this machine does.""" + ref = get_config().get_image_store_for(store_id).image_ref(REPOSITORY, "key") + return link.ON_THIS_MACHINE or not is_loopback_host(registry_host_from_ref(ref)) def settings_environment(computer: ComputerArtifact, values: dict[str, Any], *, environment_name: str, diff --git a/src/agent_env/env/envs/computer/utils/composed.py b/src/agent_env/env/envs/computer/utils/composed.py index 4d4d924e7..2aa42150e 100644 --- a/src/agent_env/env/envs/computer/utils/composed.py +++ b/src/agent_env/env/envs/computer/utils/composed.py @@ -17,6 +17,8 @@ # A pair's statuses once its image is in the store. BUILT = ("unverified", "ready") +#: The image store repository composed computers are pushed to, tagged by their keys. +REPOSITORY = "agentenv-computers" _WRITE_TRIES = 4 diff --git a/src/agent_env/env/envs/computer/utils/composer.py b/src/agent_env/env/envs/computer/utils/composer.py index d0c7a1dfb..049922d2c 100644 --- a/src/agent_env/env/envs/computer/utils/composer.py +++ b/src/agent_env/env/envs/computer/utils/composer.py @@ -37,7 +37,7 @@ from agent_env.config import get_config from agent_env.entity.export import Unwritable, checked_path from agent_env.env.envs.computer.utils import boot, builder, clashes, gate, packages, probing -from agent_env.env.envs.computer.utils.composed import BUILT, Record, failed, now +from agent_env.env.envs.computer.utils.composed import BUILT, REPOSITORY, Record, failed, now from agent_env.env.envs.computer.utils.install import Install from agent_env.env.published_images import DISK, DISK_TOOLS_IMAGE, published_disk, published_pinned from agent_env.providers.sandbox_providers.host_images import platform_tag @@ -56,8 +56,6 @@ logger = logging.getLogger(__name__) -#: The image store repository composed computers are pushed to, tagged by their keys. -REPOSITORY = "agentenv-computers" #: What the composer and its builders make of a build's inputs beyond what its key shows of them (the base, the #: Dockerfile, the build files with their executable bits, the lock): raise it with a change that alters what an image #: holds some other way, its labels or how it's built, and leave it for any other, so images aren't built again for @@ -82,17 +80,19 @@ class Composing: """How a put composes a computer env: on ``sandbox_provider``, by default the env's; emulating an arch no link runs natively unless ``emulate`` is False; each pair's stages told to ``on_progress``; the forms of the base, or for a - run (``for_run``) those its deploys on the provider boot (``forms_of``).""" + run (``for_run``) those its deploys boot (``forms_of``) on ``deploys_on``, by default the provider it composes on.""" sandbox_provider: SandboxProvider | None = None emulate: bool = True on_progress: Callable[[str], None] | None = None forms_of: Literal["base", "deploys"] = "base" + deploys_on: SandboxProvider | None = None @classmethod def for_run(cls, sandbox_provider: SandboxProvider | None, emulate: bool = True) -> Composing: - """How a run composes: on ``sandbox_provider``, by default the env's, the forms its deploys there boot.""" - return cls(sandbox_provider, emulate, forms_of="deploys") + """How a run whose deploys are on ``sandbox_provider``, by default the env's, composes: on the env's, this + machine unless the config names another, the forms its deploys boot.""" + return cls(None, emulate, forms_of="deploys", deploys_on=sandbox_provider) def provider(self) -> SandboxProvider: """``sandbox_provider``, else the env's default, looked up by each compose that needs it, once: a run with no @@ -103,7 +103,7 @@ def provider(self) -> SandboxProvider: def forms(self, env: ComputerEnv, provider: SandboxProvider) -> list[str]: """The forms of ``env``'s base this compose on ``provider``, its ``provider()``, builds.""" - return gate.forms(env.computer) if self.forms_of == "base" else booted_forms(env, provider) + return gate.forms(env.computer) if self.forms_of == "base" else booted_forms(env, self.deploys_on or provider) def compose(env: ComputerEnv, pairs: list[tuple[str, str]] | None = None, composing: Composing = Composing()) -> None: @@ -152,7 +152,8 @@ def pending(env: ComputerEnv, sandbox_provider: SandboxProvider | None = None, def left(form: str, arch: str, record: dict | None) -> bool: status = (record or {}).get("status") - return status not in BUILT or (status == "unverified" and bool(boot.native_links(provider, arch, form))) + return status not in BUILT or (status == "unverified" and bool(boot.native_links(provider, arch, form, + env.stored_id))) return _pairs(env, gate.forms(env.computer) if forms is None else forms, left) @@ -175,7 +176,7 @@ def booted_forms(env: ComputerEnv, sandbox_provider: SandboxProvider, form: str link, the one ``boot.first_native`` picks, as the deploy does.""" forms = gate.forms(env.computer) booted = {picked[0] for link in sandbox_provider.links if (picked := boot.first_native( - link, [(name, arch) for name in ([form] if form else forms) for arch in env.composed_arches()]))} + link, [(name, arch) for name in ([form] if form else forms) for arch in env.composed_arches()], env.stored_id))} return [name for name in forms if name in booted] @@ -222,7 +223,7 @@ def _probed(env: ComputerEnv, driver: ComputerDriverEnv, form: str, arch: str, r sandbox_provider: SandboxProvider, progress: _Progress) -> dict: """Gate 3: the build of ``result`` booted on the first link of ``sandbox_provider`` that runs ``arch``'s ``form`` natively, and probed at ``full`` (``probing.checked``). Still unverified, saying why, when no link does.""" - if not (links := boot.native_links(sandbox_provider, arch, form)): + if not (links := boot.native_links(sandbox_provider, arch, form, env.stored_id)): how = f"boots linux/{arch} disks" if form == "disk" else f"runs linux/{arch}" return {**result, "reason": f"no link of the sandbox provider {how} natively to probe it"} progress.at(f"{form}/{arch}", "probing", f"probing on {type(links[0]).__name__}") diff --git a/src/agent_env/env/envs/computer/utils/probing.py b/src/agent_env/env/envs/computer/utils/probing.py index c5df75333..4eeea30fc 100644 --- a/src/agent_env/env/envs/computer/utils/probing.py +++ b/src/agent_env/env/envs/computer/utils/probing.py @@ -6,6 +6,7 @@ import asyncio from typing import TYPE_CHECKING +import httpx from agentenv_protocol.computer.conformance import run_probe from agent_env.artifact.artifacts.file import FileArtifact @@ -26,6 +27,9 @@ _PROBE_TTL_S = 1800 #: How long the probe at full may take on a computer that booted: about ten seconds on a healthy one. _PROBE_TIMEOUT_S = 120 +#: How many probes, each on a fresh boot, a build gets when each ends with the connection to its driver dropped, refused +#: or timed out, which says nothing of the build, before the last one's result is recorded. +_PROBE_ATTEMPTS = 3 async def probe_first_boot(env: ComputerEnv, form: str, arch: str, sandbox_provider: SandboxProvider) -> dict: @@ -45,13 +49,36 @@ async def probe_first_boot(env: ComputerEnv, form: str, arch: str, sandbox_provi async def checked(links: list[SandboxProvider], env: ComputerEnv, driver: ComputerDriverEnv, form: str, arch: str, result: dict) -> dict: """The build of ``result`` booted on the first of ``links`` that boots it and probed at ``full``: ready, with the - driver's card, or failed at the probe; each with a summary of the probe and its report stored.""" - try: - report, card, provider = await _probe(links, env, driver, result["image"], form) - except TimeoutError: - return failed(env, form, arch, "probe", result["key"], f"the probe didn't finish within {_PROBE_TIMEOUT_S} s") - except Exception as e: - return failed(env, form, arch, "probe", result["key"], f"it didn't boot: {e}") + driver's card, or failed at the probe; each with a summary of the probe and its report stored. A probe that ended + on a transport error talking to the driver runs again on a fresh boot first, up to ``_PROBE_ATTEMPTS`` in all, the + record then keeping each attempt; one with a failed check is recorded at once.""" + attempts: list[dict] = [] + while True: + try: + report, card, provider = await _probe(links, env, driver, result["image"], form) + except TimeoutError: + outcome = failed(env, form, arch, "probe", result["key"], + f"the probe didn't finish within {_PROBE_TIMEOUT_S} s") + except Exception as e: + if not _dropped(e): + outcome = failed(env, form, arch, "probe", result["key"], f"it didn't boot: {e}") + else: + attempts.append({"status": "dropped", "reason": _transport_errors(e), "dropped_at": now()}) + if len(attempts) < _PROBE_ATTEMPTS: + continue + return {**failed(env, form, arch, "probe", result["key"], + f"each of its {_PROBE_ATTEMPTS} probes, each on a fresh boot, ended on a transport " + f"error talking to its driver, the last {attempts[-1]['reason']}"), + "attempts": attempts} + else: + outcome = _reported(env, form, arch, result, report, card, provider) + return _attempted(attempts, outcome) + + +def _reported(env: ComputerEnv, form: str, arch: str, result: dict, report: ProbeReport, card: dict, + provider: str) -> dict: + """``result`` probed: ready with the driver's card, or failed at the probe, with a summary of the probe and its + report stored.""" stored = FileArtifact.put_bytes(derive_id(env.stored_id, f"probe_report_{form}_{arch}"), description=( f"the {form}/{arch} probe of {env.stored_id} v{env.version}"), filename="probe.json", content=report.model_dump_json(indent=1).encode(), content_type="application/json") @@ -65,6 +92,28 @@ async def checked(links: list[SandboxProvider], env: ComputerEnv, driver: Comput return {**result, "status": "ready", "card": card, "probe": probe} +def _dropped(error: BaseException) -> bool: + """Whether ``error`` is only the connection to the driver dropping, refused or timing out, as httpx raises it, alone + or from a task group.""" + if isinstance(error, BaseExceptionGroup): + return all(_dropped(e) for e in error.exceptions) + return isinstance(error, httpx.TransportError) + + +def _transport_errors(error: BaseException) -> str: + if isinstance(error, BaseExceptionGroup): + return "; ".join(dict.fromkeys(_transport_errors(e) for e in error.exceptions)) + return f"{type(error).__name__}: {error}" + + +def _attempted(attempts: list[dict], result: dict) -> dict: + """``result``, with the attempts before it and itself when it was rerun.""" + if not attempts: + return result + last = {key: result[key] for key in ("status", "reason", "failed_at", "probe") if key in result} + return {**result, "attempts": [*attempts, last]} + + async def _probe(links: list[SandboxProvider], env: ComputerEnv, driver: ComputerDriverEnv, image: str, form: str = "container") -> tuple[ProbeReport, dict, str]: """``image``, of ``form``, booted at the env's settings and its apps' defaults with a control token of its own, as diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index 1d0b996c5..094b163e8 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -92,7 +92,8 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par """One attempt on one sandbox provider: the first build it runs natively, booted and readied; its record.""" from agent_env.env.envs.computer.env import ComputerEnv - if (picked := first_native(sandbox_provider, [(form, arch) for form, arch, _ in builds])) is None: + pairs = [(form, arch) for form, arch, _ in builds] + if (picked := first_native(sandbox_provider, pairs, env.stored_id)) is None: raise ValueError(f"{type(sandbox_provider).__name__} runs none of {env.id}'s ready builds natively " f"({', '.join(f'{form}/{arch}' for form, arch, _ in builds)})") form, arch = picked diff --git a/src/agent_env/providers/sandbox_providers/AGENTS.md b/src/agent_env/providers/sandbox_providers/AGENTS.md index 548bd5c56..d8d333f9d 100644 --- a/src/agent_env/providers/sandbox_providers/AGENTS.md +++ b/src/agent_env/providers/sandbox_providers/AGENTS.md @@ -57,8 +57,8 @@ with `exposed_ports=[port]` and returns the VM, and the deploy then loads and st | `create_container(...)` | inherited | containers that run an env or agent image. Override only when the platform boots images natively | | `reset_container(sandbox, *, image_name, port, env)` | inherited | a computer env's reset, which replaces the container `create_container` ran with a fresh one on the same sandbox and port. The default works on a VM backend; one whose sandbox is the container itself raises `NotImplementedError` | | `native_platforms()` | no | a computer env's deploy and gate 3, which boot a build only on a provider that runs its arch without emulation. The default is `linux/amd64`; local Docker returns this machine's | -| `disk_image_platforms()` | no | the same for a computer's disk form: the platforms whose disks `create_vm` boots without emulation. The default is none; local Docker returns this machine's where it can pass `/dev/kvm` to a container | -| `reset_vm(sandbox, *, image, env)` | no | a computer env's reset of a disk `create_vm` booted: boot it afresh on the same sandbox and ports. The default raises `NotImplementedError` | +| `disk_image_platforms()` | no | the same for a computer's disk form: the platforms whose disks `create_vm` boots without emulation. The default is none; local Docker returns this machine's where it can pass `/dev/kvm` to a container, and `modal_vm` `linux/amd64`, booting the launcher as its sandbox's image under nested virtualization | +| `reset_vm(sandbox, *, image, env)` | no | a computer env's reset of a disk `create_vm` booted: boot it afresh on the same sandbox and ports. The default raises `NotImplementedError`; local Docker and `modal_vm` run the launcher's `launch reset` in the sandbox hosting it, with `env` on its stdin | | `prepare_image(image, *, attribution)` | no | before a sandbox is asked for an image, outside any creation deadline; build a context-only image here, and pull one with an `image_digest` by it, building it from its context when it can't be pulled | | `from_config(**config)` | no | receives `[sandbox.providers..config]` with `secret:` references resolved. Override to fail early with a `ConfigError` that names the missing key (`boxd` also rejects unknown keys) | | `close()` | no | release SDK clients | diff --git a/src/agent_env/providers/sandbox_providers/disk_launcher.py b/src/agent_env/providers/sandbox_providers/disk_launcher.py new file mode 100644 index 000000000..248e0f69c --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/disk_launcher.py @@ -0,0 +1,29 @@ +"""The QEMU launcher's interface, as its launch script reads it, for any provider that boots a computer's disk form +with it: the environment that sizes its VM, forwards its ports and hands the computer its variables, and its reset.""" + +import base64 +import math + +#: How the launcher takes the computer's environment, each variable under this prefix, base64'd (see its launch script). +ENV_PREFIX = "VM_ENV_" +#: The launch script's path in the launcher image, its ENTRYPOINT, which boots the computer; a provider that clears the +#: ENTRYPOINT runs it as the command. +LAUNCH = "/usr/local/bin/launch" +#: Its command that boots the computer afresh, with the environment ``reset_input`` gives on its stdin. +RESET = (LAUNCH, "reset") + + +def launcher_variables(ports: list[int], env: dict[str, str], cpu: float, memory: int) -> dict[str, str]: + """The environment that has the launcher size its VM, forward ``ports``, refuse to boot without KVM and hand the + computer ``env``, each value base64'd, so one of several lines reaches the computer whole.""" + return {"VM_PORTS": " ".join(map(str, ports)), "VM_CPUS": str(max(1, math.ceil(cpu))), "VM_MEMORY_MB": str(memory), + "VM_REQUIRE_KVM": "1", **_computer_variables(env)} + + +def reset_input(env: dict[str, str]) -> str: + """What ``RESET`` reads on its stdin: the computer's new environment ``env``, a variable a line.""" + return "".join(f"{name}={value}\n" for name, value in _computer_variables(env).items()) + + +def _computer_variables(env: dict[str, str]) -> dict[str, str]: + return {f"{ENV_PREFIX}{name}": base64.b64encode(value.encode()).decode() for name, value in env.items()} diff --git a/src/agent_env/providers/sandbox_providers/local_docker/cli.py b/src/agent_env/providers/sandbox_providers/local_docker/cli.py new file mode 100644 index 000000000..6117fc359 --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/local_docker/cli.py @@ -0,0 +1,11 @@ +"""The docker CLI on this machine, as the local provider and its disks run it.""" + +import subprocess + + +def docker(*args: str, input: str | None = None, timeout: float | None = None) -> str: + """``docker`` with ``args``, ``input`` on its stdin: its output, stripped. A failure raises.""" + done = subprocess.run(["docker", *args], capture_output=True, text=True, timeout=timeout, input=input) + if done.returncode: + raise RuntimeError(f"docker {' '.join(args)} failed: {done.stderr.strip()}") + return done.stdout.strip() diff --git a/src/agent_env/providers/sandbox_providers/local_docker/disk.py b/src/agent_env/providers/sandbox_providers/local_docker/disk.py index a667bb367..34bc948db 100644 --- a/src/agent_env/providers/sandbox_providers/local_docker/disk.py +++ b/src/agent_env/providers/sandbox_providers/local_docker/disk.py @@ -1,24 +1,27 @@ """Booting a computer's disk form on this machine's Docker: the QEMU launcher (the published image for the sources installed, else built here from its folder) with the disk copied in, as an image this machine keeps for the next boot, -run with /dev/kvm.""" +run with /dev/kvm; and a disk exported from the image store by this machine's Docker, for a provider elsewhere to +carry.""" import asyncio -import base64 import hashlib -import math import os import subprocess import tempfile from pathlib import Path, PurePosixPath from agent_env.artifact.artifacts.computer import DISK_FILE +from agent_env.artifact.artifacts.docker_image import docker_login +from agent_env.config import get_config +from agent_env.config.paths import cache_root from agent_env.env.published_images import LAUNCHER_IMAGE, published_pinned +from agent_env.providers.sandbox_providers.disk_launcher import launcher_variables from agent_env.providers.sandbox_providers.host_images import host_platform, present +from agent_env.providers.sandbox_providers.local_docker.cli import docker +from agent_env.store.local_state import holding_locks from agent_env.utils.docker.build import build_image KVM = "/dev/kvm" -#: How the launcher takes the computer's environment, each variable under this prefix, base64'd (see its launch script). -ENV_PREFIX = "VM_ENV_" def kvm_platforms() -> frozenset[str]: @@ -37,12 +40,40 @@ async def machine_image(disk: str) -> str: def launcher_args(ports: list[int], env: dict[str, str], cpu: float, memory: int) -> list[str]: - """The ``docker run`` arguments that size the launcher's VM, forward ``ports`` and hand it ``env``, each value - base64'd, so one of several lines reaches the computer whole.""" - variables = {"VM_PORTS": " ".join(map(str, ports)), "VM_CPUS": str(max(1, math.ceil(cpu))), - "VM_MEMORY_MB": str(memory), - **{f"{ENV_PREFIX}{name}": base64.b64encode(value.encode()).decode() for name, value in env.items()}} - return ["--device", KVM, *(arg for name, value in variables.items() for arg in ("-e", f"{name}={value}"))] + """The ``docker run`` arguments that pass the launcher KVM and its ``launcher_variables``.""" + return ["--device", KVM, *(arg for name, value in launcher_variables(ports, env, cpu, memory).items() + for arg in ("-e", f"{name}={value}"))] + + +def exported_disk(image: str) -> Path: + """``image``'s ``DISK_FILE`` compressed with zstd in this machine's cache, exported from its image store by this + machine's Docker unless an export of it is there already: once for each composed disk, by its ref, whose tag is its + composer's key.""" + path = cache_root() / "disks" / f"{hashlib.sha256(image.encode()).hexdigest()[:32]}.img.zst" + with holding_locks([f"disk-export:{path.name}"]): + if not path.exists(): + _export(image, path) + return path + + +def _export(image: str, path: Path) -> None: + """Signed in to the image store first, as a boot's pull is, since ``docker create`` pulls ``image`` when it isn't + here.""" + path.parent.mkdir(parents=True, exist_ok=True) + docker_login(get_config().get_image_store().auth(image)) + container = docker("create", image, "disk") + try: + with tempfile.TemporaryDirectory(prefix="export-", dir=path.parent) as work: + raw, compressed = Path(work) / "disk.img", Path(work) / path.name + docker("cp", f"{container}:/{DISK_FILE}", str(raw)) + with raw.open("rb") as source, compressed.open("wb") as out: + done = subprocess.run(["docker", "run", "--rm", "-i", "--entrypoint", "zstd", _launcher(), "-q", "-3", + "-T0", "-c"], stdin=source, stdout=out, stderr=subprocess.PIPE, text=False) + if done.returncode: + raise RuntimeError(f"compressing {image}'s disk failed: {done.stderr.decode().strip()}") + os.replace(compressed, path) + finally: + subprocess.run(["docker", "rm", container], capture_output=True) def _launcher() -> str: diff --git a/src/agent_env/providers/sandbox_providers/local_docker/provider.py b/src/agent_env/providers/sandbox_providers/local_docker/provider.py index 84c954c0d..8ab8b04e9 100644 --- a/src/agent_env/providers/sandbox_providers/local_docker/provider.py +++ b/src/agent_env/providers/sandbox_providers/local_docker/provider.py @@ -39,6 +39,8 @@ 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, host_platform, present, pull_here +from agent_env.providers.sandbox_providers.disk_launcher import RESET, reset_input +from agent_env.providers.sandbox_providers.local_docker.cli import docker from agent_env.providers.sandbox_providers.local_docker.disk import kvm_platforms, launcher_args, machine_image from agent_env.providers.sandbox_providers.local_docker.relay import relay_image from agent_env.providers.sandbox_providers.sandbox import ( @@ -712,8 +714,9 @@ async def create_vm( return sandbox async def reset_vm(self, sandbox: Sandbox, *, image: str, env: dict[str, str]) -> None: - await sandbox.exec_script(f"docker rm -f {shlex.quote(sandbox.container_name)} > /dev/null") - await self._boot_disk(sandbox, image=image, env=env) + """The launcher in the sandbox's container booting its disk afresh with ``env``.""" + await asyncio.to_thread(docker, "exec", "-i", sandbox.container_name, *RESET, input=reset_input(env), + timeout=_RESET_SECONDS) async def _boot_disk(self, sandbox: LocalSandbox, *, image: str, env: dict[str, str]) -> None: """The launcher booting ``image`` as the sandbox's container, its ports published as a container's are.""" @@ -724,8 +727,8 @@ async def _boot_disk(self, sandbox: LocalSandbox, *, image: str, env: dict[str, publish = [arg for port in ports for spec in port_bindings(sandbox.host_ips, sandbox.host_port(port), port) for arg in ("-p", spec)] running = asyncio.ensure_future(asyncio.to_thread( - _docker, "run", "-d", "--name", sandbox.container_name, "--label", f"{SANDBOX_LABEL}={sandbox.sandbox_id}", - *publish, *launcher_args(ports, env, size["cpu"], size["memory"]), machine)) + docker, "run", "-d", "--name", sandbox.container_name, "--label", f"{SANDBOX_LABEL}={sandbox.sandbox_id}", + *publish, *launcher_args(ports, env, size["cpu"], size["memory"]), machine, timeout=_DOCKER_QUERY_SECONDS)) try: await asyncio.shield(running) except asyncio.CancelledError: # the container docker is still starting must exist before teardown removes it @@ -938,20 +941,16 @@ def _is_loopback(host: Optional[str]) -> bool: def _host_ips() -> tuple[str, ...]: """Loopback, plus on Linux the bridge gateway ``host-gateway`` resolves to, where containers reach the host.""" # Docker Desktop publishes through a host-side proxy, which can't bind the bridge address inside its VM. - if platform.system() != "Linux" or _docker("info", "--format", "{{.OperatingSystem}}") == "Docker Desktop": + if platform.system() != "Linux" or docker("info", "--format", "{{.OperatingSystem}}", + timeout=_DOCKER_QUERY_SECONDS) == "Docker Desktop": return ("127.0.0.1",) # Listing, unlike inspecting, answers a missing network with nothing rather than an error. - if not _docker("network", "ls", "--quiet", "--filter", "name=^bridge$"): + if not docker("network", "ls", "--quiet", "--filter", "name=^bridge$", timeout=_DOCKER_QUERY_SECONDS): return ("127.0.0.1",) - return ("127.0.0.1", _docker("network", "inspect", "bridge", "--format", "{{(index .IPAM.Config 0).Gateway}}")) - - -def _docker(*args: str) -> str: - """A docker CLI query's output. A failure raises, so the cache above never keeps it.""" - run = subprocess.run(["docker", *args], capture_output=True, text=True, timeout=_DOCKER_QUERY_SECONDS) - if run.returncode: - raise RuntimeError(f"docker {' '.join(args)} failed: {run.stderr.strip()}") - return run.stdout.strip() + return ("127.0.0.1", docker("network", "inspect", "bridge", "--format", "{{(index .IPAM.Config 0).Gateway}}", + timeout=_DOCKER_QUERY_SECONDS)) _DOCKER_QUERY_SECONDS = 30 +#: How long a reset may take: longer than the minute the launcher's reset waits for the computer to start again. +_RESET_SECONDS = 90 diff --git a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py index 0cec9499f..03e52d698 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -13,6 +13,11 @@ east-west concern and is irrelevant here (one sandbox, one docker bridge). Only the pure, stateless helpers are reused from modal_sandbox (no class coupling). + +With a computer's disk form as its image, ``create_vm`` makes a VM sandbox with nested virtualization whose own image +is the QEMU launcher, with no Docker in it, booting the disk mounted at create from its export on this machine +(``local_docker.disk.exported_disk``); ``reset_vm`` boots it afresh in the same sandbox. ``[sandbox.providers.modal_vm]`` +takes ``region``, where every sandbox it makes goes, and ``idle_timeout``, in seconds, after which an idle one ends. """ from __future__ import annotations @@ -38,7 +43,10 @@ ) from agent_env.attribution import Attribution from agent_env.config import get_config -from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, VmSandbox, push_object_over_stdin +from agent_env.env.published_images import LAUNCHER_IMAGE, published_digest +from agent_env.providers.sandbox_providers.disk_launcher import LAUNCH, RESET, launcher_variables, reset_input +from agent_env.providers.sandbox_providers.local_docker.disk import exported_disk +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, Sandbox, VmSandbox, push_object_over_stdin from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_VM, SandboxProvider if TYPE_CHECKING: @@ -105,6 +113,13 @@ while true; do sleep 3600; done ''' _KEEPALIVE_CMD = ("sh", "-c", _MAIN_SCRIPT) +#: Where a computer's disk is mounted on the launcher's image, compressed; the launcher expands it onto the sandbox's disk. +_MOUNTED_DISK = "/disk/disk.img.zst" +#: The memory a disk's sandbox has beyond its VM's: QEMU's own, and the sandbox's. +_DISK_HOST_MEMORY_MB = 2048 +#: How long a disk's sandbox stays up with no connection open and no command running, when the config sets no +#: ``idle_timeout``: a running QEMU isn't activity, so without one a computer nobody uses lasts its whole timeout. +_DISK_IDLE_TIMEOUT_S = 30 * 60 class ModalVmSandbox(VmSandbox): @@ -274,17 +289,27 @@ async def _download_ranged(self, url_factory: Callable[[int], Awaitable[str]], v class ModalVmSandboxProvider(SandboxProvider): + """Provisions a single Modal VM sandbox per env (the VM-mode Modal backend), or one booting a computer's disk.""" + EGRESS_HOSTS: ClassVar[tuple[str, ...]] = ("*.modal.host", "*.w.modal.host") @classmethod def supports_network_policy(cls, policy: NetworkPolicy) -> bool: return True - """Provisions a single Modal VM sandbox per env (the VM-mode Modal backend).""" - - def __init__(self, app_name: Optional[str] = None, environment_name: Optional[str] = None): + @classmethod + def disk_image_platforms(cls) -> frozenset[str]: + """Modal's VM sandboxes are amd64, and run QEMU with KVM under nested virtualization.""" + return frozenset({"linux/amd64"}) + + def __init__(self, app_name: Optional[str] = None, environment_name: Optional[str] = None, + region: Optional[str] = None, idle_timeout: Optional[int] = None): + """``region`` places every sandbox it creates, a computer's, an agent's and a gateway's alike, there; Modal + places them anywhere by default. ``idle_timeout`` ends each once it has been idle that many seconds.""" self._app_name = _resolve_app_base_name(app_name) self._environment_name = environment_name + self._region = region + self._idle_timeout = idle_timeout self._client: Optional[modal.Client] = None # Per-attribution-scope Modal App, looked up on demand (billing aggregates per App). # Self-contained copy of ModalSandboxProvider's client/app handling — intentionally @@ -326,10 +351,40 @@ async def create_vm( setup_for_gateway: bool = True, attribution: Optional[Attribution] = None, network_policy: Optional[NetworkPolicy] = None, + env: Optional[dict[str, str]] = None, ) -> ModalVmSandbox: - # image/boot_mode/disk_size_gb accepted for interface parity but inert here. - # disk_size_gb: Modal has no disk-size knob on Sandbox.create; a VM sandbox's rootfs is - # fixed at the 512 GiB max. https://modal.com/docs/guide/vm-sandboxes + """A VM sandbox running Docker, or with ``image``, a computer's disk form, one whose image is the QEMU launcher, + with nested virtualization, booting the disk, mounted at create from its export on this machine, at ``cpu`` + and ``memory`` with ``env`` in its seed. ``boot_mode`` and ``disk_size_gb`` have no effect: a VM sandbox's + root filesystem is Modal's fixed size.""" + ports = list(exposed_ports or []) + if image is None: + return await self._created(_KEEPALIVE_CMD, _VM_IMAGE, cpu=cpu, memory=memory, timeout=timeout, ports=ports, + attribution=attribution, network_policy=network_policy, + setup_for_gateway=setup_for_gateway, idle_timeout=self._idle_timeout, + experimental_options={"vm_runtime": True}) + disk = await asyncio.to_thread(exported_disk, image) + variables = {**launcher_variables(ports, env or {}, cpu, memory), "VM_DISK": _MOUNTED_DISK} + return await self._created( + (LAUNCH,), (await _launcher()).add_local_file(disk, _MOUNTED_DISK, copy=False), cpu=cpu, + memory=memory + _DISK_HOST_MEMORY_MB, timeout=timeout, ports=ports, attribution=attribution, + network_policy=network_policy, setup_for_gateway=False, runtime="vm", + experimental_options={"enable_nested_virt": True}, secrets=[modal.Secret.from_dict(variables)], + idle_timeout=self._idle_timeout or _DISK_IDLE_TIMEOUT_S) + + async def reset_vm(self, sandbox: Sandbox, *, image: str, env: dict[str, str]) -> None: + """The launcher in the sandbox booting its disk afresh with ``env``, at the same URLs.""" + code, _, stderr = await sandbox._exec_with_stdin(shlex.join(RESET), _chunks(reset_input(env))) + if code: + raise RuntimeError(f"the launcher in Modal sandbox {sandbox.sandbox_id} didn't boot its disk afresh " + f"(exit {code}): {stderr.strip()[-500:]}") + + async def _created(self, command: tuple[str, ...], image: modal.Image, *, cpu: float, memory: int, timeout: int, + ports: list[int], attribution: Optional[Attribution], network_policy: Optional[NetworkPolicy], + setup_for_gateway: bool, idle_timeout: Optional[int], **create) -> ModalVmSandbox: + """A sandbox of ``image`` running ``command``, in the configured region, its ports' tunnels up, ended once idle + ``idle_timeout`` seconds when given, set up for a gateway when ``setup_for_gateway``; one that fails before + it's returned is terminated.""" attribution = dict(attribution or {}) app_tags = _attribution_tags({}) sandbox_tags = _attribution_tags(attribution) @@ -337,11 +392,9 @@ async def create_vm( app = await self._get_app(app_name, app_tags) client = await self._get_client() effective = self.effective_network_policy(network_policy) - ports = list(exposed_ports or []) logger.info( - f"Creating Modal VM sandbox (vm_runtime; app={app_name}, tags={app_tags}, " - f"sandbox_tags={sandbox_tags}, " - f"ports={ports}, cpu={cpu}, memory={memory}MB, timeout={timeout}s, " + f"Creating Modal VM sandbox (app={app_name}, tags={app_tags}, sandbox_tags={sandbox_tags}, ports={ports}, " + f"cpu={cpu}, memory={memory}MB, timeout={timeout}s, idle_timeout={idle_timeout}, region={self._region}, " f"policy={effective.mode.value})" ) @@ -350,11 +403,14 @@ async def create_vm( # retries RESOURCE_EXHAUSTED / queueing internally, which is the real high-load case. try: sb = await modal.Sandbox._experimental_create.aio( - *_KEEPALIVE_CMD, - app=app, image=_VM_IMAGE, cpu=cpu, memory=memory, timeout=timeout, - encrypted_ports=ports, experimental_options={"vm_runtime": True}, client=client, + *command, + app=app, image=image, cpu=cpu, memory=memory, timeout=timeout, + encrypted_ports=ports, client=client, tags=sandbox_tags or None, + **({"region": self._region} if self._region else {}), + **({"idle_timeout": idle_timeout} if idle_timeout else {}), **_modal_network_kwargs(effective), + **create, ) except Exception as e: raise RuntimeError(f"Modal VM sandbox create failed [{type(e).__name__}: {e}]") from e @@ -417,6 +473,19 @@ async def terminate_sandbox(self, sandbox_id: str) -> None: await _terminate_by_id(await self._get_client(), sandbox_id) +async def _launcher() -> modal.Image: + """The QEMU launcher, its ENTRYPOINT cleared, as Modal runs an image's before the sandbox's command: the image + published for the sources installed, else built by Modal from its folder.""" + if (digest := await asyncio.to_thread(published_digest, LAUNCHER_IMAGE)) is not None: + return modal.Image.from_registry(f"{LAUNCHER_IMAGE.repository}@{digest}").entrypoint([]) + return modal.Image.from_dockerfile(str(LAUNCHER_IMAGE.dockerfile), + context_dir=str(LAUNCHER_IMAGE.dockerfile.parent)).entrypoint([]) + + +async def _chunks(text: str) -> AsyncIterator[bytes]: + yield text.encode() + + # ── Object downloads (internal) ──────────────────────────────────────────────────── # Object downloads run as one aria2c per object: a pool of HTTP range connections with small pieces, # so a single slow TCP flow only ever holds one piece instead of the whole object; resume via the diff --git a/tst/integration/computer/disk_boot_test.py b/tst/integration/computer/disk_boot_test.py index 05e0e5791..b3016941c 100644 --- a/tst/integration/computer/disk_boot_test.py +++ b/tst/integration/computer/disk_boot_test.py @@ -1,9 +1,11 @@ """The shipped desktop example's disk form, composed on this machine and booted under KVM in the QEMU launcher through ``agent-env run``. A run on local Docker, whose deploys boot the container, composes no disk; a later run whose deploy names the disk form composes it first, reusing the container, and gate 3 probes this machine's arch's disk at full as -it composes it; the deploy boots it with its settings and control token in its seed, it passes the probe at full, and a -reset boots it afresh, where it passes again. Needs Docker, the local registry and /dev/kvm.""" +it composes it, or reuses one of the same key an earlier compose made; the deploy boots it with its settings and control +token in its seed, it passes the probe at full, and a reset boots it afresh, where it passes again. Needs Docker, the +local registry and /dev/kvm.""" +import re from pathlib import Path import pytest @@ -33,9 +35,10 @@ def test_the_desktops_disk_is_composed_by_the_first_deploy_that_boots_it_then_pr ran = CliRunner().invoke(cli, ["run", str(BUNDLE), "--task", "probe"]) assert ran.exit_code == 0, ran.output - assert f"[tasks/probe.json] {DESKTOP}: composing disk/amd64, disk/arm64 before this deploy (no compose built them" in ( - ran.output) - assert f"{DESKTOP}: disk/{arch}: making the disk" in ran.output and f"{DESKTOP}: disk/{arch}: ready (" in ran.output + assert (f"[tasks/probe.json] {DESKTOP}: composing disk/amd64, disk/arm64 before this deploy on LocalSandboxProvider " + "(no compose built them") in ran.output + assert re.search(rf"{DESKTOP}: disk/{arch}: (making the disk|reusing the disk of the same key)", ran.output) + assert re.search(rf"{DESKTOP}: disk/{arch}: (reused, )?ready \(", ran.output) context = run_context(ran.output) verifications = context["metadata"]["verifications"] assert all(result["result"] for step in ("probe", "again") for result in verifications[step]["results"]) diff --git a/tst/integration/computer/disk_modal_vm_test.py b/tst/integration/computer/disk_modal_vm_test.py new file mode 100644 index 000000000..a7749975a --- /dev/null +++ b/tst/integration/computer/disk_modal_vm_test.py @@ -0,0 +1,42 @@ +"""The shipped desktop example's disk form on a Modal VM sandbox through ``agent-env run --sandbox modal_vm`` with the +default local stores, as M4.5's exit runs it: composed with this machine's Docker, exported once, compressed, and +mounted on a VM sandbox whose image is the QEMU launcher, booted under nested virtualization, probed at full, reset in +place and probed again; a second run composes and exports no amd64 disk. Needs Docker, the local registry and Modal +credentials, with ``[sandbox.attribution]`` tags for Modal's apps.""" + +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from agent_env.cli import cli +from tst.integration.computer.conftest import run_context +from tst.util.capabilities import skip_without_remote_sandbox + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_remote_sandbox("modal_vm")] + +BUNDLE = Path(__file__).parent / "disk" + + +def test_the_desktops_disk_is_carried_to_a_modal_vm_once_then_booted_probed_and_reset_in_place_there(state, + monkeypatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(state / "cache")) + + first = CliRunner().invoke(cli, ["run", str(BUNDLE), "--task", "probe", "--sandbox", "modal_vm"]) + assert first.exit_code == 0, first.output + (exported,) = (state / "cache" / "agentenv" / "disks").glob("*.img.zst") + made = exported.stat().st_mtime_ns + again = CliRunner().invoke(cli, ["run", str(BUNDLE), "--task", "probe", "--sandbox", "modal_vm"]) + + assert again.exit_code == 0, again.output + assert any("composing" in line and "disk/amd64" in line and line.endswith(" on LocalSandboxProvider") + for line in first.output.splitlines()) + assert not any("composing" in line and "disk/amd64" in line for line in again.output.splitlines()) + assert list(exported.parent.glob("*.img.zst")) == [exported] and exported.stat().st_mtime_ns == made + for ran in (first, again): + context = run_context(ran.output) + verifications = context["metadata"]["verifications"] + assert all(result["result"] for step in ("probe", "again") for result in verifications[step]["results"]) + (deployed,) = context["deployed_envs"] + assert (deployed["sandbox_type"], deployed["metadata"]["computer"]["form"]) == ("modal_vm", "disk") + assert deployed["environment_card_url"].startswith("https://") diff --git a/tst/integration/computer/launcher_test.py b/tst/integration/computer/launcher_test.py new file mode 100644 index 000000000..c5b7050ae --- /dev/null +++ b/tst/integration/computer/launcher_test.py @@ -0,0 +1,44 @@ +"""The QEMU launcher booting the shipped desktop's disk on this machine's Docker with KVM: its driver's port served +through socat, whose listen backlog takes a burst of new connections that QEMU's own forward, with a backlog of 1, +would drop, and a reset booting the disk afresh in the same container, ready again at the same URL. Needs Docker, the +local registry and /dev/kvm.""" + +import asyncio + +import httpx +import pytest + +from agent_env.env.envs.computer.utils import boot +from agent_env.env.gateway.constants import WELL_KNOWN_PATH +from agent_env.providers.sandbox_providers.host_images import host_platform +from agent_env.providers.sandbox_providers.local_docker.provider import LocalSandboxProvider +from tst.util.capabilities import skip_without_kvm + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_kvm()] + +#: New connections at once, 64 times the backlog of QEMU's forward: without socat's own retry of its connection to that +#: forward, one of them waits out the kernel's SYN backoff, past half a minute on a CI runner. +BURST = 64 + + +@pytest.mark.asyncio +async def test_the_launcher_takes_a_burst_of_new_connections_and_a_reset_boots_its_disk_afresh_at_its_url(desktop): + image = desktop.metadata["composed"]["disk"][host_platform().rsplit("/", 1)[1]]["image"] + provider, token = LocalSandboxProvider(), boot.control_token() + environment = boot.settings_environment(desktop.computer, desktop.computer_values(), + environment_name=desktop.environment_name) + booted = await boot.boot(provider, image=image, port=desktop.driver.port, environment=environment, token=token, + network_policy=None, timeout=1800, attribution=None, form="disk") + try: + statuses = await asyncio.gather(*(_status(f"{booted.url}{WELL_KNOWN_PATH}") for _ in range(BURST))) + await provider.reset_vm(booted.sandbox, image=image, env={**environment, **boot.control_environment(token)}) + again = await boot.ready(booted.url, boot.READY_TIMEOUT_S, token) + finally: + await boot.discard(booted.sandbox) + + assert statuses == [200] * BURST and again["name"] == desktop.environment_name + + +async def _status(url: str) -> int: + async with httpx.AsyncClient(timeout=60) as client: + return (await client.get(url)).status_code diff --git a/tst/unit/bundle/preflight_test.py b/tst/unit/bundle/preflight_test.py index d60fcda04..5896fe17e 100644 --- a/tst/unit/bundle/preflight_test.py +++ b/tst/unit/bundle/preflight_test.py @@ -659,6 +659,28 @@ def test_an_env_type_that_deploys_itself_is_left_to_its_own_deploy(bundle_dir, m assert dry_run_bundle(bundle_dir, sandbox="modal").infra == () +class _BootsOnlyOnModalVm(_DeploysItself): + """One that only a provider elsewhere deploys, as a computer whose disk this machine can't boot.""" + + type: ClassVar[str] = "boots_only_on_modal_vm_preflight_test" + + def deploy_refusal(self, *, sandbox_type=None, **options): + return None if sandbox_type == "modal_vm" else f"{sandbox_type or 'the default provider'} boots none of its disks" + + +def test_a_steps_own_refusal_is_checked_on_the_sandbox_the_run_deploys_it_on(bundle_dir, monkeypatch): + registry = Config.env_registry + monkeypatch.setattr(Config, "env_registry", + lambda self: {**registry(self), _BootsOnlyOnModalVm.type: _BootsOnlyOnModalVm}) + with namespace_routing(): + get_env_store().put_document(_BootsOnlyOnModalVm(id="disk", version=None)) + _task(bundle_dir, [{"id": "env", "type": "deploy_env", "env_id": "disk"}]) + + dry_run_bundle(bundle_dir, sandbox="modal_vm") + assert _problems(lambda: dry_run_bundle(bundle_dir)) == [ + "tasks/t.json: deploy_env 'env': the default provider boots none of its disks"] + + def test_a_multi_env_on_the_server_provider_is_left_to_deploy_envs_own_refusal(bundle_dir): crm = _env("crm") with namespace_routing(): diff --git a/tst/unit/env/envs/_computer.py b/tst/unit/env/envs/_computer.py index 0b8c34fb6..6198a3bba 100644 --- a/tst/unit/env/envs/_computer.py +++ b/tst/unit/env/envs/_computer.py @@ -11,6 +11,8 @@ import tarfile import urllib.parse +import httpx + from agentenv_protocol.computer import AppDeclaration, DriverDeclaration from agentenv_protocol.computer.conformance import CheckResult, ProbeReport @@ -56,6 +58,7 @@ class FakeDocker: def __init__(self, registry: Registry): self.registry, self.builds, self.fail, self.emulates, self.labels = registry, [], {}, {"linux/amd64"}, {} self.probed, self.probe_fails, self.resolved, self.timeouts = [], {}, [], {} + self.probe_drops: dict[str, int] = {} self.writes, self.saved, self.disk_builds, self.disk_sizes = {}, [], [], {} def resolve(self, request, context, source, arch, work): @@ -87,6 +90,10 @@ def resolving(on, request, source, arch, timeout): async def probe(self, links, env, driver, image, form="container"): arch = self._arch(self.registry.manifests[image.rpartition("@")[2]]["config"]["digest"]) self.probed.append((arch, image) if form == "container" else (form, arch, image)) + if self.probe_drops.get(arch): + self.probe_drops[arch] -= 1 + raise ExceptionGroup("unhandled errors in a TaskGroup", [ + httpx.RemoteProtocolError("Server disconnected without sending a response.")]) if arch in self.probe_fails: checks = [CheckResult(check="P4", name="clicks show their effect", status="failed", detail=self.probe_fails[arch])] else: diff --git a/tst/unit/env/envs/test_computer_builder.py b/tst/unit/env/envs/test_computer_builder.py index a71ca5fb5..2a4f577a4 100644 --- a/tst/unit/env/envs/test_computer_builder.py +++ b/tst/unit/env/envs/test_computer_builder.py @@ -300,6 +300,18 @@ def test_a_put_on_a_vm_provider_builds_each_arch_in_a_vm_of_its_own_and_its_imag arch: (True, r["key"]) for arch, r in composed.items()} +def test_a_run_on_a_vm_provider_composes_the_forms_its_deploys_boot_with_this_machines_docker(base, driver, docker, + signing, tmp_path): + """The run's sandbox picks the forms and the env's default provider, this machine's Docker, builds them.""" + vms = FakeVmProvider(docker, tmp_path / "vms") + + composed = Env.get(_put(base, driver, composing=Composing.for_run(vms)).id).metadata["composed"]["container"] + + assert {arch: r["builder"] for arch, r in composed.items()} == {"amd64": "LocalSandboxProvider", + "arm64": "LocalSandboxProvider"} + assert vms.vms == [] and isinstance(Composing.for_run(vms).provider(), LocalSandboxProvider) + + def test_a_chain_builds_each_arch_on_the_first_link_that_runs_it_natively(base, driver, docker, signing, tmp_path): vms = FakeVmProvider(docker, tmp_path / "vms") diff --git a/tst/unit/env/envs/test_computer_disk.py b/tst/unit/env/envs/test_computer_disk.py index a0d2b93e2..101e3e60f 100644 --- a/tst/unit/env/envs/test_computer_disk.py +++ b/tst/unit/env/envs/test_computer_disk.py @@ -150,23 +150,25 @@ def test_a_disk_is_probed_on_a_provider_that_boots_its_arch_natively(disk_base, assert env.computer_artifact().metadata["validated_boot"]["disk"] == {"arm64": disks["arm64"]["probe"]} -def _link(containers=(), disks=()): - """A sandbox provider that runs containers of ``containers`` and boots disks of ``disks`` natively.""" +def _link(containers=(), disks=(), here=True): + """A sandbox provider that runs containers of ``containers`` and boots disks of ``disks`` natively, on this machine + unless ``here`` is False.""" link = SimpleNamespace(native_platforms=lambda: frozenset(f"linux/{arch}" for arch in containers), - disk_image_platforms=lambda: frozenset(f"linux/{arch}" for arch in disks)) + disk_image_platforms=lambda: frozenset(f"linux/{arch}" for arch in disks), ON_THIS_MACHINE=here) link.links = (link,) return link _SANDBOXES = {"local": [_link(["arm64"])], "local-with-kvm": [_link(["arm64"], ["arm64"])], "vm": [_link(disks=["amd64"])], "chain": [_link(["arm64"]), _link(disks=["amd64"])], - "chain-of-containers": [_link(["arm64"]), _link(["amd64"])], "none-native": [_link(["riscv64"], ["riscv64"])]} + "chain-of-containers": [_link(["arm64"]), _link(["amd64"])], "none-native": [_link(["riscv64"], ["riscv64"])], + "vm-elsewhere": [_link(["amd64"], ["amd64"], here=False)]} @pytest.mark.parametrize("sandbox, form, forms", [ ("local", None, ["container"]), ("local-with-kvm", None, ["container"]), ("local-with-kvm", "disk", ["disk"]), ("vm", None, ["disk"]), ("chain", None, ["container", "disk"]), ("chain-of-containers", None, ["container"]), - ("none-native", None, []), + ("none-native", None, []), ("vm-elsewhere", None, ["disk"]), ]) def test_the_forms_composed_for_a_sandbox_are_those_its_deploys_boot(disk_base, driver, docker, sandbox, form, forms): """On each link the first form it boots natively, as a deploy picks, unless the deploy names one.""" @@ -175,7 +177,7 @@ def test_the_forms_composed_for_a_sandbox_are_those_its_deploys_boot(disk_base, assert composer.booted_forms(env, SimpleNamespace(links=tuple(_SANDBOXES[sandbox])), form) == forms -@pytest.mark.parametrize("sandbox", ["local", "vm", "chain"]) +@pytest.mark.parametrize("sandbox", ["local", "vm", "chain", "vm-elsewhere"]) def test_what_a_run_composes_is_what_each_link_of_its_sandbox_then_boots(disk_base, driver, docker, sandbox): """Whatever the deploy's link, its pick of the run's builds is of a form the run composed, and each form the run composed is one a link picks: the two share ``boot.first_native``.""" @@ -185,7 +187,7 @@ def test_what_a_run_composes_is_what_each_link_of_its_sandbox_then_boots(disk_ba composed = composer.pending(env, forms=Composing.for_run(sandbox).forms(env, sandbox)) with_containers = composer._with_containers(env, composed) - picked = {boot.first_native(link, with_containers) for link in links} + picked = {boot.first_native(link, with_containers, env.stored_id) for link in links} assert None not in picked and {form for form, _ in picked} == {form for form, _ in composed} @@ -448,13 +450,32 @@ async def test_a_deploy_that_composes_first_tells_the_run_why_and_each_stage_as_ context = await task.run(context=TaskStepContext()) assert context.deployed_envs[0].metadata["computer"]["form"] == "disk" - assert said[:2] == [f"{env.id}: composing disk/amd64, disk/arm64 before this deploy (no compose built them for " - "LocalSandboxProvider yet); later deploys reuse them", + assert said[:2] == [f"{env.id}: composing disk/amd64, disk/arm64 before this deploy on LocalSandboxProvider (no " + "compose built them for LocalSandboxProvider yet); later deploys reuse them", f"{env.id}: composing disk/amd64, disk/arm64 on LocalSandboxProvider"] assert any(line.startswith(f"{env.id}: disk/arm64: ready (") for line in said) assert any(line.startswith(f"{env.id}: disk/amd64: making the disk") for line in said) +@pytest.mark.asyncio +async def test_a_deploy_on_a_sandbox_that_boots_only_disks_composes_them_on_this_machine_and_says_so(disk_base, driver, + docker, + monkeypatch): + """Never on the deploy's own sandbox: the env's default provider, this machine's Docker, builds them.""" + env = _put(disk_base, driver, compose=False) + composed, said = [], [] + monkeypatch.setattr(composer, "compose", lambda env, pairs, composing: composed.append((pairs, composing))) + + with run_progress(said.append): + await Env.get(env.id).compose_unstarted(_link(disks=["amd64"])) + + ((pairs, composing),) = composed + assert pairs == [("disk", "amd64"), ("disk", "arm64")] + assert isinstance(composing.provider(), local_sandbox.LocalSandboxProvider) + assert said == [f"{env.id}: composing disk/amd64, disk/arm64 before this deploy on LocalSandboxProvider (no compose " + "built them for SimpleNamespace yet); later deploys reuse them"] + + @pytest.mark.asyncio async def test_a_deploy_of_a_disk_bigger_than_the_disk_is_refused_rather_than_booted_smaller(disk_base, driver, docker, vms): @@ -486,6 +507,22 @@ async def test_an_unverified_disk_is_probed_on_its_first_boot_with_kvm_before_it assert len(vms.created) == 1 +def test_a_sandbox_elsewhere_boots_the_disk_it_carries_not_a_container_in_this_machines_registry(base, disk_base, + driver, docker, + monkeypatch): + """The container form is in the image store's registry on this machine, which such a sandbox can't pull from; the + disk it carries itself, as a Modal VM does, so a deploy there passes its preflight with the disk form.""" + elsewhere = _link(["amd64"], ["amd64"], here=False) + monkeypatch.setattr(providers, "build_sandbox_provider", lambda spec: elsewhere) + with_disk = Env.get(_put(disk_base, driver, compose=False).id) + containers = Env.get(_put(base, driver, compose=False, id="@local/~/kit/containers", environment_name="c").id) + + assert composer.booted_forms(with_disk, elsewhere) == ["disk"] and with_disk.deploy_refusal(sandbox_type="x") is None + assert containers.deploy_refusal(sandbox_type="x") == ( + f"computer env {containers.id!r}: SimpleNamespace runs none of its arches (amd64, arm64) natively as container, " + "and SimpleNamespace can't pull its containers from the image store's registry on this machine") + + def test_a_disk_form_is_refused_at_save_where_no_provider_boots_its_disk_natively(disk_base, driver, docker): env = Env.get(_put(disk_base, driver).id) diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index 2cc8bba9b..92d5794c5 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -329,6 +329,32 @@ async def overran(*args): "failed", "probe", f"the probe didn't finish within {probing._PROBE_TIMEOUT_S} s") +def test_a_probe_ended_by_a_dropped_connection_runs_again_on_a_fresh_boot_and_one_with_a_failed_check_doesnt( + base, driver, docker): + """The record keeps every attempt; a build whose every probe drops its connection fails at the probe.""" + docker.probe_drops["arm64"] = 1 + retried = Env.get(_put(base, driver).id).metadata["composed"]["container"]["arm64"] + probes = len(docker.probed) + docker.probe_fails["arm64"] = "the click never landed" + strict = Env.get(_put(base, driver, id="@local/~/kit/strict", environment_name="strict").id).metadata[ + "composed"]["container"]["arm64"] + strict_probes = len(docker.probed) - probes + docker.probe_fails.clear() + docker.probe_drops["arm64"] = probing._PROBE_ATTEMPTS + dropped = Env.get(_put(base, driver, id="@local/~/kit/dropped", environment_name="dropped").id).metadata[ + "composed"]["container"]["arm64"] + + assert retried["status"] == "ready" and [a["status"] for a in retried["attempts"]] == ["dropped", "ready"] + assert retried["attempts"][0]["reason"] == "RemoteProtocolError: Server disconnected without sending a response." + assert retried["attempts"][1]["probe"] == retried["probe"] and probes == 2 + assert (strict["status"], strict_probes, "attempts" in strict) == ("failed", 1, False) + assert (dropped["status"], dropped["stage"]) == ("failed", "probe") + assert [a["status"] for a in dropped["attempts"]] == ["dropped"] * probing._PROBE_ATTEMPTS + assert dropped["reason"] == (f"each of its {probing._PROBE_ATTEMPTS} probes, each on a fresh boot, ended on a " + "transport error talking to its driver, the last RemoteProtocolError: Server " + "disconnected without sending a response.") + + @pytest.mark.parametrize("result, stored, replaces", [ ({"status": "unverified"}, {"status": "ready"}, False), ({"status": "failed", "stage": "build"}, {"status": "unverified"}, False), diff --git a/tst/unit/env/envs/test_disk_recipe.py b/tst/unit/env/envs/test_disk_recipe.py index e3c83e8c4..e2aecf63a 100644 --- a/tst/unit/env/envs/test_disk_recipe.py +++ b/tst/unit/env/envs/test_disk_recipe.py @@ -1,6 +1,8 @@ """The Linux disk form's recipe as the init reads what it writes: the image's start (image-env, run by the disk tools) and -the deploy's environment (the launcher's NoCloud seed), each a line per value base64'd so any bytes survive; and the -local provider booting a disk in the launcher, with KVM and the VM's ports, settings and size.""" +the deploy's environment (the launcher's NoCloud seed), each a line per value base64'd so any bytes survive; the +launcher, with QEMU and its tools stood in for, serving each port through socat in front of QEMU's forward and booting +its disk afresh in place on a reset; and the local provider booting a disk in the launcher, with KVM and the VM's ports, +settings and size, and exporting one, compressed, once.""" import asyncio import base64 @@ -9,11 +11,15 @@ import re import subprocess import threading +import time +from pathlib import Path +from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from agent_env.env.published_images import DISK +from agent_env.providers.sandbox_providers.disk_launcher import reset_input from agent_env.providers.sandbox_providers.local_docker import disk as local_disk from agent_env.providers.sandbox_providers.local_docker import provider as local_sandbox from agent_env.providers.sandbox_providers.local_docker.provider import LocalSandbox, LocalSandboxProvider @@ -85,7 +91,167 @@ def test_the_launcher_gets_kvm_the_vms_ports_and_size_and_its_environment_under_ args = local_disk.launcher_args([18765, 5900], {"ENVIRONMENT_NAME": "desk", "A": "b c"}, 1.5, 4096) assert args == ["--device", "/dev/kvm", "-e", "VM_PORTS=18765 5900", "-e", "VM_CPUS=2", "-e", "VM_MEMORY_MB=4096", - "-e", "VM_ENV_ENVIRONMENT_NAME=ZGVzaw==", "-e", "VM_ENV_A=YiBj"] + "-e", "VM_REQUIRE_KVM=1", "-e", "VM_ENV_ENVIRONMENT_NAME=ZGVzaw==", "-e", "VM_ENV_A=YiBj"] + assert reset_input({"ENVIRONMENT_NAME": "desk", "A": "b c"}) == ( + "VM_ENV_ENVIRONMENT_NAME=ZGVzaw==\nVM_ENV_A=YiBj\n") + + +#: Stand-ins for what the launcher runs, each a bash script logging its arguments as a line to $CALLS. QEMU runs until +#: it's stopped, or for $QEMU_RUNS_S seconds, logging the seed it boots with; socat exits 7 after $SOCAT_EXITS_S; zstd +#: takes $ZSTD_TAKES_S. +_STAND_INS = { + "dpkg": "echo amd64", + "qemu-img": 'echo "qemu-img $*" >> "$CALLS"', + "mkfs.vfat": ":", + "mcopy": ":", + "cp": 'touch "${@: -1}"', + "zstd": 'echo "zstd $*" >> "$CALLS"; sleep "${ZSTD_TAKES_S:-0}"; touch "${@: -1}"', + "socat": 'echo "socat $*" >> "$CALLS"; sleep "${SOCAT_EXITS_S:-30}"; exit 7', + **{qemu: f'echo "{qemu} $* seed=$(grep -vc "^#" "$VM_STATE/boot/user-data")' + f' $(sed -n 2,9p "$VM_STATE/boot/user-data" | tr "\\n" " ")" >> "$CALLS"\n' + 'trap \'echo "stopped $$" >> "$CALLS"; kill $!; exit 0\' TERM\n' + 'sleep "${QEMU_RUNS_S:-30}" & wait' for qemu in ("qemu-system-x86_64", "qemu-system-aarch64")}, +} + + +@pytest.fixture +def launcher(tmp_path): + """Run the launcher, or ``launch reset``, with the stand-ins first on its PATH: the process, its calls' log.""" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + for name, body in _STAND_INS.items(): + (bin_dir / name).write_text(f"#!/bin/bash\n{body}\n") + (bin_dir / name).chmod(0o755) + calls = tmp_path / "calls" + base = {"PATH": f"{bin_dir}:{DISK / 'launcher'}:{os.environ['PATH']}", "CALLS": str(calls), + "VM_STATE": str(tmp_path / "state"), "VM_ENV_ENVIRONMENT_NAME": "ZGVzaw=="} + started = [] + + def launch(*args: str, input: str | None = None, **variables: str) -> subprocess.Popen: + process = subprocess.Popen(["bash", str(DISK / "launcher" / "launch"), *args], env={**base, **variables}, + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + if input is not None: + process.stdin.write(input) + process.stdin.close() + started.append(process) + return process + + yield launch, calls + for process in started: + if process.poll() is None: + process.kill() + process.wait() + process.stdout.close() + process.stderr.close() + + +def _logged(calls: Path, what: str, count: int = 1) -> list[str]: + """The lines of ``calls`` starting ``what``, once there are ``count`` of them.""" + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + lines = [line for line in calls.read_text().splitlines() if line.startswith(what)] if calls.exists() else [] + if len(lines) >= count: + return lines + time.sleep(0.05) + raise AssertionError(f"{count} {what!r} lines never came: {calls.read_text() if calls.exists() else ''}") + + +def test_the_launcher_serves_each_port_through_socat_in_front_of_qemus_forward_on_loopback(launcher): + """socat's listen backlog takes a burst of connections, which QEMU's own forward, with a backlog of 1, drops, and + socat retries its connection to that forward rather than wait out the kernel's SYN backoff.""" + launch, calls = launcher + + running = launch(VM_PORTS="18765 6080") + + retry = "connect-timeout=0.5,retry=600,interval=0.05" + assert sorted(_logged(calls, "socat", 2)) == [ + f"socat TCP-LISTEN:18765,fork,reuseaddr,backlog=1024 TCP:127.0.0.1:28765,{retry}", + f"socat TCP-LISTEN:6080,fork,reuseaddr,backlog=1024 TCP:127.0.0.1:16080,{retry}"] + (qemu,) = _logged(calls, "qemu-system-x86_64") + assert "-netdev user,id=net,hostfwd=tcp:127.0.0.1:28765-:18765,hostfwd=tcp:127.0.0.1:16080-:6080 " in qemu + assert "-machine q35,accel=tcg" in qemu or "-machine q35,accel=kvm" in qemu + running.terminate() + assert running.wait(10) == 143 + + +@pytest.mark.parametrize("ports, why", [("55536", "above 65535"), ("6080 16080", "another of VM_PORTS")]) +def test_the_launcher_refuses_a_port_whose_forward_behind_the_proxy_has_no_port_of_its_own(launcher, ports, why): + launch, calls = launcher + + refused = launch(VM_PORTS=ports) + + assert refused.wait(10) == 2 and why in refused.stderr.read() + assert not calls.exists() or "qemu" not in calls.read_text() + + +@pytest.mark.parametrize("variables", [{"SOCAT_EXITS_S": "0.5"}, + {"SOCAT_EXITS_S": "0", "ZSTD_TAKES_S": "0.5", "VM_DISK": "/disk/disk.img.zst"}], + ids=["while-it-runs", "before-it-boots"]) +def test_the_launcher_stops_the_computer_when_a_proxy_exits(launcher, variables, tmp_path): + """Whenever the proxy exits: one that can't bind its port exits as the disk is expanded, before QEMU starts.""" + launch, _ = launcher + + stopped = launch(VM_PORTS="18765", **variables) + + assert stopped.wait(10) == 1 + assert "the proxy in front of a forwarded port exited (7), so the computer is stopped" in stopped.stderr.read() + assert _gone(int((tmp_path / "state" / "control" / "qemu.pid").read_text())) + + +def _gone(pid: int) -> bool: + """Whether process ``pid`` ends within 10 s.""" + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + try: + os.kill(pid, 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +def test_the_launcher_fails_rather_than_emulate_where_kvm_is_required_and_cant_run_the_disk(launcher): + launch, calls = launcher + + refused = launch(VM_ARCH="arm64", VM_REQUIRE_KVM="1") + emulated = launch(VM_ARCH="arm64", QEMU_RUNS_S="0.1") + + assert refused.wait(10) == 3 and "VM_REQUIRE_KVM refuses to emulate it" in refused.stderr.read() + assert emulated.wait(10) == 0 and "-machine virt,accel=tcg" in _logged(calls, "qemu-system-aarch64")[0] + + +def test_the_launcher_expands_a_compressed_disk_onto_its_state_folder_and_boots_that(launcher, tmp_path): + launch, calls = launcher + + done = launch(VM_DISK="/disk/disk.img.zst", QEMU_RUNS_S="0.1") + + assert done.wait(10) == 0 + assert _logged(calls, "zstd") == [f"zstd -q -d -f /disk/disk.img.zst -o {tmp_path}/state/disk.img"] + assert f"-b {tmp_path}/state/disk.img -F qcow2" in _logged(calls, "qemu-img")[0] + + +def test_a_reset_boots_the_computer_afresh_with_its_new_environment_while_its_ports_stay_up(launcher): + launch, calls = launcher + running = launch(VM_PORTS="18765", VM_ENV_A="YQ==") + (first,) = _logged(calls, "qemu-system-x86_64") + + reset = launch("reset", input="VM_ENV_ENVIRONMENT_NAME=YWdhaW4=\nnot a variable\nPATH=/nowhere\n") + + assert reset.wait(10) == 0, reset.stderr.read() + first, second = _logged(calls, "qemu-system-x86_64", 2) + assert "seed=2 " in first and "ENVIRONMENT_NAME=ZGVzaw==" in first and "A=YQ==" in first + assert "seed=1 " in second and "ENVIRONMENT_NAME=YWdhaW4=" in second and "A=YQ==" not in second + assert len(_logged(calls, "stopped")) == 1 and len(_logged(calls, "socat")) == 1 and running.poll() is None + running.terminate() + assert running.wait(10) == 143 and len(_logged(calls, "stopped", 2)) == 2 + + +def test_a_reset_with_no_computer_running_says_so(launcher): + launch, _ = launcher + + reset = launch("reset", input="") + + assert reset.wait(10) == 1 and "no computer is running here" in reset.stderr.read() def _seed(folder, env: dict[str, str]) -> None: @@ -105,7 +271,8 @@ def launched(monkeypatch): monkeypatch.setattr(local_sandbox, "machine_image", AsyncMock(return_value="agentenv-vm:0123")) monkeypatch.setattr(local_sandbox, "_free_host_port", lambda: 20001) monkeypatch.setattr(local_sandbox, "_host_ips", lambda: ("127.0.0.1",)) - monkeypatch.setattr(local_sandbox, "_docker", lambda *args: ran.append(args) or "") + monkeypatch.setattr(local_sandbox, "docker", lambda *args, input=None, timeout=None: ran.append( + args if input is None else (*args, input, timeout)) or "") monkeypatch.setattr(LocalSandbox, "pull_images", AsyncMock()) monkeypatch.setattr(LocalSandbox, "exec_script", AsyncMock(side_effect=lambda script: ran.append(script) or "")) return ran @@ -120,13 +287,14 @@ async def test_a_disk_boots_in_the_launcher_with_kvm_its_ports_published_and_a_r await provider.reset_vm(sandbox, image="localhost:5000/disk@sha256:" + "a" * 64, env={"ENVIRONMENT_NAME": "again"}) rebuilt = await provider.get_sandbox(sandbox.sandbox_id) - first, removed, second = launched + first, reset = launched assert first == ("run", "-d", "--name", sandbox.container_name, "--label", f"{SANDBOX_LABEL}={sandbox.sandbox_id}", "-p", "127.0.0.1:20001:18765", *local_disk.launcher_args([18765], {"ENVIRONMENT_NAME": "desk"}, 2, 4096), "agentenv-vm:0123") - assert removed == f"docker rm -f {sandbox.container_name} > /dev/null" - assert second[-2:] == ("VM_ENV_ENVIRONMENT_NAME=YWdhaW4=", "agentenv-vm:0123") and "VM_CPUS=2" in second - LocalSandbox.pull_images.assert_awaited_with(["localhost:5000/disk@sha256:" + "a" * 64]) + assert reset == ("exec", "-i", sandbox.container_name, "/usr/local/bin/launch", "reset", + "VM_ENV_ENVIRONMENT_NAME=YWdhaW4=\n", local_sandbox._RESET_SECONDS) + assert local_sandbox._RESET_SECONDS > 60 + LocalSandbox.pull_images.assert_awaited_once_with(["localhost:5000/disk@sha256:" + "a" * 64]) assert sandbox.tunnel_urls == {18765: "http://127.0.0.1:20001"} and rebuilt.owns_container @@ -134,13 +302,13 @@ async def test_a_disk_boots_in_the_launcher_with_kvm_its_ports_published_and_a_r async def test_a_boot_cancelled_while_docker_starts_the_launcher_waits_for_it_then_tears_it_down(launched, monkeypatch): started, release, order = threading.Event(), threading.Event(), [] - def docker(*args): + def docker(*args, **options): started.set() release.wait(5) order.append("run") return "" - monkeypatch.setattr(local_sandbox, "_docker", docker) + monkeypatch.setattr(local_sandbox, "docker", docker) monkeypatch.setattr(LocalSandbox, "terminate", AsyncMock(side_effect=lambda: order.append("terminate"))) booting = asyncio.ensure_future(LocalSandboxProvider().create_vm(image="disk", exposed_ports=[18765], env={})) await asyncio.to_thread(started.wait, 5) @@ -162,3 +330,40 @@ async def test_a_disk_isnt_booted_where_docker_cant_pass_kvm(launched, monkeypat await LocalSandboxProvider().create_vm(image="disk", exposed_ports=[18765], env={}) assert launched == [] and LocalSandboxProvider.disk_image_platforms() == frozenset() + + +def test_a_disk_is_exported_compressed_into_the_cache_once_for_each_ref(tmp_path, monkeypatch): + """By this machine's Docker, signed in to the image store as a boot's pull is: the disk copied out of a container of + its image, then compressed in the launcher.""" + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) + monkeypatch.setattr(local_disk, "_launcher", lambda: "agentenv-qemu-launcher:src-1") + ran, compressed = [], [] + store = SimpleNamespace(auth=lambda ref: f"credentials for {ref}") + monkeypatch.setattr(local_disk, "get_config", lambda: SimpleNamespace(get_image_store=lambda: store)) + monkeypatch.setattr(local_disk, "docker_login", lambda auth: ran.append(("login", auth))) + + def docker(*args): + ran.append(args) + if args[0] == "cp": + Path(args[2]).write_bytes(b"disk of " + args[1].encode()) + return "container1" if args[0] == "create" else "" + + def run(argv, stdin=None, stdout=None, **kwargs): + if argv[:2] == ["docker", "run"]: + compressed.append(argv) + stdout.write(b"zstd:" + stdin.read()) + return subprocess.CompletedProcess(argv, 0, b"", b"") + + monkeypatch.setattr(local_disk, "docker", docker) + monkeypatch.setattr(local_disk.subprocess, "run", run) + first = "localhost:5000/agentenv-computers:" + "1" * 64 + "@sha256:" + "a" * 64 + + exported = [local_disk.exported_disk(image) for image in (first, first, first.replace("1", "2"))] + + assert exported[0] == exported[1] != exported[2] and exported[0].parent == tmp_path / "cache" / "agentenv" / "disks" + assert exported[0].read_bytes() == b"zstd:disk of container1:/disk/disk.img" and exported[0].name.endswith(".img.zst") + assert [args[0] for args in ran] == ["login", "create", "cp", "login", "create", "cp"] + assert ran[0] == ("login", f"credentials for {first}") and ran[1][1:] == (first, "disk") + assert compressed[0] == ["docker", "run", "--rm", "-i", "--entrypoint", "zstd", "agentenv-qemu-launcher:src-1", + "-q", "-3", "-T0", "-c"] and len(compressed) == 2 + assert sorted(p.name for p in exported[0].parent.iterdir()) == sorted({exported[0].name, exported[2].name}) diff --git a/tst/unit/providers/sandbox_providers/modal_create_signature_test.py b/tst/unit/providers/sandbox_providers/modal_create_signature_test.py index cea72867d..1b97f2eca 100644 --- a/tst/unit/providers/sandbox_providers/modal_create_signature_test.py +++ b/tst/unit/providers/sandbox_providers/modal_create_signature_test.py @@ -14,7 +14,7 @@ from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.attribution import PIPELINE_STEP_KEY from agent_env.config import get_config, reset_config -from agent_env.providers.sandbox_providers import modal_image_build +from agent_env.providers.sandbox_providers import modal_image_build, modal_vm_sandbox from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandboxProvider from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy @@ -50,9 +50,22 @@ def _provider(cls, **kwargs): image_name="img:latest", port=8000, env={"A": "1"}, network_policy=_ALLOWLIST, attribution=_ATTRIBUTION)), "vm": (ModalVmSandboxProvider, {}, "_experimental_create", lambda p: p.create_vm( exposed_ports=[8000], network_policy=_ALLOWLIST, attribution=_ATTRIBUTION)), + "vm in a region": (ModalVmSandboxProvider, {"region": "us-west", "idle_timeout": 600}, "_experimental_create", + lambda p: p.create_vm(exposed_ports=[8000], attribution=_ATTRIBUTION)), + "disk": (ModalVmSandboxProvider, {"region": "us-west"}, "_experimental_create", lambda p: p.create_vm( + image="localhost:5000/disk@sha256:" + "a" * 64, env={"ENVIRONMENT_NAME": "desk"}, exposed_ports=[18765], + cpu=2, memory=4096, network_policy=_ALLOWLIST, attribution=_ATTRIBUTION)), } +@pytest.fixture(autouse=True) +def _exported_disk(tmp_path, monkeypatch): + """A disk's export on this machine, and the launcher's image, without Docker or a registry.""" + (tmp_path / "disk.img.zst").write_bytes(b"disk") + monkeypatch.setattr(modal_vm_sandbox, "exported_disk", lambda image: tmp_path / "disk.img.zst") + monkeypatch.setattr(modal_vm_sandbox, "published_digest", lambda image: "sha256:" + "b" * 64) + + @pytest.mark.asyncio @pytest.mark.parametrize("name", _CREATES) async def test_the_installed_modal_accepts_every_create_argument(name): diff --git a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py index 2a8ce5908..91a5c5577 100644 --- a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py @@ -412,3 +412,126 @@ async def test_terminate_sandbox_terminates_a_vm_by_id_and_counts_one_modal_no_l assert sb.terminate.aio.await_count == (1 if found else 0) sb.tunnels.aio.assert_not_called() + + +# ── a computer's disk form: the launcher as the sandbox's image, under nested virtualization ── + +from pathlib import Path # noqa: E402 +from types import SimpleNamespace # noqa: E402 + +from agent_env.providers.sandbox_providers.disk_launcher import launcher_variables # noqa: E402 + +DISK_IMAGE = "localhost:5000/agentenv-computers:" + "1" * 64 + "@sha256:" + "a" * 64 + + +@pytest.fixture +def modal_fake(monkeypatch): + """Modal's create, with each call's command and arguments, a sandbox's tunnel, the disk's export and the launcher's + image recorded; and each secret's variables.""" + made = SimpleNamespace(creates=[], secrets=[], exported=[], launcher=MagicMock(name="launcher")) + + async def create(*command, **kwargs): + made.creates.append((command, kwargs)) + sb = MagicMock() + sb.object_id = f"sb-{len(made.creates)}" + sb.tunnels.aio = AsyncMock(return_value={port: SimpleNamespace(host=f"{sb.object_id}-{port}.w.modal.host", + port=443) for port in kwargs["encrypted_ports"]}) + sb.exec.aio = AsyncMock(side_effect=AssertionError("a disk's sandbox runs nothing at create")) + return sb + + monkeypatch.setattr(mvs.modal.Sandbox, "_experimental_create", SimpleNamespace(aio=create)) + monkeypatch.setattr(mvs.modal.Secret, "from_dict", lambda variables: made.secrets.append(variables) or variables) + monkeypatch.setattr(mvs, "exported_disk", lambda image: made.exported.append(image) or Path("/cache/disk.img.zst")) + monkeypatch.setattr(mvs, "_launcher", AsyncMock(return_value=made.launcher)) + monkeypatch.setattr(mvs, "_log_sandbox_started", AsyncMock()) + return made + + +def _vm_provider(**config) -> ModalVmSandboxProvider: + provider = ModalVmSandboxProvider(app_name="agent-env-test", **config) + provider._get_app = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] + provider._get_client = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] + return provider + + +def test_modal_vm_boots_amd64_disks(): + assert ModalVmSandboxProvider.disk_image_platforms() == frozenset({"linux/amd64"}) + + +@pytest.mark.asyncio +async def test_a_disk_boots_in_a_vm_sandbox_whose_image_is_the_launcher_with_the_disk_mounted_at_create(modal_fake): + """No Docker in it, nested virtualization, KVM required, the seed and the driver's port handed to the launcher, and + the VM's size plus the launcher's own memory.""" + sandbox = await _vm_provider().create_vm(image=DISK_IMAGE, env={"ENVIRONMENT_NAME": "desk"}, exposed_ports=[18765], + cpu=2, memory=4096, timeout=3600) + + ((command, kwargs),) = modal_fake.creates + assert modal_fake.exported == [DISK_IMAGE] + modal_fake.launcher.add_local_file.assert_called_once_with(Path("/cache/disk.img.zst"), "/disk/disk.img.zst", + copy=False) + assert command == ("/usr/local/bin/launch",) and kwargs["image"] is modal_fake.launcher.add_local_file.return_value + assert {key: kwargs[key] for key in ("runtime", "experimental_options", "cpu", "memory", "timeout", "idle_timeout", + "encrypted_ports")} == { + "runtime": "vm", "experimental_options": {"enable_nested_virt": True}, "cpu": 2, "memory": 6144, + "timeout": 3600, "idle_timeout": 1800, "encrypted_ports": [18765]} + assert "region" not in kwargs and kwargs["secrets"] == modal_fake.secrets + assert modal_fake.secrets == [{**launcher_variables([18765], {"ENVIRONMENT_NAME": "desk"}, 2, 4096), + "VM_DISK": "/disk/disk.img.zst"}] + assert modal_fake.secrets[0]["VM_REQUIRE_KVM"] == "1" and modal_fake.secrets[0]["VM_PORTS"] == "18765" + assert sandbox.tunnel_urls == {18765: "https://sb-1-18765.w.modal.host"} and sandbox.type == "modal_vm" + + +@pytest.mark.asyncio +async def test_a_configured_region_and_idle_timeout_hold_for_every_sandbox_the_provider_makes(modal_fake): + """A computer's, an agent's and a gateway's alike, so they land together; neither is passed by default.""" + provider = _vm_provider(region="us-west", idle_timeout=600) + + with patch.object(ModalVmSandbox, "setup_vm_for_gateway", AsyncMock()): + await provider.create_vm(image=DISK_IMAGE, env={}, exposed_ports=[18765]) + await provider.create_sandbox(image_name="agent:1", port=8000, env={}) + await provider.create_vm(exposed_ports=[8080]) + await _vm_provider().create_vm(exposed_ports=[8080]) + + (_, disk), (_, agent), (_, gateway), (_, default) = modal_fake.creates + assert [(kwargs.get("region"), kwargs.get("idle_timeout")) for kwargs in (disk, agent, gateway, default)] == [ + ("us-west", 600), ("us-west", 600), ("us-west", 600), (None, None)] + assert agent["experimental_options"] == gateway["experimental_options"] == {"vm_runtime": True} + + +@pytest.mark.asyncio +async def test_a_reset_runs_the_launchers_reset_in_the_sandbox_with_the_new_environment_on_its_stdin(): + sandbox = _sandbox() + sent = [] + + async def with_stdin(script, stdin): + sent.append((script, b"".join([piece async for piece in stdin]))) + return 0, "", "" + + sandbox._exec_with_stdin = with_stdin + await ModalVmSandboxProvider().reset_vm(sandbox, image=DISK_IMAGE, env={"ENVIRONMENT_NAME": "again"}) + sandbox._exec_with_stdin = AsyncMock(return_value=(1, "", "launch reset: no computer is running here\n")) + + assert sent == [("/usr/local/bin/launch reset", b"VM_ENV_ENVIRONMENT_NAME=YWdhaW4=\n")] + with pytest.raises(RuntimeError, match="didn't boot its disk afresh .exit 1.: launch reset: no computer is running"): + await ModalVmSandboxProvider().reset_vm(sandbox, image=DISK_IMAGE, env={}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("digest", ["sha256:" + "b" * 64, None], ids=["published", "unpublished"]) +async def test_the_launcher_is_the_published_image_when_its_sources_are_published_else_built_from_its_folder( + monkeypatch, digest): + monkeypatch.setattr(mvs, "published_digest", lambda image: digest) + from_registry, from_dockerfile = MagicMock(), MagicMock() + monkeypatch.setattr(mvs.modal.Image, "from_registry", from_registry) + monkeypatch.setattr(mvs.modal.Image, "from_dockerfile", from_dockerfile) + + image = await mvs._launcher() + + if digest: + from_registry.assert_called_once_with(f"{mvs.LAUNCHER_IMAGE.repository}@{digest}") + else: + from_dockerfile.assert_called_once_with(str(mvs.LAUNCHER_IMAGE.dockerfile), + context_dir=str(mvs.LAUNCHER_IMAGE.dockerfile.parent)) + built = (from_registry if digest else from_dockerfile).return_value + built.entrypoint.assert_called_once_with([]) + assert image is built.entrypoint.return_value and (from_dockerfile if digest else from_registry).called is False diff --git a/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py b/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py index 5431ccac2..0b1d7004a 100644 --- a/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py +++ b/tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py @@ -15,6 +15,7 @@ from agent_env.providers.sandbox_providers.daytona.provider import DaytonaSandboxProvider from agent_env.providers.sandbox_providers.e2b.provider import E2BSandboxProvider from agent_env.providers.sandbox_providers.local_docker.provider import LocalSandboxProvider +from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandboxProvider from agent_env.providers.sandbox_providers.sail_vm.provider import SailVmSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.sandbox_providers.sandbox_provider import ( @@ -203,6 +204,19 @@ def test_modal_builtins_accept_an_environment_name(monkeypatch, tmp_path, name): assert build_sandbox_provider(name)._environment_name == "rl-prod" +def test_modal_vm_takes_a_region_and_an_idle_timeout_unset_by_default(monkeypatch, tmp_path): + cfg = _write_config(tmp_path, """ + [sandbox.providers.modal_vm.config] + region = "us-west" + idle_timeout = 600 + """) + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + provider = build_sandbox_provider("modal_vm") + + assert (provider._region, provider._idle_timeout) == ("us-west", 600) + assert (ModalVmSandboxProvider()._region, ModalVmSandboxProvider()._idle_timeout) == (None, None) + + def test_e2b_builtin_receives_interpolated_key_and_versioned_base(monkeypatch, tmp_path): cfg = _write_config(tmp_path, """ [sandbox.providers.e2b.config]