forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
c9f9b7885f
Two failure modes that cost this harness real results, both of them silent. A wrong path does not fail as a wrong path. pytest cannot collect a benchmark directory that is not there, so benchmark/conftest.py never registers --level and --mode, and the run dies with "unrecognized arguments: --level core --mode kernel" -- an error about the flags, several layers away from the missing directory. run_pytest.sh now checks the FlagGems directory, the benchmark file, SHAPE_FILE and REPLAY_FROM before starting pytest, prints "PREFLIGHT_FAILED check=<the named check>" and exits 3, distinct from pytest's own 1 and 4 so the batch layer can tell a broken environment from a broken operator. Only the file half of a node is checked; a missing ::test_<fn> already produces a clear "not found: ..." from pytest and needs no help. FlagGems keeps libtuner winners in a sqlite DB under $FLAGGEMS_CACHE_DIR (default ~/.flaggems), shared by every run on the machine. Whichever side of an A/B runs first therefore hands its winners to the second, which then compiles only those: measured on router_gemm, the second of two shared-cache runs compiled 10 variants instead of 234 and finished in 17s instead of 52s. That is how an mm A/B in this repo produced "945 variants -> 33, 136 converts -> 10" and read as a compiler result when it was cache inheritance; isolated, both runs compile 234. ZL_ISOLATE_TUNER_CACHE=1 gives a run its own cache directory, dropped afterwards with the Triton one since the chosen configs are already in autotune_records/. It defaults to 0, because cold tuning is expensive and the standard REPLAY_FROM A/B pins configs anyway; only record-against-record is exposed. The default is unchanged, but the caliber is no longer invisible: either way it goes into the run header and run_identity.json as tuner_cache, which is what was missing when the mm numbers could not be explained after the fact. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
"""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", ""))
|
|
p.add_argument("--bench-level", default=os.environ.get("BENCH_LEVEL", "core"))
|
|
# "isolated" or "shared(<path>)". Recorded because a shared libtuner cache
|
|
# lets whichever side runs first hand its winners to the other, and after
|
|
# the fact the two runs are otherwise indistinguishable.
|
|
p.add_argument("--tuner-cache", default="")
|
|
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,
|
|
bench_level=args.bench_level,
|
|
tuner_cache=args.tuner_cache,
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|