Background image service: pull-ahead + wait-barrier, unified lifecycle

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>
This commit is contained in:
sora 2026-09-18 09:45:39 +00:00
parent e9e3230e46
commit fc8ade3c03
3 changed files with 272 additions and 49 deletions

View File

@ -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(

View File

@ -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:

View File

@ -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