"""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()