Add multi-op batch screening layer; fix cudagraph fallback regressions

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.
2026-08-12 19:04:09 +00:00
parent 95895cea0e
commit da22885645
12 changed files with 14679 additions and 9 deletions
@@ -0,0 +1,294 @@
+#!/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()