"""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 — local tar batches first (es's swebench_v shipment, 500 images offline), network mirror chain second. 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 pathlib import Path from typing import Dict, Iterable, List, Optional # local tar shipments searched before the network (dir -> glob pattern) _TAR_SOURCES = [ ('/data1/sora/evalscope/docker/swebench_v', 'swebench_batch_*.tar.gz'), ('/data1/sora/evalscope/docker', 'bigcodebench-sandbox.tar.gz'), ] def _docker(args, timeout=300): return subprocess.run(['docker'] + args, capture_output=True, text=True, timeout=timeout) def _local_images() -> set: r = _docker(['images', '--format', '{{.Repository}}:{{.Tag}}'], timeout=60) return set(r.stdout.split()) if r.returncode == 0 else set() def _tar_manifest_tags(path: str) -> List[str]: """RepoTags inside a docker-save tar ( OCI layout or legacy ).""" import tarfile try: with tarfile.open(path, 'r:gz') as t: for name in ('manifest.json', 'index.json'): try: member = t.getmember(name) except KeyError: continue data = json.load(t.extractfile(member)) if name == 'manifest.json': return [tg for e in data for tg in (e.get('RepoTags') or [])] # OCI index: ref lives in annotations; best effort refs = [] for m in data.get('manifests', []): ref = (m.get('annotations') or {}).get( 'org.opencontainers.image.ref.name', '') if ref: refs.append(ref) return refs except Exception: return [] return [] 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._tar_index: Optional[Dict[str, str]] = None # tag -> tar path 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; already-local images skip the queue entirely.""" self.start() local = _local_images() for img in images: if not img: continue with self._lock: if img in self._wanted: continue self._wanted[img] = 'pulling' self._events[img] = threading.Event() if img in local: 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: # 1. local tar shipment (offline, fastest) tar = self._find_in_tars(image) if tar: self._log(f'· loading {image.split("/")[-1][:44]} from {Path(tar).parent.name}/{Path(tar).name}') r = _docker(['load', '-qi', tar], timeout=1800) if r.returncode == 0 and image in _local_images(): self._pulled.append(image) return True self._log(f' tar load failed, falling back to network') # 2. mirror chain (progress + watchdog inside) 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 # ---- tar index ---- _INDEX_CACHE = '/tmp/evalharness_tar_index.json' def _find_in_tars(self, image: str) -> Optional[str]: # the full 50-tar index costs minutes to build (gzip full-scan); # cache it on disk so only the FIRST process ever pays import glob if self._tar_index is None: cache = {} try: cache = json.load(open(self._INDEX_CACHE)) # invalidate when the tar set changes cur = sorted(glob.glob('/data1/sora/evalscope/docker/swebench_v/*.tar.gz')) if cache.get('_files') != [os.path.basename(x) for x in cur]: cache = {} except Exception: cache = {} if cache.get('tags'): self._tar_index = cache['tags'] else: self._tar_index = {} for d, pat in _TAR_SOURCES: if not os.path.isdir(d): continue for f in sorted(glob.glob(os.path.join(d, pat))): for tag in _tar_manifest_tags(f): self._tar_index.setdefault(tag, f) try: files = [os.path.basename(x) for x in glob.glob('/data1/sora/evalscope/docker/swebench_v/*.tar.gz')] json.dump({'_files': sorted(files), 'tags': self._tar_index}, open(self._INDEX_CACHE, 'w')) except Exception: pass return self._tar_index.get(image) # ---- 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