104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""推理后端统一抽象接口（PRD 5.6「⑥ 部署底座」）。
|
|||
|
|
|
|||
|
|
设计约束（对齐产品设计文档 5.6）：
|
|||
|
|
- 统一接口 `load_model / infer / health / unload`，与底层硬件无关；
|
|||
|
|
- NVIDIA 5090 实现（Triton/ONNX）与华为昇腾实现（ACL/CANN）均实现本接口；
|
|||
|
|
- 业务代码只 import 本模块，不 import 任何具体后端；
|
|||
|
|
- 切换后端 = 修改适配层配置（见 config/backends.template.yaml），业务代码零改动。
|
|||
|
|
"""
|
|||
|
|
from abc import ABC, abstractmethod
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
import json
|
|||
|
|
import time
|
|||
|
|
import urllib.request
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class InferRequest:
|
|||
|
|
"""统一的推理请求载荷（业务侧构造，与硬件无关）。"""
|
|||
|
|
|
|||
|
|
prompt: str
|
|||
|
|
max_tokens: int = 1024
|
|||
|
|
temperature: float = 0.1
|
|||
|
|
extra: dict = field(default_factory=dict)
|
|||
|
|
|
|||
|
|
def to_payload(self) -> dict:
|
|||
|
|
"""转换为后端服务 OpenAI 兼容的请求体。"""
|
|||
|
|
return {
|
|||
|
|
"model": self.extra.get("model") if self.extra.get("model") else None,
|
|||
|
|
"prompt": self.prompt,
|
|||
|
|
"max_tokens": self.max_tokens,
|
|||
|
|
"temperature": self.temperature,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class InferResult:
|
|||
|
|
"""统一的推理结果（业务侧消费，不感知硬件）。"""
|
|||
|
|
|
|||
|
|
text: str
|
|||
|
|
backend: str
|
|||
|
|
latency_ms: float = 0.0
|
|||
|
|
meta: dict = field(default_factory=dict)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class InferenceBackend(ABC):
|
|||
|
|
"""推理后端统一接口（对应 PRD 5.6：loadModel / infer / health / unload）。"""
|
|||
|
|
|
|||
|
|
#: 后端标识：gpu（NVIDIA 5090）| npu（华为昇腾）
|
|||
|
|
backend_name = "base"
|
|||
|
|
|
|||
|
|
def __init__(self, endpoint: str = "", model: str = "iaop-default",
|
|||
|
|
timeout_seconds: float = 10.0):
|
|||
|
|
self.endpoint = (endpoint or "").rstrip("/")
|
|||
|
|
self.model = model
|
|||
|
|
self.timeout = float(timeout_seconds)
|
|||
|
|
|
|||
|
|
# ---- 统一接口（业务代码仅依赖以下四个方法） ----
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
def load_model(self, model_name: str | None = None) -> dict:
|
|||
|
|
"""加载/热载模型（对应 PRD 的 loadModel）。返回加载状态。"""
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
def infer(self, request: InferRequest) -> InferResult:
|
|||
|
|
"""执行一次推理（对应 PRD 的 infer）。返回统一结果对象。"""
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
def health(self) -> dict:
|
|||
|
|
"""健康巡检（对应 PRD 的 health）。返回后端状态与版本信息。"""
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
def unload(self) -> dict:
|
|||
|
|
"""卸载模型/释放显存（对应 PRD 的 unload）。返回卸载结果。"""
|
|||
|
|
|
|||
|
|
# ---- 内部工具：HTTP JSON 调用（各后端共享） ----
|
|||
|
|
|
|||
|
|
def _post_json(self, path: str, payload: dict) -> dict:
|
|||
|
|
"""向后端服务发起 JSON POST 请求，返回解析后的响应体。"""
|
|||
|
|
url = self.endpoint + path
|
|||
|
|
data = json.dumps(payload).encode("utf-8")
|
|||
|
|
req = urllib.request.Request(
|
|||
|
|
url, data=data, headers={"Content-Type": "application/json"}
|
|||
|
|
)
|
|||
|
|
started = time.monotonic()
|
|||
|
|
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
|||
|
|
raw = resp.read().decode("utf-8")
|
|||
|
|
meta = {"http_ms": round((time.monotonic() - started) * 1000, 2)}
|
|||
|
|
body = json.loads(raw) if raw else {}
|
|||
|
|
if isinstance(body, dict):
|
|||
|
|
body.setdefault("_http_ms", meta["http_ms"])
|
|||
|
|
return body
|
|||
|
|
|
|||
|
|
def _healthz(self) -> dict:
|
|||
|
|
"""通用 /health 探针封装。"""
|
|||
|
|
if not self.endpoint:
|
|||
|
|
return {
|
|||
|
|
"status": "ok",
|
|||
|
|
"backend": self.backend_name,
|
|||
|
|
"model": self.model,
|
|||
|
|
"reason": "dry-run（未配置 endpoint，适配层就绪）",
|
|||
|
|
}
|
|||
|
|
return self._post_json("/health", {"model": self.model})
|