Skip to content

Commit da309c7

Browse files
huangruitenghuangruiteng
authored
Make SkillsBench uv bootstrap cache-first (#1741)
1 parent ae8a546 commit da309c7

3 files changed

Lines changed: 194 additions & 2 deletions

File tree

‎examples/skillsbench-uv-cache-smoke.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
#!/usr/bin/env python3
2+
from __future__ import annotations
3+
4+
import sys
5+
import tempfile
6+
from pathlib import Path
7+
8+
ROOT = Path(__file__).resolve().parents[1]
9+
if str(ROOT) not in sys.path:
10+
sys.path.insert(0, str(ROOT))
11+
12+
from loopx.benchmark_adapters import skillsbench_uv_cache as uv_cache # noqa: E402
13+
from scripts import skillsbench_automation_loop as runner # noqa: E402
14+
15+
16+
def test_staged_dockerfile_prefers_host_uv_binary_cache() -> None:
17+
with tempfile.TemporaryDirectory(prefix="skillsbench-uv-cache-smoke-") as tmp:
18+
root = Path(tmp)
19+
fake_bin = root / "bin"
20+
fake_bin.mkdir()
21+
fake_uv = fake_bin / "uv"
22+
fake_uvx = fake_bin / "uvx"
23+
fake_uv.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8")
24+
fake_uvx.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8")
25+
fake_uv.chmod(0o755)
26+
fake_uvx.chmod(0o755)
27+
28+
task = root / "tasks" / "pddl-tpp-planning"
29+
dockerfile = task / "environment" / "Dockerfile"
30+
dockerfile.parent.mkdir(parents=True)
31+
dockerfile.write_text(
32+
"FROM python:3.12-slim\n"
33+
"RUN curl -LsSf https://astral.sh/uv/0.9.22/install.sh | sh && \\\n"
34+
" install -m 0755 ${HOME}/.local/bin/uv /usr/local/bin/uv && \\\n"
35+
" install -m 0755 ${HOME}/.local/bin/uvx /usr/local/bin/uvx\n",
36+
encoding="utf-8",
37+
)
38+
(task / "task.toml").write_text("version = \"1.1\"\n", encoding="utf-8")
39+
40+
original_candidates = uv_cache.host_uv_binary_cache_candidates
41+
try:
42+
uv_cache.host_uv_binary_cache_candidates = lambda: {
43+
"uv": fake_uv,
44+
"uvx": fake_uvx,
45+
}
46+
staged_path, metadata = runner.stage_task_for_sandbox(
47+
task_path=task,
48+
jobs_dir=root / "jobs",
49+
job_name="pddl-tpp-planning-goal",
50+
sandbox="docker",
51+
include_task_skills=False,
52+
)
53+
finally:
54+
uv_cache.host_uv_binary_cache_candidates = original_candidates
55+
56+
assert metadata["dockerfile_uv_binary_cache_context_created"] is True
57+
assert metadata["dockerfile_uv_binary_cache_available"] is True
58+
assert metadata["dockerfile_uv_binary_cache_binary_count"] == 2
59+
assert metadata["dockerfile_uv_binary_cache_has_uv"] is True
60+
assert metadata["dockerfile_uv_binary_cache_has_uvx"] is True
61+
assert metadata["dockerfile_uv_binary_cache_dockerfile_patch_applied"] is True
62+
assert metadata["dockerfile_uv_binary_cache_raw_path_recorded"] is False
63+
64+
cache_dir = staged_path / "environment" / "loopx_uv_cache"
65+
assert (cache_dir / "uv").exists()
66+
assert (cache_dir / "uvx").exists()
67+
staged_text = (staged_path / "environment" / "Dockerfile").read_text(
68+
encoding="utf-8"
69+
)
70+
assert uv_cache.DOCKER_UV_BINARY_CACHE_BEGIN in staged_text
71+
assert "COPY loopx_uv_cache/ /opt/loopx_uv_cache/" in staged_text
72+
assert "uv --version >/dev/null 2>&1" in staged_text
73+
assert "python3 -m pip install ${loopx_pip_break_system_packages}" in staged_text
74+
75+
76+
if __name__ == "__main__":
77+
test_staged_dockerfile_prefers_host_uv_binary_cache()
78+
print("skillsbench-uv-cache-smoke: ok")

‎loopx/benchmark_adapters/skillsbench_uv_cache.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
from __future__ import annotations
2+
3+
import os
4+
import shutil
5+
import sys
6+
from pathlib import Path
7+
from typing import Any
8+
9+
DOCKER_UV_BINARY_CACHE_CONTEXT_DIR = "loopx_uv_cache"
10+
DOCKER_UV_BINARY_CACHE_BEGIN = "# BEGIN LOOPX_SKILLSBENCH_UV_BINARY_CACHE"
11+
DOCKER_UV_BINARY_CACHE_END = "# END LOOPX_SKILLSBENCH_UV_BINARY_CACHE"
12+
UV_BINARY_CACHE_KEYS = (
13+
"dockerfile_uv_binary_cache_context_created",
14+
"dockerfile_uv_binary_cache_available",
15+
"dockerfile_uv_binary_cache_binary_count",
16+
"dockerfile_uv_binary_cache_has_uv",
17+
"dockerfile_uv_binary_cache_has_uvx",
18+
"dockerfile_uv_binary_cache_dockerfile_patch_applied",
19+
"dockerfile_uv_binary_cache_raw_path_recorded",
20+
)
21+
22+
23+
def empty_uv_binary_cache_metadata() -> dict[str, Any]:
24+
return {
25+
"dockerfile_uv_binary_cache_context_created": False,
26+
"dockerfile_uv_binary_cache_available": False,
27+
"dockerfile_uv_binary_cache_binary_count": 0,
28+
"dockerfile_uv_binary_cache_has_uv": False,
29+
"dockerfile_uv_binary_cache_has_uvx": False,
30+
"dockerfile_uv_binary_cache_dockerfile_patch_applied": False,
31+
"dockerfile_uv_binary_cache_raw_path_recorded": False,
32+
}
33+
34+
35+
def dockerfile_uv_binary_cache_prelude() -> str:
36+
return (
37+
f"COPY {DOCKER_UV_BINARY_CACHE_CONTEXT_DIR}/ /opt/loopx_uv_cache/\n"
38+
"RUN set -eux; \\\n"
39+
f" : \"{DOCKER_UV_BINARY_CACHE_BEGIN}\"; \\\n"
40+
" if [ -x /opt/loopx_uv_cache/uv ]; then install -m 0755 /opt/loopx_uv_cache/uv /usr/local/bin/uv; fi; \\\n"
41+
" if [ -x /opt/loopx_uv_cache/uvx ]; then install -m 0755 /opt/loopx_uv_cache/uvx /usr/local/bin/uvx; fi; \\\n"
42+
" if command -v uv >/dev/null 2>&1 && command -v uvx >/dev/null 2>&1 && uv --version >/dev/null 2>&1 && uvx --version >/dev/null 2>&1; then exit 0; fi; \\\n"
43+
" rm -f /usr/local/bin/uv /usr/local/bin/uvx; \\\n"
44+
f" : \"{DOCKER_UV_BINARY_CACHE_END}\"; \\\n"
45+
)
46+
47+
48+
def host_uv_binary_cache_candidates() -> dict[str, Path]:
49+
"""Return host uv binaries that are safe to copy into Linux Docker images."""
50+
51+
if not sys.platform.startswith("linux"):
52+
return {}
53+
machine = os.uname().machine.lower()
54+
if machine not in {"x86_64", "amd64"}:
55+
return {}
56+
candidates: dict[str, Path] = {}
57+
for name in ("uv", "uvx"):
58+
raw = shutil.which(name)
59+
if not raw:
60+
continue
61+
path = Path(raw)
62+
if path.is_file() and os.access(path, os.X_OK):
63+
candidates[name] = path
64+
return candidates
65+
66+
67+
def stage_uv_binary_cache_context(environment_dir: Path) -> dict[str, Any]:
68+
"""Stage optional host uv/uvx binaries into the Docker build context."""
69+
70+
cache_dir = environment_dir / DOCKER_UV_BINARY_CACHE_CONTEXT_DIR
71+
if cache_dir.exists():
72+
shutil.rmtree(cache_dir)
73+
cache_dir.mkdir(parents=True, exist_ok=True)
74+
(cache_dir / ".loopx_keep").write_text(
75+
"Optional LoopX uv binary cache for SkillsBench Docker bootstrap.\n",
76+
encoding="utf-8",
77+
)
78+
copied: list[str] = []
79+
for name, source in host_uv_binary_cache_candidates().items():
80+
target = cache_dir / name
81+
shutil.copy2(source, target)
82+
target.chmod(0o755)
83+
copied.append(name)
84+
return {
85+
"dockerfile_uv_binary_cache_context_created": True,
86+
"dockerfile_uv_binary_cache_available": bool(copied),
87+
"dockerfile_uv_binary_cache_binary_count": len(copied),
88+
"dockerfile_uv_binary_cache_has_uv": "uv" in copied,
89+
"dockerfile_uv_binary_cache_has_uvx": "uvx" in copied,
90+
"dockerfile_uv_binary_cache_raw_path_recorded": False,
91+
}
92+
93+
94+
def discover_uv_binary_cache_metadata(
95+
prepared_task: Path,
96+
dockerfile_text: str,
97+
) -> dict[str, Any]:
98+
cache_dir = prepared_task / "environment" / DOCKER_UV_BINARY_CACHE_CONTEXT_DIR
99+
has_uv = (cache_dir / "uv").exists()
100+
has_uvx = (cache_dir / "uvx").exists()
101+
return {
102+
"dockerfile_uv_binary_cache_context_created": cache_dir.exists(),
103+
"dockerfile_uv_binary_cache_available": has_uv or has_uvx,
104+
"dockerfile_uv_binary_cache_binary_count": int(has_uv) + int(has_uvx),
105+
"dockerfile_uv_binary_cache_has_uv": has_uv,
106+
"dockerfile_uv_binary_cache_has_uvx": has_uvx,
107+
"dockerfile_uv_binary_cache_dockerfile_patch_applied": (
108+
DOCKER_UV_BINARY_CACHE_BEGIN in dockerfile_text
109+
),
110+
"dockerfile_uv_binary_cache_raw_path_recorded": False,
111+
}

‎scripts/skillsbench_automation_loop.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@
116116
from loopx.benchmark_adapters.skillsbench_task_source import ( # noqa: E402
117117
classify_missing_task_source,
118118
)
119-
from loopx.benchmark_adapters import skillsbench_runner_source as runner_source # noqa: E402
119+
from loopx.benchmark_adapters import skillsbench_runner_source as runner_source, skillsbench_uv_cache as uv_cache # noqa: E402
120120
from loopx.benchmark_adapters.skillsbench_acp_relay import ( # noqa: E402
121121
SKILLSBENCH_LOCAL_ACP_RELAY_BRIDGE_PREFLIGHT_MARKER,
122122
SKILLSBENCH_LOCAL_ACP_RELAY_BRIDGE_PREFLIGHT_PROMPT,
@@ -6277,6 +6277,7 @@ def _discover_prepared_task_staging(plan: dict[str, Any]) -> dict[str, Any]:
62776277
and "python3 -m pip install" in dockerfile_text
62786278
and "uv==${LOOPX_SKILLSBENCH_UV_VERSION}" in dockerfile_text
62796279
),
6280+
**uv_cache.discover_uv_binary_cache_metadata(prepared_task, dockerfile_text),
62806281
"dockerfile_uv_bootstrap_mirror_host": (
62816282
DEFAULT_VERIFIER_UV_RELEASE_MIRROR_HOST
62826283
if DOCKER_UV_BOOTSTRAP_MIRROR_BEGIN in dockerfile_text
@@ -7024,7 +7025,7 @@ def patch_dockerfile_uv_bootstrap_mirror(dockerfile: Path) -> dict[str, Any]:
70247025
"# official uv installer as a bounded mirror-backed fallback.\n"
70257026
f"ARG LOOPX_SKILLSBENCH_UV_RELEASE_MIRROR={DEFAULT_VERIFIER_UV_RELEASE_MIRROR_BASE}\n"
70267027
f"ARG LOOPX_SKILLSBENCH_UV_VERSION={version}\n"
7027-
"RUN set -eux; \\\n"
7028+
f"{uv_cache.dockerfile_uv_binary_cache_prelude()}"
70287029
" if command -v python3 >/dev/null 2>&1; then \\\n"
70297030
" loopx_pip_break_system_packages=''; \\\n"
70307031
" if python3 -m pip install --help 2>/dev/null | grep -q -- '--break-system-packages'; then \\\n"
@@ -7077,6 +7078,7 @@ def patch_dockerfile_uv_bootstrap_mirror(dockerfile: Path) -> dict[str, Any]:
70777078
"dockerfile_uv_bootstrap_mirror_patch_required": True,
70787079
"dockerfile_uv_bootstrap_mirror_patch_applied": True,
70797080
"dockerfile_uv_bootstrap_pip_fallback_patch_applied": True,
7081+
"dockerfile_uv_binary_cache_dockerfile_patch_applied": True,
70807082
"dockerfile_uv_bootstrap_version": version,
70817083
"dockerfile_uv_bootstrap_mirror_host": (
70827084
DEFAULT_VERIFIER_UV_RELEASE_MIRROR_HOST
@@ -8606,6 +8608,7 @@ def stage_task_for_sandbox(
86068608
)
86078609
)
86088610
),
8611+
**uv_cache.stage_uv_binary_cache_context(staged_path / "environment"),
86098612
"dockerfile_uv_bootstrap_mirror_host": (
86108613
str(
86118614
dockerfile_uv_mirror_metadata.get(

0 commit comments

Comments
 (0)