158 lines
7.1 KiB
Python
158 lines
7.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""华为昇腾 NPU 推理后端实现（ACL/CANN）。
|
||
|
||
对应 PRD 5.6 与 issue #59「昇腾 NPU 后端适配（CANN 对接）」：
|
||
昇腾实现（ACL/CANN）通过昇腾推理服务（MindIE / onnxruntime-ascend）的
|
||
OpenAI 兼容接口对外提供推理，与 NVIDIA GPU 后端实现同一
|
||
`InferenceBackend` 接口：**切换后端仅改适配层配置，业务代码零改动**。
|
||
|
||
两种运行模式（由 endpoint 是否配置决定）：
|
||
1. **服务模式**（endpoint 非空）：对接 MindIE / onnxruntime-ascend 的
|
||
OpenAI 兼容 HTTP 接口，走 ``/acl/models/load``、``/v1/chat/completions``、
|
||
``/acl/models/unload``、``/health``；
|
||
2. **直连模式**（endpoint 为空且本机可检测到 CANN）：调用本地 ACL Python
|
||
API（``acl.init`` / ``acl.rt.set_device`` / ``acl.rt.reset_device`` /
|
||
``acl.finalize``）管理昇腾设备与模型加载；
|
||
endpoint 为空且无 CANN 环境时进入 dry-run（适配层就绪，便于离线验证）。
|
||
|
||
健康巡检会附带 :func:`~inference_backend.cann_probe.probe_cann_environment`
|
||
探测到的 CANN 设备/工具链信息（``cann`` 段），便于运维确认后端就绪状态。
|
||
"""
|
||
import time
|
||
|
||
from inference_backend.base import InferenceBackend, InferRequest, InferResult
|
||
from inference_backend.cann_probe import probe_cann_environment
|
||
|
||
|
||
class AscendNpuBackend(InferenceBackend):
|
||
"""华为昇腾 NPU 后端：面向昇腾 310P/910B（CANN/MindIE）。"""
|
||
|
||
backend_name = "npu"
|
||
|
||
def __init__(self, endpoint: str = "", model: str = "iaop-default",
|
||
timeout_seconds: float = 10.0, runtime: str = "mindie",
|
||
device: str = "ascend-910b", cann_version: str = "8.0",
|
||
**kwargs):
|
||
super().__init__(endpoint, model, timeout_seconds)
|
||
self.runtime = runtime # mindie | onnx-ascend
|
||
self.device = device # ascend-310p | ascend-910b
|
||
self.cann_version = cann_version # CANN 工具链版本
|
||
self._loaded = False
|
||
self._acl_initialized = False
|
||
self._cann_cache = None
|
||
|
||
# ---- CANN 环境探测（懒加载，探测结果缓存于实例） ----
|
||
|
||
def cann_info(self) -> dict:
|
||
"""返回本机 CANN 环境探测结果（见 cann_probe.probe_cann_environment）。"""
|
||
if self._cann_cache is None:
|
||
self._cann_cache = probe_cann_environment()
|
||
return self._cann_cache
|
||
|
||
# ---- 统一接口实现 ----
|
||
|
||
def load_model(self, model_name: str | None = None) -> dict:
|
||
"""加载模型到 NPU。
|
||
|
||
服务模式：``POST /acl/models/load``（ACL aclmdlLoadFromFile 语义）；
|
||
直连模式：本地 ``acl.init()`` + ``acl.rt.set_device(device_id)``；
|
||
无 endpoint 且无 CANN 环境：dry-run，仅上报适配层就绪。
|
||
"""
|
||
model_name = model_name or self.model
|
||
cann = self.cann_info()
|
||
if self.endpoint:
|
||
body = self._post_json(
|
||
"/acl/models/load",
|
||
{"model": model_name, "device": self.device,
|
||
"cann_version": self.cann_version})
|
||
self._loaded = body.get("status") in ("ok", "loaded", "ready")
|
||
return body
|
||
if cann["available"]:
|
||
self._acl_load()
|
||
self._loaded = True
|
||
return {"status": "ok", "backend": self.backend_name,
|
||
"model": model_name, "device": self.device,
|
||
"runtime": self.runtime, "cann": cann,
|
||
"reason": "直连模式：CANN ACL 已初始化并绑定设备"}
|
||
self._loaded = True
|
||
return {"status": "ok", "backend": self.backend_name,
|
||
"model": model_name, "device": self.device,
|
||
"cann": cann,
|
||
"reason": "dry-run（未配置 endpoint 且无 CANN 环境，适配层就绪）"}
|
||
|
||
def _acl_load(self) -> None:
|
||
"""直连模式初始化 CANN ACL 并把上下文绑定到首张昇腾设备。"""
|
||
try:
|
||
import acl
|
||
except Exception as exc: # 探测与实际导入之间环境可能变化，容错
|
||
raise RuntimeError(f"CANN ACL 不可用，无法直连加载: {exc}") from exc
|
||
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):
|
||
raise RuntimeError(f"acl.init() 失败: ret={ret}")
|
||
device_id = 0 # 默认首卡；多卡资源调度由部署侧配置扩展
|
||
ret = acl.rt.set_device(device_id)
|
||
if not (ret is True or ret is None or ret == 0):
|
||
acl.finalize()
|
||
raise RuntimeError(f"acl.rt.set_device({device_id}) 失败: ret={ret}")
|
||
self._acl_initialized = True
|
||
|
||
def infer(self, request: InferRequest) -> InferResult:
|
||
"""昇腾推理（MindIE OpenAI 兼容 /v1/chat/completions）。"""
|
||
started = time.monotonic()
|
||
payload = request.to_payload()
|
||
payload["model"] = payload["model"] or self.model
|
||
body = self._post_json("/v1/chat/completions", payload)
|
||
latency_ms = round((time.monotonic() - started) * 1000, 2)
|
||
try:
|
||
text = body["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError):
|
||
text = str(body)
|
||
return InferResult(
|
||
text=text,
|
||
backend=self.backend_name,
|
||
latency_ms=latency_ms,
|
||
meta={"runtime": self.runtime, "device": self.device,
|
||
"cann_version": self.cann_version, "model": self.model,
|
||
"cann": self.cann_info(), "raw": body},
|
||
)
|
||
|
||
def health(self) -> dict:
|
||
"""健康巡检：探测 /health，返回后端/设备/CANN 环境信息。"""
|
||
base = self._healthz()
|
||
base.update({
|
||
"backend": self.backend_name,
|
||
"runtime": self.runtime,
|
||
"device": self.device,
|
||
"cann_version": self.cann_version,
|
||
"model": self.model,
|
||
"loaded": self._loaded,
|
||
"cann": self.cann_info(),
|
||
})
|
||
return base
|
||
|
||
def unload(self) -> dict:
|
||
"""卸载模型、释放 NPU 资源。
|
||
|
||
服务模式：``POST /acl/models/unload``（ACL aclmdlUnload 语义）；
|
||
直连模式：``acl.rt.reset_device`` + ``acl.finalize``；
|
||
无 endpoint 且未直连加载：dry-run。
|
||
"""
|
||
self._loaded = False
|
||
if self.endpoint:
|
||
return self._post_json("/acl/models/unload", {"model": self.model})
|
||
if self._acl_initialized:
|
||
try:
|
||
import acl
|
||
acl.rt.reset_device(0)
|
||
acl.finalize()
|
||
except Exception as exc: # noqa: BLE001
|
||
return {"status": "warn", "backend": self.backend_name,
|
||
"reason": f"ACL 资源释放失败: {exc}"}
|
||
self._acl_initialized = False
|
||
return {"status": "ok", "backend": self.backend_name,
|
||
"reason": "直连模式：ACL 资源已释放"}
|
||
return {"status": "ok", "backend": self.backend_name,
|
||
"reason": "dry-run（未配置 endpoint）"}
|