diff --git a/src/agent_env/computers/README.md b/src/agent_env/computers/README.md index 90e25095f..ce2575906 100644 --- a/src/agent_env/computers/README.md +++ b/src/agent_env/computers/README.md @@ -38,10 +38,12 @@ before reset in place started is refused a reset, to be redeployed. It boots on 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. +content and the launcher expands onto the sandbox's disk; a thin disk, a Windows one, carries its base too, exported +once for each base and expanded beside it there, so another computer on that base adds its layer alone. 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 @@ -170,10 +172,10 @@ pulls it. Gate 3 probes it at full. At the session user's sign-in, `agentenv-session`, run windowless by `pythonw`, reads the deploy's environment from the NoCloud seed, which the launcher attaches as a removable USB drive, since Windows mounts a FAT filesystem with no partition table on removable media alone, and starts each layer's services in that session, again whenever one exits. A Windows computer boots -with 4 vCPUs and 8 GB, on local Docker with KVM; it has computer-exec and data/add, through the protocol's Windows -defaults, so the probe runs its fixture there as on Linux, with Python one-liners for the commands it would give a -shell. A Windows computer composes its driver alone for now, no app, and boots on no Modal VM sandbox yet, whose export -carries a disk's top layer alone. +with 4 vCPUs and 8 GB, on local Docker with KVM and on `modal_vm`; it has computer-exec and data/add, through the +protocol's Windows defaults, so the probe runs its fixture there as on Linux, with Python one-liners for the commands it +would give a shell. On `modal_vm` it runs in your own Modal workspace, under the evaluation's terms or a licence of your +own; whether they allow it there is yours to check. A Windows computer composes its driver alone for now, no app. This bundle has no tasks of its own: other bundles pull its computers by alias, and `agent-env run computers` has nothing to run. diff --git a/src/agent_env/computers/disk/linux/launcher/launch b/src/agent_env/computers/disk/linux/launcher/launch index d27420a6b..d3acc6516 100755 --- a/src/agent_env/computers/disk/linux/launcher/launch +++ b/src/agent_env/computers/disk/linux/launcher/launch @@ -8,6 +8,8 @@ # and, under KVM, Hyper-V enlightenments. # VM_DISK the disk, /disk/disk.img by default, which may be backed by another beside it; one ending .zst is # expanded onto this machine's disk first +# VM_BASE the base of a thin VM_DISK ending .zst, compressed too, which a provider carries beside it: expanded +# beside the expanded disk, as the base.img it's backed by # VM_ENV_ the computer's variable NAME, its value base64'd, for its NoCloud seed (CIDATA; see seed) # 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; @@ -111,6 +113,9 @@ trap 'kill "$(cat "$control/qemu.pid" 2>/dev/null)" 2>/dev/null; exit 143' TERM if [[ "$disk" == *.zst ]]; then zstd -q -d -f "$disk" -o "$state/disk.img" disk=$state/disk.img + if [ -n "${VM_BASE:-}" ]; then + zstd -q -d -f "$VM_BASE" -o "$state/base.img" + fi fi boot() { diff --git a/src/agent_env/providers/sandbox_providers/AGENTS.md b/src/agent_env/providers/sandbox_providers/AGENTS.md index cd7d788b1..5c669d3d8 100644 --- a/src/agent_env/providers/sandbox_providers/AGENTS.md +++ b/src/agent_env/providers/sandbox_providers/AGENTS.md @@ -57,7 +57,7 @@ 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, and `modal_vm` `linux/amd64`, booting the launcher as its sandbox's image under nested virtualization | +| `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` and `windows/amd64`, booting the launcher as its sandbox's image under nested virtualization, a Windows disk's base carried beside it | | `vm_bundle_platforms()` | no | the same for a computer's vm_bundle form, a whole saved VM (`format = "lume"`): the platforms, as `/`, whose VM bundles `create_vm` boots. The default is none; `lume` returns `macos/arm64` on an Apple Silicon Mac. A provider that names one implements the VM-bundle methods below too | | `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 | | `has_vm_bundle(image)`, `delete_vm_bundle(image)`, `open_builder(image, name, build)`, `keep_builder(builder)`, `ssh_password` | when `vm_bundle_platforms()` names a platform | a macOS computer's compose, and a deploy on a Mac that doesn't hold the env's VM bundle: `open_builder` boots a clone of the base `image`, named `name`, with the folder `build` (named `agentenv-build`) shared read-write, and its sandbox's `exec_with_output` runs over ssh as the session user, whose sudo reads `ssh_password` from that folder; `keep_builder` shuts it down from inside and keeps it, returning `local:`; `has_vm_bundle` says whether this machine holds a `local:` VM bundle, and `delete_vm_bundle` deletes one, as a compose does one that failed its probe. Here `create_vm` boots `image`, `local:` or a registry reference pinned by digest, with `env` as the protocol's seed in a folder the VM shares, and `reset_vm` clones it afresh. See `lume/provider.py` | 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 746825943..146607dad 100644 --- a/src/agent_env/providers/sandbox_providers/local_docker/disk.py +++ b/src/agent_env/providers/sandbox_providers/local_docker/disk.py @@ -6,8 +6,10 @@ import asyncio import hashlib import os +import shutil import subprocess import tempfile +from dataclasses import dataclass from pathlib import Path, PurePosixPath from agent_env.artifact.artifacts.computer import DISK_BASE_FILE, DISK_FILE, DISK_GUEST_FILE @@ -17,11 +19,13 @@ 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.providers.sandbox_providers.local_docker.cli import DockerFailed, docker from agent_env.store.local_state import holding_locks from agent_env.utils.docker.build import build_image KVM = "/dev/kvm" +#: What ``docker cp`` says of a path its container lacks: Docker 23 and later, then earlier. +_NO_SUCH_PATH = ("Could not find the file", "No such container:path") def kvm_platforms() -> frozenset[str]: @@ -35,62 +39,153 @@ def kvm_platforms() -> frozenset[str]: async def machine_image(disk: str) -> str: """An image of the launcher holding ``disk``'s files for this machine's platform, built unless held: one for each - image ``disk`` names, by its id, so a disk built again under its tag gets a machine of its own. A disk of more than - one layer is a thin one on ``DISK_BASE_FILE``, which gets a step of its own, so every machine on that base shares - one copy of it.""" + image ``disk`` names, by its id, so a disk built again under its tag gets a machine of its own. A thin disk's base + gets a step of its own, so every machine on that base shares one copy of it.""" source = await asyncio.to_thread(launcher) - identity, layers = (await asyncio.to_thread(docker, "image", "inspect", "-f", "{{.Id}} {{len .RootFS.Layers}}", - disk)).split() - tag = "agentenv-vm:" + hashlib.sha256(f"{source} {disk} {identity}".encode()).hexdigest()[:16] + layout = await asyncio.to_thread(_layout, disk) + tag = "agentenv-vm:" + hashlib.sha256(f"{source} {disk} {layout.identity}".encode()).hexdigest()[:16] if not await asyncio.to_thread(present, tag): - layers = int(layers) - await asyncio.to_thread(_build, _machine_dockerfile(source, disk, layers), tag) + await asyncio.to_thread(_build, _machine_dockerfile(source, disk, thin=layout.base is not None), tag) return tag -def _machine_dockerfile(launcher: str, disk: str, layers: int) -> str: +def _machine_dockerfile(launcher: str, disk: str, thin: bool) -> str: folder = f"/{PurePosixPath(DISK_FILE).parent}/" - if layers == 1: + if not thin: return f"FROM {launcher}\nCOPY --from={disk} {folder} {folder}\n" return (f"FROM {launcher}\nCOPY --from={disk} /{DISK_BASE_FILE} /{DISK_BASE_FILE}\n" f"COPY --from={disk} /{DISK_FILE} /{DISK_GUEST_FILE} {folder}\n") +@dataclass(frozen=True) +class _Layout: + """A disk image's id and, for a thin disk, one of more than one layer on ``DISK_BASE_FILE``, its ``base``: the + digest of its first layer, the base's own, which every disk composed on that base shares; None for a whole one.""" + + identity: str + base: str | None + + +def _layout(disk: str) -> _Layout: + identity, layers, first = docker("image", "inspect", "-f", + "{{.Id}} {{len .RootFS.Layers}} {{index .RootFS.Layers 0}}", disk).split() + return _Layout(identity, first if int(layers) > 1 else None) + + def launcher_args(ports: list[int], env: dict[str, str], cpu: float, memory: int) -> list[str]: """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, platform: 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. ``platform`` is the disk's own, which this machine's Docker may not be.""" - 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, platform, path) - return path +@dataclass(frozen=True) +class Exported: + """A disk exported for a provider elsewhere to carry: its ``DISK_FILE``, compressed; for a thin disk, its base, + compressed, exported once for each base; and its guest file, when it names one.""" + + disk: Path + base: Path | None = None + guest: Path | None = None + + +def exported_disk(image: str, platform: str) -> Exported: + """``image``'s disk 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; a thin disk's base once for each base, by its layer, which every disk composed on it shares. ``platform`` is + the disk's own, which this machine's Docker may not be.""" + cached = _Cached(cache_root() / "disks" / f"{hashlib.sha256(image.encode()).hexdigest()[:32]}.img.zst") + with holding_locks([f"disk-export:{cached.disk.name}"]): + if not cached.disk.exists(): + _export(image, platform, cached) + return cached.exported() + + +@dataclass(frozen=True) +class _Cached: + """A disk's export in this machine's cache, named for its image's ref: the ``disk``, compressed, and beside it its + guest file, a file naming its base's export, which every disk on that base shares, and the folder an export of it + works in.""" + + disk: Path + @property + def guest(self) -> Path: + return self._beside("guest") -def _export(image: str, platform: str, path: Path) -> None: + @property + def base_link(self) -> Path: + return self._beside("base") + + @property + def work(self) -> Path: + return self._beside("work") + + def base_of(self, layer: str) -> Path: + """The export of the base whose first layer is ``layer``.""" + return self.disk.with_name(f"base-{hashlib.sha256(layer.encode()).hexdigest()[:32]}.img.zst") + + def exported(self) -> Exported: + base = self.disk.with_name(self.base_link.read_text()) if self.base_link.exists() else None + return Exported(self.disk, base, self.guest if self.guest.exists() else None) + + def _beside(self, what: str) -> Path: + return self.disk.with_name(f"{self.disk.name}.{what}") + + +def _export(image: str, platform: str, cached: _Cached) -> None: """Signed in to the image store first, as a boot's pull is, since ``docker create`` pulls ``image`` when it isn't - here; for ``platform``, as the image's index holds that one alone, which a pull for this machine's wouldn't find.""" - path.parent.mkdir(parents=True, exist_ok=True) + here; for ``platform``, as the image's index holds that one alone, which a pull for this machine's wouldn't find. + The disk is placed last, so one that's there was exported whole, with its base and guest file. Its work folder and + container are named for the export and cleared first, under its lock, so an export killed outright leaves no copy of + a base behind once the disk is exported again; the container by the folder's whole path, as another user's export + of the same disk through this machine's Docker holds a lock of its own.""" + container = f"agentenv-export-{hashlib.sha256(str(cached.work).encode()).hexdigest()[:24]}" + _cleared(cached.work, container) + cached.work.mkdir(parents=True) docker_login(get_config().get_image_store().auth(image)) - container = docker("create", "--platform", platform, image, "disk") + docker("create", "--name", container, "--platform", platform, 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) + if (guest := _copied(container, DISK_GUEST_FILE, cached.work / "guest")) is not None: + guest.replace(cached.guest) + if (layer := _layout(image).base) is not None: + base = cached.base_of(layer) + with holding_locks([f"disk-export:{base.name}"]): + if not base.exists(): + _compressed(container, DISK_BASE_FILE, base, cached.work) + cached.base_link.write_text(base.name) + _compressed(container, DISK_FILE, cached.disk, cached.work) finally: - subprocess.run(["docker", "rm", container], capture_output=True) + _cleared(cached.work, container) + + +def _cleared(work: Path, container: str) -> None: + subprocess.run(["docker", "rm", "-f", container], capture_output=True) + shutil.rmtree(work, ignore_errors=True) + + +def _copied(container: str, file: str, to: Path) -> Path | None: + """``file`` copied out of ``container`` to ``to``; None when the container has no such file. Any other failure is + raised, as an export cached without a guest file it has would boot it as the wrong OS.""" + try: + docker("cp", f"{container}:/{file}", str(to)) + except DockerFailed as e: + if not any(missing in str(e) for missing in _NO_SUCH_PATH): + raise + return None + return to + + +def _compressed(container: str, file: str, path: Path, work: Path) -> None: + """``file`` of ``container`` compressed with zstd to ``path``, through ``work``.""" + raw, compressed = work / PurePosixPath(file).name, work / path.name + docker("cp", f"{container}:/{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 a disk's {file} failed: {done.stderr.decode().strip()}") + raw.unlink() + os.replace(compressed, path) def launcher() -> str: 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 10057a1b0..c0b31491e 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -16,8 +16,9 @@ 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. +(``local_docker.disk.exported_disk``), with a thin disk's base and guest file beside it; ``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 @@ -41,6 +42,7 @@ _tunnel_url, _modal_network_kwargs, ) +from agent_env.artifact.artifacts.computer import DISK_BASE_FILE, DISK_FILE, DISK_GUEST_FILE from agent_env.attribution import Attribution from agent_env.config import get_config from agent_env.env.published_images import LAUNCHER_IMAGE, published_digest @@ -113,8 +115,12 @@ 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" +#: Where a computer's disk, compressed, a thin disk's base, compressed too, and its guest file are mounted on the +#: launcher's image: where a disk's own image holds them, as the launcher reads the guest file beside the disk. It +#: expands the disk and the base onto the sandbox's disk. +_MOUNTED_DISK, _MOUNTED_BASE, _MOUNTED_GUEST = f"/{DISK_FILE}.zst", f"/{DISK_BASE_FILE}.zst", f"/{DISK_GUEST_FILE}" +#: The platform of a disk's image in the image store, a Windows disk's too: its guest file names the OS it boots. +_DISK_IMAGE_PLATFORM = "linux/amd64" #: 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 @@ -300,8 +306,9 @@ def supports_network_policy(cls, policy: NetworkPolicy) -> bool: @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"}) + """Modal's VM sandboxes are amd64, and run QEMU with KVM under nested virtualization: a Linux disk, or a Windows + one, thin on its base.""" + return frozenset({"linux/amd64", "windows/amd64"}) def __init__(self, app_name: Optional[str] = None, environment_name: Optional[str] = None, region: Optional[str] = None, idle_timeout: Optional[int] = None): @@ -364,11 +371,16 @@ async def create_vm( attribution=attribution, network_policy=network_policy, setup_for_gateway=setup_for_gateway, idle_timeout=self._idle_timeout, experimental_options={"vm_runtime": True}) - (platform,) = self.disk_image_platforms() - disk = await asyncio.to_thread(exported_disk, image, platform) + exported = await asyncio.to_thread(exported_disk, image, _DISK_IMAGE_PLATFORM) variables = {**launcher_variables(ports, env or {}, cpu, memory), "VM_DISK": _MOUNTED_DISK} + launcher = (await _launcher()).add_local_file(exported.disk, _MOUNTED_DISK, copy=False) + if exported.base is not None: + launcher = launcher.add_local_file(exported.base, _MOUNTED_BASE, copy=False) + variables["VM_BASE"] = _MOUNTED_BASE + if exported.guest is not None: + launcher = launcher.add_local_file(exported.guest, _MOUNTED_GUEST, copy=False) return await self._created( - (LAUNCH,), (await _launcher()).add_local_file(disk, _MOUNTED_DISK, copy=False), cpu=cpu, + (LAUNCH,), launcher, 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)], diff --git a/tst/unit/env/envs/test_disk_recipe.py b/tst/unit/env/envs/test_disk_recipe.py index 213a255e0..f673c8252 100644 --- a/tst/unit/env/envs/test_disk_recipe.py +++ b/tst/unit/env/envs/test_disk_recipe.py @@ -6,6 +6,7 @@ import asyncio import base64 +import hashlib import json import os import re @@ -267,13 +268,16 @@ def test_the_launcher_fails_rather_than_emulate_where_kvm_is_required_and_cant_r 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): +@pytest.mark.parametrize("base", [None, "/disk/base.img.zst"], ids=["whole", "thin"]) +def test_the_launcher_expands_a_compressed_disk_onto_its_state_folder_and_boots_that(launcher, tmp_path, base): + """A thin disk's base, carried beside it, is expanded beside it, as the base.img it's backed by.""" launch, calls = launcher - done = launch(VM_DISK="/disk/disk.img.zst", QEMU_RUNS_S="0.1") + done = launch(VM_DISK="/disk/disk.img.zst", QEMU_RUNS_S="0.1", **({"VM_BASE": base} if base else {})) 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 _logged(calls, "zstd") == [f"zstd -q -d -f /disk/disk.img.zst -o {tmp_path}/state/disk.img"] + ( + [f"zstd -q -d -f {base} -o {tmp_path}/state/base.img"] if base else []) assert f"-b {tmp_path}/state/disk.img -F qcow2" in _logged(calls, "qemu-img")[0] @@ -532,7 +536,7 @@ async def test_a_machine_copies_a_thin_disks_base_in_a_step_of_its_own_so_machin built = [] monkeypatch.setattr(local_disk, "launcher", lambda: "launcher:1") monkeypatch.setattr(local_disk, "present", lambda tag: False) - monkeypatch.setattr(local_disk, "docker", lambda *args: f"sha256:1 {layers}") + monkeypatch.setattr(local_disk, "docker", lambda *args: f"sha256:1 {layers} sha256:first") monkeypatch.setattr(local_disk, "_build", lambda text, tag: built.append(text)) await local_disk.machine_image("disk:1") @@ -546,7 +550,7 @@ async def test_a_disk_built_again_under_its_tag_gets_a_machine_of_its_own(monkey identity = "sha256:1" monkeypatch.setattr(local_disk, "launcher", lambda: "launcher:1") monkeypatch.setattr(local_disk, "present", lambda tag: True) - monkeypatch.setattr(local_disk, "docker", lambda *args: f"{identity} 2") + monkeypatch.setattr(local_disk, "docker", lambda *args: f"{identity} 2 sha256:first") first = await local_disk.machine_image("disk:1") again = await local_disk.machine_image("disk:1") @@ -556,38 +560,114 @@ async def test_a_disk_built_again_under_its_tag_gets_a_machine_of_its_own(monkey assert first == again != rebuilt -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.""" +EXPORTED = "localhost:5000/agentenv-computers:" + "1" * 64 + "@sha256:" + "a" * 64 + + +@pytest.fixture +def exporting(tmp_path, monkeypatch): + """This machine's Docker as an export sees it, signed in to a store whose credentials name the ref, every call kept + in ``log``, ``ran`` (docker's own) and ``runs`` (the processes it starts): the disk a whole one unless ``layers`` + says more, its first layer ``first``, and its guest file missing unless ``guest`` holds its text, or the error + copying it raises; zstd as a ``zstd:`` prefix.""" + fake = SimpleNamespace(layers=1, first="sha256:linux", guest=None, log=[], ran=[], runs=[], + cache=tmp_path / "cache" / "agentenv" / "disks") 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))) + monkeypatch.setattr(local_disk, "docker_login", lambda auth: fake.ran.append(("login", auth))) def docker(*args): - ran.append(args) + fake.ran.append(args) + fake.log.append(args[0]) if args[0] == "cp": - Path(args[2]).write_bytes(b"disk of " + args[1].encode()) - return "container1" if args[0] == "create" else "" + file = args[1].partition(":")[2] + if file.endswith("/guest") and not isinstance(fake.guest, str): + raise fake.guest or DockerFailed(f"docker cp {args[1]} failed: Error response from daemon: Could not " + f"find the file {file} in container", 1) + Path(args[2]).write_bytes(fake.guest.encode() if file.endswith("/guest") else f"copy of {file}".encode()) + return f"sha256:image {fake.layers} {fake.first}" if args[0] == "image" else "" def run(argv, stdin=None, stdout=None, **kwargs): + fake.runs.append(argv) + fake.log.append(argv[1]) 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 + return fake - exported = [local_disk.exported_disk(image, "linux/amd64") 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:] == ("--platform", "linux/amd64", first, "disk") +def test_a_disk_is_exported_compressed_into_the_cache_once_for_each_ref(exporting): + """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.""" + exported = [local_disk.exported_disk(image, "linux/amd64") for image in (EXPORTED, EXPORTED, + EXPORTED.replace("1", "2"))] + + assert exported[0] == exported[1] != exported[2] and exported[0].disk.parent == exporting.cache + assert exported[0].disk.read_bytes() == b"zstd:copy of /disk/disk.img" + assert exported[0].disk.name.endswith(".img.zst") and exported[0].base is exported[0].guest is None + assert [args[0] for args in exporting.ran] == ["login", "create", "cp", "image", "cp"] * 2 + assert exporting.ran[0] == ("login", f"credentials for {EXPORTED}") and exporting.ran[1][1] == "--name" + assert exporting.ran[1][2].startswith("agentenv-export-") + assert exporting.ran[1][3:] == ("--platform", "linux/amd64", EXPORTED, "disk") + compressed = [argv for argv in exporting.runs if argv[:2] == ["docker", "run"]] 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}) + assert sorted(p.name for p in exporting.cache.iterdir()) == sorted({exported[0].disk.name, exported[2].disk.name}) + + +def test_an_export_killed_outright_leaves_nothing_behind_once_the_disk_is_exported_again(exporting, monkeypatch): + """Its work folder, holding as much of a base as it had copied, and its container are named for the disk, and + cleared under the disk's lock first.""" + left = exporting.cache / f"{hashlib.sha256(EXPORTED.encode()).hexdigest()[:32]}.img.zst.work" + left.mkdir(parents=True) + (left / "base.img").write_bytes(b"half a base") + docker = local_disk.docker + + def copying_with_no_half_base(*args): + assert args[0] != "cp" or not (left / "base.img").exists() + return docker(*args) + + monkeypatch.setattr(local_disk, "docker", copying_with_no_half_base) + + exported = local_disk.exported_disk(EXPORTED, "linux/amd64") + + assert exporting.log[:2] == ["rm", "create"] and exporting.log[-1] == "rm" + assert sorted(path.name for path in exporting.cache.iterdir()) == [exported.disk.name] + + +def test_two_users_exports_of_one_disk_through_one_docker_name_their_containers_apart(exporting, tmp_path, + monkeypatch): + """Each holds a lock of its own, so neither clears the other's container.""" + for user in ("one", "two"): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / user / "cache")) + local_disk.exported_disk(EXPORTED, "linux/amd64") + + assert len({args[2] for args in exporting.ran if args[0] == "create"}) == 2 + + +def test_an_export_that_cant_copy_a_disks_guest_file_fails_rather_than_cache_the_disk_without_it(exporting): + exporting.guest = DockerFailed("docker cp container1:/disk/guest failed: Error response from daemon: i/o timeout", + 1) + + with pytest.raises(DockerFailed, match="i/o timeout"): + local_disk.exported_disk(EXPORTED, "linux/amd64") + assert list(exporting.cache.iterdir()) == [] + + +def test_a_thin_disk_is_exported_with_its_guest_file_and_its_base_once_for_every_disk_on_that_base(exporting): + """A Windows disk is a layer on its base, whose first layer every disk composed on it shares.""" + exporting.layers, exporting.first, exporting.guest = 3, "sha256:base", "windows" + + one, two = (local_disk.exported_disk(image, "windows/amd64") for image in (EXPORTED, EXPORTED.replace("1", "2"))) + + copied = [args[1].partition(":")[2] for args in exporting.ran if args[0] == "cp"] + assert one.disk != two.disk and one.base == two.base and one.base.parent == one.disk.parent + assert one.base.read_bytes() == b"zstd:copy of /disk/base.img" + assert one.disk.read_bytes() == b"zstd:copy of /disk/disk.img" and one.guest.read_text() == "windows" + assert copied == ["/disk/guest", "/disk/base.img", "/disk/disk.img", "/disk/guest", "/disk/disk.img"] + assert local_disk.exported_disk(EXPORTED, "windows/amd64") == one + assert len([args for args in exporting.ran if args[0] == "cp"]) == 5 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 509857c4e..9995a4188 100644 --- a/tst/unit/providers/sandbox_providers/modal_create_signature_test.py +++ b/tst/unit/providers/sandbox_providers/modal_create_signature_test.py @@ -16,6 +16,7 @@ from agent_env.config import get_config, reset_config 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.local_docker.disk import Exported from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandboxProvider from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy from tst.unit.store.fakes import FakeImageStore @@ -62,7 +63,7 @@ def _provider(cls, **kwargs): 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, platform: tmp_path / "disk.img.zst") + monkeypatch.setattr(modal_vm_sandbox, "exported_disk", lambda image, platform: Exported(tmp_path / "disk.img.zst")) monkeypatch.setattr(modal_vm_sandbox, "published_digest", lambda image: "sha256:" + "b" * 64) 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 00abfa117..b138eba36 100644 --- a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py @@ -421,6 +421,7 @@ async def test_terminate_sandbox_terminates_a_vm_by_id_and_counts_one_modal_no_l from types import SimpleNamespace # noqa: E402 from agent_env.providers.sandbox_providers.disk_launcher import RESET, launcher_variables # noqa: E402 +from agent_env.providers.sandbox_providers.local_docker.disk import Exported # noqa: E402 DISK_IMAGE = "localhost:5000/agentenv-computers:" + "1" * 64 + "@sha256:" + "a" * 64 @@ -429,7 +430,8 @@ async def test_terminate_sandbox_terminates_a_vm_by_id_and_counts_one_modal_no_l 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")) + made = SimpleNamespace(creates=[], secrets=[], exported=[], launcher=MagicMock(name="launcher"), + export=Exported(Path("/cache/disk.img.zst"))) async def create(*command, **kwargs): made.creates.append((command, kwargs)) @@ -443,7 +445,7 @@ async def create(*command, **kwargs): 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, platform: made.exported.append((image, platform)) - or Path("/cache/disk.img.zst")) + or made.export) monkeypatch.setattr(mvs, "_launcher", AsyncMock(return_value=made.launcher)) monkeypatch.setattr(mvs, "_log_sandbox_started", AsyncMock()) return made @@ -456,8 +458,8 @@ def _vm_provider(**config) -> ModalVmSandboxProvider: return provider -def test_modal_vm_boots_amd64_disks(): - assert ModalVmSandboxProvider.disk_image_platforms() == frozenset({"linux/amd64"}) +def test_modal_vm_boots_amd64_disks_linux_and_windows(): + assert ModalVmSandboxProvider.disk_image_platforms() == frozenset({"linux/amd64", "windows/amd64"}) @pytest.mark.asyncio @@ -483,6 +485,26 @@ async def test_a_disk_boots_in_a_vm_sandbox_whose_image_is_the_launcher_with_the assert sandbox.tunnel_urls == {18765: "https://sb-1-18765.w.modal.host"} and sandbox.type == "modal_vm" +@pytest.mark.asyncio +async def test_a_thin_disk_boots_with_its_base_and_guest_file_mounted_beside_it_and_its_base_expanded_there(modal_fake): + modal_fake.export = Exported(Path("/cache/disk.img.zst"), Path("/cache/base-1.img.zst"), + Path("/cache/disk.img.zst.guest")) + + await _vm_provider().create_vm(image=DISK_IMAGE, env={}, exposed_ports=[18765], cpu=2, memory=4096) + + ((_, kwargs),) = modal_fake.creates + disk = modal_fake.launcher.add_local_file + base = disk.return_value.add_local_file + guest = base.return_value.add_local_file + assert [mount.call_args for mount in (disk, base, guest)] == [ + ((Path("/cache/disk.img.zst"), "/disk/disk.img.zst"), {"copy": False}), + ((Path("/cache/base-1.img.zst"), "/disk/base.img.zst"), {"copy": False}), + ((Path("/cache/disk.img.zst.guest"), "/disk/guest"), {"copy": False})] + assert kwargs["image"] is guest.return_value + assert modal_fake.secrets == [{**launcher_variables([18765], {}, 2, 4096), "VM_DISK": "/disk/disk.img.zst", + "VM_BASE": "/disk/base.img.zst"}] + + @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."""