EvalHarness/evalharness/sandbox/prefetch.py

"""Image prefetch: parallel docker pull for execution environments.

SWE-bench Verified declares ~500 per-instance images (sweb.eval.x86_64.*).
Pulling them lazily during an eval run would stall it serially; prefetch
pulls them ahead of time with bounded concurrency and progress reporting.

    from evalharness.sandbox.prefetch import prefetch_images
    done = prefetch_images(['sweb.eval.x86_64.django__django-12345', ...], workers=8)
    # CLI: evalharness sandbox prefetch swe_bench_verified --workers 8 --limit 50
"""

import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Iterable, List

from ..data.dataset import Dataset

# CN mirrors tried in order before/alongside the daemon's configured mirrors.
# Some namespaces (e.g. swebench/*) are blocked by individual CN mirrors, so we
# fall through: daemon default -> 1ms.run -> baidubce -> sjtug.
_CN_MIRROR_FALLBACKS = [
    '{img}',  # daemon default (uses its own registry-mirrors config)
    'docker.m.daocloud.io/{img}',
    'docker.1ms.run/{img}',
    'mirror.baidubce.com/{img}',
    'docker.mirrors.sjtug.sjtu.edu.cn/{img}',
    'hub.rat.dev/{img}',
]


def images_for_dataset(ds: Dataset, limit: int = 0) -> List[str]:
    """Collect distinct sandbox images declared by a dataset's samples."""
    seen: List[str] = []
    add = seen.append
    for i, s in enumerate(ds):
        if limit and i >= limit:
            break
        if s.sandbox and s.sandbox.image and s.sandbox.image not in seen:
            add(s.sandbox.image)
    return seen


def images_for_samples(samples) -> List[str]:
    """Distinct sandbox images in sample order (dedup, keeps order)."""
    seen: List[str] = []
    for s in samples:
        if s.sandbox and s.sandbox.image and s.sandbox.image not in seen:
            seen.append(s.sandbox.image)
    return seen


def local_images() -> set:
    out = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
                         capture_output=True, text=True)
    return {line for line in out.stdout.splitlines() if line}


def prefetch_images(images: Iterable[str], workers: int = 8) -> List[str]:
    """Pull images with bounded concurrency; skip ones already local.

    Returns the list newly pulled. Failures are reported and skipped (one
    missing instance must not block the rest).
    """
    have = local_images()
    todo = [img for img in images if img not in have]
    if not todo:
        print(f'prefetch: all {len(images)} images already local')
        return []
    print(f'prefetch: {len(todo)} to pull ({len(images) - len(todo)} already local), '
          f'workers={workers}')
    pulled: List[str] = []
    failed: List[str] = []
    t0 = time.time()
    done = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futs = {pool.submit(_pull_one, img): img for img in todo}
        for fut in as_completed(futs):
            img = futs[fut]
            done += 1
            try:
                fut.result()
                pulled.append(img)
            except Exception as e:
                failed.append(img)
                print(f'  FAIL {img}: {str(e)[:120]}', flush=True)
            if done % 10 == 0 or done == len(todo):
                rate = done / max(time.time() - t0, 1e-6)
                print(f'  [{done}/{len(todo)}] {rate:.2f} imgs/s, '
                      f'{len(failed)} failed', flush=True)
    print(f'prefetch done: {len(pulled)} pulled, {len(failed)} failed '
          f'in {time.time() - t0:.0f}s')
    return pulled


def _pull_one(image: str, verbose: bool = True) -> None:
    """Pull via CN-mirror fallback chain; retag to the canonical name on hit.

    A multi-GB pull runs silently for many minutes with capture_output --
    which reads as a hang -- so say WHICH source is being tried and how
    long it took."""
    import sys as _sys
    import time as _time
    last_err = None
    for template in _CN_MIRROR_FALLBACKS:
        ref = template.format(img=image)
        if verbose:
            src = 'docker.io(daemon mirrors)' if '{img}' in template else template.split('/{img}')[0]
            print(f'· pulling sandbox image {image} via {src} ...',
                  file=_sys.stderr, flush=True)
        t0 = _time.monotonic()
        # stream progress: a multi-GB pull with captured output is minutes
        # of silence that reads as a hang. Non-TTY docker pull emits clean
        # per-layer lines; forward the informative ones (throttled).
        proc = subprocess.Popen(['docker', 'pull', ref],
                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                text=True, bufsize=1)
        last_line = ''
        last_show = 0.0
        layer_prog = {}   # layer id -> (done_bytes, total_bytes): the raw
        # per-layer events give no sense of OVERALL progress -- sum them
        import os as _os
        import re as _re

        def _bytes(s):
            s = s.strip()
            mult = {'kB': 1e3, 'KB': 1e3, 'MB': 1e6, 'GB': 1e9, 'B': 1}
            for suf, m in mult.items():
                if s.endswith(suf):
                    return int(float(s[:-len(suf)]) * m)
            return int(float(s or 0))

        for line in proc.stdout:
            last_line = line.strip()
            m = _re.match(r'^([0-9a-f]{12}): Downloading.*?([\d.]+[kKMG]?B)/([\d.]+[kKMG]?B)', line)
            if m:
                layer_prog[m.group(1)] = (_bytes(m.group(2)), _bytes(m.group(3)))
            if 'Pull complete' in line:
                lid = line.split(':')[0]
                if lid in layer_prog:
                    layer_prog[lid] = (layer_prog[lid][1], layer_prog[lid][1])
            now = _time.monotonic()
            key_evt = any(k in line for k in
                          ('Pulling from', 'Status', 'error', 'denied'))
            if verbose and key_evt:
                print(f'  {last_line[:100]}', file=_sys.stderr, flush=True)
                last_show = now
            elif verbose and now - last_show > 3 and layer_prog:
                done = sum(v[0] for v in layer_prog.values())
                tot = sum(v[1] for v in layer_prog.values())
                if tot:
                    pct = done / tot * 100
                    big = max(layer_prog.items(), key=lambda kv: kv[1][1] - kv[1][0])
                    print(f'  ⏳ {done / 1e9:.2f}/{tot / 1e9:.2f} GB ({pct:.0f}%)'
                          f' · 最大层 {big[0]}: {big[1][0] / 1e9:.2f}/{big[1][1] / 1e9:.2f} GB',
                          file=_sys.stderr, flush=True)
                last_show = now
        proc.wait(timeout=3600)
        r = subprocess.CompletedProcess(proc.args, proc.returncode,
                                        stdout='', stderr=last_line)
        if r.returncode == 0 and verbose:
            print(f'✓ image ready in {_time.monotonic() - t0:.0f}s',
                  file=_sys.stderr, flush=True)
        if r.returncode == 0:
            if ref != image:
                # retag BEFORE rmi and VERIFY: a silent tag failure followed
                # by rmi deleted the only tag and lost a 342s multi-GB pull
                tag = subprocess.run(['docker', 'tag', ref, image],
                                     capture_output=True, text=True)
                if tag.returncode != 0:
                    print(f'!! docker tag {ref} -> {image} failed: '
                          f'{(tag.stderr or "")[:120]} (keeping the mirror tag)',
                          file=_sys.stderr, flush=True)
                    return   # do NOT rmi: the mirror tag is the only handle
                subprocess.run(['docker', 'rmi', ref], check=False)
            return
        last_err = (ref, (r.stderr or '').strip()[:150])
    raise RuntimeError(f'all mirrors failed for {image}: {last_err}')