"""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, 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}')