zhoulin
da22885645
run_batch.py drives run_pytest.sh per operator across GPUs (stable-hash sharding, per-op subprocess isolation, process-group timeouts, retry with deterministic-failure cutoff, two-level dtype fallback, .complete resume, per-op REPLAY_FROM). batch_summary.py aggregates run.log tables into summary.csv. ops/ holds the curated assets: dual-repo inventories rebuilt via AST scan + pytest collect verification, shape sets migrated from the old regression harness and merged with upstream core_shapes class-name keys (upstream's set_shapes falls back op_name -> MRO class name -> 1-D DEFAULT_SHAPES, so replacing the shape file without class keys crashes the BLAS family), and a dismiss list where all 76 entries carry verified reasons. Validated end to end: 1036-op full screen with zero failures. Also fix two cudagraph plugin regressions: newer torch appends "enable device-side assertions" to every CUDA error, so the loose fatal-error marker disabled the documented do_bench fallback entirely; and an aborted graph capture can leave the default CUDA RNG generator stuck in capturing state, poisoning every later torch.randn - captures now run under a throwaway RNG state. run_pytest.sh gains an optional DTYPES passthrough.
46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
"""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 hanging with no output, GPU at 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()
|