Aharrypotter/zl_bench
1
0 Fork 0
forked from zhoulin/zl_bench
Code Pull Requests Activity

feat: make the harness backend-agnostic

Browse Source 
Everything about which accelerator a run targets was previously spelled CUDA:
device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES,
the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in
run_pytest.sh with no way to drop the CUDA-only members. Running the same
operator corpus on another vendor's device therefore meant editing the drivers.

The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with
the usual aliases):

- backends/ owns per-backend device enumeration, visibility env keys and
  capability flags, so nothing above it names a vendor;
- plugin_plan.py derives the -p list from those capabilities instead of a
  hardcoded array, which is what lets a non-CUDA run omit the graph-capture and
  cluster-fix plugins rather than watch them fail;
- _timing_plugin selects the timer from the same capabilities: CUDA keeps
  graph capture with its per-shape fallback, other backends stay on upstream
  do_bench and record timer=direct, because capture-then-fallback on a device
  without graph support only produces noise;
- _device_setup_plugin replaces the NVIDIA-only vendor guard;
- providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE
  directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to
  FlagGems' own core_shapes;
- _run_identity.py writes run_identity.json next to each run, recording backend,
  device, timer, autotune mode and the resolved test node. Two runs whose
  calibers differ are otherwise indistinguishable after the fact.

run_batch gains --backend and --devices (--gpus kept as an alias) and carries
backend/device/timer into ops_status.csv. test_backends.py covers the registry,
alias resolution, visibility env and plugin planning.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
AharrypotterandClaude Opus 5 committed 2026-09-09 16:39:14 +08:00
1 parent 35a928a516
commit c8ea586d0f
24 files changed
+1448 -159
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

run_batch.py
+84 -47
View File
@@ -1,13 +1,16 @@
#!/usr/bin/env python3
"""多算子批量测试驱动：按算子清单逐个调 run_pytest.sh，多卡分片、串行测量。
每个算子一个独立子进程（CUDA crash 只废单个算子）、独立产物目录
每个算子一个独立子进程（设备 crash 只废单个算子）、独立产物目录
<batch>/<repo>/<op>/（与单算子跑 run_pytest.sh 的产物结构完全一致，含
run.log / shapes.yaml / autotune_records / ttgir）。批量层只做编排：
run.log / shapes.yaml / autotune_records / ttgir / run_identity.json）。
批量层只做编排：
- 分片：按算子名稳定哈希到 GPU（每卡内部串行，测量口径与单算子一致；
- 分片：按算子名稳定哈希到设备（每卡内部串行，测量口径与单算子一致；
跨卡并行测量存在功耗耦合噪声，本模式定位为筛查口径，可疑算子回
单算子串行流程确认）；
- 后端：--backend / ZL_BACKEND 选择 cuda/musa/rocm/ascend，设备可见性
由 backends/ 写入，不再写死 CUDA_VISIBLE_DEVICES；
- 超时：--op-timeout 杀死挂死算子（记 TIMEOUT，不重试，继续跑）；
- 重试：pytest rc=1（用例失败/flaky）与信号杀（rc>=128）最多重试 2 次，
连续两次同 rc 视为确定性失败提前止损；--dtypes 下两级降级——遇上游
@@ -23,7 +26,7 @@ run.log / shapes.yaml / autotune_records / ttgir）。批量层只做编排：
示例：
# 全量筛查（两仓库全部可收集算子，排除 marked_skip 与 dismiss）
python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --op-timeout 1800
python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --op-timeout 1800
# 指定算子子集（repo 前缀默认 main；重名用 @文件名 消歧）
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -50,6 +53,8 @@ from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple
import batch_summary
from backends import UnknownBackendError, get_backend
from _run_identity import load_identity
ZL_BENCH = Path(__file__).resolve().parent
OPS_DIR = ZL_BENCH / "ops"
@@ -62,7 +67,9 @@ REPO_DIRS = {
}
STATUS_FIELDS = [
"repo", "op", "op_file", "run_name", "gpu", "status", "return_code",
"repo", "op", "op_file", "run_name",
"backend", "device", "gpu", "timer",
"status", "return_code",
"elapsed_ms", "rows_total", "rows_success", "notes", "out_dir",
]
@@ -194,48 +201,44 @@ def resolve_selection(args: argparse.Namespace,
return cases
# ---------------- GPU 分片 ----------------
# ---------------- 设备分片 ----------------
def visible_gpus(arg: Optional[str]) -> List[str]:
if arg:
return [g.strip() for g in arg.split(",") if g.strip()]
env = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
if env:
return [g.strip() for g in env.split(",") if g.strip()]
try:
out = subprocess.run(["nvidia-smi", "--list-gpus"], text=True,
capture_output=True, timeout=30).stdout
n = sum(1 for line in out.splitlines() if line.startswith("GPU "))
if n:
return [str(i) for i in range(n)]
except Exception:
pass
return ["0"]
def resolve_device_arg(devices: Optional[str], gpus: Optional[str]) -> Optional[str]:
"""Prefer --devices; --gpus is an alias. Passing both is an error."""
if devices is not None and gpus is not None:
raise ValueError("do not pass both --devices and --gpus")
return devices if devices is not None else gpus
def partition(cases: List[Case], gpus: List[str]) -> Dict[str, List[Case]]:
def visible_devices(backend, explicit: Optional[str] = None) -> List[str]:
return list(backend.enumerate_devices(explicit))
def partition(cases: List[Case], devices: List[str]) -> Dict[str, List[Case]]:
"""按 (repo, run_name) 稳定哈希分片：与选择顺序无关，resume 时同卡复用。"""
parts: Dict[str, List[Case]] = {g: [] for g in gpus}
parts: Dict[str, List[Case]] = {d: [] for d in devices}
for c in cases:
digest = hashlib.sha256(f"{c.repo}/{c.run_name}".encode()).digest()
parts[gpus[int.from_bytes(digest[:8], "big") % len(gpus)]].append(c)
parts[devices[int.from_bytes(digest[:8], "big") % len(devices)]].append(c)
return parts
# ---------------- 单算子执行 ----------------
def _spawn(case: Case, out_dir: Path, gpu: str, args: argparse.Namespace,
def _spawn(case: Case, out_dir: Path, device: str, args: argparse.Namespace,
dtypes: str) -> subprocess.Popen:
env = os.environ.copy()
env.pop("PARALLEL_WARMUP_GPUS", None) # 批量层已分片，算子内不再并行
env.pop("REPLAY_FROM", None)
env.pop("OUT_DIR", None)
env = args.backend_obj.apply_visibility(env, device)
env.update({
"OP": case.op,
"OP_FILE": case.op_file,
"FLAGGEMS_DIR": REPO_DIRS[case.repo],
"OUT_DIR": str(out_dir),
"CUDA_VISIBLE_DEVICES": gpu,
"ZL_BACKEND": args.backend,
"ZL_DEVICE": device,
"FLAGGEMS_PERF_COLOR": "never",
"NO_COLOR": "1",
"USE_FLAGTUNE": str(args.use_flagtune),
@@ -271,13 +274,14 @@ def _wait(proc: subprocess.Popen, timeout: int) -> Tuple[int, bool]:
return 124, True
def run_case(case: Case, batch_dir: Path, gpu: str,
def run_case(case: Case, batch_dir: Path, device: str,
args: argparse.Namespace) -> Dict[str, str]:
out_dir = batch_dir / case.repo / case.run_name
marker = out_dir / ".complete"
if args.resume and marker.is_file() and (out_dir / "run.log").is_file():
row = _classify(case, out_dir, gpu, 0, 0, "", timed_out=False)
row = _classify(case, out_dir, device, 0, 0, "", timed_out=False,
args=args)
row["notes"] = _join("resumed", row["notes"])
return row
@@ -295,7 +299,7 @@ def run_case(case: Case, batch_dir: Path, gpu: str,
log = out_dir / "run.log"
if log.is_file():
log.rename(out_dir / f"run.attempt{attempt - 1}.log")
rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, dtypes),
rc, timed_out = _wait(_spawn(case, out_dir, device, args, dtypes),
args.op_timeout)
rcs.append(rc)
if timed_out:
@@ -315,20 +319,22 @@ def run_case(case: Case, batch_dir: Path, gpu: str,
# 覆盖 baseline 对受限 dtype 编译失败（torch jiterator 的 special 函数族等）
# 但报错文案不是上游 "can't be supported by this op" 的场景。
if dtypes and not timed_out:
probe = _classify(case, out_dir, gpu, rcs[-1], 0, "", False)
probe = _classify(case, out_dir, device, rcs[-1], 0, "", False,
args=args)
if probe["status"] == "FAIL" and probe["rows_success"] == "0":
log = out_dir / "run.log"
if log.is_file():
log.rename(out_dir / f"run.attempt{len(rcs) - 1}.log")
notes = _join(notes, "dtype_rescue")
rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, ""),
rc, timed_out = _wait(_spawn(case, out_dir, device, args, ""),
args.op_timeout)
rcs.append(rc)
elapsed_ms = int((time.time() - t0) * 1000)
if len(rcs) > 1:
notes = _join(notes, "retry_rc=" + "->".join(map(str, rcs)))
row = _classify(case, out_dir, gpu, rcs[-1], elapsed_ms, notes, timed_out)
row = _classify(case, out_dir, device, rcs[-1], elapsed_ms, notes, timed_out,
args=args)
if row["status"] in ("PASS", "SKIP"):
tmp = marker.with_name(".complete.tmp")
tmp.write_text(row["status"] + "\n")
@@ -345,8 +351,21 @@ def _join(*notes: str) -> str:
return ";".join(n for n in notes if n)
def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
notes: str, timed_out: bool) -> Dict[str, str]:
def _identity_fields(out_dir: Path, device: str, backend: str) -> Dict[str, str]:
ident = load_identity(str(out_dir))
backend_v = str(ident.get("backend") or backend or "")
device_v = str(ident.get("device") or device or "")
timer_v = str(ident.get("timer") or "")
return {
"backend": backend_v,
"device": device_v,
"gpu": device_v,
"timer": timer_v,
}
def _classify(case: Case, out_dir: Path, device: str, rc: int, elapsed_ms: int,
notes: str, timed_out: bool, args: argparse.Namespace) -> Dict[str, str]:
text = _read_log(out_dir)
rows = batch_summary.parse_run_log(text)
if timed_out:
@@ -356,12 +375,13 @@ def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
n_skipped = batch_summary.pytest_skipped_count(text)
if status == "PASS" and n_skipped:
notes = _join(notes, f"partial_skip={n_skipped}")
ident = _identity_fields(out_dir, device, args.backend)
return {
"repo": case.repo,
"op": case.op,
"op_file": case.op_file,
"run_name": case.run_name,
"gpu": gpu,
**ident,
"status": status,
"return_code": str(rc),
"elapsed_ms": str(elapsed_ms),
@@ -393,8 +413,12 @@ def parse_args() -> argparse.Namespace:
help="--all 时不剔除带无条件 skip 标记的算子")
run = p.add_argument_group("执行")
run.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"),
help="设备后端：cuda/musa/rocm/ascend（默认 ZL_BACKEND 或 cuda）")
run.add_argument("--devices", default=None,
help="逗号分隔设备列表（默认该后端可见设备；优先于 --gpus）")
run.add_argument("--gpus", default=None,
help="逗号分隔 GPU 列表（默认 CUDA_VISIBLE_DEVICES 或全部）")
help="--devices 的别名（两者同时给出则报错）")
run.add_argument("--batch-dir", type=Path, default=None,
help="产物目录（默认 runs/batch_<时间戳>；resume 时必须指定）")
run.add_argument("--op-timeout", type=int, default=1800,
@@ -429,19 +453,30 @@ def parse_args() -> argparse.Namespace:
def main() -> None:
args = parse_args()
try:
backend = get_backend(args.backend)
except UnknownBackendError as exc:
raise SystemExit(f"{LOG} {exc}") from exc
args.backend_obj = backend
args.backend = backend.name
try:
explicit = resolve_device_arg(args.devices, args.gpus)
except ValueError as exc:
raise SystemExit(f"{LOG} {exc}") from exc
devices = visible_devices(backend, explicit)
inv = load_inventory()
cases = resolve_selection(args, inv)
if not cases:
raise SystemExit(f"{LOG} 选择结果为空")
gpus = visible_gpus(args.gpus)
batch_dir = (args.batch_dir or
ZL_BENCH / "runs" / f"batch_{datetime.now():%Y%m%d-%H%M%S}")
batch_dir = batch_dir.resolve()
batch_dir.mkdir(parents=True, exist_ok=True)
parts = {g: cs for g, cs in partition(cases, gpus).items() if cs}
print(f"{LOG} {len(cases)} 个算子 -> {len(parts)} 张卡 "
parts = {d: cs for d, cs in partition(cases, devices).items() if cs}
print(f"{LOG} backend={args.backend} {len(cases)} 个算子 -> {len(parts)} 台设备 "
f"(shards: {[len(cs) for cs in parts.values()]}) out={batch_dir}")
if args.replay_root:
print(f"{LOG} replay 模式: REPLAY_FROM={args.replay_root}/<repo>/<op>")
@@ -461,17 +496,19 @@ def main() -> None:
writer.writerows(rows)
os.replace(tmp, status_csv)
def run_partition(gpu: str, part: List[Case]) -> None:
def run_partition(device: str, part: List[Case]) -> None:
for case in part:
try:
row = run_case(case, batch_dir, gpu, args)
row = run_case(case, batch_dir, device, args)
except Exception as exc: # 编排层异常不拖垮整卡分片
out_dir = batch_dir / case.repo / case.run_name
ident = _identity_fields(out_dir, device, args.backend)
row = {
"repo": case.repo, "op": case.op, "op_file": case.op_file,
"run_name": case.run_name, "gpu": gpu, "status": "FAIL",
"run_name": case.run_name, **ident, "status": "FAIL",
"return_code": "EXC", "elapsed_ms": "0", "rows_total": "0",
"rows_success": "0", "notes": f"harness_exc={exc}",
"out_dir": str(batch_dir / case.repo / case.run_name),
"out_dir": str(out_dir),
}
with lock:
done[0] += 1
@@ -480,16 +517,16 @@ def main() -> None:
print(f" ({done[0]:>4}/{len(cases)}) {row['status']:<7} "
f"{row['elapsed_ms']:>8}ms rows={row['rows_success']}/"
f"{row['rows_total']:<4} {case.repo}:{case.run_name:<{wname}} "
f"[gpu{gpu}]"
f"[device{device}]"
+ (f" {row['notes']}" if row["notes"] else ""),
flush=True)
if len(parts) == 1:
gpu, part = next(iter(parts.items()))
run_partition(gpu, part)
device, part = next(iter(parts.items()))
run_partition(device, part)
else:
with ThreadPoolExecutor(max_workers=len(parts)) as pool:
futures = [pool.submit(run_partition, g, p) for g, p in parts.items()]
futures = [pool.submit(run_partition, d, p) for d, p in parts.items()]
for f in futures:
f.result()