EvalHarness/evalharness/sandbox/image_service.py

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