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
    # 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