From fa3e45cc54452ad3b0523c0fae27cb402d74e74b Mon Sep 17 00:00:00 2001 From: codex Date: Sat, 18 Jul 2026 21:18:50 +0000 Subject: [PATCH] fix(firecracker): honor cached image policy --- bot_bottle/backend/firecracker/backend.py | 11 ++-- bot_bottle/backend/firecracker/freezer.py | 2 +- .../backend/firecracker/image_builder.py | 9 ++- bot_bottle/backend/firecracker/launch.py | 47 ++++++++++----- tests/unit/test_firecracker_image_builder.py | 9 +++ tests/unit/test_stale_checks.py | 57 +++++++++++++++++++ 6 files changed, 115 insertions(+), 20 deletions(-) diff --git a/bot_bottle/backend/firecracker/backend.py b/bot_bottle/backend/firecracker/backend.py index e5cb44e..a8a9df9 100644 --- a/bot_bottle/backend/firecracker/backend.py +++ b/bot_bottle/backend/firecracker/backend.py @@ -93,16 +93,17 @@ class FirecrackerBottleBackend( ) def _build_or_load_images(self, plan: FirecrackerBottlePlan) -> BottleImages: - # Firecracker builds its rootfs image inside _launch_impl (via buildah in - # a builder VM), so there is nothing to pre-resolve here. - return BottleImages(agent=plan.image) + return BottleImages(agent=_launch.build_or_load_agent_base(plan)) + + def prelaunch_checks(self, plan: FirecrackerBottlePlan) -> None: + _launch.stale_checks(plan) @contextmanager def _launch_impl( self, plan: FirecrackerBottlePlan, images: BottleImages, ) -> Generator[FirecrackerBottle, None, None]: - del images # rootfs built inside _launch.launch; images unused here - with _launch.launch(plan, provision=self.provision) as bottle: + assert isinstance(images.agent, Path) + with _launch.launch(plan, images.agent, provision=self.provision) as bottle: yield bottle def prepare_cleanup(self) -> FirecrackerBottleCleanupPlan: diff --git a/bot_bottle/backend/firecracker/freezer.py b/bot_bottle/backend/firecracker/freezer.py index 4ffbaed..988eb0f 100644 --- a/bot_bottle/backend/firecracker/freezer.py +++ b/bot_bottle/backend/firecracker/freezer.py @@ -5,7 +5,7 @@ backend — we stream the guest root filesystem out over the control channel (SSH here). Unlike the other backends this needs no Docker: the tar *is* the resumable artifact. `resume` extracts it and rebuilds a fresh per-bottle ext4 with `mke2fs -d` (see `util.build_committed_rootfs_dir` -and `launch._build_agent_base`). The bottle keeps running after the +and `launch.build_or_load_agent_base`). The bottle keeps running after the snapshot. """ diff --git a/bot_bottle/backend/firecracker/image_builder.py b/bot_bottle/backend/firecracker/image_builder.py index 176ce31..3b18e0b 100644 --- a/bot_bottle/backend/firecracker/image_builder.py +++ b/bot_bottle/backend/firecracker/image_builder.py @@ -45,6 +45,13 @@ def _dockerfile_hash(dockerfile: Path) -> str: return hashlib.sha256(dockerfile.read_bytes()).hexdigest()[:16] +def cached_agent_rootfs_dir(dockerfile: Path) -> Path | None: + """Return the ready cached rootfs for ``dockerfile``, if one exists.""" + digest = _dockerfile_hash(dockerfile) + base = util.cache_dir() / "rootfs" / f"agent-{digest}" + return base if (base / ".bb-ready").is_file() else None + + def build_agent_rootfs_dir( dockerfile: Path, *, image_tag: str, smoke_test: tuple[str, ...] = (), ) -> Path: @@ -58,7 +65,7 @@ def build_agent_rootfs_dir( silent-failure image at build time rather than at first agent use.""" digest = _dockerfile_hash(dockerfile) base = util.cache_dir() / "rootfs" / f"agent-{digest}" - if (base / ".bb-ready").is_file(): + if cached_agent_rootfs_dir(dockerfile) is not None: info(f"using cached agent rootfs {base.name}") return base diff --git a/bot_bottle/backend/firecracker/launch.py b/bot_bottle/backend/firecracker/launch.py index 26d3779..8a22867 100644 --- a/bot_bottle/backend/firecracker/launch.py +++ b/bot_bottle/backend/firecracker/launch.py @@ -45,7 +45,8 @@ from ...git_gate import ( provision_git_gate_dynamic_keys, revoke_git_gate_provisioned_keys, ) -from ...log import info, warn +from ...image_cache import check_stale_path +from ...log import die, info, warn from ...supervise import SUPERVISE_PORT from ..docker.egress import EGRESS_PORT from ..util import AGENT_CA_BUNDLE, AGENT_CA_PATH @@ -64,6 +65,7 @@ _GIT_HTTP_PORT = 9420 @contextmanager def launch( plan: FirecrackerBottlePlan, + agent_base: Path, *, provision: Callable[[FirecrackerBottlePlan, "FirecrackerBottle"], str | None], ) -> Generator[FirecrackerBottle, None, None]: @@ -85,11 +87,9 @@ def launch( raise teardown_exc try: - # Step 1: agent rootfs. Built from the Dockerfile inside a Firecracker - # builder VM (buildah, no host docker); a committed snapshot is reused - # when present. Returns the base dir the per-bottle ext4 is made from. - plan, agent_base = _build_agent_base(plan) - + # Step 1 (rootfs resolution/build) runs in BottleBackend.launch before + # this context starts resources. ``agent_base`` is the selected cache, + # fresh build, or committed snapshot. # Step 2: mint the git-gate dynamic (gitea) deploy keys, if any. git_gate_plan = plan.git_gate_plan if git_gate_plan.upstreams: @@ -206,9 +206,7 @@ def launch( teardown() -def _build_agent_base( - plan: FirecrackerBottlePlan, -) -> tuple[FirecrackerBottlePlan, Path]: +def build_or_load_agent_base(plan: FirecrackerBottlePlan) -> Path: """Produce the agent's base rootfs dir. Primary path: build the Dockerfile inside a Firecracker builder VM (buildah, no host docker), smoke-testing the image before export. A committed snapshot (freeze/migrate) is resumed @@ -217,13 +215,36 @@ def _build_agent_base( committed_tar = committed_rootfs_path(plan.slug) if committed and committed_tar.is_file(): info(f"resuming from committed rootfs {committed_tar}") - return plan, util.build_committed_rootfs_dir(committed_tar) - base = image_builder.build_agent_rootfs_dir( - Path(plan.dockerfile_path), + return util.build_committed_rootfs_dir(committed_tar) + dockerfile = Path(plan.dockerfile_path) + if plan.spec.image_policy == "cached": + cached = image_builder.cached_agent_rootfs_dir(dockerfile) + if cached is None: + die( + f"cached agent rootfs for {plan.image!r} not found; " + "run without --cached-images to build it" + ) + info(f"using cached agent rootfs {cached.name}") + return cached + return image_builder.build_agent_rootfs_dir( + dockerfile, image_tag=plan.image, smoke_test=runtime_for(plan.agent_provider_template).smoke_test, ) - return plan, base + + +def stale_checks(plan: FirecrackerBottlePlan) -> None: + """Raise when the cached rootfs selected by this plan is stale.""" + if plan.spec.image_policy != "cached": + return + committed = read_committed_image(plan.slug) + committed_tar = committed_rootfs_path(plan.slug) + 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)) + if cached is not None: + check_stale_path(f"agent rootfs {cached}", cached / ".bb-ready") # --- agent guest env ------------------------------------------------- diff --git a/tests/unit/test_firecracker_image_builder.py b/tests/unit/test_firecracker_image_builder.py index 40e132c..7ee483b 100644 --- a/tests/unit/test_firecracker_image_builder.py +++ b/tests/unit/test_firecracker_image_builder.py @@ -35,6 +35,15 @@ class TestBuildAgentRootfsDir(unittest.TestCase): build.assert_not_called() self.assertEqual(base, out) + def test_cached_lookup_requires_ready_marker(self): + digest = image_builder._dockerfile_hash(self.dockerfile) + base = self.cache / "rootfs" / f"agent-{digest}" + base.mkdir(parents=True) + with patch.object(image_builder.util, "cache_dir", return_value=self.cache): + self.assertIsNone(image_builder.cached_agent_rootfs_dir(self.dockerfile)) + (base / ".bb-ready").write_text("ok\n") + self.assertEqual(base, image_builder.cached_agent_rootfs_dir(self.dockerfile)) + def test_cache_miss_builds_injects_and_marks_ready(self): with patch.object(image_builder.util, "cache_dir", return_value=self.cache), \ patch.object(image_builder, "_build_in_infra") as build, \ diff --git a/tests/unit/test_stale_checks.py b/tests/unit/test_stale_checks.py index d1efcf1..9aa0187 100644 --- a/tests/unit/test_stale_checks.py +++ b/tests/unit/test_stale_checks.py @@ -7,6 +7,7 @@ calls are mocked at the module boundary.""" from __future__ import annotations import unittest +from pathlib import Path from types import SimpleNamespace from typing import Any, cast from unittest.mock import MagicMock, patch @@ -204,5 +205,61 @@ class TestMacosContainerStaleChecks(unittest.TestCase): sc.assert_called_once_with(plan) +# --------------------------------------------------------------------------- +# Firecracker cached rootfs selection and stale checks +# --------------------------------------------------------------------------- + +class TestFirecrackerCachedRootfs(unittest.TestCase): + def _plan(self, policy: str = "cached") -> Any: + return cast(Any, SimpleNamespace( + spec=SimpleNamespace(image_policy=policy), + slug="dev-abc", + image="bot-bottle-agent:latest", + dockerfile_path="/repo/Dockerfile", + agent_provider_template="claude", + )) + + def test_cached_policy_reuses_ready_rootfs_without_building(self) -> None: + from bot_bottle.backend.firecracker import launch as mod + cached = Path("/cache/rootfs/agent-deadbeef") + with patch.object(mod, "read_committed_image", return_value=""), \ + patch.object( + mod.image_builder, "cached_agent_rootfs_dir", return_value=cached, + ), \ + patch.object(mod.image_builder, "build_agent_rootfs_dir") as build: + self.assertEqual(cached, mod.build_or_load_agent_base(self._plan())) + build.assert_not_called() + + def test_cached_policy_fails_when_rootfs_is_missing(self) -> None: + from bot_bottle.backend.firecracker import launch as mod + with patch.object(mod, "read_committed_image", return_value=""), \ + patch.object( + mod.image_builder, "cached_agent_rootfs_dir", return_value=None, + ), \ + patch.object(mod.image_builder, "build_agent_rootfs_dir") as build, \ + self.assertRaises(SystemExit): + mod.build_or_load_agent_base(self._plan()) + build.assert_not_called() + + def test_stale_checks_use_ready_marker_timestamp(self) -> None: + from bot_bottle.backend.firecracker import launch as mod + cached = Path("/cache/rootfs/agent-deadbeef") + with patch.object(mod, "read_committed_image", return_value=""), \ + patch.object( + mod.image_builder, "cached_agent_rootfs_dir", return_value=cached, + ), \ + patch.object(mod, "check_stale_path") as check: + mod.stale_checks(self._plan()) + check.assert_called_once_with(f"agent rootfs {cached}", cached / ".bb-ready") + + def test_backend_prelaunch_checks_delegates(self) -> None: + from bot_bottle.backend.firecracker.backend import FirecrackerBottleBackend + from bot_bottle.backend.firecracker import launch as mod + plan = self._plan() + with patch.object(mod, "stale_checks") as checks: + FirecrackerBottleBackend().prelaunch_checks(plan) + checks.assert_called_once_with(plan) + + if __name__ == "__main__": unittest.main()