fc8ade3c03
One ImageService per process. When tasks arrive (runner) or a sample starts (env), its images are REGISTERED; a background worker pool delivers them -- local tar shipments first (es's swebench_v 500-image batch set, disk-cached index), network mirror chain second. Sample execution waits on a readiness barrier instead of the old failing timings (docker-run implicit pull killed at 120s; score-phase batch pull ran after generation had already failed). - runner registers every pending sample's image up front (pull-ahead overlaps generation) - env blocks on wait_ready(1800s) before docker run -- a slow pull delays that sample, never fails it - tar index cached at /tmp/evalharness_tar_index.json (full scan costs minutes; only the first process pays) - images the service loaded are released at exit (atexit; opt out with EVALHARNESS_KEEP_SWE_IMAGES); pre-existing local images never touched - EVALHARNESS_IMAGE_WORKERS (default 2) tunes the pool Verified E2E: register -> background load from swebench_batch_001.tar.gz -> image present locally (matplotlib-14623); wait barrier semantics confirmed (blocks until load completes). Co-Authored-By: Claude <noreply@anthropic.com>
249 lines
9.0 KiB
Python
249 lines
9.0 KiB
Python
"""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
|