forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
c8ea586d0f
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>
549 lines
22 KiB
Python
549 lines
22 KiB
Python
#!/usr/bin/env python3
|
||
"""多算子批量测试驱动：按算子清单逐个调 run_pytest.sh，多卡分片、串行测量。
|
||
|
||
每个算子一个独立子进程（设备 crash 只废单个算子）、独立产物目录
|
||
<batch>/<repo>/<op>/（与单算子跑 run_pytest.sh 的产物结构完全一致，含
|
||
run.log / shapes.yaml / autotune_records / ttgir / run_identity.json）。
|
||
批量层只做编排：
|
||
|
||
- 分片：按算子名稳定哈希到设备（每卡内部串行，测量口径与单算子一致；
|
||
跨卡并行测量存在功耗耦合噪声，本模式定位为筛查口径，可疑算子回
|
||
单算子串行流程确认）；
|
||
- 后端：--backend / ZL_BACKEND 选择 cuda/musa/rocm/ascend，设备可见性
|
||
由 backends/ 写入，不再写死 CUDA_VISIBLE_DEVICES；
|
||
- 超时：--op-timeout 杀死挂死算子（记 TIMEOUT，不重试，继续跑）；
|
||
- 重试：pytest rc=1（用例失败/flaky）与信号杀（rc>=128）最多重试 2 次，
|
||
连续两次同 rc 视为确定性失败提前止损；--dtypes 下两级降级——遇上游
|
||
"can't be supported by this op" 即去掉限制重跑，其余失败且无成功行时
|
||
也去掉限制最后救一次（dtype_rescue）；
|
||
- 断点续跑：--resume 复用带 .complete 标记的算子（PASS/SKIP 才打标记，
|
||
FAIL/TIMEOUT 不打、下次重跑）；
|
||
- replay：--replay-root 指向上一次批量目录，逐算子传 REPLAY_FROM
|
||
<root>/<repo>/<op>，语义与单算子 REPLAY_FROM 相同。
|
||
|
||
产物：<batch>/ops_status.csv（每算子一行）+ summary.csv（每表行一行，
|
||
由 batch_summary.py 生成）。
|
||
|
||
示例：
|
||
# 全量筛查（两仓库全部可收集算子，排除 marked_skip 与 dismiss）
|
||
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_
|
||
|
||
# 锁 config 复测（口径同单算子 REPLAY_FROM）
|
||
python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import hashlib
|
||
import os
|
||
import signal
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from collections import Counter
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from datetime import datetime
|
||
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"
|
||
LOG = "[batch]"
|
||
MAX_RETRIES = 2
|
||
|
||
REPO_DIRS = {
|
||
"main": "/workspace/dev/FlagGems",
|
||
"vllm": "/workspace/dev/FlagGems-vllm",
|
||
}
|
||
|
||
STATUS_FIELDS = [
|
||
"repo", "op", "op_file", "run_name",
|
||
"backend", "device", "gpu", "timer",
|
||
"status", "return_code",
|
||
"elapsed_ms", "rows_total", "rows_success", "notes", "out_dir",
|
||
]
|
||
|
||
DTYPE_UNSUPPORTED_MARK = "can't be supported by this op"
|
||
|
||
|
||
# ---------------- 清单与选择 ----------------
|
||
|
||
class Case:
|
||
__slots__ = ("repo", "op", "op_file", "run_name")
|
||
|
||
def __init__(self, repo: str, op: str, op_file: str, run_name: str):
|
||
self.repo, self.op, self.op_file, self.run_name = repo, op, op_file, run_name
|
||
|
||
@property
|
||
def key(self) -> str:
|
||
return f"{self.repo}:{self.op}@{self.op_file}"
|
||
|
||
|
||
def load_inventory() -> Dict[str, List[Dict[str, str]]]:
|
||
inv: Dict[str, List[Dict[str, str]]] = {}
|
||
for label in REPO_DIRS:
|
||
path = OPS_DIR / f"inventory_{label}.csv"
|
||
if path.is_file():
|
||
with path.open(newline="") as f:
|
||
inv[label] = list(csv.DictReader(f))
|
||
if not inv:
|
||
raise SystemExit(f"{LOG} ops/inventory_*.csv 不存在，先跑 ops/gen_inventory.py")
|
||
return inv
|
||
|
||
|
||
def _run_name(row: Dict[str, str], dup_ops: Set[Tuple[str, str]]) -> str:
|
||
# 目录名默认用 op；仓库内重名时追加 @op_file 消歧（按 inventory 全局判定，
|
||
# 与选择子集无关，保证 resume/replay 路径稳定）。op_file 可能带子目录
|
||
# （test_FLA/ 等），做目录名时把 "/" 压成 "_"。
|
||
op = row["op"]
|
||
if (row["repo"], op) not in dup_ops:
|
||
return op
|
||
return f"{op}@{row['op_file'].replace('/', '_')}"
|
||
|
||
|
||
def _dup_ops(inv: Dict[str, List[Dict[str, str]]]) -> Set[Tuple[str, str]]:
|
||
dup = set()
|
||
for label, rows in inv.items():
|
||
counts = Counter(r["op"] for r in rows)
|
||
dup |= {(label, op) for op, n in counts.items() if n > 1}
|
||
return dup
|
||
|
||
|
||
def load_dismiss(path: Path) -> Set[Tuple[str, str]]:
|
||
if not path.is_file():
|
||
return set()
|
||
out = set()
|
||
for line in path.read_text().splitlines():
|
||
line = line.strip()
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
repo, _, op = line.rpartition(":")
|
||
out.add((repo or "main", op))
|
||
return out
|
||
|
||
|
||
def resolve_selection(args: argparse.Namespace,
|
||
inv: Dict[str, List[Dict[str, str]]]) -> List[Case]:
|
||
dup = _dup_ops(inv)
|
||
index: Dict[Tuple[str, str, str], Dict[str, str]] = {}
|
||
by_op: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
|
||
for label, rows in inv.items():
|
||
for r in rows:
|
||
index[(label, r["op"], r["op_file"])] = r
|
||
by_op.setdefault((label, r["op"]), []).append(r)
|
||
|
||
dismissed = load_dismiss(args.dismiss)
|
||
|
||
def make_case(row: Dict[str, str]) -> Case:
|
||
return Case(row["repo"], row["op"], row["op_file"], _run_name(row, dup))
|
||
|
||
if args.all:
|
||
cases = []
|
||
for label, rows in inv.items():
|
||
if args.repos and label not in args.repos:
|
||
continue
|
||
for r in rows:
|
||
if r.get("collected") == "no":
|
||
continue
|
||
if r.get("marked_skip") == "yes" and not args.include_marked_skip:
|
||
continue
|
||
if (label, r["op"]) in dismissed:
|
||
continue
|
||
cases.append(make_case(r))
|
||
return cases
|
||
|
||
specs: List[str] = list(args.ops or [])
|
||
if args.ops_file:
|
||
for line in args.ops_file.read_text().splitlines():
|
||
line = line.strip()
|
||
if line and not line.startswith("#"):
|
||
specs.append(line)
|
||
if not specs:
|
||
raise SystemExit(f"{LOG} 未选择算子：用 --all / --ops / --ops-file 之一")
|
||
|
||
cases, errors = [], []
|
||
for spec in specs:
|
||
repo, _, rest = spec.rpartition(":")
|
||
repo = repo or "main"
|
||
op, _, op_file = rest.partition("@")
|
||
if repo not in inv:
|
||
errors.append(f"{spec}: 未知仓库 {repo!r}")
|
||
continue
|
||
if op_file:
|
||
row = index.get((repo, op, op_file))
|
||
if row is None:
|
||
errors.append(f"{spec}: inventory 中无此条目")
|
||
continue
|
||
cases.append(make_case(row))
|
||
continue
|
||
rows = by_op.get((repo, op)) or []
|
||
if not rows:
|
||
errors.append(f"{spec}: inventory 中无此算子")
|
||
elif len(rows) > 1:
|
||
files = ", ".join(r["op_file"] for r in rows)
|
||
errors.append(f"{spec}: 重名，需用 @op_file 消歧（候选: {files}）")
|
||
else:
|
||
cases.append(make_case(rows[0]))
|
||
if errors:
|
||
for e in errors:
|
||
print(f"{LOG} error: {e}", file=sys.stderr)
|
||
raise SystemExit(1)
|
||
return cases
|
||
|
||
|
||
# ---------------- 设备分片 ----------------
|
||
|
||
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 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]] = {d: [] for d in devices}
|
||
for c in cases:
|
||
digest = hashlib.sha256(f"{c.repo}/{c.run_name}".encode()).digest()
|
||
parts[devices[int.from_bytes(digest[:8], "big") % len(devices)]].append(c)
|
||
return parts
|
||
|
||
|
||
# ---------------- 单算子执行 ----------------
|
||
|
||
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),
|
||
"ZL_BACKEND": args.backend,
|
||
"ZL_DEVICE": device,
|
||
"FLAGGEMS_PERF_COLOR": "never",
|
||
"NO_COLOR": "1",
|
||
"USE_FLAGTUNE": str(args.use_flagtune),
|
||
})
|
||
if args.shape_file:
|
||
env["SHAPE_FILE"] = str(args.shape_file)
|
||
if dtypes:
|
||
env["DTYPES"] = dtypes
|
||
if args.replay_root:
|
||
env["REPLAY_FROM"] = str(args.replay_root / case.repo / case.run_name)
|
||
return subprocess.Popen(
|
||
["bash", str(ZL_BENCH / "run_pytest.sh")],
|
||
cwd=str(ZL_BENCH), env=env, start_new_session=True,
|
||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||
)
|
||
|
||
|
||
def _wait(proc: subprocess.Popen, timeout: int) -> Tuple[int, bool]:
|
||
"""返回 (rc, timed_out)；超时对整个进程组 SIGTERM→SIGKILL。"""
|
||
try:
|
||
return proc.wait(timeout=timeout or None), False
|
||
except subprocess.TimeoutExpired:
|
||
try:
|
||
pgid = os.getpgid(proc.pid)
|
||
os.killpg(pgid, signal.SIGTERM)
|
||
try:
|
||
proc.wait(timeout=10)
|
||
except subprocess.TimeoutExpired:
|
||
os.killpg(pgid, signal.SIGKILL)
|
||
except ProcessLookupError:
|
||
pass # 组内进程恰好在超时瞬间退出
|
||
proc.wait()
|
||
return 124, True
|
||
|
||
|
||
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, device, 0, 0, "", timed_out=False,
|
||
args=args)
|
||
row["notes"] = _join("resumed", row["notes"])
|
||
return row
|
||
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
marker.unlink(missing_ok=True)
|
||
|
||
dtypes = args.dtypes or ""
|
||
t0 = time.time()
|
||
rcs: List[int] = []
|
||
notes = ""
|
||
timed_out = False
|
||
for attempt in range(MAX_RETRIES + 1):
|
||
if attempt:
|
||
# 保留上一轮日志再重试
|
||
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, device, args, dtypes),
|
||
args.op_timeout)
|
||
rcs.append(rc)
|
||
if timed_out:
|
||
break
|
||
log_text = _read_log(out_dir)
|
||
# dtype 不支持：去掉 --dtypes 降级重试（只降一次，不计入重试预算）
|
||
if rc != 0 and dtypes and DTYPE_UNSUPPORTED_MARK in log_text:
|
||
dtypes = ""
|
||
notes = _join(notes, "dtype_fallback")
|
||
continue
|
||
if rc == 0 or not (rc == 1 or rc >= 128):
|
||
break
|
||
if len(rcs) >= 2 and rcs[-1] == rcs[-2]:
|
||
break # 确定性失败：同 rc 连续两次，不再烧预算
|
||
|
||
# dtype 营救：限制 dtype 时整体失败且没有任何成功行，去掉限制最后救一次。
|
||
# 覆盖 baseline 对受限 dtype 编译失败（torch jiterator 的 special 函数族等）
|
||
# 但报错文案不是上游 "can't be supported by this op" 的场景。
|
||
if dtypes and not timed_out:
|
||
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, 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, 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")
|
||
os.replace(tmp, marker)
|
||
return row
|
||
|
||
|
||
def _read_log(out_dir: Path) -> str:
|
||
log = out_dir / "run.log"
|
||
return log.read_text(errors="replace") if log.is_file() else ""
|
||
|
||
|
||
def _join(*notes: str) -> str:
|
||
return ";".join(n for n in notes if n)
|
||
|
||
|
||
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:
|
||
status = "TIMEOUT"
|
||
else:
|
||
status = batch_summary.classify(rc, rows, text)
|
||
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,
|
||
**ident,
|
||
"status": status,
|
||
"return_code": str(rc),
|
||
"elapsed_ms": str(elapsed_ms),
|
||
"rows_total": str(len(rows)),
|
||
"rows_success": str(sum(1 for r in rows if r["row_status"] == "SUCCESS")),
|
||
"notes": _join(notes, batch_summary.marker_notes(text)),
|
||
"out_dir": str(out_dir),
|
||
}
|
||
|
||
|
||
# ---------------- 主流程 ----------------
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
p = argparse.ArgumentParser(
|
||
description=__doc__.splitlines()[0],
|
||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
sel = p.add_argument_group("算子选择")
|
||
sel.add_argument("--all", action="store_true",
|
||
help="跑 inventory 全量（排除 collected=no / marked_skip / dismiss）")
|
||
sel.add_argument("--ops", nargs="*", default=None, metavar="[repo:]op[@op_file]",
|
||
help="直接指定算子（默认 main 仓；重名需 @op_file 消歧）")
|
||
sel.add_argument("--ops-file", type=Path, default=None,
|
||
help="算子清单文件，每行一个 [repo:]op[@op_file]，# 注释")
|
||
sel.add_argument("--repos", nargs="*", default=None, choices=list(REPO_DIRS),
|
||
help="--all 时限定仓库（默认全部）")
|
||
sel.add_argument("--dismiss", type=Path, default=OPS_DIR / "dismiss.txt",
|
||
help="排除清单（默认 ops/dismiss.txt，仅 --all 生效）")
|
||
sel.add_argument("--include-marked-skip", action="store_true",
|
||
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="--devices 的别名（两者同时给出则报错）")
|
||
run.add_argument("--batch-dir", type=Path, default=None,
|
||
help="产物目录（默认 runs/batch_<时间戳>；resume 时必须指定）")
|
||
run.add_argument("--op-timeout", type=int, default=1800,
|
||
help="单算子超时秒数，0=不限（默认 %(default)s）")
|
||
run.add_argument("--resume", action="store_true",
|
||
help="复用 --batch-dir 里带 .complete 标记的算子结果")
|
||
run.add_argument("--replay-root", type=Path, default=None,
|
||
help="上次批量目录：逐算子 REPLAY_FROM <root>/<repo>/<op>")
|
||
run.add_argument("--shape-file", type=Path, default=OPS_DIR / "shapes_single.yaml",
|
||
help="shape yaml（默认 %(default)s；深查用 ops/shapes_multi.yaml；"
|
||
"传空串则用各算子上游默认 shape）")
|
||
run.add_argument("--dtypes", default="",
|
||
help='限制 dtype 集（空格分隔，如 "bfloat16"）；空=上游默认扫描')
|
||
run.add_argument("--use-flagtune", type=int, default=0, choices=(0, 1),
|
||
help="透传 USE_FLAGTUNE（默认 0；批量 A/B 两侧必须同值）")
|
||
args = p.parse_args()
|
||
|
||
if args.resume and not args.batch_dir:
|
||
raise SystemExit(f"{LOG} --resume 需要显式 --batch-dir")
|
||
if args.shape_file and str(args.shape_file).strip():
|
||
args.shape_file = args.shape_file.resolve()
|
||
if not args.shape_file.is_file():
|
||
raise SystemExit(f"{LOG} shape 文件不存在: {args.shape_file}")
|
||
else:
|
||
args.shape_file = None
|
||
if args.replay_root:
|
||
args.replay_root = args.replay_root.resolve()
|
||
if not args.replay_root.is_dir():
|
||
raise SystemExit(f"{LOG} replay 目录不存在: {args.replay_root}")
|
||
return args
|
||
|
||
|
||
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} 选择结果为空")
|
||
|
||
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 = {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>")
|
||
|
||
lock = threading.Lock()
|
||
done = [0]
|
||
results: Dict[str, Dict[str, str]] = {}
|
||
status_csv = batch_dir / "ops_status.csv"
|
||
wname = max((len(c.run_name) for c in cases), default=8)
|
||
|
||
def flush_status() -> None:
|
||
rows = [results[c.key] for c in cases if c.key in results]
|
||
tmp = status_csv.with_name(status_csv.name + ".tmp")
|
||
with tmp.open("w", newline="") as f:
|
||
writer = csv.DictWriter(f, fieldnames=STATUS_FIELDS)
|
||
writer.writeheader()
|
||
writer.writerows(rows)
|
||
os.replace(tmp, status_csv)
|
||
|
||
def run_partition(device: str, part: List[Case]) -> None:
|
||
for case in part:
|
||
try:
|
||
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, **ident, "status": "FAIL",
|
||
"return_code": "EXC", "elapsed_ms": "0", "rows_total": "0",
|
||
"rows_success": "0", "notes": f"harness_exc={exc}",
|
||
"out_dir": str(out_dir),
|
||
}
|
||
with lock:
|
||
done[0] += 1
|
||
results[case.key] = row
|
||
flush_status()
|
||
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"[device{device}]"
|
||
+ (f" {row['notes']}" if row["notes"] else ""),
|
||
flush=True)
|
||
|
||
if len(parts) == 1:
|
||
device, part = next(iter(parts.items()))
|
||
run_partition(device, part)
|
||
else:
|
||
with ThreadPoolExecutor(max_workers=len(parts)) as pool:
|
||
futures = [pool.submit(run_partition, d, p) for d, p in parts.items()]
|
||
for f in futures:
|
||
f.result()
|
||
|
||
summary = batch_summary.summarize(batch_dir)
|
||
rows = [results[c.key] for c in cases if c.key in results]
|
||
counts = {s: sum(1 for r in rows if r["status"] == s)
|
||
for s in ("PASS", "SKIP", "FAIL", "TIMEOUT")}
|
||
print(f"{LOG} 完成: " + " ".join(f"{k}={v}" for k, v in counts.items()))
|
||
print(f"{LOG} 状态表: {status_csv}")
|
||
print(f"{LOG} 总表: {summary}")
|
||
bad = [r for r in rows if r["status"] in ("FAIL", "TIMEOUT")]
|
||
if bad:
|
||
print(f"{LOG} 失败算子: " +
|
||
", ".join(f"{r['repo']}:{r['run_name']}" for r in bad))
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|