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:
@@ -0,0 +1,45 @@
|
||||
"""Guard against FlagGems' device-detection hang (benchmark-side fix).
|
||||
|
||||
FlagGems' `runtime.backend.device_finder` falls back to
|
||||
`subprocess.run(nvidia-smi)` **without a timeout**. Under this conda env's
|
||||
forked/inconsistent subprocess (`_posixsubprocess` symbol mismatch), that
|
||||
child can hang indefinitely, leaving `import flag_gems` stuck in `wait4`
|
||||
(seen as run_pytest "卡住 with no result", GPU 0%).
|
||||
|
||||
This plugin runs at import time — before any test module imports flag_gems —
|
||||
detects the vendor via torch (no subprocess), and sets GEMS_VENDOR so
|
||||
device_finder takes its env fast-path (`_get_vendor_from_env`) and never
|
||||
reaches the hang-prone subprocess probe. No change to the FlagGems repo.
|
||||
|
||||
Side effect: get_device_properties initializes the CUDA context very early in
|
||||
the pytest process. Fine for the current single-process runs; revisit if
|
||||
fork-based parallelism (e.g. pytest-xdist) is ever introduced.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
_VENDOR_ENV_KEYS = ("GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND")
|
||||
|
||||
|
||||
def _guard_device_vendor():
|
||||
# Respect an explicit choice if the user already set one.
|
||||
if any(k in os.environ for k in _VENDOR_ENV_KEYS):
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
name = torch.cuda.get_device_properties(0).name.upper()
|
||||
if "NVIDIA" in name:
|
||||
os.environ["GEMS_VENDOR"] = "nvidia"
|
||||
from _term_style import tag
|
||||
print(f"{tag('[device-guard-plugin]')} set GEMS_VENDOR=nvidia "
|
||||
"(skip flag_gems nvidia-smi subprocess probe, avoids import hang)",
|
||||
file=sys.stderr, flush=True)
|
||||
except Exception as e: # pragma: no cover
|
||||
print(f"[device-guard-plugin] torch vendor probe failed ({e}); "
|
||||
"leaving detection to flag_gems",
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
_guard_device_vendor()
|
||||
Reference in New Issue
Block a user