zhoulin
da22885645
run_batch.py drives run_pytest.sh per operator across GPUs (stable-hash sharding, per-op subprocess isolation, process-group timeouts, retry with deterministic-failure cutoff, two-level dtype fallback, .complete resume, per-op REPLAY_FROM). batch_summary.py aggregates run.log tables into summary.csv. ops/ holds the curated assets: dual-repo inventories rebuilt via AST scan + pytest collect verification, shape sets migrated from the old regression harness and merged with upstream core_shapes class-name keys (upstream's set_shapes falls back op_name -> MRO class name -> 1-D DEFAULT_SHAPES, so replacing the shape file without class keys crashes the BLAS family), and a dismiss list where all 76 entries carry verified reasons. Validated end to end: 1036-op full screen with zero failures. Also fix two cudagraph plugin regressions: newer torch appends "enable device-side assertions" to every CUDA error, so the loose fatal-error marker disabled the documented do_bench fallback entirely; and an aborted graph capture can leave the default CUDA RNG generator stuck in capturing state, poisoning every later torch.randn - captures now run under a throwaway RNG state. run_pytest.sh gains an optional DTYPES passthrough.
295 lines
14 KiB
Python
295 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""重建批量测试的算子清单（inventory），并可校验迁移 shape yaml。
|
||
|
||
对每个 FlagGems 仓库静态扫描 benchmark/ 下的 test_*.py（含 test_ 前缀的
|
||
子目录，如 test_FLA/——op_file 记相对路径去前后缀，run_pytest.sh 的
|
||
`benchmark/test_${OP_FILE}.py` 拼接对其天然成立；models_benchmark/ 等
|
||
非 test_ 前缀路径不在性能批量范围）：
|
||
|
||
- 每个模块级 test_* 函数出一行 inventory：op（函数名去 test_ 前缀，即
|
||
run_pytest.sh 的 OP）、op_file（文件名去 test_ 前缀/.py 后缀，即 OP_FILE）；
|
||
- 标注 marked_skip（函数带无条件 pytest.mark.skip，跑了必 SKIP，批量驱动默认剔除）；
|
||
- 提取 op_name 字符串常量（op_name= 关键字实参，以及 *Benchmark 类
|
||
__init__ 里 super().__init__ 的首个字符串位置实参），用于 shape yaml 键校验；
|
||
- --verify-collect 时额外跑一次 pytest --collect-only 核实函数确实可被收集
|
||
（AST 见到 ≠ pytest 收得到，import 失败/条件定义都会导致差异）。
|
||
|
||
产物：
|
||
ops/inventory_<label>.csv # repo,op,op_file,file,marked_skip,collected,op_names
|
||
--migrate-shapes 时按 op_name 并集过滤输入 yaml、合并上游 core_shapes 底座：
|
||
<out>.yaml + <out>.unmatched.yaml（未匹配任何 op_name 的键，供人工复核后删除）
|
||
|
||
用法示例（默认扫描两个新仓库）：
|
||
python ops/gen_inventory.py --verify-collect
|
||
# 上游 op_name/core_shapes 变动后，对现有 shape 集重新校验（输入=输出即原地刷新）
|
||
python ops/gen_inventory.py --migrate-shapes ops/shapes_single.yaml \
|
||
--shapes-out ops/shapes_single.yaml
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import ast
|
||
import csv
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Set, Tuple
|
||
|
||
import yaml
|
||
|
||
OPS_DIR = Path(__file__).resolve().parent
|
||
|
||
DEFAULT_REPOS = [
|
||
("main", Path("/workspace/dev/FlagGems")),
|
||
("vllm", Path("/workspace/dev/FlagGems-vllm")),
|
||
]
|
||
|
||
INVENTORY_FIELDS = ["repo", "op", "op_file", "file", "marked_skip", "collected", "op_names"]
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
||
p.add_argument("--repo", action="append", default=None, metavar="LABEL=PATH",
|
||
help="要扫描的仓库（可多次），默认 main=/workspace/dev/FlagGems "
|
||
"与 vllm=/workspace/dev/FlagGems-vllm")
|
||
p.add_argument("--verify-collect", action="store_true",
|
||
help="用 pytest --collect-only 核实每个测试函数可被收集（慢，"
|
||
"需要完整运行环境；结果写进 collected 列）")
|
||
p.add_argument("--out-dir", type=Path, default=OPS_DIR,
|
||
help="inventory CSV 输出目录（默认 ops/）")
|
||
p.add_argument("--migrate-shapes", type=Path, default=None,
|
||
help="旧 shape yaml 路径；按扫描出的 op_name 并集过滤后迁移")
|
||
p.add_argument("--shapes-out", type=Path, default=None,
|
||
help="迁移后 yaml 输出路径（--migrate-shapes 时必填）")
|
||
return p.parse_args()
|
||
|
||
|
||
def _resolve_repos(raw: Optional[List[str]]) -> List[Tuple[str, Path]]:
|
||
if not raw:
|
||
return [(label, path) for label, path in DEFAULT_REPOS if path.is_dir()]
|
||
repos = []
|
||
for item in raw:
|
||
label, _, path = item.partition("=")
|
||
if not path:
|
||
raise SystemExit(f"--repo 需要 LABEL=PATH 形式，收到: {item!r}")
|
||
repos.append((label, Path(path)))
|
||
return repos
|
||
|
||
|
||
def _is_unconditional_skip(dec: ast.AST) -> bool:
|
||
"""pytest.mark.skip（非 skipif）——带不带 reason 实参都算无条件。"""
|
||
target = dec.func if isinstance(dec, ast.Call) else dec
|
||
if not (isinstance(target, ast.Attribute) and target.attr == "skip"):
|
||
return False
|
||
mark = target.value
|
||
return (isinstance(mark, ast.Attribute) and mark.attr == "mark"
|
||
and isinstance(mark.value, ast.Name) and mark.value.id == "pytest")
|
||
|
||
|
||
def _op_names_in(node: ast.AST) -> Set[str]:
|
||
"""节点范围内的 op_name 字符串常量，覆盖三种上游写法：
|
||
1. 任意调用的 op_name= 关键字实参；
|
||
2. super().__init__ 的首个字符串位置实参（FusedDeepseekV4... 模式）；
|
||
3. *Benchmark 类实例化的首个字符串位置实参（ScaledMMBenchmark("scaled_mm",
|
||
...) 这类经形参转发、常量提取够不到 super() 调用的模式）。
|
||
误报只会让 shape yaml 多留一个无人读取的键，无害。"""
|
||
names: Set[str] = set()
|
||
for sub in ast.walk(node):
|
||
if not isinstance(sub, ast.Call):
|
||
continue
|
||
for kw in sub.keywords:
|
||
if kw.arg == "op_name" and isinstance(kw.value, ast.Constant) \
|
||
and isinstance(kw.value.value, str):
|
||
names.add(kw.value.value)
|
||
first_str = (sub.args[0].value
|
||
if sub.args and isinstance(sub.args[0], ast.Constant)
|
||
and isinstance(sub.args[0].value, str) else None)
|
||
if first_str is None:
|
||
continue
|
||
func = sub.func
|
||
if (isinstance(func, ast.Attribute) and func.attr == "__init__"
|
||
and isinstance(func.value, ast.Call)
|
||
and isinstance(func.value.func, ast.Name)
|
||
and func.value.func.id == "super"):
|
||
names.add(first_str)
|
||
callee = func.id if isinstance(func, ast.Name) else (
|
||
func.attr if isinstance(func, ast.Attribute) else "")
|
||
if callee.endswith("Benchmark"):
|
||
names.add(first_str)
|
||
return names
|
||
|
||
|
||
def scan_repo(label: str, root: Path) -> Tuple[List[Dict[str, str]], Set[str]]:
|
||
"""返回 (inventory 行, 该仓库全部 op_name 集合)。"""
|
||
bench = root / "benchmark"
|
||
rows: List[Dict[str, str]] = []
|
||
op_names_all: Set[str] = set()
|
||
for path in sorted(bench.rglob("test_*.py")):
|
||
rel = path.relative_to(bench)
|
||
# 仅收 test_ 前缀链路（根级 test_x.py 与 test_XXX/ 子目录），保证
|
||
# "benchmark/test_" + op_file + ".py" 能原样重建路径。
|
||
if not str(rel).startswith("test_"):
|
||
continue
|
||
try:
|
||
tree = ast.parse(path.read_text())
|
||
except SyntaxError as exc:
|
||
print(f"[gen] warning: {path} 解析失败，跳过: {exc}", file=sys.stderr)
|
||
continue
|
||
op_names_all |= _op_names_in(tree)
|
||
op_file = str(rel)[len("test_"):-len(".py")]
|
||
for node in tree.body:
|
||
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||
continue
|
||
if not node.name.startswith("test_"):
|
||
continue
|
||
marked = any(_is_unconditional_skip(d) for d in node.decorator_list)
|
||
func_ops = sorted(_op_names_in(node))
|
||
rows.append({
|
||
"repo": label,
|
||
"op": node.name[len("test_"):],
|
||
"op_file": op_file,
|
||
"file": f"benchmark/{rel}",
|
||
"marked_skip": "yes" if marked else "",
|
||
"collected": "",
|
||
"op_names": ";".join(func_ops),
|
||
})
|
||
return rows, op_names_all
|
||
|
||
|
||
def verify_collect(label: str, root: Path, rows: List[Dict[str, str]]) -> None:
|
||
"""跑 pytest --collect-only 核实测试函数可被收集，写 collected 列。"""
|
||
env = os.environ.copy()
|
||
env.setdefault("GEMS_VENDOR", "nvidia") # 跳过 nvidia-smi 子进程探测（可能挂死）
|
||
proc = subprocess.run(
|
||
[sys.executable, "-m", "pytest", "--collect-only", "-q",
|
||
"--continue-on-collection-errors", "benchmark/"],
|
||
cwd=str(root), text=True, capture_output=True, env=env, timeout=1200,
|
||
)
|
||
collected_funcs: Set[Tuple[str, str]] = set()
|
||
for line in proc.stdout.splitlines():
|
||
line = line.strip()
|
||
if "::" not in line or not line.startswith("benchmark/"):
|
||
continue
|
||
file_part, _, rest = line.partition("::")
|
||
func = rest.split("::")[0].split("[")[0]
|
||
collected_funcs.add((file_part, func))
|
||
if not collected_funcs:
|
||
print(f"[gen] warning: {label} collect 无结果 (rc={proc.returncode})，"
|
||
f"stderr 尾部: {proc.stderr[-500:]}", file=sys.stderr)
|
||
return
|
||
for row in rows:
|
||
key = (row["file"], f"test_{row['op']}")
|
||
row["collected"] = "yes" if key in collected_funcs else "no"
|
||
n_missing = sum(1 for r in rows if r["collected"] == "no")
|
||
if n_missing:
|
||
print(f"[gen] {label}: {n_missing} 个 AST 可见但 pytest 未收集到的函数"
|
||
f"（import 失败/条件定义），inventory 中 collected=no")
|
||
|
||
|
||
def carry_over_collected(out_dir: Path, label: str,
|
||
rows: List[Dict[str, str]]) -> None:
|
||
"""未跑 --verify-collect 时，从已有 CSV 继承 collected 列，避免重写清单时
|
||
把上次核实的结果冲掉。"""
|
||
path = out_dir / f"inventory_{label}.csv"
|
||
if not path.is_file():
|
||
return
|
||
with path.open(newline="") as f:
|
||
prev = {(r["op"], r["op_file"]): r.get("collected", "")
|
||
for r in csv.DictReader(f)}
|
||
for row in rows:
|
||
row["collected"] = prev.get((row["op"], row["op_file"]), "")
|
||
|
||
|
||
def write_inventory(out_dir: Path, label: str, rows: List[Dict[str, str]]) -> Path:
|
||
out = out_dir / f"inventory_{label}.csv"
|
||
tmp = out.with_name(out.name + ".tmp")
|
||
with tmp.open("w", newline="") as f:
|
||
writer = csv.DictWriter(f, fieldnames=INVENTORY_FIELDS)
|
||
writer.writeheader()
|
||
writer.writerows(rows)
|
||
os.replace(tmp, out)
|
||
return out
|
||
|
||
|
||
def migrate_shapes(old_path: Path, out_path: Path, valid_op_names: Set[str],
|
||
repos: List[Tuple[str, Path]]) -> None:
|
||
"""迁移旧 shape yaml，并以上游 core_shapes.yaml 为底座合并。
|
||
|
||
上游 set_shapes 的回退链是 op_name 键 → MRO 类名键（BlasBenchmark 等）→
|
||
基类 DEFAULT_SHAPES（1 维）。--shape_file 是整体替换而非叠加，若我们的
|
||
yaml 缺少类名级条目，BLAS 这类需要 (B,M,N,K) 的算子会跌到 1 维默认值上
|
||
解包崩溃。因此把各仓库 core_shapes.yaml 里我们没有的键（含类名键与
|
||
op 级键）原样并入，迁移条目优先，先并 main 后并 vllm（同名键先到先得）。"""
|
||
data = yaml.safe_load(old_path.read_text()) or {}
|
||
core_maps = []
|
||
for label, root in repos:
|
||
core = root / "benchmark" / "core_shapes.yaml"
|
||
if core.is_file():
|
||
core_maps.append((label, yaml.safe_load(core.read_text()) or {}))
|
||
core_keys = {k for _, m in core_maps for k in m}
|
||
kept = {k: v for k, v in data.items() if k in valid_op_names}
|
||
# core_shapes 来源的键（类名键等）不算"未匹配"——它们每次都从上游取新版，
|
||
# 这样输入=输出的在地刷新也能同步上游 core_shapes 的变动。
|
||
dropped = {k: v for k, v in data.items()
|
||
if k not in valid_op_names and k not in core_keys}
|
||
n_curated = len(kept)
|
||
merged_from = []
|
||
for label, m in core_maps:
|
||
n_before = len(kept)
|
||
for key, val in m.items():
|
||
kept.setdefault(key, val)
|
||
merged_from.append(f"{label}:+{len(kept) - n_before}")
|
||
header = (f"# 由 gen_inventory.py 从 {old_path.name} 迁移：保留新上游仍存在的"
|
||
f" op_name 键 {n_curated}/{len(data)}，"
|
||
f"并合并上游 core_shapes.yaml 缺失键（{' '.join(merged_from)}）。\n"
|
||
"# 顶层键 = FlagGems op_name 或 Benchmark 类名（上游回退链需要）；\n"
|
||
"# 两者都无键的算子使用其基类 DEFAULT_SHAPES。\n")
|
||
out_path.write_text(header + yaml.safe_dump(
|
||
kept, sort_keys=True, default_flow_style=None, allow_unicode=True))
|
||
if dropped:
|
||
unmatched = out_path.with_suffix(".unmatched.yaml")
|
||
unmatched.write_text(
|
||
f"# {old_path.name} 中未匹配新上游任何 op_name 的键（{len(dropped)} 个），"
|
||
"供人工复核后手动挪回。\n"
|
||
+ yaml.safe_dump(dropped, sort_keys=True, default_flow_style=None,
|
||
allow_unicode=True))
|
||
print(f"[gen] shapes: 保留 {len(kept)}，剔除 {len(dropped)} -> {unmatched}")
|
||
else:
|
||
print(f"[gen] shapes: 全部 {len(kept)} 个键有效")
|
||
# 覆盖率：有效 op_name 里有多少没有 shape 条目（用上游默认 shape，仅提示）
|
||
uncovered = sorted(valid_op_names - set(kept))
|
||
print(f"[gen] shapes: 新上游 {len(valid_op_names)} 个 op_name 中 "
|
||
f"{len(uncovered)} 个无 shape 条目（将用上游默认 shape）")
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
repos = _resolve_repos(args.repo)
|
||
if not repos:
|
||
raise SystemExit("没有可扫描的仓库")
|
||
args.out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
all_op_names: Set[str] = set()
|
||
for label, root in repos:
|
||
rows, op_names = scan_repo(label, root)
|
||
all_op_names |= op_names
|
||
if args.verify_collect:
|
||
verify_collect(label, root, rows)
|
||
else:
|
||
carry_over_collected(args.out_dir, label, rows)
|
||
out = write_inventory(args.out_dir, label, rows)
|
||
n_skip = sum(1 for r in rows if r["marked_skip"])
|
||
print(f"[gen] {label}: {len(rows)} 个测试函数 -> {out}"
|
||
f"（其中 {n_skip} 个带无条件 skip 标记）")
|
||
|
||
if args.migrate_shapes:
|
||
if not args.shapes_out:
|
||
raise SystemExit("--migrate-shapes 需要同时给 --shapes-out")
|
||
migrate_shapes(args.migrate_shapes, args.shapes_out, all_op_names, repos)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|