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})
|