diff --git a/evalharness/agent/envs/swe_agentic.py b/evalharness/agent/envs/swe_agentic.py index c840d0a..3700ecd 100644 --- a/evalharness/agent/envs/swe_agentic.py +++ b/evalharness/agent/envs/swe_agentic.py @@ -20,48 +20,7 @@ from typing import Any, Dict from ...data.sample import ChatMessage, Sample from ..loop import Environment, register_env -# ---- image lifecycle tracking ------------------------------------------- -# per-instance swe images are BIG (1-4GB each, 500 total ~2TB): every -# image THIS process pulled is registered here and released at exit -# (normal end, Ctrl+C kill, crash) so a run never strands its footprint. -import atexit as _atexit -_IMAGES_PULLED: list = [] - - -def _release_pulled_images() -> None: - import os as _os - - if not _IMAGES_PULLED or _os.environ.get('EVALHARNESS_KEEP_SWE_IMAGES'): - return - # in-use containers first (docker refuses rmi otherwise) - try: - _docker(['ps', '-aq', '--filter', 'name=eh-swe-'], timeout=30) - except Exception: - pass - try: - _docker(['ps', '-aq', '--filter', 'ancestor=none'], timeout=30) - except Exception: - pass - import subprocess as _sp - - try: - ids = _sp.run(['docker', 'ps', '-aq', '--filter', 'name=eh-swe'], - capture_output=True, text=True, timeout=30).stdout.split() - if ids: - _docker(['rm', '-f'] + ids, timeout=120) - except Exception: - pass - for img in _IMAGES_PULLED: - try: - _docker(['rmi', '-f', img], timeout=120) - print(f'· released swe image {img}', file=sys.stderr, flush=True) - except Exception: - pass - - -_atexit.register(_release_pulled_images) -# ------------------------------------------------------------------------ SENTINEL = 'COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT' @@ -165,15 +124,13 @@ class SWEAgenticEnvironment(Environment): # per-instance image: ensure (local check -> mirror-chain pull with # progress) BEFORE docker run -- `docker run` auto-pulls with zero # output and its 120s timeout kills runs on slow mirrors - from ...sandbox.docker import ensure_image + # unified image service: pulls happen in the BACKGROUND from task + # arrival; a not-yet-ready image BLOCKS here instead of failing + # (the old inline pull timed-out docker run at 120s and killed runs) + from ...sandbox.image_service import get_image_service - had_it = _docker(['image', 'inspect', image]).returncode == 0 - try: - ensure_image(image) - except RuntimeError as e: - raise RuntimeError(f'image {image} unavailable: {str(e)[:150]}') - if not had_it and image not in _IMAGES_PULLED: - _IMAGES_PULLED.append(image) # we pulled it -> we release it + if not get_image_service().wait_ready(image, timeout_s=1800): + raise RuntimeError(f'image {image} unavailable after all sources') name = f'eh-swe-{uuid.uuid4().hex[:10]}' r = _docker(['run', '-d', '--name', name, '-w', '/testbed', @@ -208,6 +165,11 @@ class SWEAgenticEnvironment(Environment): if not image: raise RuntimeError('swe_agentic: sample has no sandbox image ' '(dataset plugin must declare it)') + # register on ARRIVAL: the background pool starts pulling while + # earlier samples are still generating + from ...sandbox.image_service import get_image_service + + get_image_service().register([image]) self.container = self._start(image) ps = (sample.metadata or {}).get('problem_statement') or sample.input_text messages = [ChatMessage(role='user', content=INSTANCE_TEMPLATE.format( diff --git a/evalharness/model/runner.py b/evalharness/model/runner.py index 3c8d6cf..c057d6d 100644 --- a/evalharness/model/runner.py +++ b/evalharness/model/runner.py @@ -433,6 +433,19 @@ async def generate_predictions( if progress_reporter is not None: progress_reporter.reset_samples(len(work), dataset_name, completed=len(restored)) + + # sandbox-image pull-ahead: register every pending sample's image the + # moment the task list is known -- the background pool starts pulling + # while generation is still in flight (swe: 1 image per sample) + try: + _imgs = {s.sandbox.image for s in work + if getattr(s, 'sandbox', None) and s.sandbox.image} + if _imgs: + from ..sandbox.image_service import get_image_service + + get_image_service().register(sorted(_imgs)) + except Exception: + pass # let the adapter surface retry attempts to the bar members = getattr(adapter, 'adapters', [adapter]) for m_ in members: diff --git a/evalharness/sandbox/image_service.py b/evalharness/sandbox/image_service.py new file mode 100644 index 0000000..4b2d82f --- /dev/null +++ b/evalharness/sandbox/image_service.py @@ -0,0 +1,248 @@ +"""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