From 6ffa8d8843a1ccb222fe10ef823f4d5bb1a5dc52 Mon Sep 17 00:00:00 2001 From: codex Date: Mon, 27 Jul 2026 17:14:14 +0000 Subject: [PATCH] feat(firecracker): pull pinned agent images --- .../backend/firecracker/image_builder.py | 77 +++++++++++++++++++ bot_bottle/backend/firecracker/launch.py | 18 ++++- tests/unit/test_firecracker_image_builder.py | 14 ++++ 3 files changed, 107 insertions(+), 2 deletions(-) diff --git a/bot_bottle/backend/firecracker/image_builder.py b/bot_bottle/backend/firecracker/image_builder.py index 57615453..13cf9ef5 100644 --- a/bot_bottle/backend/firecracker/image_builder.py +++ b/bot_bottle/backend/firecracker/image_builder.py @@ -131,6 +131,46 @@ def cached_agent_rootfs_dir(dockerfile: Path) -> Path | None: return base if (base / ".bb-ready").is_file() else None +def _image_rootfs_digest(image: str) -> str: + h = hashlib.sha256() + h.update(image.encode()) + h.update(b"\0") + h.update(util._GUEST_INIT.encode()) + return h.hexdigest()[:16] + + +def cached_agent_image_rootfs_dir(image: str) -> Path | None: + """Return a ready rootfs exported from an immutable OCI image.""" + base = util.cache_dir() / "rootfs" / f"agent-image-{_image_rootfs_digest(image)}" + return base if (base / ".bb-ready").is_file() else None + + +def acquire_agent_image_rootfs_dir( + image: str, *, smoke_test: tuple[str, ...] = (), +) -> Path: + """Pull a digest-pinned image in the infra VM and export its rootfs.""" + if "@sha256:" not in image: + die(f"prebuilt Firecracker agent image is not digest-pinned: {image}") + digest = _image_rootfs_digest(image) + base = util.cache_dir() / "rootfs" / f"agent-image-{digest}" + cached = cached_agent_image_rootfs_dir(image) + if cached is not None: + info(f"using cached agent rootfs {cached.name}") + return cached + with _build_lock(): + if (base / ".bb-ready").is_file(): + return base + staging = util.cache_dir() / "rootfs" / f".pulling-{digest}" + shutil.rmtree(staging, ignore_errors=True) + staging.mkdir(parents=True) + _pull_in_infra(image, staging, smoke_test, digest) + util.inject_guest_boot(staging) + (staging / ".bb-ready").write_text("ok\n") + shutil.rmtree(base, ignore_errors=True) + os.rename(staging, base) + return base + + def build_agent_rootfs_dir( dockerfile: Path, *, image_tag: str, smoke_test: tuple[str, ...] = (), ) -> Path: @@ -227,6 +267,43 @@ def _build_in_infra( _cleanup() +def _pull_in_infra( + image: str, base: Path, smoke_test: tuple[str, ...], digest: str, +) -> None: + """Pull and export a published agent image inside the orchestrator VM.""" + service = FirecrackerInfraService() + service.ensure_running() + key, ip = service.orchestrator().ssh_target() + tag = f"bot-bottle-agent-pull-{digest}" + smoke_ctr, export_ctr = f"{tag}-smoke", f"{tag}-export" + quoted_image = shlex.quote(image) + + def cleanup() -> None: + _ssh( + key, + ip, + f"buildah rm {smoke_ctr} {export_ctr} >/dev/null 2>&1; " + f"buildah rmi {_STORE_FLAG} {tag} >/dev/null 2>&1", + timeout=60, + ) + + cleanup() + try: + result = _ssh( + key, + ip, + f"buildah pull {_STORE_FLAG} {quoted_image} && " + f"buildah tag {_STORE_FLAG} {quoted_image} {tag}", + timeout=_BUILD_TIMEOUT_SECONDS, + ) + if result.returncode != 0: + die(f"pulling pinned agent image failed: {result.stderr.strip()}") + _smoke_test(key, ip, tag, smoke_ctr, smoke_test) + _stream_rootfs(key, ip, tag, export_ctr, base) + finally: + cleanup() + + def _ssh(private_key: Path, guest_ip: str, script: str, *, timeout: float = 60.0) -> subprocess.CompletedProcess[str]: return subprocess.run( diff --git a/bot_bottle/backend/firecracker/launch.py b/bot_bottle/backend/firecracker/launch.py index 6fa0d0b8..05ae62cc 100644 --- a/bot_bottle/backend/firecracker/launch.py +++ b/bot_bottle/backend/firecracker/launch.py @@ -228,8 +228,13 @@ def build_or_load_agent_base(plan: FirecrackerBottlePlan) -> Path: info(f"resuming from committed rootfs {committed_tar}") return util.build_committed_rootfs_dir(committed_tar) dockerfile = Path(plan.dockerfile_path) + prebuilt = "@sha256:" in plan.image if plan.spec.image_policy == "cached": - cached = image_builder.cached_agent_rootfs_dir(dockerfile) + cached = ( + image_builder.cached_agent_image_rootfs_dir(plan.image) + if prebuilt + else image_builder.cached_agent_rootfs_dir(dockerfile) + ) if cached is None: die( f"cached agent rootfs for {plan.image!r} not found; " @@ -237,6 +242,11 @@ def build_or_load_agent_base(plan: FirecrackerBottlePlan) -> Path: ) info(f"using cached agent rootfs {cached.name}") return cached + if prebuilt: + return image_builder.acquire_agent_image_rootfs_dir( + plan.image, + smoke_test=runtime_for(plan.agent_provider_template).smoke_test, + ) return image_builder.build_agent_rootfs_dir( dockerfile, image_tag=plan.image, @@ -253,7 +263,11 @@ def stale_checks(plan: FirecrackerBottlePlan) -> None: if committed and committed_tar.is_file(): check_stale_path(f"agent rootfs {committed_tar}", committed_tar) return - cached = image_builder.cached_agent_rootfs_dir(Path(plan.dockerfile_path)) + cached = ( + image_builder.cached_agent_image_rootfs_dir(plan.image) + if "@sha256:" in plan.image + else image_builder.cached_agent_rootfs_dir(Path(plan.dockerfile_path)) + ) if cached is not None: check_stale_path(f"agent rootfs {cached}", cached / ".bb-ready") diff --git a/tests/unit/test_firecracker_image_builder.py b/tests/unit/test_firecracker_image_builder.py index b3087fbd..8fe9d10c 100644 --- a/tests/unit/test_firecracker_image_builder.py +++ b/tests/unit/test_firecracker_image_builder.py @@ -91,6 +91,20 @@ class TestBuildAgentRootfsDir(unittest.TestCase): second = image_builder._rootfs_digest(self.dockerfile) self.assertNotEqual(first, second) + def test_pinned_image_is_pulled_and_cached(self): + image = f"registry.example/agent@sha256:{'a' * 64}" + with patch.object(image_builder.util, "cache_dir", return_value=self.cache), \ + patch.object(image_builder, "_pull_in_infra") as pull, \ + patch.object(image_builder.util, "inject_guest_boot"): + first = image_builder.acquire_agent_image_rootfs_dir(image) + second = image_builder.acquire_agent_image_rootfs_dir(image) + pull.assert_called_once() + self.assertEqual(first, second) + + def test_mutable_image_cannot_use_prebuilt_path(self): + with self.assertRaises(SystemExit): + image_builder.acquire_agent_image_rootfs_dir("agent:latest") + class TestBuildContext(unittest.TestCase): """The COPY-source parsing + context shipping that lets a Dockerfile pin an