diff --git a/cookbook/capability_routing_pilot/poc/run_profiling.py b/cookbook/capability_routing_pilot/poc/run_profiling.py index 13b20441b96..a4220f31a56 100644 --- a/cookbook/capability_routing_pilot/poc/run_profiling.py +++ b/cookbook/capability_routing_pilot/poc/run_profiling.py @@ -573,10 +573,13 @@ def _tag_image(source: str, target: str) -> None: subprocess.run(["docker", "tag", source, target], check=True, capture_output=True, text=True) -def _build_image(root_dir: Path, tag: str, dockerfile: str) -> None: - subprocess.run( - ["docker", "build", "--tag", tag, "--file", "-", "."], - cwd=root_dir, +def _build_image( + tag: str, + dockerfile: str, + runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run, +) -> None: + runner( + ["docker", "build", "--tag", tag, "-"], input=dockerfile, check=True, capture_output=True, @@ -819,7 +822,7 @@ def _pull_batch( task: Final = batch.tasks[0] remaining: Final = ProfileBatch(batch.index, batch.tasks[1:]) if preflight_rows.get(task.task_id, {}).get("record_type") == "preflight_failed": - return _pull_batch(config, remaining, ops, preflight_rows) + return _pull_batch(config, remaining, ops, preflight_rows, prepared, owned) if task.image_ref is None: if task.task_id not in preflight_rows: _append_preflight(config, task, "preflight_failed", "env", None, None, 0.0) @@ -984,7 +987,7 @@ def _default_ops( inspect_image=_inspect_image, pull_image=_pull_image, tag_image=_tag_image, - build_image=lambda tag, dockerfile: _build_image(config.root_dir, tag, dockerfile), + build_image=lambda tag, dockerfile: _build_image(tag, dockerfile), evaluate=lambda task_id, patch, label, paths: _evaluate(task_id, patch, label, paths, env), remove_image=_remove_image, image_exists=lambda reference: _inspect_image(reference) is not None, diff --git a/cookbook/capability_routing_pilot/tests/test_run_profiling.py b/cookbook/capability_routing_pilot/tests/test_run_profiling.py index 1ea6eea2b42..20558a67cfc 100644 --- a/cookbook/capability_routing_pilot/tests/test_run_profiling.py +++ b/cookbook/capability_routing_pilot/tests/test_run_profiling.py @@ -1,6 +1,7 @@ import hashlib import json from pathlib import Path +from unittest.mock import Mock import pytest @@ -14,6 +15,7 @@ from poc.run_profiling import ( ProfileTask, RuntimeOps, TransientInfraError, + _build_image, _cleanup_plan, _pull_batch, _guard_sha256, @@ -385,6 +387,55 @@ def test_transient_pull_leaves_rows_absent_and_retry_prepares_task(tmp_path: Pat assert plan.prepared[0].final_image_id == "upstream-id" +def test_preflight_failed_task_preserves_later_pull_plan_and_cleanup(tmp_path: Path) -> None: + first = ProfileTask(0, "task-0", "repository_debugging", model_order_for_task(0), "repo/first@sha256:a", "sha256:a") + later = ProfileTask(1, "task-1", "repository_debugging", model_order_for_task(1), "repo/later@sha256:b", "sha256:b") + batch = ProfileBatch(1, (first, later)) + config = ProfileConfig(tmp_path, tmp_path / "manifest.json", "manifest-sha") + images: dict[str, str] = {} + + def tag(source: str, target: str) -> None: + images[target] = images[source] + + operations = RuntimeOps( + disk_stats=lambda: {"free_bytes": 100 * 1024**3}, + inspect_image=lambda reference: images.get(reference), + pull_image=lambda reference: images.__setitem__(reference, reference), + tag_image=tag, + build_image=lambda _tag, _dockerfile: None, + evaluate=lambda *_args: ("unresolved", 0.1), + remove_image=lambda reference: not images.pop(reference, None), + image_exists=lambda reference: reference in images, + attempt=lambda *_args: {}, + ) + plan = _pull_batch( + config, + batch, + operations, + {"task-0": {"record_type": "preflight_failed"}}, + ) + + assert tuple(task_image.task_id for task_image in plan.prepared) == ("task-1",) + assert plan.owned_refs == ("repo/later@sha256:b", "repo/later:latest") + _cleanup_plan(config, batch, plan, operations) + assert not images + + +def test_overlay_build_uses_empty_docker_context() -> None: + runner = Mock() + dockerfile = "FROM repo/image@sha256:digest\nRUN true\n" + + _build_image("routing_poc/profiling-overlay:task-0", dockerfile, runner) + + runner.assert_called_once_with( + ["docker", "build", "--tag", "routing_poc/profiling-overlay:task-0", "-"], + input=dockerfile, + check=True, + capture_output=True, + text=True, + ) + + def test_profile_guard_enforces_phase_cap_and_pre_pull_disk_reserve(tmp_path: Path) -> None: config = _config(tmp_path, ("task-0",)) guard = profile_guard(