diff --git a/src/agent_env/computers/README.md b/src/agent_env/computers/README.md index c5da91ac8..7874544b8 100644 --- a/src/agent_env/computers/README.md +++ b/src/agent_env/computers/README.md @@ -150,19 +150,20 @@ the desktop signs in as at every boot; `setup.ps1` runs once at `agentadmin`'s f setup that fails is read back over ssh at once (`setup-failed`), then turns off what would get in an agent's way (updates, the lock screen, notifications, Widgets, Copilot, first-run pages, the page file and the firewall, behind QEMU's network, which reaches the guest at the ports the launcher forwards alone), installs the pinned Python with Tk -and `agentenv-session`, and switches the sign-in to `user`. Setup takes about 13 minutes. The installed disk is kept -in this machine's Docker as the base's image, `disk/base.img` in a layer of its own under an empty overlay of it, with -the date its 90-day evaluation expires; a compose installs a base afresh once fewer than 7 days are left, and every -input of an install (the ISO's sha256 and the recipe) is part of its key, so a newer ISO installs a new base. Name it -as `windows-11`, or by its id, `@local/agentenv-framework/computers/windows-11`. +and `agentenv-session`, and switches the sign-in to `user`. Setup takes about 13 minutes. The installed disk, trimmed, +becomes the base's image, `disk/base.img` in a layer of its own under an empty overlay of it, with the date its 90-day +evaluation expires and the key to `agentadmin`; it's pushed to the env's image store, as Linux disks +are, under a tag of its recipe, so another machine composing on it pulls it rather than installing Windows, and fetches +no ISO. A compose installs a base afresh once fewer than 7 days of its evaluation are left, from the ISO Microsoft's +link serves then. Name it as `windows-11`, or by its id, `@local/agentenv-framework/computers/windows-11`. A Windows computer env is composed on the same machine (`env/envs/computer/utils/windows_disk.py`): the base booted on an overlay kept in a folder of its own, each layer installed over ssh as `agentadmin` by a PowerShell script composing writes, as its `[install.windows]` says (its files under `C:\agentenv\drivers\`, its downloads checked by sha256, its venv on the pinned Python, its PowerShell run scripts, elevated, and its services), then shut down and -packed onto the base's image, so every computer composed on a base shares one copy of it. A composed disk stays in this -machine's Docker, as a VM bundle stays on its Mac, under a tag of its key, which a compose of the same inputs reuses -while it's there; gate 3 probes it at full, and one that fails is removed. At the session user's sign-in, +packed onto the base's image, so every computer composed on a base shares one copy of it, and pushed to the env's image +store under its key, pinned, as a Linux disk is; a compose of the same inputs on any machine reuses it, and a deploy +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 diff --git a/src/agent_env/env/envs/computer/utils/composer.py b/src/agent_env/env/envs/computer/utils/composer.py index 6caa6d94c..021d672be 100644 --- a/src/agent_env/env/envs/computer/utils/composer.py +++ b/src/agent_env/env/envs/computer/utils/composer.py @@ -54,7 +54,8 @@ from agent_env.store.ids import derive_id from agent_env.store.image_store.image_contents import image_config, labels from agent_env.store.image_store.oci_registry_credentials import names_registry -from agent_env.store.image_store.registry_api import NoSuchImageError, pin_digest, write_index +from agent_env.store.image_store.image_store import ImageStore +from agent_env.store.image_store.registry_api import pin_digest, pin_if_served, write_index from agent_env.store.local_state import holding_locks from agent_env.utils.docker.build import DockerBuildError @@ -176,7 +177,7 @@ def left(form: str, arch: str, record: dict | None) -> bool: if status == "unverified": return bool(boot.native_links(provider, arch, form, env.stored_id, env.computer.os.type)) if form == "disk" and env.computer.os.type == "windows": - return not windows_disk.held(record["image"]) or windows_disk.expiring(record.get("expires")) + return windows_disk.expiring(record.get("expires")) if form != "vm_bundle": return False links = boot.native_links(provider, arch, form, env.stored_id) @@ -336,11 +337,11 @@ def _composed_vm_bundle(env: ComputerEnv, base: ComputerArtifact, driver: Comput def _composed_windows_disk(env: ComputerEnv, base: ComputerArtifact, driver: ComputerDriverEnv, arch: str, provider: SandboxProvider, progress: _Progress) -> dict: - """The Windows disk of ``arch`` composed with this machine's Docker and KVM (``windows_disk``), which keeps it: its - base installed from Microsoft's media first unless held, then the driver installed over ssh as a layer on it and - packed under a tag of its key; or one of that key reused, held here already. Then probed at full (gate 3), and - removed when that fails, so no deploy here boots it. All of it holds a lock of its tag, which keeps two composes of - one key from breaking each other: the other waits, then reuses it.""" + """The Windows disk of ``arch``: one of its key the env's image store holds, reused; else composed with this + machine's Docker and KVM (``windows_disk``), on its base, this machine's, the store's or installed from Microsoft's + media, the driver installed over ssh as a layer on it, packed and pushed to the store under its key, as Linux disks + are. Then probed at full (gate 3). All of it holds a lock of its key, which keeps two composes of one key on this + machine from breaking each other: the other waits, then reuses it.""" pair = f"disk/{arch}" with (_attempting(env, "disk", arch) as attempt, tempfile.TemporaryDirectory(prefix="agentenv-compose-") as work): @@ -369,33 +370,29 @@ def say(line: str) -> None: progress.say(f"{pair}: {line}") attempt.stage = "base" - iso = windows_disk.media(say) + registry = windows_disk.Registry(_image_store(env), REPOSITORY) disk_gb, user = base.forms.disk.min_disk_gb, base.session.user - made = windows_disk.base(_key(iso=windows_disk.sha256(iso), recipe=_digest(windows_disk.RECIPE), - disk_gb=disk_gb, user=user, form="windows-base", format=FORMAT), - iso, disk_gb, user, say) + made = windows_disk.base(_key(recipe=_digest(windows_disk.RECIPE), disk_gb=disk_gb, user=user, + form="windows-base", format=FORMAT), + lambda: windows_disk.media(say), disk_gb, user, say, registry) script = windows_disk.write(build, [part], settings, {"base": made.image, "layers": [part.folder]}) - key = attempt.key = _key(base=windows_disk.image_id(made.image), provision=script, context=_digest(context), + expires = made.expires.isoformat(timespec="seconds") + key = attempt.key = _key(base=made.image, base_expires=expires, provision=script, context=_digest(context), build_settings=settings, form="disk", arch=arch, format=FORMAT) - image = f"agentenv-windows:{key[7:23]}" attempt.stage = "build" - with holding_locks([f"windows-disk:{image}"], lambda: progress.at( - pair, "waiting", f"waiting for the compose or probe of {image} under way on this machine")): - if windows_disk.held(image): - progress.at(pair, "reusing", "reusing the disk of the same key, composed on this machine earlier") - result = {"status": "unverified", "key": key, "image": image, "reused": True, - "expires": made.expires.isoformat(timespec="seconds"), "built_at": now()} + with holding_locks([f"windows-disk:{key}"], lambda: progress.at( + pair, "waiting", "waiting for the compose of the same disk under way on this machine")): + if (pinned := registry.served(windows_disk.disk_names(key).tag)) is not None: + progress.at(pair, "reusing", "reusing the disk of the same key, composed earlier") + result = {"status": "unverified", "key": key, "image": pinned, "reused": True, "expires": expires, + "built_at": now()} else: progress.at(pair, "composing", "composing on this machine's Docker") - windows_disk.compose(made, build, image, {KEY_LABEL: key, windows_disk.EXPIRES_LABEL: - made.expires.isoformat(timespec="seconds")}, say) - result = {"status": "unverified", "key": key, "image": image, "reused": False, - "builder": "LocalSandboxProvider", "expires": made.expires.isoformat(timespec="seconds"), - "built_at": now()} - probed = _probed(env, driver, "disk", arch, result, provider, progress) - if probed["status"] == "failed": - windows_disk.remove(image) - return probed + pinned = windows_disk.compose(made, build, key, + {KEY_LABEL: key, windows_disk.EXPIRES_LABEL: expires}, say, registry) + result = {"status": "unverified", "key": key, "image": pinned, "reused": False, + "builder": "LocalSandboxProvider", "expires": expires, "built_at": now()} + return _probed(env, driver, "disk", arch, result, provider, progress) return attempt.result @@ -654,13 +651,19 @@ def _key(**inputs: Any) -> str: return "sha256:" + hashlib.sha256(json.dumps(inputs, sort_keys=True).encode()).hexdigest() -def _stored(env: ComputerEnv, key: str) -> tuple[Any, str, str | None]: +def _stored(env: ComputerEnv, key: str) -> tuple[ImageStore, str, str | None]: """The env's image store, the ref there of the image of ``key``, and that ref pinned when the store serves one already, for any env that built it.""" + store = _image_store(env) + ref = store.image_ref(REPOSITORY, key.removeprefix("sha256:")) + return store, ref, pin_if_served(ref, store.auth(ref)) + + +def _image_store(env: ComputerEnv) -> ImageStore: + """The env's image store, with the repository its computers' images are kept in.""" store = get_config().get_image_store_for(env.stored_id) store.ensure_repository(REPOSITORY) - ref = store.image_ref(REPOSITORY, key.removeprefix("sha256:")) - return store, ref, _served(ref, store) + return store def _base_disk(base: ComputerArtifact, arch: str) -> tuple[str | None, str]: @@ -912,14 +915,6 @@ def _wants(layers: list[_Layer], protocol: bool) -> list[packages.Want]: return wants -def _served(ref: str, store) -> str | None: - """``ref`` pinned to the digest the image store serves for it, or None when it serves none.""" - try: - return pin_digest(ref, store.auth(ref)) - except NoSuchImageError: - return None - - def _computer(env: ComputerEnv, base: ComputerArtifact, built: dict[str, str], disks: dict[str, str], bundles: dict[str, str]) -> dict | None: """Write the env's computer over the images ``built``, by arch, joined into one index in the env's image store, signed diff --git a/src/agent_env/env/envs/computer/utils/windows_disk.py b/src/agent_env/env/envs/computer/utils/windows_disk.py index c598c8850..652ea7f0a 100644 --- a/src/agent_env/env/envs/computer/utils/windows_disk.py +++ b/src/agent_env/env/envs/computer/utils/windows_disk.py @@ -12,7 +12,6 @@ import asyncio import base64 import contextlib -import hashlib import io import json import re @@ -26,15 +25,19 @@ from collections.abc import Callable, Iterator from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from pathlib import Path, PureWindowsPath -from typing import Any +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import Any, NamedTuple from agent_env.artifact.artifacts.computer import DISK_BASE_FILE, DISK_FILE, DISK_OWNER +from agent_env.artifact.artifacts.docker_image import docker_login, push_local_image from agent_env.config.paths import cache_root from agent_env.env.envs.computer.utils.install import WindowsInstall -from agent_env.providers.sandbox_providers.host_images import SOURCES_LABEL, present +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.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 from agent_env.store.local_state import holding_locks from agent_env.utils.docker.build import build_image @@ -51,6 +54,9 @@ #: The fewest days of its evaluation a base may have left, as a deploy needs some. RENEW_DAYS = 7 EXPIRES_LABEL, BUILD_LABEL, KEY_LABEL = "dev.agentenv.windows.expires", "dev.agentenv.windows.build", "dev.agentenv.key" +#: Where a base's image keeps the key to its admin, which composing a layer on it takes over ssh. The image is in the +#: user's own image store, and the guest's ssh is reached only on this machine's loopback while a layer is composed. +ADMIN_KEY = PurePosixPath("/disk/admin.key") #: How long an install, and a compose, may take, and how long the session user's first sign-in is left to settle. INSTALL_TIMEOUT_S, COMPOSE_TIMEOUT_S, SETTLE_S = 3600, 3600, 60 #: How long a guest may take to power off once told. @@ -67,13 +73,62 @@ "once; it's installed for evaluation under Microsoft's terms and runs 90 days from each install") +@dataclass(frozen=True) +class Registry: + """The env's image store as Windows disks are kept in it: by tag in ``repository``, pinned by digest, as Linux + disks are.""" + + store: ImageStore + repository: str + + def served(self, tag: str) -> str | None: + """``tag``'s ref, pinned to the digest the store serves for it now; None when it serves none.""" + ref = self.store.image_ref(self.repository, tag) + return pin_if_served(ref, self.store.auth(ref)) + + def labels(self, pinned: str) -> dict[str, str]: + return image_labels(pinned, "amd64", self.store.auth(pinned)) + + def push(self, local: str, tag: str) -> str: + """``local`` pushed as ``tag``; its ref there, pinned.""" + ref = self.store.image_ref(self.repository, tag) + push_local_image(local, ref, self.store) + return pin_digest(ref, self.store.auth(ref)) + + def pull(self, pinned: str) -> None: + docker_login(self.store.auth(pinned)) + docker("pull", "-q", pinned) + + +class Names(NamedTuple): + """What a base or a disk is called: ``local`` in this machine's Docker, ``tag`` in the image store.""" + + local: str + tag: str + + +def base_names(key: str) -> Names: + return Names(f"agentenv-windows-base:{_short(key)}", f"windows-base-{_short(key)}") + + +def disk_names(key: str) -> Names: + """A disk's names, its tag in the image store being its key, as a Linux disk's is.""" + return Names(f"agentenv-windows:{_short(key)}", key.removeprefix("sha256:")) + + +def _short(key: str) -> str: + return key.removeprefix("sha256:")[:16] + + @dataclass(frozen=True) class Base: - """A base on this machine: its image's tag, the folder of its admin's key, and when its evaluation expires.""" + """A base: its image's tag on this machine, the folder of its admin's key, when its evaluation expires, and, for one + the image store holds that this machine doesn't yet, its ref there, pulled when a layer is composed on it.""" image: str keys: Path expires: datetime + pinned: str | None = None @dataclass(frozen=True) @@ -108,43 +163,72 @@ def media(say: Callable[[str], None]) -> Path: return iso -def sha256(path: Path) -> str: - """``path``'s sha256, kept beside it for its size and mtime, since an ISO takes a while to read.""" - stat, kept = path.stat(), path.with_name(path.name + ".sha256") - stamp = f"{stat.st_size} {stat.st_mtime_ns}" - if kept.exists() and (lines := kept.read_text().split("\n"))[0] == stamp: - return lines[1] - hashed = hashlib.sha256() - with path.open("rb") as source: - while chunk := source.read(1 << 20): - hashed.update(chunk) - kept.write_text(f"{stamp}\n{hashed.hexdigest()}\n") - return hashed.hexdigest() - - -def base(key: str, iso: Path, disk_gb: int, user: str, say: Callable[[str], None]) -> Base: - """The base of ``key`` on this machine, installed from ``iso`` first when there's none, its key is gone, or its - evaluation has fewer than ``RENEW_DAYS`` left; one install of a key runs at once, and another waits for it.""" - image, keys = f"agentenv-windows-base:{key[7:23]}", cache_root() / "windows" / "bases" / key[7:23] - with holding_locks([f"windows-base:{key}"], lambda: say("waiting for the install of the same base under way")): +def base(key: str, iso: Callable[[], Path], disk_gb: int, user: str, say: Callable[[str], None], + registry: Registry) -> Base: + """The base of ``key``: this machine's, pushed to the image store unless it holds one of ``key``; else the one the + store holds, pulled, with the key to its admin it carries, when a layer is composed on it; else Windows installed + from the ISO ``iso`` gives, and pushed there. One with fewer than ``RENEW_DAYS`` of its evaluation left is installed + afresh. One install of a key runs at once on a machine, and another waits for it.""" + (image, tag), keys = base_names(key), cache_root() / "windows" / "bases" / _short(key) + with holding_locks([_base_lock(image)], lambda: say("waiting for the install of the same base under way")): if present(image) and (keys / "id").exists(): expires = datetime.fromisoformat(_label(image, EXPIRES_LABEL)) - if expires - datetime.now(timezone.utc) >= timedelta(days=RENEW_DAYS): + if not expiring(expires.isoformat()): + if registry.served(tag) is None: + say("pushing the base to the image store, which doesn't hold it") + registry.push(image, tag) return Base(image, keys, expires) say(f"its base's evaluation expires {expires:%Y-%m-%d}; installing it afresh") + elif (pinned := registry.served(tag)) is not None and not expiring( + stored := registry.labels(pinned).get(EXPIRES_LABEL)): + return Base(image, keys, datetime.fromisoformat(stored), pinned) shutil.rmtree(keys, ignore_errors=True) keys.mkdir(parents=True) subprocess.run(["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-C", f"agentenv {image}", "-f", str(keys / "id")], check=True) - with _working(f"base-{key[7:23]}") as work: - facts = _installed(work, iso, keys, disk_gb, user, say) + media = iso() + with _working(f"base-{_short(key)}") as work: + facts = _installed(work, media, keys, disk_gb, user, say) say("packing the base") expires = datetime.now(timezone.utc) + timedelta(minutes=facts["grace_minutes"]) + shutil.copy(keys / "id", work / ADMIN_KEY.name) _pack(work, image, {KEY_LABEL: key, EXPIRES_LABEL: expires.isoformat(timespec="seconds"), BUILD_LABEL: facts["build"]}, base_image=None) + say("pushing the base to the image store") + registry.push(image, tag) return Base(image, keys, expires) +def _here(base_: Base, registry: Registry, say: Callable[[str], None]) -> None: + """``base_`` on this machine, with the key to its admin: pulled from the image store when it's the store's, under + the base's lock, so a compose of another disk on it waits, then finds it.""" + if base_.pinned is None: + return + with holding_locks([_base_lock(base_.image)], lambda: say("waiting for the pull of the same base under way")): + if present(base_.image) and (base_.keys / "id").exists(): + return + say(f"pulling its base from the image store, {base_.pinned}") + registry.pull(base_.pinned) + docker("tag", base_.pinned, base_.image) + _take_admin_key(base_.image, base_.keys) + + +def _base_lock(image: str) -> str: + return f"windows-base:{image}" + + +def _take_admin_key(image: str, keys: Path) -> None: + """The key to the admin of the base ``image`` holds, written to ``keys``, as composing a layer on it needs.""" + shutil.rmtree(keys, ignore_errors=True) + keys.mkdir(parents=True) + container = docker("create", image, "none") + try: + docker("cp", f"{container}:{ADMIN_KEY}", str(keys / "id")) + finally: + subprocess.run(["docker", "rm", container], capture_output=True) + (keys / "id").chmod(0o600) + + def _installed(work: Path, iso: Path, keys: Path, disk_gb: int, user: str, say: Callable[[str], None]) -> dict: """Windows installed from ``iso`` onto ``work/disk.img`` by the launcher, answered as the recipe says: then booted until its session user has signed in and settled, and shut down; what it reports of itself.""" @@ -187,13 +271,18 @@ def _installed(work: Path, iso: Path, keys: Path, disk_gb: int, user: str, say: return facts -def compose(base_: Base, build: Path, image: str, labels: dict[str, str], say: Callable[[str], None]) -> None: - """``image``, a disk composed on ``base_``: the base booted on an overlay kept in a folder of this machine, ``build`` - copied in and its ``provision.ps1`` run as the admin over ssh, then shut down and packed onto the base's image.""" - with _working(f"disk-{image.rpartition(':')[2]}") as work: +def compose(base_: Base, build: Path, key: str, labels: dict[str, str], say: Callable[[str], None], + registry: Registry) -> str: + """The disk of ``key`` composed on ``base_``: the base booted on an overlay kept in a folder of this machine, + ``build`` copied in and its ``provision.ps1`` run as the admin over ssh, then shut down, packed onto the base's + image with ``labels`` and pushed to the image store; its ref there, pinned. A base the store holds is pulled + first.""" + image, tag = disk_names(key) + _here(base_, registry, say) + with _working(f"disk-{_short(key)}") as work: out = work / "out" out.mkdir() - name = f"agentenv-windows-compose-{image.rpartition(':')[2]}" + 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) @@ -221,11 +310,8 @@ def compose(base_: Base, build: Path, image: str, labels: dict[str, str], say: C (out / "guest").write_text("windows") say("packing the disk") _pack(out, image, labels, base_image=base_.image) - - -def held(image: str) -> bool: - """Whether this machine's Docker holds ``image``.""" - return present(image) + say("pushing the disk to the image store") + return registry.push(image, tag) def expiring(expires: str | None) -> bool: @@ -234,16 +320,6 @@ def expiring(expires: str | None) -> bool: return expires is None or datetime.fromisoformat(expires) - datetime.now(timezone.utc) < timedelta(days=RENEW_DAYS) -def image_id(image: str) -> str: - """The id of ``image`` in this machine's Docker, which a base installed afresh under its tag changes.""" - return docker("image", "inspect", "-f", "{{.Id}}", image) - - -def remove(image: str) -> None: - """``image`` removed from this machine's Docker, so no deploy here boots it.""" - subprocess.run(["docker", "image", "rm", "-f", image], capture_output=True) - - def write(build: Path, parts: list[Part], build_settings: dict[str, Any], image: dict) -> str: """Write ``build``'s ``provision.ps1``, which installs ``parts`` from ``build/context``, their services for agentenv-session, ``build_settings`` as every service's ``AGENTENV_SETTING_`` and ``image`` as what the disk @@ -398,14 +474,14 @@ def _pack(folder: Path, image: str, labels: dict[str, str], base_image: str | No check=True, capture_output=True) (folder / "guest").write_text("windows") dockerfile = (f"FROM scratch\nCOPY {owner} base.img /{DISK_BASE_FILE}\n" + f"COPY --chmod=600 {ADMIN_KEY.name} {ADMIN_KEY}\n" f"COPY {owner} disk.img guest /{Path(DISK_FILE).parent}/\n") else: dockerfile = f"FROM {base_image}\nCOPY {owner} disk.img guest /{Path(DISK_FILE).parent}/\n" for name in ("answers", "out"): shutil.rmtree(folder / name, ignore_errors=True) (folder / "Dockerfile").write_text(dockerfile) - # Labelled as built here, so a deploy boots this machine's image and pulls none of its name. - build_image(folder / "Dockerfile", folder, image, platform="linux/amd64", labels={**labels, SOURCES_LABEL: image}) + build_image(folder / "Dockerfile", folder, image, platform="linux/amd64", labels=labels) @contextlib.contextmanager diff --git a/src/agent_env/store/image_store/registry_api.py b/src/agent_env/store/image_store/registry_api.py index 61ff38bdb..4544cfed8 100644 --- a/src/agent_env/store/image_store/registry_api.py +++ b/src/agent_env/store/image_store/registry_api.py @@ -68,6 +68,14 @@ def pin_digest(ref: str, auth: RegistryAuth | None, *, timeout: float = _TIMEOUT return ref +def pin_if_served(ref: str, auth: RegistryAuth | None) -> str | None: + """``ref`` pinned as ``pin_digest`` pins it, or None when its registry serves no such image.""" + try: + return pin_digest(ref, auth) + except NoSuchImageError: + return None + + def pinned_from(pinned: str, ref: str) -> bool: """Whether ``pinned``, a ref as ``pin_digest`` returns it, pins ``ref``: ``ref`` itself when it gives a digest, else the same name and tag (``latest`` when none is given) with a digest.""" diff --git a/tst/unit/env/envs/conftest.py b/tst/unit/env/envs/conftest.py index 3f561951f..4096f25b0 100644 --- a/tst/unit/env/envs/conftest.py +++ b/tst/unit/env/envs/conftest.py @@ -10,6 +10,7 @@ from agent_env.artifact.store import get_artifact_store from agent_env.env.envs.computer.utils import builder, composer, probing from agent_env.providers.sandbox_providers.local_docker import provider as local_sandbox +from agent_env.store.image_store import registry_api from agent_env.store.image_store.image_store import OciRegistryImageStore from agent_env.store.image_store.local_registry_image_store import LocalRegistryImageStore from agent_env.store.object_store.local.store import LocalFilesystemObjectStore @@ -27,6 +28,7 @@ def docker(local_stores, cli_routing, monkeypatch): monkeypatch.setattr(builder, "docker_login", lambda auth: None) monkeypatch.setattr(composer, "write_index", fake.joined) monkeypatch.setattr(composer, "pin_digest", registry.pin_digest) + monkeypatch.setattr(registry_api, "pin_digest", registry.pin_digest) monkeypatch.setattr(docker_image, "pin_digest", registry.pin_digest) monkeypatch.setattr(builder, "host_platform", lambda: "linux/arm64") monkeypatch.setattr(local_sandbox, "host_platform", lambda: "linux/arm64") diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index a8a363820..9789c25bf 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -630,7 +630,7 @@ def test_builds_of_one_env_at_once_keep_each_others_pairs_and_a_failure_doesnt_r composer.compose(first, [("container", "amd64")]) composer.compose(second, [("container", "arm64")]) - monkeypatch.setattr(composer, "_served", lambda ref, store: None) + monkeypatch.setattr(composer, "pin_if_served", lambda ref, auth: None) docker.fail["linux/arm64"] = "boom" composer.compose(third, [("container", "arm64")]) diff --git a/tst/unit/env/envs/test_computer_windows.py b/tst/unit/env/envs/test_computer_windows.py index 04b263a7a..a2b584be9 100644 --- a/tst/unit/env/envs/test_computer_windows.py +++ b/tst/unit/env/envs/test_computer_windows.py @@ -4,6 +4,9 @@ the env's computer is a disk of the base's form. A Windows base is amd64, disk-only and started by agentenv-session.""" import json +import threading +import time +from types import SimpleNamespace from datetime import datetime, timedelta, timezone from pathlib import Path, PureWindowsPath @@ -43,37 +46,49 @@ def disk_image_platforms(self): return self.platforms +class FakeRegistry: + """The env's image store as Windows disks are kept in it: what it serves, by tag, pinned.""" + + def __init__(self): + self.held = {} + + def served(self, tag): + return self.held.get(tag) + + def push(self, local, tag): + self.held[tag] = f"registry.test/agentenv-computers:{tag}@sha256:{len(self.held):064x}" + return self.held[tag] + + class FakeDocker: - """windows_disk on this machine's Docker: bases installed and disks composed, as images it holds.""" + """windows_disk on one machine's Docker: bases installed and disks composed and pushed to the image store.""" - def __init__(self, tmp_path: Path): - self.tmp, self.images, self.installs, self.composed = tmp_path, set(), [], [] + def __init__(self, tmp_path: Path, registry: FakeRegistry): + self.tmp, self.registry, self.bases, self.installs, self.composed = tmp_path, registry, {}, [], [] - def base(self, key, iso, disk_gb, user, say): - image = f"agentenv-windows-base:{key[7:23]}" - if image not in self.images: - self.installs.append((key, iso, disk_gb, user)) - self.images.add(image) - return windows_disk.Base(image, self.tmp / "keys", datetime.now(timezone.utc) + timedelta(days=89)) + def base(self, key, iso, disk_gb, user, say, registry): + if key not in self.bases: + self.installs.append((key, iso(), disk_gb, user)) + self.bases[key] = datetime.now(timezone.utc) + timedelta(days=89) + return windows_disk.Base(f"agentenv-windows-base:{key[7:23]}", self.tmp / "keys", self.bases[key]) - def compose(self, base, build, image, labels, say): + def compose(self, base, build, key, labels, say, registry): self.composed.append((base.image, (build / "provision.ps1").read_bytes().decode(), sorted(p.relative_to(build).as_posix() for p in build.rglob("*") if p.is_file()), labels)) - self.images.add(image) + return registry.push(*windows_disk.disk_names(key)) @pytest.fixture def docker(local_stores, cli_routing, monkeypatch, tmp_path): - fake = FakeDocker(tmp_path) + fake = FakeDocker(tmp_path, FakeRegistry()) iso = tmp_path / "windows.iso" iso.write_bytes(b"iso") monkeypatch.setattr(windows_disk, "media", lambda say: iso) - monkeypatch.setattr(windows_disk, "sha256", lambda path: "ab" * 32) monkeypatch.setattr(windows_disk, "base", fake.base) monkeypatch.setattr(windows_disk, "compose", fake.compose) - monkeypatch.setattr(windows_disk, "held", lambda image: image in fake.images) - monkeypatch.setattr(windows_disk, "image_id", lambda image: f"sha256:{image}") - monkeypatch.setattr(windows_disk, "remove", fake.images.discard) + monkeypatch.setattr(windows_disk, "Registry", lambda store, repository: fake.registry) + monkeypatch.setattr(composer, "get_config", lambda: SimpleNamespace( + get_image_store_for=lambda stored_id: SimpleNamespace(ensure_repository=lambda repository: None))) monkeypatch.setattr(composer, "_probed", lambda env, driver, form, arch, result, provider, progress: { **result, "status": "ready"}) return fake @@ -118,7 +133,8 @@ def test_a_windows_env_composes_a_thin_disk_on_its_base_here_the_driver_installe env = _put(_driver(tmp_path / "desk"), FakeLinux(), theme="dark") built = env.metadata["composed"]["disk"]["amd64"] - assert built["status"] == "ready" and built["image"].startswith("agentenv-windows:") and not built["reused"] + assert built["status"] == "ready" and not built["reused"] + assert built["image"] == docker.registry.held[built["key"].removeprefix("sha256:")] (installed,) = docker.installs assert installed[1:] == (tmp_path / "windows.iso", 64, "user") ((base, script, files, labels),) = docker.composed @@ -143,18 +159,16 @@ def test_a_windows_env_composes_a_thin_disk_on_its_base_here_the_driver_installe assert computer.forms.disk.image is None and env.metadata["computer"]["disks"] == {"amd64": built["image"]} -def test_a_compose_of_the_same_inputs_reuses_the_disk_this_machine_holds_and_composes_one_it_lost(docker, tmp_path): +def test_a_compose_of_the_same_inputs_reuses_the_disk_the_image_store_holds(docker, tmp_path): driver = _driver(tmp_path / "desk") first = _put(driver, FakeLinux()) image = first.metadata["composed"]["disk"]["amd64"]["image"] again = _put(driver, FakeLinux()) - assert again.metadata["composed"]["disk"]["amd64"]["reused"] and len(docker.composed) == 1 - assert composer.pending(again, FakeLinux()) == [] - docker.images.discard(image) - assert composer.pending(again, FakeLinux()) == [("disk", "amd64")] - assert len(docker.installs) == 1 + built = again.metadata["composed"]["disk"]["amd64"] + assert built["reused"] and built["image"] == image and len(docker.composed) == 1 and len(docker.installs) == 1 + assert composer.pending(again, FakeLinux()) == [] def test_a_machine_that_boots_no_windows_disk_composes_none(docker, tmp_path): @@ -216,6 +230,7 @@ def test_a_disk_whose_evaluation_nears_its_end_is_composed_again(docker, tmp_pat @pytest.mark.parametrize("days_left, installs", [(30, 0), (3, 1)]) def test_a_base_is_reused_until_its_evaluation_nears_its_end_then_installed_afresh(monkeypatch, tmp_path, days_left, installs): + """Either way the image store gets it, which a push that failed after an install, or another env's store, lacks.""" monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) key = "sha256:" + "12" * 32 keys = tmp_path / "agentenv" / "windows" / "bases" / key[7:23] @@ -229,11 +244,75 @@ def test_a_base_is_reused_until_its_evaluation_nears_its_end_then_installed_afre monkeypatch.setattr(windows_disk, "_installed", lambda *args: installed.append(args) or { "grace_minutes": 90 * 24 * 60, "build": "26200.6584"}) monkeypatch.setattr(windows_disk, "_pack", lambda folder, image, labels, base_image: None) + registry = FakeRegistry() - base = windows_disk.base(key, tmp_path / "windows.iso", 64, "user", lambda line: None) + base = windows_disk.base(key, lambda: tmp_path / "windows.iso", 64, "user", lambda line: None, registry) assert len(installed) == installs and base.image == f"agentenv-windows-base:{key[7:23]}" assert (base.expires - datetime.now(timezone.utc)).days == (days_left - 1 if not installs else 89) + assert list(registry.held) == [f"windows-base-{key[7:23]}"] + + +def test_a_base_this_machine_holds_goes_to_an_image_store_lacking_it_once(monkeypatch, tmp_path): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + key = "sha256:" + "56" * 32 + keys = tmp_path / "agentenv" / "windows" / "bases" / key[7:23] + keys.mkdir(parents=True) + (keys / "id").write_text("key") + expires = datetime.now(timezone.utc) + timedelta(days=30) + monkeypatch.setattr(windows_disk, "present", lambda image: True) + monkeypatch.setattr(windows_disk, "_label", lambda image, name: expires.isoformat(timespec="seconds")) + registry, pushed = FakeRegistry(), [] + push = registry.push + registry.push = lambda local, tag: pushed.append(local) or push(local, tag) + + for _ in range(2): + windows_disk.base(key, lambda: tmp_path / "windows.iso", 64, "user", lambda line: None, registry) + + assert pushed == [f"agentenv-windows-base:{key[7:23]}"] + + +def test_composes_of_two_disks_on_a_base_the_store_holds_pull_it_and_take_its_key_once(monkeypatch, tmp_path): + """The second waits for the first's pull, then finds the base here.""" + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + keys = tmp_path / "keys" + made, pulls = threading.Event(), [] + + def pull(pinned): + pulls.append(pinned) + time.sleep(0.2) + + monkeypatch.setattr(windows_disk, "present", lambda image: made.is_set()) + monkeypatch.setattr(windows_disk, "docker", lambda *args: None) + monkeypatch.setattr(windows_disk, "_take_admin_key", lambda image, at: (at.mkdir(), (at / "id").write_text("key"), + made.set())) + registry = SimpleNamespace(pull=pull) + base = windows_disk.Base("agentenv-windows-base:1", keys, datetime.now(timezone.utc), "registry.test/base@sha256:1") + + composes = [threading.Thread(target=windows_disk._here, args=(base, registry, lambda line: None)) for _ in range(2)] + for compose in composes: + compose.start() + for compose in composes: + compose.join(10) + + assert pulls == ["registry.test/base@sha256:1"] and (keys / "id").read_text() == "key" + + +def test_a_machine_without_the_base_takes_the_image_stores_unless_its_evaluation_nears_its_end(monkeypatch, tmp_path): + """Found without the ISO, which only an install fetches; pulled only when a layer is composed on it.""" + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + key = "sha256:" + "34" * 32 + registry = FakeRegistry() + pinned = registry.push("agentenv-windows-base:x", f"windows-base-{key[7:23]}") + expires = datetime.now(timezone.utc) + timedelta(days=60) + registry.labels = lambda ref: {windows_disk.EXPIRES_LABEL: expires.isoformat(timespec="seconds")} + monkeypatch.setattr(windows_disk, "present", lambda image: False) + fetched = [] + + base = windows_disk.base(key, lambda: fetched.append("iso"), 64, "user", lambda line: None, registry) + + assert base.pinned == pinned and base.image == f"agentenv-windows-base:{key[7:23]}" and fetched == [] + assert (base.expires - datetime.now(timezone.utc)).days == 59 def test_a_windows_computers_display_is_1920x1080_for_now(local_stores, cli_routing): diff --git a/tst/unit/store/registry_api_test.py b/tst/unit/store/registry_api_test.py index 10bd8fcde..aec0d54fe 100644 --- a/tst/unit/store/registry_api_test.py +++ b/tst/unit/store/registry_api_test.py @@ -14,6 +14,7 @@ NoSuchImageError, RegistryRefusedError, pin_digest, + pin_if_served, pinned_from, read_blob, read_manifest, @@ -187,6 +188,17 @@ def test_a_missing_image_and_a_refusal_each_have_their_own_error(fake_registry, assert type(raised.value) is (kind or ValueError) +@pytest.mark.parametrize("status, kind", [(404, None), (401, RegistryRefusedError), (500, ValueError)]) +def test_pin_if_served_is_none_for_a_missing_image_and_raises_anything_else(fake_registry, status, kind): + fake_registry(lambda request: httpx.Response(status)) + + if kind is None: + assert pin_if_served("ghcr.io/org/tool:v9", None) is None + else: + with pytest.raises(kind): + pin_if_served("ghcr.io/org/tool:v9", None) + + def test_an_unreachable_registry_is_named(fake_registry): def down(request): raise httpx.ConnectError("connection refused")