5e98f6b701
The preflight did a bare 'docker pull' -- docker.io is slow/unreachable from CN without luck; it now walks the same fallback chain as the SWE prefetch (daemon mirrors -> daocloud -> 1ms.run -> baidubce -> sjtu -> rat.dev), retagging the hit to the canonical name. Co-Authored-By: Claude <noreply@anthropic.com>
110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
"""Image prefetch: parallel docker pull for execution environments.
|
|
|
|
SWE-bench Verified declares ~500 per-instance images (sweb.eval.x86_64.*).
|
|
Pulling them lazily during an eval run would stall it serially; prefetch
|
|
pulls them ahead of time with bounded concurrency and progress reporting.
|
|
|
|
from evalharness.sandbox.prefetch import prefetch_images
|
|
done = prefetch_images(['sweb.eval.x86_64.django__django-12345', ...], workers=8)
|
|
# CLI: evalharness sandbox prefetch swe_bench_verified --workers 8 --limit 50
|
|
"""
|
|
|
|
import subprocess
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Iterable, List
|
|
|
|
from ..data.dataset import Dataset
|
|
|
|
# CN mirrors tried in order before/alongside the daemon's configured mirrors.
|
|
# Some namespaces (e.g. swebench/*) are blocked by individual CN mirrors, so we
|
|
# fall through: daemon default -> 1ms.run -> baidubce -> sjtug.
|
|
_CN_MIRROR_FALLBACKS = [
|
|
'{img}', # daemon default (uses its own registry-mirrors config)
|
|
'docker.m.daocloud.io/{img}',
|
|
'docker.1ms.run/{img}',
|
|
'mirror.baidubce.com/{img}',
|
|
'docker.mirrors.sjtug.sjtu.edu.cn/{img}',
|
|
'hub.rat.dev/{img}',
|
|
]
|
|
|
|
|
|
def images_for_dataset(ds: Dataset, limit: int = 0) -> List[str]:
|
|
"""Collect distinct sandbox images declared by a dataset's samples."""
|
|
seen: List[str] = []
|
|
add = seen.append
|
|
for i, s in enumerate(ds):
|
|
if limit and i >= limit:
|
|
break
|
|
if s.sandbox and s.sandbox.image and s.sandbox.image not in seen:
|
|
add(s.sandbox.image)
|
|
return seen
|
|
|
|
|
|
def images_for_samples(samples) -> List[str]:
|
|
"""Distinct sandbox images in sample order (dedup, keeps order)."""
|
|
seen: List[str] = []
|
|
for s in samples:
|
|
if s.sandbox and s.sandbox.image and s.sandbox.image not in seen:
|
|
seen.append(s.sandbox.image)
|
|
return seen
|
|
|
|
|
|
def local_images() -> set:
|
|
out = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
|
|
capture_output=True, text=True)
|
|
return {line for line in out.stdout.splitlines() if line}
|
|
|
|
|
|
def prefetch_images(images: Iterable[str], workers: int = 8) -> List[str]:
|
|
"""Pull images with bounded concurrency; skip ones already local.
|
|
|
|
Returns the list newly pulled. Failures are reported and skipped (one
|
|
missing instance must not block the rest).
|
|
"""
|
|
have = local_images()
|
|
todo = [img for img in images if img not in have]
|
|
if not todo:
|
|
print(f'prefetch: all {len(images)} images already local')
|
|
return []
|
|
print(f'prefetch: {len(todo)} to pull ({len(images) - len(todo)} already local), '
|
|
f'workers={workers}')
|
|
pulled: List[str] = []
|
|
failed: List[str] = []
|
|
t0 = time.time()
|
|
done = 0
|
|
with ThreadPoolExecutor(max_workers=workers) as pool:
|
|
futs = {pool.submit(_pull_one, img): img for img in todo}
|
|
for fut in as_completed(futs):
|
|
img = futs[fut]
|
|
done += 1
|
|
try:
|
|
fut.result()
|
|
pulled.append(img)
|
|
except Exception as e:
|
|
failed.append(img)
|
|
print(f' FAIL {img}: {str(e)[:120]}', flush=True)
|
|
if done % 10 == 0 or done == len(todo):
|
|
rate = done / max(time.time() - t0, 1e-6)
|
|
print(f' [{done}/{len(todo)}] {rate:.2f} imgs/s, '
|
|
f'{len(failed)} failed', flush=True)
|
|
print(f'prefetch done: {len(pulled)} pulled, {len(failed)} failed '
|
|
f'in {time.time() - t0:.0f}s')
|
|
return pulled
|
|
|
|
|
|
def _pull_one(image: str) -> None:
|
|
"""Pull via CN-mirror fallback chain; retag to the canonical name on hit."""
|
|
last_err = None
|
|
for template in _CN_MIRROR_FALLBACKS:
|
|
ref = template.format(img=image)
|
|
r = subprocess.run(['docker', 'pull', ref], capture_output=True, text=True,
|
|
timeout=3600)
|
|
if r.returncode == 0:
|
|
if ref != image: # retag the mirrored pull to the canonical name
|
|
subprocess.run(['docker', 'tag', ref, image], check=False)
|
|
subprocess.run(['docker', 'rmi', ref], check=False)
|
|
return
|
|
last_err = (ref, (r.stderr or '').strip()[:150])
|
|
raise RuntimeError(f'all mirrors failed for {image}: {last_err}')
|