mirror of
https://github.com/microsoft/SkillOpt.git
synced 2026-07-06 07:32:34 +08:00
681 lines
24 KiB
Python
Executable File
681 lines
24 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Launch the SearchQA / SpreadsheetBench ablation matrix.
|
|
|
|
By default this script prints commands only. Pass --execute to actually start
|
|
runs. Every run writes to a unique out_root under the run root and logs stdout
|
|
/ stderr to logs/<run_id>.log.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
PYTHON_BIN = Path("/home/azureuser/workspace-gzy/miniconda3/envs/reflact/bin/python")
|
|
|
|
T2_ENDPOINT = "https://t2vgoaigpt4o3.openai.azure.com/"
|
|
SEARCHAGENT5_ENDPOINT = "https://searchagent5.cognitiveservices.azure.com/"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Experiment:
|
|
run_id: str
|
|
benchmark: str
|
|
config: str
|
|
overrides: tuple[str, ...]
|
|
|
|
|
|
BENCH_CONFIG = {
|
|
"searchqa": "configs/searchqa/default.yaml",
|
|
"spreadsheetbench": "configs/spreadsheetbench/default.yaml",
|
|
"livemathematicianbench": "configs/livemathematicianbench/default.yaml",
|
|
"alfworld": "configs/alfworld/default.yaml",
|
|
"docvqa": "configs/docvqa/default.yaml",
|
|
}
|
|
|
|
DEFAULT_SPLIT = {
|
|
"searchqa": "data/ablation_splits/searchqa/2-1-7_seed42",
|
|
"spreadsheetbench": "data/ablation_splits/spreadsheetbench/2-1-7_seed42",
|
|
"livemathematicianbench": "data/ablation_splits/livemathematicianbench/2-1-7_seed42",
|
|
"alfworld": "data/ablation_splits/alfworld/2-1-7_seed42",
|
|
"docvqa": "/home/azureuser/zisu/SkillReflection/data/docvqa/splits",
|
|
}
|
|
|
|
DEFAULT_TRAIN_SIZE = {
|
|
"searchqa": 400,
|
|
"spreadsheetbench": 80,
|
|
"livemathematicianbench": 35,
|
|
"alfworld": 39,
|
|
"docvqa": 1070,
|
|
}
|
|
|
|
BATCH_SIZE_VALUES: tuple[int | str, ...] = (8, 24, 40, 56, "full")
|
|
|
|
SPLITS = {
|
|
"searchqa": {
|
|
"1shot": ("data/ablation_splits/searchqa/1shot_seed42", ("optimizer.slow_update_samples=1",)),
|
|
"1-1-8": ("data/ablation_splits/searchqa/1-1-8_seed42", ()),
|
|
"2-1-7": ("data/ablation_splits/searchqa/2-1-7_seed42", ()),
|
|
"4-1-5": ("data/ablation_splits/searchqa/4-1-5_seed42", ()),
|
|
},
|
|
"spreadsheetbench": {
|
|
"1shot": ("data/ablation_splits/spreadsheetbench/1shot_seed42", ("optimizer.slow_update_samples=1",)),
|
|
"1-1-8": ("data/ablation_splits/spreadsheetbench/1-1-8_seed42", ()),
|
|
"2-1-7": ("data/ablation_splits/spreadsheetbench/2-1-7_seed42", ()),
|
|
"4-1-5": ("data/ablation_splits/spreadsheetbench/4-1-5_seed42", ()),
|
|
},
|
|
"livemathematicianbench": {
|
|
"1shot": ("data/ablation_splits/livemathematicianbench/1shot_seed42", ("optimizer.slow_update_samples=1",)),
|
|
"1-1-8": ("data/ablation_splits/livemathematicianbench/1-1-8_seed42", ()),
|
|
"2-1-7": ("data/ablation_splits/livemathematicianbench/2-1-7_seed42", ()),
|
|
"4-1-5": ("data/ablation_splits/livemathematicianbench/4-1-5_seed42", ()),
|
|
},
|
|
"alfworld": {
|
|
"1shot": ("data/ablation_splits/alfworld/1shot_seed42", ("optimizer.slow_update_samples=1",)),
|
|
"1-1-8": ("data/ablation_splits/alfworld/1-1-8_seed42", ()),
|
|
"2-1-7": ("data/ablation_splits/alfworld/2-1-7_seed42", ()),
|
|
"4-1-5": ("data/ablation_splits/alfworld/4-1-5_seed42", ()),
|
|
},
|
|
"docvqa": {
|
|
"1shot": ("data/ablation_splits/docvqa/1shot_seed42", ("optimizer.slow_update_samples=1",)),
|
|
"1-1-8": ("data/ablation_splits/docvqa/1-1-8_seed42", ()),
|
|
"2-1-7": ("/home/azureuser/zisu/SkillReflection/data/docvqa/splits", ()),
|
|
"4-1-5": ("data/ablation_splits/docvqa/4-1-5_seed42", ()),
|
|
},
|
|
}
|
|
|
|
|
|
def common_overrides(benchmark: str, out_root: Path) -> list[str]:
|
|
return [
|
|
"model.teacher_backend=openai_chat",
|
|
"model.student_backend=openai_chat",
|
|
"model.teacher=gpt-5.5",
|
|
"model.student=gpt-5.5",
|
|
f"model.teacher_azure_openai_endpoint={T2_ENDPOINT}",
|
|
"model.teacher_azure_openai_api_version=2024-12-01-preview",
|
|
"model.teacher_azure_openai_auth_mode=azure_cli",
|
|
f"model.student_azure_openai_endpoint={T2_ENDPOINT}",
|
|
"model.student_azure_openai_api_version=2024-12-01-preview",
|
|
"model.student_azure_openai_auth_mode=azure_cli",
|
|
"model.reasoning_effort=medium",
|
|
"train.num_epochs=4",
|
|
"train.train_size=0",
|
|
"train.batch_size=40",
|
|
"train.accumulation=1",
|
|
"train.seed=42",
|
|
"gradient.minibatch_size=8",
|
|
"gradient.merge_batch_size=8",
|
|
"gradient.analyst_workers=16",
|
|
"gradient.use_deep_reflect=false",
|
|
"optimizer.learning_rate=4",
|
|
"optimizer.min_learning_rate=2",
|
|
"optimizer.lr_scheduler=cosine",
|
|
"optimizer.skill_update_mode=patch",
|
|
"optimizer.use_slow_update=true",
|
|
"optimizer.slow_update_samples=20",
|
|
"optimizer.use_meta_skill=true",
|
|
"optimizer.use_meta_reflect=false",
|
|
"evaluation.use_gate=true",
|
|
"evaluation.eval_test=true",
|
|
"env.split_mode=split_dir",
|
|
f"env.split_dir={DEFAULT_SPLIT[benchmark]}",
|
|
f"env.out_root={out_root}",
|
|
]
|
|
|
|
|
|
def make_experiment(
|
|
group: str,
|
|
benchmark: str,
|
|
suffix: str,
|
|
run_root: Path,
|
|
overrides: list[str],
|
|
) -> Experiment:
|
|
run_id = f"{group}-{benchmark}-{suffix}"
|
|
out_root = run_root / run_id
|
|
all_overrides = common_overrides(benchmark, out_root)
|
|
all_overrides.extend(overrides)
|
|
return Experiment(
|
|
run_id=run_id,
|
|
benchmark=benchmark,
|
|
config=BENCH_CONFIG[benchmark],
|
|
overrides=tuple(all_overrides),
|
|
)
|
|
|
|
|
|
def build_matrix(
|
|
groups: set[str],
|
|
benchmarks: list[str],
|
|
run_root: Path,
|
|
*,
|
|
include_duplicate_defaults: bool = False,
|
|
) -> list[Experiment]:
|
|
exps: list[Experiment] = []
|
|
group_order = [
|
|
"default",
|
|
"split",
|
|
"batch",
|
|
"mbs",
|
|
"lr",
|
|
"sched",
|
|
"slown",
|
|
"mod",
|
|
"smodel",
|
|
"longpair",
|
|
"lrctrl",
|
|
]
|
|
|
|
for group in group_order:
|
|
if group not in groups:
|
|
continue
|
|
for benchmark in benchmarks:
|
|
if group == "default":
|
|
exps.append(make_experiment("DEFAULT", benchmark, "5.5", run_root, []))
|
|
continue
|
|
|
|
if group == "split":
|
|
for tag, (split_dir, extra) in SPLITS[benchmark].items():
|
|
if not include_duplicate_defaults and tag == "2-1-7":
|
|
continue
|
|
exps.append(make_experiment(
|
|
"SPLIT",
|
|
benchmark,
|
|
tag,
|
|
run_root,
|
|
[f"env.split_dir={split_dir}", *extra],
|
|
))
|
|
continue
|
|
|
|
if group == "mbs":
|
|
for value in (1, 2, 4, 8, 16, 32):
|
|
if not include_duplicate_defaults and value == 8:
|
|
continue
|
|
exps.append(make_experiment(
|
|
"MBS",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[f"gradient.minibatch_size={value}"],
|
|
))
|
|
continue
|
|
|
|
if group == "batch":
|
|
for value in BATCH_SIZE_VALUES:
|
|
if not include_duplicate_defaults and value == 40:
|
|
continue
|
|
batch_size = DEFAULT_TRAIN_SIZE[benchmark] if value == "full" else int(value)
|
|
exps.append(make_experiment(
|
|
"BATCH",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[
|
|
f"train.batch_size={batch_size}",
|
|
"gradient.minibatch_size=8",
|
|
],
|
|
))
|
|
continue
|
|
|
|
if group == "lr":
|
|
for value in (1, 2, 4, 8, 16):
|
|
exps.append(make_experiment(
|
|
"LR",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[
|
|
"optimizer.lr_scheduler=constant",
|
|
"optimizer.min_learning_rate=1",
|
|
f"optimizer.learning_rate={value}",
|
|
],
|
|
))
|
|
continue
|
|
|
|
if group == "sched":
|
|
for value in ("constant", "cosine", "linear"):
|
|
if not include_duplicate_defaults and value == "cosine":
|
|
continue
|
|
exps.append(make_experiment(
|
|
"SCHED",
|
|
benchmark,
|
|
value,
|
|
run_root,
|
|
[f"optimizer.lr_scheduler={value}"],
|
|
))
|
|
continue
|
|
|
|
if group == "slown":
|
|
for value in (5, 10, 20, 40):
|
|
if not include_duplicate_defaults and value == 20:
|
|
continue
|
|
exps.append(make_experiment(
|
|
"SLOWN",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[f"optimizer.slow_update_samples={value}"],
|
|
))
|
|
continue
|
|
|
|
if group == "mod":
|
|
settings = {
|
|
"slow-meta": ("true", "true"),
|
|
"slow-only": ("true", "false"),
|
|
"meta-only": ("false", "true"),
|
|
"none": ("false", "false"),
|
|
}
|
|
for tag, (slow, meta) in settings.items():
|
|
if not include_duplicate_defaults and tag == "slow-meta":
|
|
continue
|
|
exps.append(make_experiment(
|
|
"MOD",
|
|
benchmark,
|
|
tag,
|
|
run_root,
|
|
[
|
|
f"optimizer.use_slow_update={slow}",
|
|
f"optimizer.use_meta_skill={meta}",
|
|
],
|
|
))
|
|
continue
|
|
|
|
if group == "smodel":
|
|
student_settings = {
|
|
"5.4": [
|
|
"model.student=gpt-5.4-pro",
|
|
f"model.student_azure_openai_endpoint={T2_ENDPOINT}",
|
|
"model.student_azure_openai_api_version=2025-03-01-preview",
|
|
"model.student_azure_openai_auth_mode=azure_cli",
|
|
],
|
|
"5.4-mini": [
|
|
"model.student=gpt-5.4-mini",
|
|
f"model.student_azure_openai_endpoint={SEARCHAGENT5_ENDPOINT}",
|
|
"model.student_azure_openai_api_version=2024-12-01-preview",
|
|
"model.student_azure_openai_auth_mode=azure_cli",
|
|
],
|
|
"5.5": [],
|
|
}
|
|
for tag, overrides in student_settings.items():
|
|
if not include_duplicate_defaults and tag == "5.5":
|
|
continue
|
|
exps.append(make_experiment("SMODEL", benchmark, tag, run_root, overrides))
|
|
continue
|
|
|
|
if group == "longpair":
|
|
for value in ("changed", "unchanged"):
|
|
exps.append(make_experiment(
|
|
"LONGPAIR",
|
|
benchmark,
|
|
value,
|
|
run_root,
|
|
[f"optimizer.longitudinal_pair_policy={value}"],
|
|
))
|
|
continue
|
|
|
|
if group == "lrctrl":
|
|
settings = {
|
|
"autonomous": ["optimizer.lr_control_mode=autonomous"],
|
|
"full-rewrite": [
|
|
"optimizer.lr_control_mode=none",
|
|
"optimizer.skill_update_mode=full_rewrite_minibatch",
|
|
],
|
|
}
|
|
for tag, overrides in settings.items():
|
|
exps.append(make_experiment("LRCTRL", benchmark, tag, run_root, overrides))
|
|
continue
|
|
|
|
return exps
|
|
|
|
|
|
def _build_matrix_legacy(
|
|
groups: set[str],
|
|
benchmarks: list[str],
|
|
run_root: Path,
|
|
*,
|
|
include_duplicate_defaults: bool = False,
|
|
) -> list[Experiment]:
|
|
exps: list[Experiment] = []
|
|
for benchmark in benchmarks:
|
|
if "default" in groups:
|
|
exps.append(make_experiment("DEFAULT", benchmark, "5.5", run_root, []))
|
|
|
|
if "split" in groups:
|
|
for tag, (split_dir, extra) in SPLITS[benchmark].items():
|
|
if not include_duplicate_defaults and tag == "2-1-7":
|
|
continue
|
|
exps.append(make_experiment(
|
|
"SPLIT",
|
|
benchmark,
|
|
tag,
|
|
run_root,
|
|
[f"env.split_dir={split_dir}", *extra],
|
|
))
|
|
|
|
if "mbs" in groups:
|
|
for value in (1, 2, 4, 8, 16, 32):
|
|
if not include_duplicate_defaults and value == 8:
|
|
continue
|
|
exps.append(make_experiment(
|
|
"MBS",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[f"gradient.minibatch_size={value}"],
|
|
))
|
|
|
|
if "batch" in groups:
|
|
for value in BATCH_SIZE_VALUES:
|
|
if not include_duplicate_defaults and value == 40:
|
|
continue
|
|
batch_size = DEFAULT_TRAIN_SIZE[benchmark] if value == "full" else int(value)
|
|
exps.append(make_experiment(
|
|
"BATCH",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[
|
|
f"train.batch_size={batch_size}",
|
|
"gradient.minibatch_size=8",
|
|
],
|
|
))
|
|
|
|
if "lr" in groups:
|
|
for value in (1, 2, 4, 8, 16):
|
|
exps.append(make_experiment(
|
|
"LR",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[
|
|
"optimizer.lr_scheduler=constant",
|
|
"optimizer.min_learning_rate=1",
|
|
f"optimizer.learning_rate={value}",
|
|
],
|
|
))
|
|
|
|
if "sched" in groups:
|
|
for value in ("constant", "cosine", "linear"):
|
|
if not include_duplicate_defaults and value == "cosine":
|
|
continue
|
|
exps.append(make_experiment(
|
|
"SCHED",
|
|
benchmark,
|
|
value,
|
|
run_root,
|
|
[f"optimizer.lr_scheduler={value}"],
|
|
))
|
|
|
|
if "slown" in groups:
|
|
for value in (5, 10, 20, 40):
|
|
if not include_duplicate_defaults and value == 20:
|
|
continue
|
|
exps.append(make_experiment(
|
|
"SLOWN",
|
|
benchmark,
|
|
str(value),
|
|
run_root,
|
|
[f"optimizer.slow_update_samples={value}"],
|
|
))
|
|
|
|
if "mod" in groups:
|
|
settings = {
|
|
"slow-meta": ("true", "true"),
|
|
"slow-only": ("true", "false"),
|
|
"meta-only": ("false", "true"),
|
|
"none": ("false", "false"),
|
|
}
|
|
for tag, (slow, meta) in settings.items():
|
|
if not include_duplicate_defaults and tag == "slow-meta":
|
|
continue
|
|
exps.append(make_experiment(
|
|
"MOD",
|
|
benchmark,
|
|
tag,
|
|
run_root,
|
|
[
|
|
f"optimizer.use_slow_update={slow}",
|
|
f"optimizer.use_meta_skill={meta}",
|
|
],
|
|
))
|
|
|
|
if "smodel" in groups:
|
|
student_settings = {
|
|
"5.4": [
|
|
"model.student=gpt-5.4-pro",
|
|
f"model.student_azure_openai_endpoint={T2_ENDPOINT}",
|
|
"model.student_azure_openai_api_version=2025-03-01-preview",
|
|
"model.student_azure_openai_auth_mode=azure_cli",
|
|
],
|
|
"5.4-mini": [
|
|
"model.student=gpt-5.4-mini",
|
|
f"model.student_azure_openai_endpoint={SEARCHAGENT5_ENDPOINT}",
|
|
"model.student_azure_openai_api_version=2024-12-01-preview",
|
|
"model.student_azure_openai_auth_mode=azure_cli",
|
|
],
|
|
"5.5": [],
|
|
}
|
|
for tag, overrides in student_settings.items():
|
|
if not include_duplicate_defaults and tag == "5.5":
|
|
continue
|
|
exps.append(make_experiment("SMODEL", benchmark, tag, run_root, overrides))
|
|
|
|
if "longpair" in groups:
|
|
for value in ("changed", "unchanged"):
|
|
exps.append(make_experiment(
|
|
"LONGPAIR",
|
|
benchmark,
|
|
value,
|
|
run_root,
|
|
[f"optimizer.longitudinal_pair_policy={value}"],
|
|
))
|
|
|
|
if "lrctrl" in groups:
|
|
settings = {
|
|
"autonomous": ["optimizer.lr_control_mode=autonomous"],
|
|
"full-rewrite": [
|
|
"optimizer.lr_control_mode=none",
|
|
"optimizer.skill_update_mode=full_rewrite_minibatch",
|
|
],
|
|
}
|
|
for tag, overrides in settings.items():
|
|
exps.append(make_experiment("LRCTRL", benchmark, tag, run_root, overrides))
|
|
|
|
return exps
|
|
|
|
|
|
def command_for(exp: Experiment) -> list[str]:
|
|
return [
|
|
str(PYTHON_BIN),
|
|
"scripts/train.py",
|
|
"--config",
|
|
exp.config,
|
|
"--cfg-options",
|
|
*exp.overrides,
|
|
]
|
|
|
|
|
|
def active_run_ids(run_root: Path, valid_run_ids: set[str] | None = None) -> set[str]:
|
|
try:
|
|
raw = subprocess.check_output(["pgrep", "-af", "scripts/train.py"], text=True)
|
|
except subprocess.CalledProcessError:
|
|
return set()
|
|
pattern = re.compile(re.escape(str(run_root)) + r"/([^\s]+)")
|
|
active: set[str] = set()
|
|
for line in raw.splitlines():
|
|
for match in pattern.finditer(line):
|
|
run_id = match.group(1).strip("'\"")
|
|
if run_id.endswith(".log") or "/" in run_id:
|
|
continue
|
|
if valid_run_ids is not None and run_id not in valid_run_ids:
|
|
continue
|
|
active.add(run_id)
|
|
return active
|
|
|
|
|
|
def completed_run_ids(run_root: Path) -> set[str]:
|
|
return {
|
|
path.parent.name
|
|
for path in run_root.glob("*/summary.json")
|
|
if path.is_file()
|
|
}
|
|
|
|
|
|
def print_commands(exps: list[Experiment]) -> None:
|
|
for exp in exps:
|
|
cmd = command_for(exp)
|
|
print(f"\n# {exp.run_id}")
|
|
print(" ".join(subprocess.list2cmdline([part]) for part in cmd))
|
|
|
|
|
|
def run_commands(
|
|
exps: list[Experiment],
|
|
run_root: Path,
|
|
max_parallel: int,
|
|
run_retries: int,
|
|
) -> int:
|
|
logs_dir = run_root / "logs"
|
|
logs_dir.mkdir(parents=True, exist_ok=True)
|
|
active: list[tuple[Experiment, subprocess.Popen, object]] = []
|
|
valid_run_ids = {exp.run_id for exp in exps}
|
|
skipped_completed = completed_run_ids(run_root)
|
|
skipped_active = active_run_ids(run_root, valid_run_ids)
|
|
pending: list[tuple[Experiment, int]] = [
|
|
(exp, 0)
|
|
for exp in exps
|
|
if exp.run_id not in skipped_completed and exp.run_id not in skipped_active
|
|
]
|
|
for run_id in sorted(skipped_completed):
|
|
print(f"[SKIP_COMPLETED] {run_id}", flush=True)
|
|
for run_id in sorted(skipped_active):
|
|
print(f"[SKIP_ACTIVE] {run_id}", flush=True)
|
|
failures = 0
|
|
|
|
while pending or active:
|
|
external_active = active_run_ids(run_root, valid_run_ids) - {exp.run_id for exp, _, _ in active}
|
|
while pending and len(active) + len(external_active) < max_parallel:
|
|
exp, attempt = pending.pop(0)
|
|
log_path = logs_dir / f"{exp.run_id}.log"
|
|
if attempt:
|
|
log_path = logs_dir / f"{exp.run_id}.retry{attempt}.log"
|
|
log_f = open(log_path, "w", encoding="utf-8")
|
|
print(f"[START] {exp.run_id} attempt={attempt + 1} log={log_path}", flush=True)
|
|
proc = subprocess.Popen(
|
|
command_for(exp),
|
|
cwd=PROJECT_ROOT,
|
|
stdout=log_f,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
setattr(proc, "_attempt", attempt)
|
|
active.append((exp, proc, log_f))
|
|
|
|
time.sleep(5)
|
|
still_active: list[tuple[Experiment, subprocess.Popen, object]] = []
|
|
for exp, proc, log_f in active:
|
|
rc = proc.poll()
|
|
if rc is None:
|
|
still_active.append((exp, proc, log_f))
|
|
continue
|
|
log_f.close()
|
|
if rc == 0:
|
|
print(f"[DONE] {exp.run_id}", flush=True)
|
|
else:
|
|
if getattr(proc, "_attempt", 0) < run_retries:
|
|
next_attempt = getattr(proc, "_attempt", 0) + 1
|
|
pending.append((exp, next_attempt))
|
|
print(f"[RETRY] {exp.run_id} rc={rc} next_attempt={next_attempt + 1}", flush=True)
|
|
else:
|
|
failures += 1
|
|
print(f"[FAIL] {exp.run_id} rc={rc}", flush=True)
|
|
active = still_active
|
|
|
|
return failures
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--groups",
|
|
nargs="+",
|
|
default=["default"],
|
|
choices=[
|
|
"default",
|
|
"split",
|
|
"batch",
|
|
"mbs",
|
|
"lr",
|
|
"sched",
|
|
"slown",
|
|
"mod",
|
|
"smodel",
|
|
"longpair",
|
|
"lrctrl",
|
|
"all",
|
|
],
|
|
help="Experiment groups to include. Default: default.",
|
|
)
|
|
parser.add_argument(
|
|
"--bench",
|
|
nargs="+",
|
|
default=["searchqa", "spreadsheetbench"],
|
|
choices=["searchqa", "spreadsheetbench", "livemathematicianbench", "alfworld", "docvqa"],
|
|
)
|
|
parser.add_argument("--run-root", default="", help="Output root. Default: outputs/ablation_<UTC timestamp>.")
|
|
parser.add_argument("--max-parallel", type=int, default=1)
|
|
parser.add_argument("--run-retries", type=int, default=1, help="Retry failed runs this many times. Default: 1.")
|
|
parser.add_argument(
|
|
"--include-duplicate-defaults",
|
|
action="store_true",
|
|
help="Also run ablation points that are exactly the default setting.",
|
|
)
|
|
parser.add_argument("--execute", action="store_true", help="Actually start runs. Without this, print commands only.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
groups = set(args.groups)
|
|
if "all" in groups:
|
|
groups = {"default", "split", "batch", "mbs", "lr", "sched", "slown", "mod", "smodel"}
|
|
|
|
ts = time.strftime("%Y%m%d_%H%M%S", time.gmtime())
|
|
run_root = Path(args.run_root) if args.run_root else PROJECT_ROOT / "outputs" / f"ablation_{ts}"
|
|
if not run_root.is_absolute():
|
|
run_root = PROJECT_ROOT / run_root
|
|
run_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
exps = build_matrix(
|
|
groups,
|
|
args.bench,
|
|
run_root,
|
|
include_duplicate_defaults=args.include_duplicate_defaults,
|
|
)
|
|
print(f"run_root={run_root}")
|
|
print(f"num_experiments={len(exps)}")
|
|
print(f"groups={','.join(sorted(groups))}")
|
|
print(f"bench={','.join(args.bench)}")
|
|
|
|
if not args.execute:
|
|
print_commands(exps)
|
|
return
|
|
|
|
max_parallel = max(1, int(args.max_parallel))
|
|
failures = run_commands(
|
|
exps,
|
|
run_root,
|
|
max_parallel=max_parallel,
|
|
run_retries=max(0, int(args.run_retries)),
|
|
)
|
|
if failures:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|