zl_bench/backends/base.py

"""Device backend contract for zl_bench screening runs.

This is the orchestration-layer adapter: device enumeration, process visibility
env, and capability flags that decide which pytest plugins are safe to load.
It does not implement kernel timing receipts or compiler-artifact identity;
those stay with the caller (FlagGems pytest for screening, a separate harness
for qualification).
"""

from __future__ import annotations

import os
import subprocess
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

# Every vendor visibility variable zl_bench knows about. apply_visibility()
# clears the whole set before writing the active backend's keys so a parent
# CUDA_VISIBLE_DEVICES cannot leak onto a MUSA/Ascend child.
ALL_VISIBILITY_KEYS: Tuple[str, ...] = (
    "CUDA_VISIBLE_DEVICES",
    "MUSA_VISIBLE_DEVICES",
    "HIP_VISIBLE_DEVICES",
    "ROCR_VISIBLE_DEVICES",
    "ASCEND_RT_VISIBLE_DEVICES",
    "ASCEND_VISIBLE_DEVICES",
)


def parse_device_list(raw: Optional[str]) -> List[str]:
    if not raw:
        return []
    return [tok.strip() for tok in raw.split(",") if tok.strip()]


class DeviceBackend:
    """One hardware runtime (CUDA, MUSA, ROCm, Ascend, ...).

    Subclasses only fill class attributes and optional probe hooks. Callers
    should not branch on ``backend.name`` for visibility or plugin policy;
    they should read these attributes / methods.
    """

    name: str = ""
    # First key is the canonical one written into run_identity.json.
    visibility_env_keys: Tuple[str, ...] = ()
    # torch.<namespace> used for optional device_count / identity probes.
    torch_namespace: str = ""
    # FlagGems GEMS_VENDOR value. Lives on the backend object because zl_bench
    # still only drives FlagGems, but the mapping is owned by providers.flaggems
    # and copied here so DeviceBackend.prepare_environment() can stay generic.
    flaggems_vendor: str = ""
    supports_cuda_graph: bool = False
    supports_mm_cluster_fix: bool = False
    supports_triton_hooks: bool = True

    # --- graph-capture failure classification -----------------------------
    # A kernel that refuses capture is normal and falls back. A device error is
    # not, and the two are told apart by vendor error text, so the wording lives
    # on the backend rather than in the timing plugin.
    #
    # fatal    : the kernel itself failed on the device. Re-raise; recording a
    #            fallback latency here would report a broken kernel as a result.
    # poisoning: the capture ATTEMPT left the context unusable. Falling back runs
    #            the measurement inside a dead process, which is how an operator
    #            ends up emitting a fallback marker and then still dying at
    #            teardown with zero rows. Re-raise too, but under its own marker,
    #            because the operator is capture-incompatible rather than broken:
    #            the actionable response is to re-run it with ZL_ENABLE_CUDAGRAPH=0.
    fatal_error_markers: Tuple[str, ...] = ()
    poisoning_error_markers: Tuple[str, ...] = ()
    # Operators whose capture attempt is known to poison the context on this
    # backend. Skipped up front so one of them cannot cost a whole run.
    capture_deny_ops: Tuple[str, ...] = ()

    def classify_capture_error(self, exc: BaseException) -> str:
        """-> "fatal" | "poisoned" | "recoverable".

        Matched against the exception text itself, never a generic footer: torch
        appends "Compile with `TORCH_USE_CUDA_DSA` to enable device-side
        assertions." to every AcceleratorError, so a loose "device-side assert"
        marker would classify every capture refusal as fatal and disable the
        documented fallback path outright.
        """
        msg = f"{type(exc).__name__}: {exc}".lower()
        if any(m in msg for m in self.fatal_error_markers):
            return "fatal"
        if any(m in msg for m in self.poisoning_error_markers):
            return "poisoned"
        return "recoverable"

    def denies_capture(self, op: str) -> bool:
        return bool(op) and op in self.capture_deny_ops

    # Plugins that patch around upstream assumptions this backend violates, e.g.
    # FlagGems benchmarks that gate themselves on torch.cuda.is_available() even
    # though the operator under test runs on whatever device FlagGems dispatched
    # to. Loaded first so collection-time overrides land before anything reads
    # the item list. Each such plugin must scope itself to its own vendor.
    compat_plugins: Tuple[str, ...] = ()
    # Optional listing command; (argv, line_prefix). prefix="" means count
    # non-empty lines. Must be a real vendor CLI; do not invent flags.
    list_command: Optional[Tuple[Sequence[str], str]] = None

    def visibility_env(self, device: str) -> Dict[str, str]:
        return {key: device for key in self.visibility_env_keys}

    def apply_visibility(self, env: Dict[str, str], device: str) -> Dict[str, str]:
        out = dict(env)
        for key in ALL_VISIBILITY_KEYS:
            out.pop(key, None)
        out.update(self.visibility_env(device))
        return out

    def enumerate_devices(self, explicit: Optional[str] = None) -> List[str]:
        listed = parse_device_list(explicit)
        if listed:
            return listed
        for key in self.visibility_env_keys:
            listed = parse_device_list(os.environ.get(key, ""))
            if listed:
                return listed
        n = self.probe_device_count()
        if n > 0:
            return [str(i) for i in range(n)]
        return ["0"]

    def probe_device_count(self) -> int:
        n = self._probe_list_command()
        if n > 0:
            return n
        return self._probe_torch_count()

    def _probe_list_command(self) -> int:
        if not self.list_command:
            return 0
        argv, prefix = self.list_command
        try:
            proc = subprocess.run(
                list(argv), text=True, capture_output=True, timeout=30,
            )
        except Exception:
            return 0
        if proc.returncode != 0:
            return 0
        lines = [ln for ln in proc.stdout.splitlines() if ln.strip()]
        if prefix:
            lines = [ln for ln in lines if ln.startswith(prefix)]
        return len(lines)

    def _probe_torch_count(self) -> int:
        if not self.torch_namespace:
            return 0
        try:
            import torch
            api = getattr(torch, self.torch_namespace, None)
            if api is None or not hasattr(api, "device_count"):
                return 0
            return int(api.device_count())
        except Exception:
            return 0

    def device_identity(self) -> Dict[str, str]:
        """Best-effort identity for screening logs. Empty strings are allowed."""
        info = {
            "backend": self.name,
            "device_name": "",
            "device_arch": "",
            "timer_api": f"torch.{self.torch_namespace}" if self.torch_namespace else "",
            "flaggems_vendor": self.flaggems_vendor,
        }
        if not self.torch_namespace:
            return info
        try:
            import torch
            api = getattr(torch, self.torch_namespace, None)
            if api is None:
                return info
            count = int(api.device_count()) if hasattr(api, "device_count") else 0
            if count <= 0:
                return info
            index = int(api.current_device()) if hasattr(api, "current_device") else 0
            if hasattr(api, "get_device_name"):
                info["device_name"] = str(api.get_device_name(index))
            if hasattr(api, "get_device_capability"):
                cap = api.get_device_capability(index)
                info["device_arch"] = f"{cap[0]}.{cap[1]}"
            elif hasattr(api, "get_device_properties"):
                props = api.get_device_properties(index)
                arch = getattr(props, "name", "") or getattr(props, "arch", "")
                info["device_arch"] = str(arch)
        except Exception:
            return info
        return info