75 lines
3.1 KiB
Python
75 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""NVIDIA GPU 推理后端实现（5090，Triton/ONNX）。
|
||
|
||
对应 PRD 5.6：5090 实现（Triton/ONNX）——通过 OpenAI 兼容接口
|
||
（vLLM / Triton Inference Server）对外提供推理，业务侧统一走
|
||
`InferenceBackend` 接口，不感知硬件。
|
||
"""
|
||
from inference_backend.base import InferenceBackend, InferRequest, InferResult
|
||
|
||
|
||
class NvidiaGpuBackend(InferenceBackend):
|
||
"""NVIDIA GPU 后端：面向 5090 服务器（Triton/ONNX，OpenAI 兼容）。"""
|
||
|
||
backend_name = "gpu"
|
||
|
||
def __init__(self, endpoint: str = "", model: str = "iaop-default",
|
||
timeout_seconds: float = 10.0, runtime: str = "vllm",
|
||
device: str = "nvidia-5090", **kwargs):
|
||
super().__init__(endpoint, model, timeout_seconds)
|
||
self.runtime = runtime # vllm | triton
|
||
self.device = device # 硬件型号（默认 5090）
|
||
self._loaded = False
|
||
|
||
def load_model(self, model_name: str | None = None) -> dict:
|
||
"""加载模型到 GPU。无 endpoint 时仅维护本地状态（适配层就绪）。"""
|
||
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("/v1/models/load",
|
||
{"model": model_name, "device": self.device})
|
||
self._loaded = body.get("status") in ("ok", "loaded", "ready")
|
||
return body
|
||
|
||
def infer(self, request: InferRequest) -> InferResult:
|
||
"""OpenAI 兼容生成（/v1/chat/completions）。"""
|
||
started = __import__("time").monotonic()
|
||
payload = request.to_payload()
|
||
payload["model"] = payload["model"] or self.model
|
||
body = self._post_json("/v1/chat/completions", payload)
|
||
latency_ms = round((__import__("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,
|
||
"model": self.model, "raw": body},
|
||
)
|
||
|
||
def health(self) -> dict:
|
||
"""健康巡检：探测 /health，返回后端/设备/运行时信息。"""
|
||
base = self._healthz()
|
||
base.update({
|
||
"backend": self.backend_name,
|
||
"runtime": self.runtime,
|
||
"device": self.device,
|
||
"model": self.model,
|
||
"loaded": self._loaded,
|
||
})
|
||
return base
|
||
|
||
def unload(self) -> dict:
|
||
"""卸载模型、释放显存。"""
|
||
self._loaded = False
|
||
if not self.endpoint:
|
||
return {"status": "ok", "backend": self.backend_name,
|
||
"reason": "dry-run（未配置 endpoint）"}
|
||
return self._post_json("/v1/models/unload", {"model": self.model})
|