"""Background sandbox-image service: pull-ahead + wait-if-missing. One service per process. Benches declare the images they will need (the runner registers them when a task arrives); a small worker pool pulls them in the background from the CN mirror chain. Sample execution calls ``ensure``/``wait_ready`` BEFORE ``docker run``: an image still in flight BLOCKS that sample (no failure, no timeout at 120s) until the pull service delivers it or exhausts all sources. This replaces the two broken timings: - gen-phase `docker run` implicit pull (zero output, 120s timeout kill) - score-phase batch pull that ran AFTER generation already failed. """ import json import os import queue import subprocess import threading from typing import Dict, Iterable, List, Optional def _docker(args, timeout=300): return subprocess.run(['docker'] + args, capture_output=True, text=True, timeout=timeout) _LOCAL_TTL_S = 60.0 # `docker images` on a daemon busy with saves _LOCAL_CACHE = [0.0, set()] # [expires_at, names] -- one query per TTL def _local_images() -> set: import time now = time.monotonic() if now < _LOCAL_CACHE[0]: return _LOCAL_CACHE[1] r = _docker(['images', '--format', '{{.Repository}}:{{.Tag}}'], timeout=60) names = set(r.stdout.split()) if r.returncode == 0 else set() _LOCAL_CACHE[0] = now + _LOCAL_TTL_S _LOCAL_CACHE[1] = names return names class ImageService: """Pull-ahead pool + readiness barrier, one per process.""" def __init__(self, workers: int = 2, verbose: bool = True): self.workers = workers self.verbose = verbose self._wanted: Dict[str, str] = {} # image -> 'pulling'|'ready'|'failed' self._lock = threading.Lock() self._events: Dict[str, threading.Event] = {} self._queue: 'queue.Queue[Optional[str]]' = queue.Queue() self._threads: List[threading.Thread] = [] self._started = False self._pulled: List[str] = [] # for atexit release # ---- lifecycle ---- def start(self): if self._started: return self._started = True for i in range(self.workers): t = threading.Thread(target=self._worker, daemon=True, name=f'img-pull-{i}') t.start() self._threads.append(t) def register(self, images: Iterable[str]): """Declare upcoming images (called when tasks arrive). Idempotent and MEMORY-FIRST: the run bulk-registers everything once, so the per-task register must be a dict hit with ZERO docker calls -- `docker images` per task (x100) against a daemon busy with concurrent `docker save`s timed out at 60s and cascaded whole runs into failure. Only genuinely UNKNOWN images trigger one (TTL-cached) docker query for the batch; if even that fails under daemon pressure, mark ready optimistically and let `docker run` be the referee -- it has its own timeout and a clear error.""" self.start() unknown = [] with self._lock: for img in images: if not img: continue if img in self._wanted: continue self._wanted[img] = 'pulling' self._events[img] = threading.Event() unknown.append(img) if not unknown: return try: local = _local_images() except Exception: local = None # daemon query failed: optimistic path below for img in unknown: if local is None or img in local: if local is None: self._log(f'· docker images unavailable -- assuming ' f'local: {img.split("/")[-1][:44]}') with self._lock: self._wanted[img] = 'ready' self._events[img].set() continue self._queue.put(img) self._log(f'· image queued: {img}') def wait_ready(self, image: str, timeout_s: float = 3600) -> bool: """Block until the image is local (True) or all sources failed / timeout (False). Register-on-demand so a wait without register still works.""" self.register([image]) ev = self._events.get(image) if ev is None: return True # never registered -> assume caller handled it ok = ev.wait(timeout_s) if not ok: self._log(f'· image wait TIMEOUT {image} ({timeout_s:.0f}s)') return False with self._lock: return self._wanted.get(image) == 'ready' # ---- worker ---- def _worker(self): while True: img = self._queue.get() if img is None: return try: ok = self._pull_one(img) except Exception: ok = False with self._lock: self._wanted[img] = 'ready' if ok else 'failed' self._events[img].set() self._queue.task_done() def _pull_one(self, image: str) -> bool: # fast path: a concurrent worker / previous run may have delivered # it already (also flips any other queued-now-local images ready) if image in _local_images(): self._sweep_local() return True # 1. OUR ModelScope dataset shipment (CN CDN, per-image tars). The # swebench/* namespace is 403 on every public CN mirror, so this # is the only network source for those images. Opt out with # EVALHARNESS_MS_IMAGE_REPO='' . repo = os.environ.get('EVALHARNESS_MS_IMAGE_REPO', 'SoraAmami/swebench-images') if repo: try: from .ms_images import ms_pull if ms_pull(image, repo): self._pulled.append(image) self._sweep_local() return True except Exception: pass # 2. CN mirror chain (progress + idle watchdog inside); on hit the # puller retags to the canonical name so callers see it directly try: from .prefetch import _pull_one as net_pull rc, _ = net_pull(image) if rc == 0: self._pulled.append(image) return True except Exception: pass return False def _sweep_local(self): """A tar load makes ~10 images local at once: flip every queued 'pulling' image that is now local to 'ready' (and drop it from the release list duty) -- otherwise each waits its own full queue turn just to discover it already arrived.""" with self._lock: waiting = [img for img, st in self._wanted.items() if st == 'pulling'] if not waiting: return local = _local_images() with self._lock: for img in waiting: if img in local: self._wanted[img] = 'ready' self._events[img].set() self._pulled.append(img) # we caused it; release at exit # ---- misc ---- def _log(self, msg): if self.verbose: print(msg, flush=True) def stats(self) -> Dict[str, int]: with self._lock: vals = list(self._wanted.values()) return {'wanted': len(vals), 'ready': vals.count('ready'), 'pulling': vals.count('pulling'), 'failed': vals.count('failed')} def release_all(self): """Remove images this service pulled (atexit / bench end).""" import atexit ids = _docker(['ps', '-aq', '--filter', 'name=eh-swe-'], timeout=60).stdout.split() if ids: try: _docker(['rm', '-f'] + ids, timeout=180) except Exception: pass for img in list(self._pulled): try: _docker(['rmi', '-f', img], timeout=180) self._log(f'· released {img}') except Exception: pass self._pulled.clear() _SERVICE: Optional[ImageService] = None def get_image_service() -> ImageService: """Process-wide singleton; release hook auto-registered once.""" global _SERVICE if _SERVICE is None: _SERVICE = ImageService(workers=int(os.environ.get( 'EVALHARNESS_IMAGE_WORKERS', '2'))) if not os.environ.get('EVALHARNESS_KEEP_SWE_IMAGES'): import atexit atexit.register(_SERVICE.release_all) return _SERVICE