225 lines
7.6 KiB
Python
225 lines
7.6 KiB
Python
# -*- 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 可执行但未解析到昇腾设备",
|
||
}
|