3afd57b5a4
- Allow --dataset swe_bench_verified,swe_bench_lite to pre-pull multiple SWE-bench variants in one run. - Validate dataset keys before loading. - Update myread.md with multi-dataset example.
354 lines
12 KiB
Python
Executable File
354 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
预拉取 SWE-bench 评测所需的 Docker 镜像，支持断点续传。
|
||
|
||
用法:
|
||
# 默认从上次进度继续
|
||
python bash/pull_swe_bench_images.py --dataset swe_bench_verified --max-workers 4
|
||
|
||
# 强制重新开始（删除进度文件）
|
||
rm -f output/.swe_bench_pull_state.json
|
||
python bash/pull_swe_bench_images.py --dataset swe_bench_verified --max-workers 4
|
||
|
||
# 测试前 3 个镜像
|
||
python bash/pull_swe_bench_images.py --dry-run 3 --max-workers 1
|
||
|
||
进度文件:
|
||
output/.swe_bench_pull_state.json
|
||
记录已拉取/跳过/失败的镜像，Ctrl+C 后下次自动跳过。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
import subprocess
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
|
||
# 国内网络加速（必须在导入 evalscope/huggingface 之前设置）
|
||
os.environ.setdefault('USE_MODELSCOPE_HUB', '1')
|
||
os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com')
|
||
os.environ.setdefault('PYTHONUNBUFFERED', '1')
|
||
|
||
from tqdm import tqdm
|
||
|
||
# 让 evalscope 可导入
|
||
PROJECT_ROOT = Path(__file__).parent.parent.resolve()
|
||
sys.path.insert(0, str(PROJECT_ROOT / 'evalscope'))
|
||
|
||
from evalscope.api.dataset import RemoteDataLoader, FieldSpec
|
||
|
||
|
||
DATASET_IDS = {
|
||
'swe_bench_verified': 'princeton-nlp/SWE-bench_Verified',
|
||
# 'swe_bench_lite': 'princeton-nlp/SWE-bench_Lite',
|
||
# 'swe_bench_verified_mini': 'evalscope/swe-bench-verified-mini',
|
||
# 'swe_bench_pro': 'evalscope/swe-bench-pro',
|
||
}
|
||
|
||
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / 'output'
|
||
STATE_FILENAME = '.swe_bench_pull_state.json'
|
||
|
||
|
||
def get_state_path(output_dir: Path) -> Path:
|
||
return Path(output_dir) / STATE_FILENAME
|
||
|
||
|
||
def load_state(output_dir: Path) -> dict:
|
||
"""Load previous pull state; return empty state if none exists."""
|
||
state_path = get_state_path(output_dir)
|
||
if not state_path.exists():
|
||
return {'done': [], 'skipped': [], 'failed': []}
|
||
try:
|
||
with state_path.open('r', encoding='utf-8') as f:
|
||
state = json.load(f)
|
||
# Ensure required keys exist
|
||
for key in ('done', 'skipped', 'failed'):
|
||
state.setdefault(key, [])
|
||
return state
|
||
except Exception as e:
|
||
print(f'WARNING: failed to load state file {state_path}: {e}')
|
||
return {'done': [], 'skipped': [], 'failed': []}
|
||
|
||
|
||
def save_state(output_dir: Path, state: dict):
|
||
"""Atomically save current pull state."""
|
||
state_path = get_state_path(output_dir)
|
||
state_path.parent.mkdir(parents=True, exist_ok=True)
|
||
tmp_path = state_path.with_suffix('.tmp')
|
||
try:
|
||
with tmp_path.open('w', encoding='utf-8') as f:
|
||
json.dump(state, f, ensure_ascii=False, indent=2)
|
||
tmp_path.replace(state_path)
|
||
except Exception as e:
|
||
print(f'WARNING: failed to save state file {state_path}: {e}')
|
||
|
||
|
||
def load_swe_samples(dataset_id: str):
|
||
"""Load SWE-bench samples and wrap them into the format expected by build_images."""
|
||
print(f'正在加载数据集元数据: {dataset_id} ...')
|
||
loader = RemoteDataLoader(
|
||
data_id_or_path=dataset_id,
|
||
split='test',
|
||
sample_fields=FieldSpec(
|
||
input='problem_statement',
|
||
metadata=[
|
||
'instance_id',
|
||
'repo',
|
||
'base_commit',
|
||
'patch',
|
||
'PASS_TO_PASS',
|
||
'FAIL_TO_PASS',
|
||
'test_patch',
|
||
'version',
|
||
'environment_setup_commit',
|
||
'hints_text',
|
||
'created_at',
|
||
],
|
||
),
|
||
)
|
||
return loader.load()
|
||
|
||
|
||
def load_multiple_datasets(dataset_keys: list[str]) -> list:
|
||
"""Load samples from one or more dataset keys and tag each sample with its source."""
|
||
all_samples = []
|
||
for key in dataset_keys:
|
||
dataset_id = DATASET_IDS[key]
|
||
samples = load_swe_samples(dataset_id)
|
||
# Tag each sample so get_image_name can still work and state stays per-dataset.
|
||
for sample in samples:
|
||
if isinstance(sample, dict):
|
||
sample.setdefault('dataset_key', key)
|
||
else:
|
||
metadata = getattr(sample, 'metadata', None)
|
||
if metadata is not None:
|
||
metadata.setdefault('dataset_key', key)
|
||
all_samples.extend(samples)
|
||
print(f' {key}: {len(samples)} 个样本')
|
||
return all_samples
|
||
|
||
|
||
def get_image_name(instance, namespace: str, arch: str) -> str:
|
||
"""根据 swebench 规则生成远程镜像名称。"""
|
||
if isinstance(instance, dict):
|
||
metadata = instance
|
||
else:
|
||
metadata = getattr(instance, 'metadata', {}) or {}
|
||
|
||
instance_id = metadata.get('instance_id', '')
|
||
if not instance_id and not isinstance(instance, dict):
|
||
instance_id = getattr(instance, 'instance_id', '')
|
||
|
||
if not instance_id:
|
||
return None
|
||
|
||
image_name = f'{namespace}/sweb.eval.{arch}.{instance_id.lower()}:latest'
|
||
image_name = image_name.replace('__', '_1776_')
|
||
return image_name
|
||
|
||
|
||
def get_local_images() -> set:
|
||
"""Return a set of all locally available Docker image names."""
|
||
result = subprocess.run(
|
||
['docker', 'images', '--format', '{{.Repository}}:{{.Tag}}'],
|
||
capture_output=True, text=True
|
||
)
|
||
if result.returncode != 0:
|
||
return set()
|
||
return {line.strip() for line in result.stdout.splitlines() if line.strip()}
|
||
|
||
|
||
_print_lock = threading.Lock()
|
||
|
||
|
||
def _log(msg: str):
|
||
with _print_lock:
|
||
print(msg, flush=True)
|
||
|
||
|
||
def _stream_pull(image_name: str) -> tuple:
|
||
"""Run docker pull and stream output in real time.
|
||
|
||
Returns (returncode, last_lines_of_output).
|
||
"""
|
||
process = subprocess.Popen(
|
||
['docker', 'pull', image_name],
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
bufsize=1,
|
||
)
|
||
|
||
lines = []
|
||
if process.stdout:
|
||
for line in process.stdout:
|
||
line = line.rstrip('\n')
|
||
lines.append(line)
|
||
_log(f' {line}')
|
||
# Keep only the last 50 lines for error reporting
|
||
if len(lines) > 50:
|
||
lines.pop(0)
|
||
|
||
process.wait()
|
||
return process.returncode, '\n'.join(lines)
|
||
|
||
|
||
def pull_one_image(args_tuple, state: dict, output_dir: Path, local_images: set):
|
||
"""Pull a single image and update shared state.
|
||
|
||
Returns (idx, image_name, status, message).
|
||
"""
|
||
idx, total, instance, namespace, arch = args_tuple
|
||
image_name = get_image_name(instance, namespace, arch)
|
||
if not image_name:
|
||
return idx, None, 'skip', 'missing instance_id'
|
||
|
||
# Resume: skip if already recorded in state
|
||
if image_name in state['done'] or image_name in state['skipped']:
|
||
return idx, image_name, 'skip', 'already recorded in state'
|
||
|
||
# Also check local docker images (in case state was deleted)
|
||
if image_name in local_images:
|
||
state['skipped'].append(image_name)
|
||
save_state(output_dir, state)
|
||
return idx, image_name, 'skip', 'already exists locally'
|
||
|
||
_log(f'⬇️ [{idx+1}/{total}] START {image_name}')
|
||
|
||
rc, output = _stream_pull(image_name)
|
||
|
||
if rc == 0:
|
||
state['done'].append(image_name)
|
||
save_state(output_dir, state)
|
||
return idx, image_name, 'success', ''
|
||
else:
|
||
msg = output.strip()[-500:] if output.strip() else 'docker pull failed'
|
||
state['failed'].append(image_name)
|
||
save_state(output_dir, state)
|
||
return idx, image_name, 'fail', msg
|
||
|
||
|
||
def main():
|
||
import argparse
|
||
parser = argparse.ArgumentParser(description='Pre-pull SWE-bench Docker images with resume support')
|
||
parser.add_argument(
|
||
'--dataset',
|
||
default='swe_bench_verified',
|
||
type=lambda s: [x.strip() for x in s.split(',') if x.strip()],
|
||
help='SWE-bench dataset variant(s), comma-separated, e.g. swe_bench_verified or swe_bench_verified,swe_bench_lite',
|
||
)
|
||
parser.add_argument(
|
||
'--max-workers',
|
||
type=int,
|
||
default=4,
|
||
help='Max concurrent docker pulls (default: 4)',
|
||
)
|
||
parser.add_argument(
|
||
'--force-arch',
|
||
default='x86_64',
|
||
choices=['', 'arm64', 'x86_64'],
|
||
help='Image architecture (default: x86_64)',
|
||
)
|
||
parser.add_argument(
|
||
'--dockerhub-username',
|
||
default='swebench',
|
||
help='DockerHub namespace for remote images (default: swebench)',
|
||
)
|
||
parser.add_argument(
|
||
'--output-dir',
|
||
default=str(DEFAULT_OUTPUT_DIR),
|
||
help='Directory to store progress state (default: output/)',
|
||
)
|
||
parser.add_argument(
|
||
'--dry-run',
|
||
type=int,
|
||
default=0,
|
||
help='Only test the first N images',
|
||
)
|
||
parser.add_argument(
|
||
'--retry-failed',
|
||
action='store_true',
|
||
help='Retry images that failed in previous runs',
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
output_dir = Path(args.output_dir)
|
||
state = load_state(output_dir)
|
||
|
||
# Optionally retry previously failed images
|
||
if args.retry_failed:
|
||
failed = state['failed']
|
||
state['failed'] = []
|
||
print(f'🔄 重试 {len(failed)} 个之前失败的镜像')
|
||
|
||
# Validate dataset keys
|
||
for key in args.dataset:
|
||
if key not in DATASET_IDS:
|
||
print(f'ERROR: unknown dataset "{key}". Supported: {list(DATASET_IDS.keys())}')
|
||
sys.exit(1)
|
||
|
||
samples = load_multiple_datasets(args.dataset)
|
||
|
||
if args.dry_run > 0:
|
||
samples = samples[:args.dry_run]
|
||
print(f'⚠️ 测试模式：仅处理前 {args.dry_run} 个样本以验证镜像名和网络...')
|
||
else:
|
||
print(f'✅ 成功加载 {len(samples)} 个样本（数据集: {args.dataset}）。')
|
||
|
||
already_done = len(state['done']) + len(state['skipped'])
|
||
print(f'📂 进度文件: {get_state_path(output_dir)}')
|
||
print(f' 已下载/已存在: {len(state["done"])} | 已跳过: {len(state["skipped"])} | 历史失败: {len(state["failed"])}')
|
||
|
||
# Cache local image names once to avoid 500+ `docker images -q` calls.
|
||
print('🔄 正在扫描本地镜像...')
|
||
local_images = get_local_images()
|
||
print(f' 本地已有 {len(local_images)} 个镜像')
|
||
print('-' * 60)
|
||
|
||
tasks = [
|
||
(idx, len(samples), sample, args.dockerhub_username, args.force_arch)
|
||
for idx, sample in enumerate(samples)
|
||
]
|
||
|
||
success_count = 0
|
||
skip_count = 0
|
||
fail_count = 0
|
||
|
||
with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
|
||
futures = {executor.submit(pull_one_image, t, state, output_dir, local_images): t[0] for t in tasks}
|
||
for future in tqdm(as_completed(futures), total=len(futures), desc='Pulling images', unit='img', mininterval=1.0):
|
||
idx = futures[future]
|
||
try:
|
||
_, image_name, status, msg = future.result()
|
||
except Exception as e:
|
||
image_name, status, msg = None, 'fail', str(e)
|
||
|
||
if status == 'success':
|
||
success_count += 1
|
||
_log(f'✅ [{idx+1}/{len(samples)}] DONE {image_name}')
|
||
elif status == 'skip':
|
||
skip_count += 1
|
||
_log(f'⏭️ [{idx+1}/{len(samples)}] SKIP {image_name or msg}')
|
||
else:
|
||
fail_count += 1
|
||
_log(f'❌ [{idx+1}/{len(samples)}] FAILED {image_name}')
|
||
if msg:
|
||
_log(f' {msg}')
|
||
|
||
print('=' * 60)
|
||
print('🎉 预拉取任务结束！')
|
||
print(f' 总计: {len(samples)} | 已存在/跳过: {skip_count} | 本次成功: {success_count} | 本次失败: {fail_count}')
|
||
print(f' 进度文件: {get_state_path(output_dir)}')
|
||
print('=' * 60)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
try:
|
||
main()
|
||
except KeyboardInterrupt:
|
||
print('\n\n⚠️ 用户手动中断。进度已保存，下次运行会自动跳过已完成的镜像。')
|
||
sys.exit(0)
|
||
|