zl_bench/_device_setup_plugin.py

"""Set FlagGems vendor + record backend identity before flag_gems is imported.

Replaces the NVIDIA-only `_device_guard_plugin`. Setting GEMS_VENDOR from the
zl_bench backend name skips FlagGems' hang-prone nvidia-smi subprocess probe
on CUDA, and maps musa/rocm/ascend onto the matching FlagGems vendor.

Import-time side effect is intentional: FlagGems reads GEMS_VENDOR while the
benchmark module is being collected. Device name / arch are filled later in
pytest_configure so import does not initialize a device context.
"""

from __future__ import annotations

import os
import sys

from backends import get_backend
from providers.flaggems import apply_vendor_env
from _run_identity import update_identity


def _setup_at_import() -> None:
    try:
        backend = get_backend()
    except Exception as exc:
        print(f"[device-setup-plugin] cannot resolve backend ({exc}); "
              "leaving GEMS_VENDOR to flag_gems",
              file=sys.stderr, flush=True)
        return
    vendor = apply_vendor_env(backend)
    os.environ["ZL_BACKEND"] = backend.name
    device = os.environ.get("ZL_DEVICE", "").strip()
    if not device:
        for key in backend.visibility_env_keys:
            raw = os.environ.get(key, "").strip()
            if raw:
                device = raw.split(",")[0].strip()
                break
    update_identity(
        backend=backend.name,
        device=device,
        flaggems_vendor=vendor or "",
        visibility_env=backend.visibility_env_keys[0] if backend.visibility_env_keys else "",
    )
    if vendor:
        from _term_style import tag
        print(f"{tag('[device-setup-plugin]')} backend={backend.name} "
              f"GEMS_VENDOR={vendor}",
              file=sys.stderr, flush=True)


def pytest_configure(config):
    try:
        backend = get_backend()
        ident = backend.device_identity()
        update_identity(
            backend=backend.name,
            device_name=ident.get("device_name", ""),
            device_arch=ident.get("device_arch", ""),
            timer_api=ident.get("timer_api", ""),
            flaggems_vendor=ident.get("flaggems_vendor", ""),
        )
    except Exception as exc:
        from _term_style import tag
        print(f"{tag('[device-setup-plugin]')} identity probe failed ({exc})",
              file=sys.stderr, flush=True)


_setup_at_import()

# FlagGems 的厂商实现按 flag_gems.runtime.backend._<vendor>.ops.* 组织。benchmark
# 常写成 `try: import <某厂商实现>; AVAILABLE = True`——但 import 只证明那个 .py 在
# 磁盘上，它不碰设备，所以在任何后端上都会成功。于是 gate 形同虚设，测出来的是别家
# 的 kernel 跑在本机上。router_gemm 就是这样：test_router_gemm.py 无条件导入
# _nvidia.hopper 版，在 MUSA 上 gate 照过，最后死在 JIT 深处的
# "dynamic_func() missing BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K"——那个报错完全指不
# 向真正的原因，而 _mthreads 下明明有一份能用的实现。
#
# 这里不拦，只报：跑完看一眼载入过谁家的实现。
_VENDOR_MOD_PREFIX = "flag_gems.runtime.backend._"


def foreign_vendor_modules(active_vendor: str, modules=None) -> list:
    """本次进程里载入过的、不属于当前厂商的 FlagGems 后端模块。"""
    if modules is None:
        modules = sys.modules
    out = set()
    for name in list(modules):
        if not name.startswith(_VENDOR_MOD_PREFIX):
            continue
        # 只保留厂商根，_nvidia.hopper.ops.mm 与 _nvidia.ops.mm 算同一条。
        # 按段比而不是按前缀：前缀会让 _<vendor>_something 被误认成自家。
        root = name[len(_VENDOR_MOD_PREFIX):].split(".")[0]
        if root != active_vendor:
            out.add(root)
    return sorted(out)


def pytest_sessionfinish(session, exitstatus):
    from _term_style import tag
    try:
        vendor = get_backend().flaggems_vendor
    except Exception:
        return
    foreign = foreign_vendor_modules(vendor)
    if not foreign:
        return
    print(f"{tag('[device-setup-plugin]')} FOREIGN_VENDOR_MODULE active={vendor} "
          f"loaded={','.join(foreign)}", file=sys.stderr, flush=True)
    print(f"{tag('[device-setup-plugin]')} 本次载入了别家厂商的 FlagGems 实现。"
          "如果这个算子的结果可疑，先查 benchmark 的可用性判据是不是只做了 import "
          "而没判 vendor_name", file=sys.stderr, flush=True)