feat(pull_swe_bench_images): support comma-separated --dataset

- 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.
This commit is contained in:
sora 2026-07-23 02:34:49 +00:00
parent 2ee0af0728
commit 3afd57b5a4
2 changed files with 35 additions and 5 deletions

View File

@ -111,6 +111,25 @@ def load_swe_samples(dataset_id: str):
return loader.load() 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: def get_image_name(instance, namespace: str, arch: str) -> str:
"""根据 swebench 规则生成远程镜像名称。""" """根据 swebench 规则生成远程镜像名称。"""
if isinstance(instance, dict): if isinstance(instance, dict):
@ -217,8 +236,8 @@ def main():
parser.add_argument( parser.add_argument(
'--dataset', '--dataset',
default='swe_bench_verified', default='swe_bench_verified',
choices=list(DATASET_IDS.keys()), type=lambda s: [x.strip() for x in s.split(',') if x.strip()],
help='SWE-bench dataset variant', help='SWE-bench dataset variant(s), comma-separated, e.g. swe_bench_verified or swe_bench_verified,swe_bench_lite',
) )
parser.add_argument( parser.add_argument(
'--max-workers', '--max-workers',
@ -264,14 +283,19 @@ def main():
state['failed'] = [] state['failed'] = []
print(f'🔄 重试 {len(failed)} 个之前失败的镜像') print(f'🔄 重试 {len(failed)} 个之前失败的镜像')
dataset_id = DATASET_IDS[args.dataset] # Validate dataset keys
samples = load_swe_samples(dataset_id) 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: if args.dry_run > 0:
samples = samples[:args.dry_run] samples = samples[:args.dry_run]
print(f'⚠️ 测试模式:仅处理前 {args.dry_run} 个样本以验证镜像名和网络...') print(f'⚠️ 测试模式:仅处理前 {args.dry_run} 个样本以验证镜像名和网络...')
else: else:
print(f'✅ 成功加载 {len(samples)} 个样本') print(f'✅ 成功加载 {len(samples)} 个样本(数据集: {args.dataset}')
already_done = len(state['done']) + len(state['skipped']) already_done = len(state['done']) + len(state['skipped'])
print(f'📂 进度文件: {get_state_path(output_dir)}') print(f'📂 进度文件: {get_state_path(output_dir)}')

View File

@ -243,6 +243,11 @@ export PYTHONUNBUFFERED=1
python bash/pull_swe_bench_images.py \ python bash/pull_swe_bench_images.py \
--dataset swe_bench_verified \ --dataset swe_bench_verified \
--max-workers 1 --max-workers 1
# 也可以同时预拉多个数据集(逗号分隔)
python bash/pull_swe_bench_images.py \
--dataset swe_bench_verified,swe_bench_lite \
--max-workers 1
``` ```
脚本特性: 脚本特性:
@ -250,6 +255,7 @@ python bash/pull_swe_bench_images.py \
- 断点续传:状态文件 `output/.swe_bench_pull_state.json`,中断后再次运行会自动跳过已成功的镜像。 - 断点续传:状态文件 `output/.swe_bench_pull_state.json`,中断后再次运行会自动跳过已成功的镜像。
- 本地已有镜像自动跳过。 - 本地已有镜像自动跳过。
- 单 worker 模式稳定,减少限流概率;若网络非常好,可适当提高 `--max-workers` - 单 worker 模式稳定,减少限流概率;若网络非常好,可适当提高 `--max-workers`
- `--dataset` 支持逗号分隔的多个数据集。
镜像会存到 Docker `data-root` 配置的路径: 镜像会存到 Docker `data-root` 配置的路径: