82 lines
3.5 KiB
Python
82 lines
3.5 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""华为昇腾 NPU 推理后端实现（ACL/CANN）。
|
|||
|
|
|
|||
|
|
对应 PRD 5.6：昇腾实现（ACL/CANN）——通过昇腾推理服务
|
|||
|
|
（MindIE / onnxruntime-ascend）的 OpenAI 兼容接口对外提供推理，
|
|||
|
|
与 NVIDIA GPU 后端实现同一 `InferenceBackend` 接口：
|
|||
|
|
**切换后端仅改适配层配置，业务代码零改动**。
|
|||
|
|
"""
|
|||
|
|
from inference_backend.base import InferenceBackend, InferRequest, InferResult
|
|||
|
|
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
def load_model(self, model_name: str | None = None) -> dict:
|
|||
|
|
"""加载模型到 NPU（ACL aclmdlLoadFromFile 语义）。"""
|
|||
|
|
model_name = model_name or self.model
|
|||
|
|
if not self.endpoint:
|
|||
|
|
self._loaded = True
|
|||
|
|
return {"status": "ok", "backend": self.backend_name,
|
|||
|
|
"model": model_name, "device": self.device,
|
|||
|
|
"reason": "dry-run（未配置 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
|
|||
|
|
|
|||
|
|
def infer(self, request: InferRequest) -> InferResult:
|
|||
|
|
"""昇腾推理（MindIE OpenAI 兼容 /v1/chat/completions）。"""
|
|||
|
|
import time
|
|||
|
|
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,
|
|||
|
|
"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,
|
|||
|
|
})
|
|||
|
|
return base
|
|||
|
|
|
|||
|
|
def unload(self) -> dict:
|
|||
|
|
"""卸载模型、释放 NPU 资源（ACL aclmdlUnload 语义）。"""
|
|||
|
|
self._loaded = False
|
|||
|
|
if not self.endpoint:
|
|||
|
|
return {"status": "ok", "backend": self.backend_name,
|
|||
|
|
"reason": "dry-run（未配置 endpoint）"}
|
|||
|
|
return self._post_json("/acl/models/unload", {"model": self.model})
|