c23665fed7
Single-image benches keep fail-fast; >8 distinct images (swe's 500 per-instance) now pull concurrently with resume state inside the run itself -- exactly like the old auto-pull behavior, just resilient: scoring proceeds with whatever images landed, missing ones score 0 and only a total wipeout fails the bench. Ctrl+C-safe (state file), network recovery resumes automatically on the next run. Co-Authored-By: Claude <noreply@anthropic.com>
203 lines
6.8 KiB
Python
203 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
||
"""SWE-bench 逐题镜像预拉取（断点续传、并发、失败重试）。
|
||
|
||
移植自 es 的 pull_swe_bench_images.py，去掉 evalscope 依赖（用我们自己的
|
||
数据集注册表），加进 sandbox 子命令：
|
||
|
||
evalharness sandbox pull swe_bench_verified --max-workers 4
|
||
evalharness sandbox pull swe_bench_verified --dry-run 3 # 只试前 3 个
|
||
evalharness sandbox pull swe_bench_verified --retry-failed
|
||
|
||
进度文件: <cache-dir>/swe_pull_state.json（done/failed/skipped），
|
||
Ctrl+C 后重跑自动跳过已完成的。
|
||
"""
|
||
|
||
import json
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
|
||
_print_lock = threading.Lock()
|
||
|
||
|
||
def _log(msg):
|
||
with _print_lock:
|
||
print(msg, flush=True)
|
||
|
||
|
||
def swe_image_names(dataset_name: str):
|
||
"""数据集 → 官方逐题镜像名列表（与 dataset 插件的命名完全一致）。"""
|
||
from evalharness import get_dataset
|
||
|
||
ds = get_dataset(dataset_name)
|
||
ds.materialize()
|
||
names = []
|
||
for s in ds:
|
||
if s.sandbox and s.sandbox.image:
|
||
names.append(s.sandbox.image)
|
||
return names
|
||
|
||
|
||
def _state_path(cache_dir: Path) -> Path:
|
||
return cache_dir / 'swe_pull_state.json'
|
||
|
||
|
||
def _load_state(cache_dir: Path) -> dict:
|
||
p = _state_path(cache_dir)
|
||
if not p.exists():
|
||
return {'done': [], 'failed': [], 'skipped': []}
|
||
try:
|
||
state = json.load(open(p, encoding='utf-8'))
|
||
for k in ('done', 'failed', 'skipped'):
|
||
state.setdefault(k, [])
|
||
return state
|
||
except Exception:
|
||
return {'done': [], 'failed': [], 'skipped': []}
|
||
|
||
|
||
def _save_state(cache_dir: Path, state: dict):
|
||
p = _state_path(cache_dir)
|
||
p.parent.mkdir(parents=True, exist_ok=True)
|
||
tmp = p.with_suffix('.tmp')
|
||
tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding='utf-8')
|
||
tmp.replace(p)
|
||
|
||
|
||
def _pull_one(image: str) -> tuple:
|
||
"""docker pull 单个镜像，实时透传输出。返回 (returncode, 尾部输出)。"""
|
||
proc = subprocess.Popen(['docker', 'pull', image],
|
||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||
text=True, bufsize=1)
|
||
lines = []
|
||
for line in proc.stdout:
|
||
line = line.rstrip('\n')
|
||
lines.append(line)
|
||
_log(f' {line}')
|
||
if len(lines) > 50:
|
||
lines.pop(0)
|
||
proc.wait()
|
||
return proc.returncode, '\n'.join(lines)
|
||
|
||
|
||
def pull_many(names, cache_dir, max_workers: int = 4,
|
||
verbose: bool = True) -> int:
|
||
"""批量拉取一组镜像名（断点续传复用 <cache-dir>/swe_pull_state.json）。
|
||
|
||
返回失败数。给判分预检用：多镜像 bench（swe 500 个）逐个 fail-fast
|
||
会把整个 bench 卡死在第一个失败上；这里并发拉、不中断、返回失败数。"""
|
||
cdir = Path(cache_dir)
|
||
state = _load_state(cdir)
|
||
local = set()
|
||
r = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
|
||
capture_output=True, text=True)
|
||
if r.returncode == 0:
|
||
local = {l.strip() for l in r.stdout.splitlines() if l.strip()}
|
||
todo = [nm for nm in names
|
||
if nm not in state['done'] and nm not in local]
|
||
if verbose:
|
||
_log(f'· {len(names)} 个沙箱镜像: 本地/已完成 {len(names) - len(todo)}, '
|
||
f'待拉 {len(todo)} (并发 {max_workers}, 进度存 {_state_path(cdir)})')
|
||
|
||
def work(nm):
|
||
rc, out = _pull_one(nm)
|
||
if rc == 0:
|
||
state['done'].append(nm)
|
||
else:
|
||
state['failed'].append(nm)
|
||
_save_state(cdir, state)
|
||
return nm, rc
|
||
|
||
fails = 0
|
||
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
||
for i, (nm, rc) in enumerate(
|
||
pool.map(work, todo), 1):
|
||
if rc:
|
||
fails += 1
|
||
if verbose and (i % 10 == 0 or i == len(todo)):
|
||
_log(f' 镜像进度 [{i}/{len(todo)}] 失败 {fails}')
|
||
return fails
|
||
|
||
|
||
def run_pull(dataset: str, cache_dir: str, max_workers: int = 4,
|
||
dry_run: int = 0, retry_failed: bool = False):
|
||
cdir = Path(cache_dir)
|
||
state = _load_state(cdir)
|
||
if retry_failed:
|
||
n = len(state['failed'])
|
||
state['failed'] = []
|
||
_log(f'🔄 重试 {n} 个之前失败的镜像')
|
||
|
||
names = swe_image_names(dataset)
|
||
if dry_run:
|
||
names = names[:dry_run]
|
||
local = set()
|
||
r = subprocess.run(['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
|
||
capture_output=True, text=True)
|
||
if r.returncode == 0:
|
||
local = {l.strip() for l in r.stdout.splitlines() if l.strip()}
|
||
|
||
todo = []
|
||
for nm in names:
|
||
if nm in state['done'] or nm in state['skipped']:
|
||
continue
|
||
if nm in local:
|
||
state['skipped'].append(nm)
|
||
_save_state(cdir, state)
|
||
continue
|
||
todo.append(nm)
|
||
_log(f'共 {len(names)} 个镜像: 已完成/本地 {len(names) - len(todo)}, '
|
||
f'待拉 {len(todo)}, 并发 {max_workers}')
|
||
|
||
def work(nm):
|
||
_log(f'⬇ START {nm}')
|
||
rc, out = _pull_one(nm)
|
||
if rc == 0:
|
||
state['done'].append(nm)
|
||
else:
|
||
state['failed'].append(nm)
|
||
_save_state(cdir, state)
|
||
return nm, rc, out[-300:] if rc else ''
|
||
|
||
fails = []
|
||
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
||
futs = {pool.submit(work, nm): nm for nm in todo}
|
||
for i, fut in enumerate(as_completed(futs), 1):
|
||
nm, rc, out = fut.result()
|
||
mark = '✓' if rc == 0 else '✗'
|
||
_log(f'[{i}/{len(todo)}] {mark} {nm}')
|
||
if rc:
|
||
fails.append((nm, out))
|
||
|
||
_log(f'\n完成: {len(state["done"])} 成功, {len(state["skipped"])} 跳过, '
|
||
f'{len(state["failed"])} 失败 (进度存 {_state_path(cdir)})')
|
||
for nm, out in fails[:5]:
|
||
_log(f' 失败样例 {nm}: {out[:150]}')
|
||
return 0 if not fails else 1
|
||
|
||
|
||
def main(argv) -> int:
|
||
import argparse
|
||
|
||
p = argparse.ArgumentParser(description='Pre-pull SWE-bench per-instance '
|
||
'images with resume')
|
||
p.add_argument('dataset', nargs='?', default='swe_bench_verified')
|
||
p.add_argument('--max-workers', type=int, default=4)
|
||
p.add_argument('--dry-run', type=int, default=0,
|
||
help='only pull the first N images')
|
||
p.add_argument('--retry-failed', action='store_true')
|
||
p.add_argument('--cache-dir', default='')
|
||
a = p.parse_args(argv)
|
||
|
||
cache = a.cache_dir
|
||
if not cache:
|
||
import os
|
||
cache = os.environ.get('EVALHARNESS_CACHE',
|
||
str(Path.home() / '.cache/evalharness'))
|
||
return run_pull(a.dataset, cache, a.max_workers, a.dry_run, a.retry_failed)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
sys.exit(main(sys.argv[1:]))
|