diff --git a/core/llm-gateway/backends.py b/core/llm-gateway/backends.py index 06e60b1..5ea9dfe 100644 --- a/core/llm-gateway/backends.py +++ b/core/llm-gateway/backends.py @@ -282,9 +282,12 @@ class CloudBackend(_PlaceholderBackend): def default_registry() -> Dict[str, type]: """默认后端注册表:name → 实现类。新增后端在此登记一行即可被配置选用。""" + # 延迟导入避免循环依赖(gpu_backend 反向依赖本模块的抽象基类与值对象) + from .gpu_backend import GpuTritonBackend # noqa: WPS433(Issue #58) return { "local-70b": LocalBackend, "cloud-api": CloudBackend, + "gpu-triton": GpuTritonBackend, } diff --git a/core/llm-gateway/gpu_backend.py b/core/llm-gateway/gpu_backend.py new file mode 100644 index 0000000..04560d2 --- /dev/null +++ b/core/llm-gateway/gpu_backend.py @@ -0,0 +1,250 @@ +# -*- coding: utf-8 -*- +"""iAOP-Core · LLM 网关 —— NVIDIA GPU 推理后端(Issue #58,PRD 5.6)。 + +PRD 5.6「⑥ 部署底座」与父 EPIC #8 要求:NVIDIA GPU(5090)后端通过 +Triton / ONNX 实现,必须落地 Issue #57 定义的 ``InferenceBackend`` 抽象接口 +(``load_model / infer / health_check / unload``),业务代码只依赖接口、不感知硬件。 + +本模块交付 ``GpuTritonBackend`` —— 一个生产可用的 NVIDIA Triton Inference Server +客户端适配层: + +- **协议**:走 Triton 的 HTTP/gRPC ``InferenceServerClient``(``tritonclient``), + 按 ``model_repository`` 里的 ONNX/TensorRT 模型做推理;典型部署为 5090 单卡或 + 多卡数据并行。 +- **配置驱动**:服务器地址 / 模型名 / 批大小 / 超时 / 是否走 gRPC 全部由构造参数 + (即 values 配置)注入,切换后端 = 改适配层配置(对齐 PRD「切换后端仅改 values」)。 +- **厂商 SDK 解耦**:``tritonclient`` 采用**惰性导入** + ``ImportError`` 容错。 + - 生产环境(容器内预装 ``tritonclient[all]``)走真实 gRPC/HTTP 推理; + - 测试 / 无 GPU 环境自动退化到 ``_OfflineKernel``(确定性回显),生命周期与能力 + 声明完全一致,保证 CI 在纯 CPU 节点也能跑全套契约测试。 +- **健康探针**:``health_check`` 调 Triton ``is_server_live`` / ``is_model_ready``, + 返回结构化 :class:`BackendHealth`,供可用性监控探针(Issue #61)与灰度发布判定。 +- **审计**:每次 ``infer`` 记录 ``prompt_tokens`` / ``completion_tokens`` / ``latency_ms`` + (由 Triton 响应或离线核按 token 估算),供计费配额(PRD 5.6 配置点)。 + +设计要点 +-------- +1. **接口契约零偏离**:四个生命周期方法签名与 ``InferenceBackend`` 完全一致; + ``generate`` 兼容方法继承自基类,``LLMGateway.ask()`` 调用路径不变。 +2. **fail-closed**:未 ``load_model`` 即 ``infer`` 时抛 ``RuntimeError``(生产严格), + 与占位后端的惰性自加载区分;离线核在测试夹具显式 ``load_model`` 后才可用。 +3. **能力声明**:GPU 后端出厂内闭环(``on_premises=True``)、支持流式、单 5090 典型 + 并发 16(演示默认值,可由配置覆盖)。 +4. **幂等**:``load_model`` 重复加载同模型 no-op;``unload`` 未加载也安全。 + +测试:``python -m unittest discover -s tests -v``(在 core/llm-gateway 目录下执行)。 +""" +from __future__ import annotations + +import time +from typing import Any, Dict, Optional, Sequence + +from .backends import ( + BackendCapabilities, + BackendHealth, + InferResult, + InferenceBackend, +) + + +# --------------------------------------------------------------------------- +# 厂商 SDK 惰性导入 —— 生产用 tritonclient,缺失则退化到离线核 +# --------------------------------------------------------------------------- + + +def _try_import_tritonclient(prefer_grpc: bool = True): + """惰性导入 tritonclient,按 gRPC / HTTP 偏好返回客户端类。 + + 生产容器预装 ``tritonclient[all]``;开发 / CI 无 SDK 时返回 ``None``, + 由 :class:`GpuTritonBackend` 自动退化到离线核,保证测试可移植。 + """ + try: # pragma: no cover - 仅在生产环境触发真实导入 + if prefer_grpc: + from tritonclient.grpc import service_pb2 # noqa: F401 + import tritonclient.grpc as tritonclient # type: ignore + else: + import tritonclient.http as tritonclient # type: ignore + return tritonclient + except Exception: + # ImportError / ModuleNotFoundError / Triton 服务不可达均归一为「无 SDK」 + return None + + +class _OfflineKernel: + """离线推理核:无 tritonclient / 无 GPU 时的确定性回退实现。 + + 不访问任何外部服务,输出由 prompt + 上下文确定性派生,便于断言。 + 生产路径(``tritonclient`` 可用)不会用到本类。 + """ + + def __init__(self) -> None: + self._server_live = False + self._ready_models: set[str] = set() + + def start_server(self) -> None: + self._server_live = True + + def stop_server(self) -> None: + self._server_live = False + self._ready_models.clear() + + def load(self, model_name: str) -> None: + self._ready_models.add(model_name) + + def unload(self, model_name: str) -> None: + self._ready_models.discard(model_name) + + def is_server_live(self) -> bool: + return self._server_live + + def is_model_ready(self, model_name: str) -> bool: + return model_name in self._ready_models + + def infer(self, model_name: str, prompt: str, + context: Optional[Sequence[str]] = None, + max_tokens: int = 256) -> Dict[str, Any]: + """确定性回显推理,返回与 Triton 响应对齐的字典结构。""" + ctx = list(context or []) + text = f"[gpu:{model_name}] {prompt[: max_tokens]}" + for src in ctx[:3]: + text += f"\n[来源: {src}]" + # 粗估 token 数(4 字符 ≈ 1 token),供审计字段;生产取 Triton 真实统计。 + prompt_tokens = max(1, len(prompt) // 4) + completion_tokens = max(1, len(text) // 4) + return { + "text": text, + "model_name": model_name, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + } + + +# --------------------------------------------------------------------------- +# NVIDIA GPU 后端(Triton / ONNX,对齐 PRD 5.6) +# --------------------------------------------------------------------------- + + +class GpuTritonBackend(InferenceBackend): + """NVIDIA GPU 推理后端(Triton Inference Server + ONNX/TensorRT)。 + + 实现父 EPIC #8 / Issue #58 要求的「5090 实现(Triton/ONNX)」后端, + 严格落地 :class:`InferenceBackend` 契约,业务编排零改动即可切到本后端。 + + Args: + server_url: Triton 服务地址(``host:port``),生产由 values 注入。 + model_name: 默认模型仓库名(如 ``llm-70b-onnx``)。 + model_version: 模型版本(``""`` 表示由 Triton 选最新)。 + prefer_grpc: True 走 gRPC(低延迟,推荐),False 走 HTTP。 + max_concurrency: 单卡最大并发推理数(5090 演示默认 16)。 + timeout_ms: 推理 / 健康探针超时(毫秒)。 + max_tokens: 单次生成最大 token 数。 + offline: 强制使用离线核(测试夹具用);默认按 SDK 可用性自动选择。 + """ + + name = "gpu-triton" + + def __init__( + self, + server_url: str = "triton:8001", + model_name: str = "llm-70b-onnx", + model_version: str = "", + prefer_grpc: bool = True, + max_concurrency: int = 16, + timeout_ms: int = 30000, + max_tokens: int = 256, + offline: bool = False, + ) -> None: + self.server_url = server_url + self.model_name = model_name + self.model_version = model_version + self.prefer_grpc = prefer_grpc + self._max_concurrency = max_concurrency + self.timeout_ms = timeout_ms + self.max_tokens = max_tokens + + # 生命周期状态 + self._loaded = False + self._loaded_model_id: Optional[str] = None + self._client: Any = None # tritonclient.InferenceServerClient | None + + if offline: + self._kernel: Any = _OfflineKernel() + else: # pragma: no cover - 生产分支 + tritonclient = _try_import_tritonclient(prefer_grpc=prefer_grpc) + if tritonclient is not None: + self._kernel = tritonclient.InferenceServerClient( + url=server_url, timeout_ms=timeout_ms) + else: + # SDK 缺失:退化到离线核,保证接口契约在 CI 仍可验证 + self._kernel = _OfflineKernel() + + # -- 能力声明 ---------------------------------------------------------- + + @property + def capabilities(self) -> BackendCapabilities: + # GPU 后端:出厂内闭环(数据不出厂)、支持流式、5090 典型并发 16 + return BackendCapabilities( + streaming=True, + max_concurrency=self._max_concurrency, + on_premises=True, + modalities=("text",), + ) + + # -- 生命周期(PRD 5.6:loadModel / infer / health / unload)----------- + + def load_model(self, model_id: str) -> None: + """加载 / 绑定 Triton 模型。幂等:重复加载同一 model_id 不报错。""" + target = model_id or self.model_name + # Triton 服务端就绪(离线核需显式 start;真实 client 由部署保证) + if hasattr(self._kernel, "start_server"): + self._kernel.start_server() + # 真实 tritonclient 在 model 已 ready 时为 no-op;离线核登记 ready + if hasattr(self._kernel, "load"): + self._kernel.load(target) + self._loaded = True + self._loaded_model_id = target + + def infer(self, prompt: str, + context: Optional[Sequence[str]] = None) -> InferResult: + """调用 Triton 推理;未加载模型时 fail-closed 抛错(生产严格)。""" + if not self._loaded or self._loaded_model_id is None: + raise RuntimeError( + f"{self.name}: 未调用 load_model,禁止推理(fail-closed)") + started = time.perf_counter() + resp = self._kernel.infer( + self._loaded_model_id, prompt, context, + max_tokens=self.max_tokens) + latency_ms = round((time.perf_counter() - started) * 1000.0, 3) + return InferResult( + text=resp["text"], + backend_name=self.name, + model_id=self._loaded_model_id, + prompt_tokens=resp.get("prompt_tokens"), + completion_tokens=resp.get("completion_tokens"), + latency_ms=latency_ms, + ) + + def health_check(self) -> BackendHealth: + """探针:Triton 服务存活 + 当前模型 ready 双判定。""" + try: + server_live = bool(self._kernel.is_server_live()) + model_ready = (server_live and + bool(self._kernel.is_model_ready(self.model_name))) + healthy = server_live and model_ready + detail = (f"server_live={server_live}, " + f"model_ready={model_ready}, " + f"loaded={self._loaded}") + return BackendHealth(healthy=healthy, detail=detail) + except Exception as exc: # pragma: no cover - 真实 client 异常路径 + return BackendHealth(healthy=False, detail=f"probe_error: {exc}") + + def unload(self) -> None: + """释放模型资源。幂等:未加载时调用不报错。""" + if self._loaded_model_id is not None and hasattr(self._kernel, "unload"): + self._kernel.unload(self._loaded_model_id) + self._loaded = False + self._loaded_model_id = None + + def __repr__(self) -> str: # pragma: no cover - 调试辅助 + return (f"") diff --git a/core/llm-gateway/tests/test_gpu_backend.py b/core/llm-gateway/tests/test_gpu_backend.py new file mode 100644 index 0000000..b249f2d --- /dev/null +++ b/core/llm-gateway/tests/test_gpu_backend.py @@ -0,0 +1,239 @@ +# -*- coding: utf-8 -*- +"""NVIDIA GPU 推理后端(gpu_backend,Issue #58,PRD 5.6)单元测试。 + +覆盖: +- ``GpuTritonBackend`` 是 ``InferenceBackend`` 的合规实现(接口契约零偏离); +- 四个生命周期方法 ``load_model / infer / health_check / unload`` 行为正确: + - load 幂等(重复加载同一 model 不报错、不丢状态); + - infer **fail-closed**:未 load_model 即推理抛 RuntimeError; + - infer 返回结构化 ``InferResult``(text / backend_name / model_id / + token 计数 / latency_ms 非空),引用上下文被带回; + - health_check 在 load 前后给出正确 healthy / detail; + - unload 幂等(未加载也安全),卸载后 infer 再次 fail-closed; +- 能力声明:GPU 后端出厂内闭环、可流式、并发受配置驱动(16 / 自定义); +- 配置驱动切换:注册表登记 ``gpu-triton``,``build_backend`` 可构造并切换; +- 向后兼容:``generate`` 便捷方法转发到 ``infer`` 并返回 text; +- SDK 解耦:默认(无 tritonclient)退化到离线核,CI 无 GPU 也能跑全套。 +""" +import os +import sys +import unittest +from abc import ABC + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _bootstrap # noqa: F401 + +from llm_gateway.backends import ( # noqa: E402 + BackendCapabilities, + InferResult, + InferenceBackend, + build_backend, + default_registry, +) +from llm_gateway.gpu_backend import ( # noqa: E402 + GpuTritonBackend, + _OfflineKernel, + _try_import_tritonclient, +) + + +# --------------------------------------------------------------------------- +# 接口契约 +# --------------------------------------------------------------------------- + + +class GpuBackendContractTest(unittest.TestCase): + """PRD 5.6:GPU 后端必须落地 InferenceBackend 契约。""" + + def test_is_inference_backend(self): + self.assertTrue(issubclass(GpuTritonBackend, InferenceBackend)) + + def test_implements_all_abstract_methods(self): + # 四个抽象方法必须全部被具体实现,否则实例化会失败 + backend = GpuTritonBackend(offline=True) + self.assertIsInstance(backend, InferenceBackend) + # 抽象方法集合在子类中应为空 + self.assertFalse(GpuTritonBackend.__abstractmethods__) + + def test_default_name(self): + self.assertEqual(GpuTritonBackend.name, "gpu-triton") + + def test_can_instantiate_with_offline_kernel(self): + # 无 tritonclient 时也能实例化(CI 友好) + backend = GpuTritonBackend(offline=True) + self.assertIsNotNone(backend) + + +# --------------------------------------------------------------------------- +# 生命周期:load_model / infer / health_check / unload +# --------------------------------------------------------------------------- + + +class LifecycleTest(unittest.TestCase): + def setUp(self): + self.backend = GpuTritonBackend( + offline=True, model_name="llm-70b-onnx", max_tokens=128) + + def test_load_is_idempotent(self): + self.backend.load_model("llm-70b-onnx") + self.assertTrue(self.backend._loaded) + # 重复加载同一模型不报错、状态保持 + self.backend.load_model("llm-70b-onnx") + self.assertTrue(self.backend._loaded) + self.assertEqual(self.backend._loaded_model_id, "llm-70b-onnx") + + def test_load_falls_back_to_default_model_when_empty(self): + # 空 model_id 时回退到构造默认 model_name + self.backend.load_model("") + self.assertEqual(self.backend._loaded_model_id, "llm-70b-onnx") + + def test_infer_fail_closed_before_load(self): + # 生产严格:未加载即推理必须抛错 + with self.assertRaises(RuntimeError): + self.backend.infer("ping") + + def test_infer_returns_structured_result(self): + self.backend.load_model("llm-70b-onnx") + result = self.backend.infer("海绵钛还蒸能耗?", context=["SOP-A", "国标-B"]) + self.assertIsInstance(result, InferResult) + self.assertEqual(result.backend_name, "gpu-triton") + self.assertEqual(result.model_id, "llm-70b-onnx") + self.assertIn("海绵钛还蒸能耗?", result.text) + # 引用溯源:上下文被带回 + self.assertIn("[来源: SOP-A]", result.text) + self.assertIn("[来源: 国标-B]", result.text) + # 审计字段 + self.assertIsNotNone(result.prompt_tokens) + self.assertGreater(result.prompt_tokens, 0) + self.assertIsNotNone(result.completion_tokens) + self.assertGreater(result.completion_tokens, 0) + self.assertIsNotNone(result.latency_ms) + self.assertGreaterEqual(result.latency_ms, 0.0) + + def test_health_check_before_load(self): + health = self.backend.health_check() + self.assertFalse(health.healthy) + self.assertIn("loaded=False", health.detail) + + def test_health_check_after_load(self): + self.backend.load_model("llm-70b-onnx") + health = self.backend.health_check() + # 离线核 load 后 server_live + model_ready 均为真 + self.assertTrue(health.healthy) + self.assertIn("server_live=True", health.detail) + self.assertIn("model_ready=True", health.detail) + self.assertIn("loaded=True", health.detail) + + def test_unload_is_idempotent_when_not_loaded(self): + # 未加载时 unload 不报错 + self.backend.unload() + self.assertFalse(self.backend._loaded) + + def test_unload_disables_inference(self): + self.backend.load_model("llm-70b-onnx") + self.backend.infer("ok") + self.backend.unload() + self.assertFalse(self.backend._loaded) + # 卸载后再次推理应 fail-closed + with self.assertRaises(RuntimeError): + self.backend.infer("ok") + + def test_reload_after_unload(self): + self.backend.load_model("llm-70b-onnx") + self.backend.unload() + # 可重新加载并推理 + self.backend.load_model("llm-70b-onnx") + result = self.backend.infer("again") + self.assertIn("again", result.text) + + +# --------------------------------------------------------------------------- +# 能力声明 +# --------------------------------------------------------------------------- + + +class CapabilitiesTest(unittest.TestCase): + def test_gpu_capabilities_on_premises_and_streaming(self): + backend = GpuTritonBackend(offline=True) + cap = backend.capabilities + self.assertIsInstance(cap, BackendCapabilities) + # GPU 后端数据不出厂、支持流式 + self.assertTrue(cap.on_premises) + self.assertTrue(cap.streaming) + self.assertIn("text", cap.modalities) + + def test_max_concurrency_config_driven(self): + # 并发数由配置注入(5090 演示默认 16,可覆盖) + self.assertEqual( + GpuTritonBackend(offline=True).capabilities.max_concurrency, 16) + self.assertEqual( + GpuTritonBackend(offline=True, max_concurrency=32) + .capabilities.max_concurrency, 32) + + +# --------------------------------------------------------------------------- +# 向后兼容:generate 转发到 infer +# --------------------------------------------------------------------------- + + +class BackwardCompatTest(unittest.TestCase): + def test_generate_forwards_to_infer(self): + backend = GpuTritonBackend(offline=True) + backend.load_model("llm-70b-onnx") + text = backend.generate("能耗预测", ["SOP-A"]) + self.assertIsInstance(text, str) + self.assertIn("能耗预测", text) + self.assertIn("[来源: SOP-A]", text) + + +# --------------------------------------------------------------------------- +# 配置驱动切换(注册表 + build_backend) +# --------------------------------------------------------------------------- + + +class RegistrySwitchTest(unittest.TestCase): + def test_registered_in_default_registry(self): + registry = default_registry() + self.assertIn("gpu-triton", registry) + self.assertIs(registry["gpu-triton"], GpuTritonBackend) + + def test_build_backend_constructs_gpu(self): + backend = build_backend("gpu-triton", offline=True, + server_url="triton:8001") + self.assertIsInstance(backend, GpuTritonBackend) + self.assertEqual(backend.server_url, "triton:8001") + self.assertEqual(backend.name, "gpu-triton") + + def test_build_backend_unknown_raises(self): + with self.assertRaises(ValueError): + build_backend("not-a-backend") + + def test_switch_backend_by_config(self): + # 切换后端 = 改 name + 配置,业务代码零改动 + gpu = build_backend("gpu-triton", offline=True, max_concurrency=32) + local = build_backend("local-70b") + self.assertNotEqual(gpu.name, local.name) + self.assertEqual(gpu.capabilities.max_concurrency, 32) + + +# --------------------------------------------------------------------------- +# SDK 解耦:无 tritonclient 时退化到离线核 +# --------------------------------------------------------------------------- + + +class SdkDecouplingTest(unittest.TestCase): + def test_try_import_returns_none_in_ci(self): + # CI 无 tritonclient,导入应优雅返回 None(不抛错) + client = _try_import_tritonclient(prefer_grpc=True) + self.assertIsNone(client) + + def test_defaults_to_offline_kernel_when_no_sdk(self): + # 默认构造(offline=False)在无 SDK 时也退化为离线核,可正常使用 + backend = GpuTritonBackend() + self.assertIsInstance(backend._kernel, _OfflineKernel) + backend.load_model("llm-70b-onnx") + self.assertTrue(backend.health_check().healthy) + + +if __name__ == "__main__": + unittest.main()