Stabilize first-run cudagraph timing and colorize terminal output
- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture (internal 5-iter warmup is too short), tag fallback markers with the failing phase; first-run latency no longer jitters run-to-run. - _device_guard_plugin (new): set GEMS_VENDOR via torch probe before importing flag_gems, avoiding its timeout-less nvidia-smi subprocess probe that can hang import in fork-broken environments. - _pretty_report_plugin (new) + _term_style (new): fold inputs identical across all result rows into a legend line, color status/plugin tags/markers on the live terminal; run.log is ANSI-stripped and keeps upstream SUCCESS/column wording for grep compatibility. - run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a config section with Chinese comments, add start/end banners. - README: document the warmup semantics, new plugins/env vars, and the A/B rule that both sides must use the same USE_FLAGTUNE.
This commit is contained in:
+13
-6
@@ -114,7 +114,8 @@ def _compile_listener(*, src, metadata, metadata_group, times, cache_hit):
|
||||
f, indent=1, sort_keys=True)
|
||||
os.replace(tmp, path)
|
||||
except Exception as exc: # never break compilation over a metadata dump
|
||||
print(f"[ir-meta-plugin] sidecar dump failed: {exc}", file=sys.stderr)
|
||||
from _term_style import tag
|
||||
print(f"{tag('[ir-meta-plugin]')} sidecar dump failed: {exc}", file=sys.stderr)
|
||||
|
||||
|
||||
# --- shape + launch tracking ----------------------------------------------
|
||||
@@ -202,7 +203,8 @@ def pytest_collection_finish(session):
|
||||
if own is not None:
|
||||
cls.get_input_iter = _wrap_input_iter(own)
|
||||
wrapped += 1
|
||||
print(f"[ir-meta-plugin] shape tracking wrapped on {wrapped} Benchmark classes",
|
||||
from _term_style import tag
|
||||
print(f"{tag('[ir-meta-plugin]')} shape tracking wrapped on {wrapped} Benchmark classes",
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
@@ -312,7 +314,9 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
|
||||
|
||||
cache, dest = pathlib.Path(cache_dir), pathlib.Path(dump_dir)
|
||||
if not cache.is_dir():
|
||||
print(f"[ir-meta-plugin] no cache dir {cache}; nothing to dump", file=sys.stderr)
|
||||
from _term_style import tag
|
||||
print(f"{tag('[ir-meta-plugin]')} no cache dir {cache}; nothing to dump",
|
||||
file=sys.stderr)
|
||||
return
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@@ -386,11 +390,12 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
|
||||
|
||||
shapes = sorted({r[0] for r in rows if r[0] != "-"})
|
||||
losers = sum(1 for r in rows if r[0] == "-")
|
||||
from _term_style import DIM, paint
|
||||
print(f">>> [Dump] ttgir -> {dest} ({copied} files across {len(shapes)} shapes; "
|
||||
f"{losers} unused variants index-only; legend: naming.md)", flush=True)
|
||||
for s in shapes:
|
||||
n = sum(1 for r in rows if r[0] == s)
|
||||
print(f">>> [Dump] {s}: {n}", flush=True)
|
||||
print(paint(f">>> [Dump] {s}: {n}", DIM), flush=True)
|
||||
|
||||
|
||||
# --- registration -----------------------------------------------------------
|
||||
@@ -419,10 +424,12 @@ def pytest_configure(config):
|
||||
_autotuner.Autotuner._bench = _wrap_bench(_autotuner.Autotuner._bench)
|
||||
atexit.register(_dump_ttgir, cache_dir, dump_dir)
|
||||
else:
|
||||
print("[ir-meta-plugin] warning: dump dir set but TRITON_CACHE_DIR "
|
||||
from _term_style import tag
|
||||
print(f"{tag('[ir-meta-plugin]')} warning: dump dir set but TRITON_CACHE_DIR "
|
||||
"is not; ttgir dump disabled", file=sys.stderr, flush=True)
|
||||
|
||||
print("[ir-meta-plugin] compilation listener registered"
|
||||
from _term_style import tag
|
||||
print(f"{tag('[ir-meta-plugin]')} compilation listener registered"
|
||||
+ (f"; shape/launch tracking on, ttgir dump -> {dump_dir}"
|
||||
if _dump_enabled else ""),
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
Reference in New Issue
Block a user