feat: 完成 issue #59 ⑥ 昇腾 NPU 后端适配（CANN 对接）

2026-08-05 04:13:23 +08:00
parent 5d9a76c3ff
commit 95dace5ef5
4 changed files with 584 additions and 23 deletions
@@ -0,0 +1,224 @@
+# -*- coding: utf-8 -*-
+"""昇腾 CANN 运行时环境探测（issue #59「昇腾 NPU 后端适配（CANN 对接）」）。
+
+目标：
+- 在**无 CANN 硬件的开发/CI 环境**下安全降级（available=False，绝不抛异常）；
+- 在**有 CANN 的生产环境**正确上报设备数量/名称与工具链版本，
+  供 :class:`~inference_backend.npu_backend.AscendNpuBackend` 的健康巡检
+  与直连模式（本地 ACL 生命周期管理）使用。
+
+探测优先级（取第一个可用的运行时）：
+1. ``torch_npu`` —— PyTorch 昇腾插件（``torch.npu.*``）；
+2. ``acl`` —— CANN 基础 ACL Python API（``acl.init`` / ``acl.rt.*``）；
+3. ``npu-smi`` —— 昇腾命令行工具（``npu-smi info`` 输出解析）。
+
+ACL API 各 CANN 版本存在差异（部分接口返回 ``(ret, value)`` 元组、
+部分直接返回值），本模块统一用 ``_unwrap`` 兼容两种风格，并逐调用
+try/except 兜底 —— 探测逻辑在任何版本/异常下都不会让调用方崩溃。
+"""
+from __future__ import annotations
+
+import shutil
+import subprocess
+from typing import Optional
+
+#: 探测失败兜底结果（永不抛异常）
+_NO_CANN = {
+    "available": False,
+    "runtime": "none",
+    "cann_version": "",
+    "device_count": 0,
+    "device_names": [],
+    "detail": "未检测到 CANN 运行时（torch_npu / acl / npu-smi 均不可用）",
+}
+
+
+def probe_cann_environment() -> dict:
+    """探测当前环境的昇腾 CANN 运行时，返回结构化结果。
+
+    Returns:
+        dict，固定键：``available`` / ``runtime`` / ``cann_version`` /
+        ``device_count`` / ``device_names`` / ``detail``。
+        任一探测路径失败均不抛异常，仅将 ``available`` 置为 False。
+    """
+    for probe in (_probe_torch_npu, _probe_acl, _probe_npu_smi):
+        info = probe()
+        if info is not None:
+            return info
+    return dict(_NO_CANN)
+
+
+def _unwrap(ret_value):
+    """兼容 ACL API 两种返回风格：``(ret, value)`` 元组或直接返回 value。"""
+    if isinstance(ret_value, tuple):
+        if len(ret_value) >= 2:
+            return ret_value[1]
+        return None
+    return ret_value
+
+
+def _version_of(module, name: str) -> str:
+    """尽力取模块版本号（torch_npu.__version__ / acl.__version__）。"""
+    try:
+        ver = getattr(module, "__version__", "")
+        return str(ver) if ver else ""
+    except Exception:
+        return ""
+
+
+def _probe_torch_npu() -> Optional[dict]:
+    """探测 torch_npu（PyTorch 昇腾插件）路径。"""
+    try:
+        import torch  # noqa: F401
+        import torch_npu  # noqa: F401
+    except Exception:
+        return None  # 未安装该运行时，交给下一个探测路径
+    base = {
+        "runtime": "torch_npu",
+        "cann_version": _version_of(torch_npu, "torch_npu"),
+    }
+    try:
+        if not torch_npu.npu.is_available():
+            return {
+                **base,
+                "available": False,
+                "device_count": 0,
+                "device_names": [],
+                "detail": "torch_npu 可导入但 NPU 不可用（未检测到昇腾设备）",
+            }
+        count = int(torch.npu.device_count())
+        names = []
+        for i in range(count):
+            try:
+                names.append(str(torch.npu.get_device_name(i)))
+            except Exception:
+                names.append("ascend-device-%d" % i)
+        return {
+            **base,
+            "available": True,
+            "device_count": count,
+            "device_names": names,
+            "detail": "torch_npu 探测成功（%d 张昇腾设备）" % count,
+        }
+    except Exception as exc:  # noqa: BLE001 —— 探测必须容错
+        return {
+            **base,
+            "available": False,
+            "device_count": 0,
+            "device_names": [],
+            "detail": "torch_npu 探测失败: %s" % exc,
+        }
+
+
+def _probe_acl() -> Optional[dict]:
+    """探测 CANN ACL Python API 路径。"""
+    try:
+        import acl
+    except Exception:
+        return None  # 未安装该运行时，交给下一个探测路径
+    base = {
+        "runtime": "acl",
+        "cann_version": _version_of(acl, "acl"),
+    }
+    try:
+        ret = acl.init()
+        # 注意：不能用 `ret not in (0, None, True)` —— Python 中 1 == True，
+        # 会把错误码 1 误判为成功；ACL 约定 ret=0（ACL_SUCCESS）为成功。
+        if not (ret is True or ret is None or ret == 0):
+            return {
+                **base,
+                "available": False,
+                "device_count": 0,
+                "device_names": [],
+                "detail": "acl.init() 返回错误码 %s" % ret,
+            }
+        try:
+            count = int(_unwrap(acl.rt.get_device_count()))
+        except Exception:
+            count = 0
+        names = []
+        for i in range(count):
+            name = None
+            for method in ("get_device_name", "get_soc_name"):
+                fn = getattr(acl.rt, method, None)
+                if fn is None:
+                    continue
+                try:
+                    name = _unwrap(fn(i))
+                    break
+                except Exception:
+                    name = None
+            names.append(str(name) if name else "ascend-device-%d" % i)
+        # CANN 版本：优先 acl.rt.get_version()，其次模块 __version__
+        try:
+            ver = _unwrap(acl.rt.get_version()) or base["cann_version"]
+        except Exception:
+            ver = base["cann_version"]
+        return {
+            **base,
+            "cann_version": str(ver) if ver else "",
+            "available": True,
+            "device_count": count,
+            "device_names": names,
+            "detail": "ACL 探测成功（%d 张昇腾设备）" % count,
+        }
+    except Exception as exc:  # noqa: BLE001
+        return {
+            **base,
+            "available": False,
+            "device_count": 0,
+            "device_names": [],
+            "detail": "ACL 探测失败: %s" % exc,
+        }
+    finally:
+        try:
+            acl.finalize()  # 探测后立即释放，避免占用设备
+        except Exception:
+            pass
+
+
+def _probe_npu_smi() -> Optional[dict]:
+    """探测 npu-smi 命令行工具路径。"""
+    exe = shutil.which("npu-smi")
+    if not exe:
+        return None  # 未安装该工具
+    try:
+        proc = subprocess.run(
+            [exe, "info", "-l"],
+            capture_output=True, text=True, timeout=10,
+            check=False,
+        )
+        text = proc.stdout + "\n" + proc.stderr
+    except Exception as exc:
+        return {
+            "available": False,
+            "runtime": "npu-smi",
+            "cann_version": "",
+            "device_count": 0,
+            "device_names": [],
+            "detail": "npu-smi 执行失败: %s" % exc,
+        }
+    # 解析 "Device Count : N" 与 "Name : xxx"（npu-smi info -l 常见输出）
+    device_count = 0
+    names = []
+    for line in text.splitlines():
+        low = line.lower()
+        if "device count" in low and ":" in line:
+            try:
+                device_count = int(line.split(":")[-1].strip())
+            except ValueError:
+                device_count = 0
+        elif "name" in low and ":" in line:
+            name = line.split(":", 1)[-1].strip()
+            if name:
+                names.append(name)
+    available = device_count > 0 or bool(names)
+    return {
+        "available": available,
+        "runtime": "npu-smi",
+        "cann_version": "",
+        "device_count": device_count,
+        "device_names": names[: max(device_count, len(names))],
+        "detail": "npu-smi 探测成功（%d 张昇腾设备）" % device_count
+        if available else "npu-smi 可执行但未解析到昇腾设备",
+    }