adb46d4927
16 total. The intl entries need real egress (usually blocked on this host) but cost only 10s each with the idle watchdog on networks that have them. Also noted: a 10s watchdog can kill sources mid slow negotiation (rat.dev emitted 'Pulling fs layer' then went quiet) -- EVALHARNESS_PULL_IDLE_S=30 softens that. Co-Authored-By: Claude <noreply@anthropic.com>
246 lines
11 KiB
Python
246 lines
11 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 + international mainstream fallbacks. History on this host: daocloud
|
|
# (139s) and 1ms.run (342s + trickle) have delivered this 9GB image; the
|
|
# international entries (Google's hub cache, AWS ECR public) need real
|
|
# egress -- usually blocked here, kept for networks that have it.
|
|
_CN_MIRROR_FALLBACKS = [
|
|
'{img}', # daemon default (uses its own registry-mirrors config)
|
|
'docker.m.daocloud.io/{img}',
|
|
'docker.1ms.run/{img}',
|
|
'docker.1panel.live/{img}',
|
|
'hub.rat.dev/{img}',
|
|
'docker.xuanyuan.me/{img}',
|
|
'dockerhub.icu/{img}',
|
|
'dockerproxy.com/{img}',
|
|
'hub-mirror.c.163.com/{img}',
|
|
'mirror.baidubce.com/{img}',
|
|
'docker.mirrors.ustc.edu.cn/{img}',
|
|
'docker.mirrors.sjtug.sjtu.edu.cn/{img}',
|
|
'docker.nju.edu.cn/{img}',
|
|
'docker.mirrors.tuna.tsinghua.edu.cn/{img}',
|
|
'mirror.gcr.io/{img}', # Google's docker-hub cache (intl)
|
|
'public.ecr.aws/docker/{img}', # AWS ECR public (library/* coverage)
|
|
]
|
|
# 用户/环境可整体覆盖（逗号分隔模板，{img} 占位）
|
|
import os as _os
|
|
|
|
_env = _os.environ.get('EVALHARNESS_DOCKER_MIRRORS', '').strip()
|
|
if _env:
|
|
_CN_MIRROR_FALLBACKS = [t.strip() for t in _env.split(',') if t.strip()]
|
|
|
|
|
|
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, verbose: bool = True) -> None:
|
|
"""Pull via CN-mirror fallback chain; retag to the canonical name on hit.
|
|
|
|
A multi-GB pull runs silently for many minutes with capture_output --
|
|
which reads as a hang -- so say WHICH source is being tried and how
|
|
long it took."""
|
|
import sys as _sys
|
|
import time as _time
|
|
last_err = None
|
|
for template in _CN_MIRROR_FALLBACKS:
|
|
ref = template.format(img=image)
|
|
if verbose:
|
|
src = 'docker.io(daemon mirrors)' if '{img}' in template else template.split('/{img}')[0]
|
|
print(f'· pulling sandbox image {image} via {src} ...',
|
|
file=_sys.stderr, flush=True)
|
|
t0 = _time.monotonic()
|
|
# stream progress: a multi-GB pull with captured output is minutes
|
|
# of silence that reads as a hang. Non-TTY docker pull emits clean
|
|
# per-layer lines; forward the informative ones (throttled).
|
|
proc = subprocess.Popen(['docker', 'pull', ref],
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
bufsize=0)
|
|
# idle watchdog: a dead mirror emits NOTHING (0 bytes/s for hours);
|
|
# the old fixed 3600s timeout let one dead source stall the whole
|
|
# chain. 5 minutes of silence = kill and try the next mirror.
|
|
import os as _os
|
|
import selectors as _sel
|
|
|
|
_selr = _sel.DefaultSelector()
|
|
_selr.register(proc.stdout, _sel.EVENT_READ)
|
|
_fd = proc.stdout.fileno()
|
|
last_line = ''
|
|
last_show = 0.0
|
|
layer_prog = {} # layer id -> (done_bytes, total_bytes): the raw
|
|
# per-layer events give no sense of OVERALL progress -- sum them
|
|
import os as _os
|
|
import re as _re
|
|
|
|
def _bytes(s):
|
|
s = s.strip()
|
|
mult = {'kB': 1e3, 'KB': 1e3, 'MB': 1e6, 'GB': 1e9, 'B': 1}
|
|
for suf, m in mult.items():
|
|
if s.endswith(suf):
|
|
return int(float(s[:-len(suf)]) * m)
|
|
return int(float(s or 0))
|
|
|
|
# docker progress lines end with \r (carriage-return refresh), not
|
|
# \n -- readline() buffered them until a rare \n arrived, so the
|
|
# Downloading snapshots (and the ⏳ summary fed by them) never
|
|
# surfaced. Split on BOTH terminators, chunked reads.
|
|
buf = ''
|
|
_lines_iter = iter(lambda: None, 1) # placeholder, replaced below
|
|
|
|
# no output at all -> dead mirror, switch source. ANY output resets
|
|
# the timer (manifest/layer announcements/DOWNLOADING refreshes), so
|
|
# this only fires on sources that print nothing. 10s default;
|
|
# override with EVALHARNESS_PULL_IDLE_S for slow-negotiating mirrors.
|
|
import os as _os2
|
|
|
|
_IDLE_S = float(_os2.environ.get('EVALHARNESS_PULL_IDLE_S', '10'))
|
|
|
|
def _iter_lines():
|
|
nonlocal buf
|
|
while True:
|
|
if not _selr.select(timeout=_IDLE_S):
|
|
proc.kill()
|
|
if verbose:
|
|
print(f'· no data for {_IDLE_S:.0f}s -- dead mirror, '
|
|
'trying next source ...',
|
|
file=_sys.stderr, flush=True)
|
|
return
|
|
chunk = _os.read(_fd, 4096)
|
|
if not chunk:
|
|
if buf:
|
|
yield buf
|
|
return
|
|
buf += chunk.decode('utf-8', 'replace')
|
|
parts = buf.split('\r') if '\r' in buf else buf.split('\n')
|
|
if len(parts) > 1:
|
|
for p in parts[:-1]:
|
|
yield p
|
|
buf = parts[-1]
|
|
|
|
for line in _iter_lines():
|
|
line = line.strip('\n ')
|
|
last_line = line
|
|
m = _re.match(r'^([0-9a-f]{12}): Downloading.*?([\d.]+[kKMG]?B)/([\d.]+[kKMG]?B)', line)
|
|
if m:
|
|
layer_prog[m.group(1)] = (_bytes(m.group(2)), _bytes(m.group(3)))
|
|
if 'Pull complete' in line:
|
|
lid = line.split(':')[0]
|
|
if lid in layer_prog:
|
|
layer_prog[lid] = (layer_prog[lid][1], layer_prog[lid][1])
|
|
now = _time.monotonic()
|
|
key_evt = any(k in line for k in
|
|
('Pulling from', 'Status', 'error', 'denied'))
|
|
if verbose and key_evt:
|
|
print(f' {last_line[:100]}', file=_sys.stderr, flush=True)
|
|
last_show = now
|
|
elif verbose and now - last_show > 3 and layer_prog:
|
|
done = sum(v[0] for v in layer_prog.values())
|
|
tot = sum(v[1] for v in layer_prog.values())
|
|
if tot:
|
|
pct = done / tot * 100
|
|
big = max(layer_prog.items(), key=lambda kv: kv[1][1] - kv[1][0])
|
|
print(f' ⏳ {done / 1e9:.2f}/{tot / 1e9:.2f} GB ({pct:.0f}%)'
|
|
f' · 最大层 {big[0]}: {big[1][0] / 1e9:.2f}/{big[1][1] / 1e9:.2f} GB',
|
|
file=_sys.stderr, flush=True)
|
|
last_show = now
|
|
proc.wait(timeout=3600)
|
|
r = subprocess.CompletedProcess(proc.args, proc.returncode,
|
|
stdout='', stderr=last_line)
|
|
if r.returncode == 0 and verbose:
|
|
print(f'✓ image ready in {_time.monotonic() - t0:.0f}s',
|
|
file=_sys.stderr, flush=True)
|
|
if r.returncode == 0:
|
|
if ref != image:
|
|
# retag BEFORE rmi and VERIFY: a silent tag failure followed
|
|
# by rmi deleted the only tag and lost a 342s multi-GB pull
|
|
tag = subprocess.run(['docker', 'tag', ref, image],
|
|
capture_output=True, text=True)
|
|
if tag.returncode != 0:
|
|
print(f'!! docker tag {ref} -> {image} failed: '
|
|
f'{(tag.stderr or "")[:120]} (keeping the mirror tag)',
|
|
file=_sys.stderr, flush=True)
|
|
return # do NOT rmi: the mirror tag is the only handle
|
|
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}')
|