Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 12 additions & 10 deletions src/agent_env/computers/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -34,16 +34,18 @@ partition kept, so nothing is emulated or privileged. A disk boots where `/dev/k
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; a container an agent-env from
before reset in place started is refused a reset, to be redeployed. 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; 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.
before reset in place started is refused a reset, to be redeployed. It boots on local Docker where Docker can pass it
`/dev/kvm`, the disk image's files mounted read-only into the launcher where Docker mounts images (28 and later, on the
containerd image store), so no copy of a disk, nor of a Windows disk's base, is made for each launcher, else copied into
an image of the launcher it keeps for the next boot; 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; 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
Expand Down
18 changes: 9 additions & 9 deletions src/agent_env/env/envs/computer/utils/windows_disk.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
from agent_env.env.envs.computer.utils.install import WindowsInstall
from agent_env.providers.sandbox_providers.host_images import present
from agent_env.providers.sandbox_providers.local_docker.cli import docker
from agent_env.providers.sandbox_providers.local_docker.disk import KVM, launcher, machine_image
from agent_env.providers.sandbox_providers.local_docker.disk import KVM, launcher, machine_for
from agent_env.store.image_store.image_contents import labels as image_labels
from agent_env.store.image_store.image_store import ImageStore
from agent_env.store.image_store.registry_api import pin_digest, pin_if_served
Expand Down Expand Up @@ -284,8 +284,8 @@ def compose(base_: Base, build: Path, key: str, labels: dict[str, str], say: Cal
out.mkdir()
name = f"agentenv-windows-compose-{_short(key)}"
say(f"booting its base in {name}")
machine = asyncio.run(machine_image(base_.image))
container = _launched(name, {"VM_KEEP": "/out/disk.img"}, ["-v", f"{out}:/out"], machine)
machine = asyncio.run(machine_for(base_.image))
container = _launched(name, {"VM_KEEP": "/out/disk.img"}, ["-v", f"{out}:/out", *machine.args], machine.image)
try:
port, deadline = _ssh_port(container), time.monotonic() + COMPOSE_TIMEOUT_S
while ssh(port, base_.keys, "echo ready", check=False, timeout=60).returncode != 0:
Expand All @@ -303,7 +303,7 @@ def compose(base_: Base, build: Path, key: str, labels: dict[str, str], say: Cal
_shut_down(container, port, base_.keys)
finally:
subprocess.run(["docker", "rm", "-f", container], capture_output=True)
_converted(out, machine, backing=f"/{DISK_BASE_FILE}")
_converted(out, machine.image, backing=f"/{DISK_BASE_FILE}", args=machine.args)
subprocess.run(["docker", "run", "--rm", "-v", f"{out}:/out", "--entrypoint", "qemu-img", launcher(), "rebase",
"-u", "-b", PureWindowsPath(DISK_BASE_FILE).name, "-F", "qcow2", "/out/disk.img"],
check=True, capture_output=True)
Expand Down Expand Up @@ -419,12 +419,12 @@ def _compacted(port: int, keys: Path, say: Callable[[str], None]) -> None:
say(f"trimming exited {done.returncode}, so its image keeps what it held: {(done.stderr + done.stdout)[-300:]}")


def _converted(folder: Path, image: str, backing: str | None = None) -> None:
"""``folder``'s ``disk.img`` rewritten by ``qemu-img convert`` in ``image``, which holds its backing file, keeping
only the clusters it uses: a base's whole, a layer's over ``backing``, the backing file ``image`` holds."""
def _converted(folder: Path, image: str, backing: str | None = None, args: tuple[str, ...] = ()) -> None:
"""``folder``'s ``disk.img`` rewritten by ``qemu-img convert`` in ``image`` run with ``args``, which give it its
backing file, keeping only the clusters it uses: a base's whole, a layer's over ``backing``."""
over = ["-B", backing, "-F", "qcow2"] if backing else []
subprocess.run(["docker", "run", "--rm", "-v", f"{folder}:/out", "--entrypoint", "qemu-img", image, "convert", "-O",
"qcow2", *over, "/out/disk.img", "/out/compact.img"], check=True, capture_output=True)
subprocess.run(["docker", "run", "--rm", "-v", f"{folder}:/out", *args, "--entrypoint", "qemu-img", image, "convert",
"-O", "qcow2", *over, "/out/disk.img", "/out/compact.img"], check=True, capture_output=True)
(folder / "compact.img").replace(folder / "disk.img")


Expand Down
52 changes: 49 additions & 3 deletions src/agent_env/providers/sandbox_providers/local_docker/disk.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
"""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; and a disk exported from the image store by this machine's Docker, for a provider elsewhere to
carry."""
installed, else built here from its folder) run with /dev/kvm, the disk image's files mounted into it where this
machine's Docker mounts images, else copied into an image of it this machine keeps for the next boot; and a disk
exported from the image store by this machine's Docker, for a provider elsewhere to carry."""

import asyncio
import functools
import hashlib
import os
import shutil
Expand All @@ -24,6 +25,8 @@
from agent_env.utils.docker.build import build_image

KVM = "/dev/kvm"
#: How long asking this machine's Docker whether it mounts images may take, before a boot copies the disk instead.
_DOCKER_INFO_SECONDS = 30
#: 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")

Expand All @@ -37,6 +40,49 @@ def kvm_platforms() -> frozenset[str]:
return frozenset({here, "windows/amd64"} if here == "linux/amd64" else {here})


@dataclass(frozen=True)
class Machine:
"""How the launcher boots a disk on this machine: ``image`` run with ``args``. Where Docker mounts images (28 and
later, on the containerd image store), the launcher itself with the disk image's folder of files mounted read-only
where the launcher finds them, so no copy of a disk, nor of its base, is made for each launcher; else a machine
image of the launcher with the disk copied in."""

image: str
args: tuple[str, ...] = ()


async def machine_for(disk: str) -> Machine:
if await asyncio.to_thread(mounts_images):
folder = PurePosixPath(DISK_FILE).parent
return Machine(await asyncio.to_thread(launcher),
("--mount", f"type=image,source={disk},target=/{folder},image-subpath={folder}"))
return Machine(await machine_image(disk))


def mounts_images() -> bool:
"""Whether this machine's Docker mounts an image's files into a container (``--mount type=image``, with its
``image-subpath``): its client, which parses the option, and its server both 28 or later, on the containerd image
store. A Docker that fails to say, or doesn't within ``_DOCKER_INFO_SECONDS``, gets the copy this time and is asked
again next time; an answer is kept for the process."""
try:
return _docker_mounts_images()
except (DockerFailed, subprocess.TimeoutExpired):
return False


@functools.cache
def _docker_mounts_images() -> bool:
client, server, status = docker("info", "--format",
"{{.ClientInfo.Version}} {{.ServerVersion}} {{json .DriverStatus}}",
timeout=_DOCKER_INFO_SECONDS).split(" ", 2)
return _major(client) >= 28 and _major(server) >= 28 and "io.containerd.snapshotter.v1" in status


def _major(version: str) -> int:
head = version.split(".", 1)[0]
return int(head) if head.isdigit() else 0


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 thin disk's base
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
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, UNRESETTABLE, reset_input, unresettable
from agent_env.providers.sandbox_providers.local_docker.cli import DockerFailed, 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.disk import kvm_platforms, launcher_args, machine_for
from agent_env.providers.sandbox_providers.local_docker.relay import relay_image
from agent_env.providers.sandbox_providers.sandbox import (
SANDBOX_LABEL,
Expand Down Expand Up @@ -729,13 +729,14 @@ async def _boot_disk(self, sandbox: LocalSandbox, *, image: str, env: dict[str,
"""The launcher booting ``image`` as the sandbox's container, its ports published as a container's are."""
size = json.loads((sandbox.work_dir / _VM_SIZE_FILE).read_text())
await sandbox.pull_images([image])
machine = await machine_image(image)
machine = await machine_for(image)
ports = list(sandbox._port_map)
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, timeout=_DOCKER_QUERY_SECONDS))
*publish, *machine.args, *launcher_args(ports, env, size["cpu"], size["memory"]), machine.image,
timeout=_DOCKER_QUERY_SECONDS))
try:
await asyncio.shield(running)
except asyncio.CancelledError: # the container docker is still starting must exist before teardown removes it
Expand Down
14 changes: 8 additions & 6 deletions tst/unit/env/envs/test_computer_windows.py
Original file line number Diff line number Diff line change
Expand Up @@ -325,9 +325,11 @@ def test_a_windows_computers_display_is_1920x1080_for_now(local_stores, cli_rout
"display: a Windows computer's display is 1920x1080 for now"]


@pytest.mark.parametrize("backing, over", [(None, []), ("/disk/base.img", ["-B", "/disk/base.img", "-F", "qcow2"])],
ids=["a-base", "a-layer-over-its-base"])
def test_a_disk_is_rewritten_keeping_only_the_clusters_it_uses(monkeypatch, tmp_path, backing, over):
@pytest.mark.parametrize("backing, over, args", [
(None, [], ()), ("/disk/base.img", ["-B", "/disk/base.img", "-F", "qcow2"], ()),
("/disk/base.img", ["-B", "/disk/base.img", "-F", "qcow2"], ("--mount", "type=image,source=base:1,target=/disk"))],
ids=["a-base", "a-layer-over-its-base", "a-layer-over-its-base-mounted"])
def test_a_disk_is_rewritten_keeping_only_the_clusters_it_uses(monkeypatch, tmp_path, backing, over, args):
ran = []

def run(argv, **kwargs):
Expand All @@ -337,8 +339,8 @@ def run(argv, **kwargs):
(tmp_path / "disk.img").write_text("as Windows left it")
monkeypatch.setattr(windows_disk.subprocess, "run", run)

windows_disk._converted(tmp_path, "machine:1", backing=backing)
windows_disk._converted(tmp_path, "machine:1", backing=backing, args=args)

assert ran == [["docker", "run", "--rm", "-v", f"{tmp_path}:/out", "--entrypoint", "qemu-img", "machine:1", "convert",
"-O", "qcow2", *over, "/out/disk.img", "/out/compact.img"]]
assert ran == [["docker", "run", "--rm", "-v", f"{tmp_path}:/out", *args, "--entrypoint", "qemu-img", "machine:1",
"convert", "-O", "qcow2", *over, "/out/disk.img", "/out/compact.img"]]
assert (tmp_path / "disk.img").read_text() == "compacted" and not (tmp_path / "compact.img").exists()
78 changes: 77 additions & 1 deletion tst/unit/env/envs/test_disk_recipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -441,7 +441,7 @@ def launched(monkeypatch):
ran = []
monkeypatch.setattr(local_sandbox, "kvm_platforms", lambda: frozenset({"linux/arm64"}))
monkeypatch.setattr(local_sandbox, "host_platform", lambda: "linux/arm64")
monkeypatch.setattr(local_sandbox, "machine_image", AsyncMock(return_value="agentenv-vm:0123"))
monkeypatch.setattr(local_sandbox, "machine_for", AsyncMock(return_value=local_disk.Machine("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, input=None, timeout=None: ran.append(
Expand Down Expand Up @@ -526,6 +526,82 @@ async def test_a_disk_isnt_booted_where_docker_cant_pass_kvm(launched, monkeypat
assert launched == [] and LocalSandboxProvider.disk_image_platforms() == frozenset()


@pytest.mark.asyncio
async def test_a_disk_boots_with_its_image_mounted_into_the_launcher_where_docker_mounts_images(launched, monkeypatch):
mounted = local_disk.Machine("launcher:1", ("--mount", "type=image,source=disk:1,target=/disk,image-subpath=disk"))
monkeypatch.setattr(local_sandbox, "machine_for", AsyncMock(return_value=mounted))

await LocalSandboxProvider().create_vm(image="disk:1", exposed_ports=[18765], env={}, cpu=2, memory=4096)

(run,) = launched
assert run[-1] == "launcher:1" and run[run.index("--mount"):run.index("--mount") + 2] == mounted.args


@pytest.mark.asyncio
@pytest.mark.parametrize("mounts", [True, False], ids=["mounted", "copied"])
async def test_a_disk_is_mounted_into_the_launcher_where_docker_mounts_images_else_copied_into_a_machine_image(
monkeypatch, mounts):
"""No copy of a disk, nor of a Windows disk's base, is made for each launcher where Docker can mount it."""
monkeypatch.setattr(local_disk, "mounts_images", lambda: mounts)
monkeypatch.setattr(local_disk, "launcher", lambda: "launcher:1")
monkeypatch.setattr(local_disk, "machine_image", AsyncMock(return_value="agentenv-vm:1"))

machine = await local_disk.machine_for("disk:1")

assert machine == (local_disk.Machine("launcher:1", (
"--mount", "type=image,source=disk:1,target=/disk,image-subpath=disk")) if mounts else local_disk.Machine(
"agentenv-vm:1"))
assert local_disk.machine_image.await_count == (0 if mounts else 1)


@pytest.mark.parametrize("info, mounts", [
('29.9.3 29.9.3 [["driver-type","io.containerd.snapshotter.v1"]]', True),
('28.0.0 28.0.0 [["driver-type","io.containerd.snapshotter.v1"]]', True),
('27.5.1 27.5.1 [["driver-type","io.containerd.snapshotter.v1"]]', False),
('27.5.1 29.9.3 [["driver-type","io.containerd.snapshotter.v1"]]', False),
('29.9.3 29.9.3 [["Backing Filesystem","extfs"],["Supports d_type","true"]]', False),
("4.9.3 4.9.3 null", False),
(DockerFailed("docker info failed: template: :1: can't evaluate field ClientInfo", 1), False),
(subprocess.TimeoutExpired(["docker", "info"], 30), False)],
ids=["29-containerd", "28-containerd", "27", "27-client-on-a-29-server", "29-overlay2", "podman",
"a-client-too-old-to-say", "no-answer"])
def test_docker_mounts_images_from_28_on_both_ends_on_the_containerd_image_store(monkeypatch, info, mounts):
"""A Docker that fails to say, or doesn't in time, gets the copy."""
local_disk._docker_mounts_images.cache_clear()
asked = []

def docker(*args, timeout=None):
asked.append(timeout)
if isinstance(info, Exception):
raise info
return info

monkeypatch.setattr(local_disk, "docker", docker)

assert local_disk.mounts_images() is mounts and asked == [local_disk._DOCKER_INFO_SECONDS]
local_disk._docker_mounts_images.cache_clear()


def test_a_docker_that_failed_to_say_is_asked_again_and_an_answer_is_kept(monkeypatch):
"""A daemon busy for a moment doesn't leave a long-lived process copying disks for good."""
local_disk._docker_mounts_images.cache_clear()
answers = [subprocess.TimeoutExpired(["docker", "info"], 30),
'29.9.3 29.9.3 [["driver-type","io.containerd.snapshotter.v1"]]']
asked = []

def docker(*args, timeout=None):
asked.append(args)
answer = answers.pop(0)
if isinstance(answer, Exception):
raise answer
return answer

monkeypatch.setattr(local_disk, "docker", docker)

assert [local_disk.mounts_images() for _ in range(3)] == [False, True, True] and len(asked) == 2
local_disk._docker_mounts_images.cache_clear()


@pytest.mark.asyncio
@pytest.mark.parametrize("layers, dockerfile", [
("1", "FROM launcher:1\nCOPY --from=disk:1 /disk/ /disk/\n"),
Expand Down
Loading