forked from zhoulin/zl_bench
feat: make the harness backend-agnostic
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:
1 parent
35a928a516
commit
c8ea586d0f
24 files changed
+1448
-159
No files matched your search
@@ -0,0 +1,89 @@
|
||||
"""Merge-update `$OUT_DIR/run_identity.json` so plugins can record backend facts.
|
||||
|
||||
Shell writes a stub before pytest; plugins add device_name / timer / vendor
|
||||
without clobbering earlier keys. Missing OUT_DIR is a silent no-op so unit
|
||||
tests and ad-hoc pytest invocations do not fail.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
IDENTITY_NAME = "run_identity.json"
|
||||
|
||||
|
||||
def identity_path(out_dir: Optional[str] = None) -> Optional[Path]:
|
||||
root = out_dir or os.environ.get("OUT_DIR", "").strip()
|
||||
if not root:
|
||||
return None
|
||||
return Path(root) / IDENTITY_NAME
|
||||
|
||||
|
||||
def load_identity(out_dir: Optional[str] = None) -> Dict[str, Any]:
|
||||
path = identity_path(out_dir)
|
||||
if path is None or not path.is_file():
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def update_identity(out_dir: Optional[str] = None, **fields: Any) -> Dict[str, Any]:
|
||||
path = identity_path(out_dir)
|
||||
if path is None:
|
||||
return dict(fields)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
current = load_identity(out_dir)
|
||||
for key, value in fields.items():
|
||||
if value is None:
|
||||
continue
|
||||
current[key] = value
|
||||
tmp = path.with_name(path.name + ".tmp")
|
||||
tmp.write_text(json.dumps(current, indent=2, ensure_ascii=False) + "\n",
|
||||
encoding="utf-8")
|
||||
os.replace(tmp, path)
|
||||
return current
|
||||
|
||||
|
||||
def init_identity(out_dir: str, **fields: Any) -> Dict[str, Any]:
|
||||
payload = {"schema_version": 1, "provider": "flaggems"}
|
||||
payload.update(fields)
|
||||
os.environ["OUT_DIR"] = str(out_dir)
|
||||
return update_identity(out_dir, **payload)
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("out_dir")
|
||||
p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
|
||||
p.add_argument("--op", default=os.environ.get("OP", ""))
|
||||
p.add_argument("--op-file", default=os.environ.get("OP_FILE", ""))
|
||||
p.add_argument("--flaggems-dir", default=os.environ.get("FLAGGEMS_DIR", ""))
|
||||
p.add_argument("--autotune-mode", default="")
|
||||
p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
|
||||
p.add_argument("--test-file", default=os.environ.get("TEST_FILE", ""))
|
||||
p.add_argument("--shape-mode", default=os.environ.get("SHAPE_MODE", ""))
|
||||
args = p.parse_args(argv)
|
||||
init_identity(
|
||||
args.out_dir,
|
||||
backend=args.backend,
|
||||
op=args.op,
|
||||
op_file=args.op_file,
|
||||
flaggems_dir=args.flaggems_dir,
|
||||
autotune_mode=args.autotune_mode,
|
||||
device=args.device,
|
||||
test_file=args.test_file,
|
||||
shape_mode=args.shape_mode,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in new issue
Block a user