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:
@@ -115,8 +115,10 @@ def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> O
|
||||
|
||||
def _emit_marker(reason: str) -> None:
|
||||
# Marker lands in run.log (grep it to audit replay coverage); flush so it
|
||||
# appears before the kernel runs.
|
||||
print(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", flush=True)
|
||||
# appears before the kernel runs. Yellow on the live terminal only —
|
||||
# run_pytest.sh strips ANSI from run.log.
|
||||
from _term_style import YELLOW, paint
|
||||
print(paint(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", YELLOW), flush=True)
|
||||
|
||||
|
||||
def _record_run(original):
|
||||
@@ -235,16 +237,19 @@ def _dump_record(record_dir: str) -> None:
|
||||
path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n")
|
||||
except Exception as exc:
|
||||
# Don't break the run if dump fails; the perf row is still produced.
|
||||
print(f"[autotune-record-plugin] dump failed: {exc}", file=sys.stderr, flush=True)
|
||||
from _term_style import tag
|
||||
print(f"{tag('[autotune-record-plugin]')} dump failed: {exc}",
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
from _term_style import tag
|
||||
record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
|
||||
replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
|
||||
if record_dir and replay_dir:
|
||||
# Mutually exclusive: recording while replaying is pointless. Surface loudly.
|
||||
print("[autotune-record-plugin] error: both RECORD_DIR and REPLAY_DIR set; "
|
||||
"ignoring both (no-op)", file=sys.stderr, flush=True)
|
||||
print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
|
||||
"REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
|
||||
return
|
||||
if not record_dir and not replay_dir:
|
||||
return
|
||||
@@ -265,11 +270,11 @@ def pytest_configure(config):
|
||||
# atexit (not sessionfinish): persist whatever was recorded even if an op
|
||||
# crash kills the session; a later replay run falls back for missing keys.
|
||||
atexit.register(_dump_record, record_dir)
|
||||
print(f"[autotune-record-plugin] recording autotune configs to "
|
||||
print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
|
||||
f"{_record_path(record_dir)} ({'+'.join(patched)})",
|
||||
file=sys.stderr, flush=True)
|
||||
else:
|
||||
loaded = _load_replay_dir(replay_dir)
|
||||
print(f"[autotune-record-plugin] replaying {loaded} recorded entries from "
|
||||
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
|
||||
f"{_record_path(replay_dir)} ({'+'.join(patched)})",
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
Reference in New Issue
Block a user