| """ |
| 推理模块单元测试 |
| |
| 测试推理模块的所有功能，包括： |
| - llama-cpp 后端 |
| - vLLM 后端 |
| - SGLang 后端 |
| - 统一推理接口 |
| - 性能监控 |
| - 数据结构 |
| |
| 使用 mock 避免实际模型加载和推理过程。 |
| """ |
|
|
| import os |
| import sys |
| import pytest |
| from pathlib import Path |
| from unittest.mock import MagicMock, patch, Mock |
| import time |
|
|
| |
| sys.path.insert(0, str(Path(__file__).parent.parent)) |
|
|
| from hos_optimizer.inference import ( |
| InferenceRequest, |
| InferenceResult, |
| PerformanceStats, |
| PerformanceMonitor, |
| InferenceBackend, |
| LlamaCppBackend, |
| VLLMBackend, |
| SGLangBackend, |
| UnifiedInferenceEngine, |
| get_gpu_memory_usage_mb, |
| get_total_gpu_memory_mb, |
| detect_best_backend, |
| ) |
|
|
|
|
| class TestDataStructures: |
| """数据结构测试""" |
|
|
| def test_inference_request_creation(self): |
| """测试推理请求数据结构创建""" |
| request = InferenceRequest( |
| prompt="测试提示", |
| max_tokens=100, |
| temperature=0.8, |
| top_p=0.95, |
| top_k=50, |
| stop=["END"], |
| json_schema={"type": "object"}, |
| extra={"custom": "param"} |
| ) |
| |
| assert request.prompt == "测试提示" |
| assert request.max_tokens == 100 |
| assert request.temperature == 0.8 |
| assert request.top_p == 0.95 |
| assert request.top_k == 50 |
| assert request.stop == ["END"] |
| assert request.json_schema == {"type": "object"} |
| assert request.extra == {"custom": "param"} |
|
|
| def test_inference_request_defaults(self): |
| """测试推理请求默认值""" |
| request = InferenceRequest(prompt="测试") |
| |
| assert request.max_tokens == 256 |
| assert request.temperature == 0.7 |
| assert request.top_p == 0.9 |
| assert request.top_k == 50 |
| assert request.stop is None |
| assert request.json_schema is None |
| assert request.extra is None |
|
|
| def test_inference_result_creation(self): |
| """测试推理结果数据结构创建""" |
| result = InferenceResult( |
| text="生成的文本", |
| token_ids=[1, 2, 3], |
| prompt="提示", |
| latency_ms=150.5, |
| tokens_per_second=33.3, |
| prompt_tokens=10, |
| completion_tokens=50, |
| metadata={"key": "value"} |
| ) |
| |
| assert result.text == "生成的文本" |
| assert result.token_ids == [1, 2, 3] |
| assert result.prompt == "提示" |
| assert result.latency_ms == 150.5 |
| assert result.tokens_per_second == 33.3 |
| assert result.prompt_tokens == 10 |
| assert result.completion_tokens == 50 |
| assert result.metadata == {"key": "value"} |
|
|
| def test_performance_stats_properties(self): |
| """测试性能统计属性计算""" |
| stats = PerformanceStats( |
| total_requests=10, |
| total_prompt_tokens=100, |
| total_completion_tokens=500, |
| total_latency_ms=1000.0, |
| peak_vram_mb=2048.0, |
| wall_time_s=5.0 |
| ) |
| |
| assert stats.avg_latency_ms == 100.0 |
| assert stats.throughput_tokens_per_s == 100.0 |
| assert stats.requests_per_s == 2.0 |
|
|
| def test_performance_stats_zero_requests(self): |
| """测试零请求时的性能统计""" |
| stats = PerformanceStats() |
| |
| assert stats.avg_latency_ms == 0.0 |
| assert stats.throughput_tokens_per_s == 0.0 |
| assert stats.requests_per_s == 0.0 |
|
|
| def test_performance_stats_summary(self): |
| """测试性能统计摘要生成""" |
| stats = PerformanceStats( |
| total_requests=5, |
| total_prompt_tokens=50, |
| total_completion_tokens=250, |
| total_latency_ms=500.0, |
| peak_vram_mb=1024.0, |
| wall_time_s=2.5 |
| ) |
| |
| summary = stats.summary() |
| |
| assert "性能统计摘要" in summary |
| assert "5" in summary |
| assert "50" in summary |
| assert "250" in summary |
| assert "1024.0" in summary |
|
|
|
|
| class TestPerformanceMonitor: |
| """性能监控器测试""" |
|
|
| def test_monitor_start_end(self): |
| """测试监控器启动和结束""" |
| monitor = PerformanceMonitor() |
| |
| monitor.start() |
| time.sleep(0.1) |
| monitor.end() |
| |
| assert monitor.stats.wall_time_s > 0 |
| assert monitor.stats.wall_time_s >= 0.1 |
|
|
| def test_monitor_record_request(self): |
| """测试记录请求""" |
| monitor = PerformanceMonitor() |
| |
| monitor.record_request(prompt_tokens=10, completion_tokens=50, latency_ms=100) |
| monitor.record_request(prompt_tokens=20, completion_tokens=60, latency_ms=150) |
| |
| assert monitor.stats.total_requests == 2 |
| assert monitor.stats.total_prompt_tokens == 30 |
| assert monitor.stats.total_completion_tokens == 110 |
| assert monitor.stats.total_latency_ms == 250 |
|
|
| def test_monitor_update_peak_vram(self): |
| """测试更新峰值显存""" |
| monitor = PerformanceMonitor() |
| |
| with patch("hos_optimizer.inference.get_gpu_memory_usage_mb") as mock_vram: |
| mock_vram.return_value = 2048.0 |
| monitor.update_peak_vram() |
| |
| assert monitor.stats.peak_vram_mb == 2048.0 |
| |
| |
| mock_vram.return_value = 3072.0 |
| monitor.update_peak_vram() |
| assert monitor.stats.peak_vram_mb == 3072.0 |
| |
| |
| mock_vram.return_value = 1024.0 |
| monitor.update_peak_vram() |
| assert monitor.stats.peak_vram_mb == 3072.0 |
|
|
| def test_monitor_reset(self): |
| """测试重置监控器""" |
| monitor = PerformanceMonitor() |
| monitor.start() |
| monitor.record_request(10, 50, 100) |
| monitor.end() |
| |
| monitor.reset() |
| |
| assert monitor.stats.total_requests == 0 |
| assert monitor.stats.total_latency_ms == 0.0 |
| assert monitor._start_time is None |
| assert monitor._end_time is None |
|
|
|
|
| class TestGPUUtilities: |
| """GPU 工具函数测试""" |
|
|
| def test_get_gpu_memory_usage_no_cuda(self): |
| """测试无 CUDA 时的显存使用""" |
| with patch("torch.cuda.is_available", return_value=False): |
| usage = get_gpu_memory_usage_mb() |
| assert usage == 0.0 |
|
|
| def test_get_gpu_memory_usage_with_cuda(self): |
| """测试有 CUDA 时的显存使用""" |
| with patch("torch.cuda.is_available", return_value=True), \ |
| patch("torch.cuda.device_count", return_value=1), \ |
| patch("torch.cuda.max_memory_allocated", return_value=2 * 1024 ** 3): |
| |
| usage = get_gpu_memory_usage_mb() |
| assert usage > 0 |
|
|
| def test_get_total_gpu_memory_no_cuda(self): |
| """测试无 CUDA 时的总显存""" |
| with patch("torch.cuda.is_available", return_value=False): |
| total = get_total_gpu_memory_mb() |
| assert total == 0.0 |
|
|
| def test_get_total_gpu_memory_with_cuda(self): |
| """测试有 CUDA 时的总显存""" |
| mock_props = MagicMock() |
| mock_props.total_mem = 8 * 1024 ** 3 |
| |
| with patch("torch.cuda.is_available", return_value=True), \ |
| patch("torch.cuda.get_device_properties", return_value=mock_props): |
| |
| total = get_total_gpu_memory_mb() |
| assert total > 0 |
|
|
|
|
| class TestBackendDetection: |
| """后端检测测试""" |
|
|
| def test_detect_best_backend_vllm(self): |
| """测试检测到 vLLM 作为最优后端""" |
| with patch("hos_optimizer.inference.get_total_gpu_memory_mb", return_value=8192), \ |
| patch("hos_optimizer.inference._check_import") as mock_check: |
| |
| mock_check.side_effect = lambda name: name == "vllm" |
| |
| backend = detect_best_backend() |
| assert backend == "vllm" |
|
|
| def test_detect_best_backend_sglang(self): |
| """测试检测到 SGLang 作为最优后端""" |
| with patch("hos_optimizer.inference.get_total_gpu_memory_mb", return_value=4096), \ |
| patch("hos_optimizer.inference._check_import") as mock_check: |
| |
| mock_check.side_effect = lambda name: name == "sglang" |
| |
| backend = detect_best_backend() |
| assert backend == "sglang" |
|
|
| def test_detect_best_backend_llama_cpp(self): |
| """测试检测到 llama-cpp 作为最优后端""" |
| with patch("hos_optimizer.inference.get_total_gpu_memory_mb", return_value=0), \ |
| patch("hos_optimizer.inference._check_import") as mock_check: |
| |
| mock_check.side_effect = lambda name: name == "llama_cpp" |
| |
| backend = detect_best_backend() |
| assert backend == "llama_cpp" |
|
|
| def test_detect_best_backend_no_backend_available(self): |
| """测试没有可用后端时抛出异常""" |
| with patch("hos_optimizer.inference.get_total_gpu_memory_mb", return_value=0), \ |
| patch("hos_optimizer.inference._check_import", return_value=False): |
| |
| with pytest.raises(RuntimeError) as exc_info: |
| detect_best_backend() |
| |
| assert "没有可用的推理后端" in str(exc_info.value) |
|
|
|
|
| class TestLlamaCppBackend: |
| """llama-cpp 后端测试""" |
|
|
| def test_llama_cpp_initialization(self): |
| """测试 llama-cpp 后端初始化""" |
| backend = LlamaCppBackend( |
| model_path="/path/to/model.gguf", |
| n_gpu_layers=32, |
| n_ctx=2048, |
| n_threads=4 |
| ) |
| |
| assert backend.model_path == "/path/to/model.gguf" |
| assert backend.n_gpu_layers == 32 |
| assert backend.n_ctx == 2048 |
| assert backend.n_threads == 4 |
| assert backend._loaded is False |
|
|
| def test_llama_cpp_load_success(self): |
| """测试 llama-cpp 模型加载成功""" |
| with patch("llama_cpp.Llama") as mock_llama: |
| mock_model = MagicMock() |
| mock_llama.return_value = mock_model |
| |
| backend = LlamaCppBackend("/path/to/model.gguf") |
| |
| with patch("hos_optimizer.inference.get_total_gpu_memory_mb", return_value=8192): |
| backend.load() |
| |
| assert backend._loaded is True |
| assert backend._model is not None |
|
|
| def test_llama_cpp_load_import_error(self): |
| """测试 llama-cpp 导入失败""" |
| backend = LlamaCppBackend("/path/to/model.gguf") |
| |
| with patch.dict("sys.modules", {"llama_cpp": None}): |
| with pytest.raises(ImportError) as exc_info: |
| backend.load() |
| |
| assert "llama-cpp-python 未安装" in str(exc_info.value) |
|
|
| def test_llama_cpp_generate_not_loaded(self): |
| """测试未加载时生成抛出异常""" |
| backend = LlamaCppBackend("/path/to/model.gguf") |
| request = InferenceRequest(prompt="测试") |
| |
| with pytest.raises(RuntimeError) as exc_info: |
| backend.generate(request) |
| |
| assert "模型未加载" in str(exc_info.value) |
|
|
| def test_llama_cpp_generate_success(self): |
| """测试 llama-cpp 生成成功""" |
| backend = LlamaCppBackend("/path/to/model.gguf") |
| backend._loaded = True |
| backend._model = MagicMock() |
| |
| |
| backend._model.return_value = { |
| "choices": [{"text": "生成的文本"}], |
| "usage": {"prompt_tokens": 10, "completion_tokens": 20} |
| } |
| |
| request = InferenceRequest(prompt="测试", max_tokens=50) |
| result = backend.generate(request) |
| |
| assert isinstance(result, InferenceResult) |
| assert result.text == "生成的文本" |
| assert result.prompt_tokens == 10 |
| assert result.completion_tokens == 20 |
|
|
| def test_llama_cpp_shutdown(self): |
| """测试 llama-cpp 关闭""" |
| backend = LlamaCppBackend("/path/to/model.gguf") |
| backend._loaded = True |
| backend._model = MagicMock() |
| |
| backend.shutdown() |
| |
| assert backend._loaded is False |
| assert backend._model is None |
|
|
|
|
| class TestVLLMBackend: |
| """vLLM 后端测试""" |
|
|
| def test_vllm_initialization(self): |
| """测试 vLLM 后端初始化""" |
| backend = VLLMBackend( |
| model_path="/path/to/model", |
| gpu_memory_utilization=0.85, |
| max_model_len=2048, |
| max_num_seqs=64 |
| ) |
| |
| assert backend.model_path == "/path/to/model" |
| assert backend.gpu_memory_utilization == 0.85 |
| assert backend.max_model_len == 2048 |
| assert backend.max_num_seqs == 64 |
| assert backend._loaded is False |
|
|
| def test_vllm_load_success(self): |
| """测试 vLLM 模型加载成功""" |
| with patch("vllm.LLM") as mock_llm: |
| mock_engine = MagicMock() |
| mock_llm.return_value = mock_engine |
| |
| backend = VLLMBackend("/path/to/model") |
| backend.load() |
| |
| assert backend._loaded is True |
| assert backend._llm is not None |
|
|
| def test_vllm_load_import_error(self): |
| """测试 vLLM 导入失败""" |
| backend = VLLMBackend("/path/to/model") |
| |
| with patch.dict("sys.modules", {"vllm": None}): |
| with pytest.raises(ImportError) as exc_info: |
| backend.load() |
| |
| assert "vLLM 未安装" in str(exc_info.value) |
|
|
| def test_vllm_generate_not_loaded(self): |
| """测试未加载时生成抛出异常""" |
| backend = VLLMBackend("/path/to/model") |
| request = InferenceRequest(prompt="测试") |
| |
| with pytest.raises(RuntimeError) as exc_info: |
| backend.generate(request) |
| |
| assert "vLLM 引擎未加载" in str(exc_info.value) |
|
|
| def test_vllm_generate_success(self): |
| """测试 vLLM 生成成功""" |
| backend = VLLMBackend("/path/to/model") |
| backend._loaded = True |
| backend._llm = MagicMock() |
| |
| |
| mock_output = MagicMock() |
| mock_output.outputs = [MagicMock(text="生成的文本", token_ids=[1, 2, 3])] |
| mock_output.prompt_token_ids = [10, 20, 30] |
| backend._llm.generate.return_value = [mock_output] |
| |
| request = InferenceRequest(prompt="测试", max_tokens=50) |
| |
| with patch("vllm.SamplingParams"): |
| result = backend.generate(request) |
| |
| assert isinstance(result, InferenceResult) |
| assert result.text == "生成的文本" |
| assert result.token_ids == [1, 2, 3] |
| assert result.prompt_tokens == 3 |
|
|
| def test_vllm_batch_generate(self): |
| """测试 vLLM 批量生成""" |
| backend = VLLMBackend("/path/to/model") |
| backend._loaded = True |
| backend._llm = MagicMock() |
| |
| |
| mock_output1 = MagicMock() |
| mock_output1.outputs = [MagicMock(text="文本1", token_ids=[1, 2])] |
| mock_output1.prompt_token_ids = [10, 20] |
| mock_output1.prompt = "提示1" |
| |
| mock_output2 = MagicMock() |
| mock_output2.outputs = [MagicMock(text="文本2", token_ids=[3, 4])] |
| mock_output2.prompt_token_ids = [30, 40] |
| mock_output2.prompt = "提示2" |
| |
| backend._llm.generate.return_value = [mock_output1, mock_output2] |
| |
| requests = [ |
| InferenceRequest(prompt="提示1"), |
| InferenceRequest(prompt="提示2") |
| ] |
| |
| with patch("vllm.SamplingParams"): |
| results = backend.generate_batch(requests) |
| |
| assert len(results) == 2 |
| assert results[0].text == "文本1" |
| assert results[1].text == "文本2" |
|
|
| def test_vllm_shutdown(self): |
| """测试 vLLM 关闭""" |
| backend = VLLMBackend("/path/to/model") |
| backend._loaded = True |
| backend._llm = MagicMock() |
| |
| with patch("torch.cuda.is_available", return_value=True), \ |
| patch("torch.cuda.empty_cache"): |
| backend.shutdown() |
| |
| assert backend._loaded is False |
| assert backend._llm is None |
|
|
|
|
| class TestSGLangBackend: |
| """SGLang 后端测试""" |
|
|
| def test_sglang_initialization(self): |
| """测试 SGLang 后端初始化""" |
| backend = SGLangBackend( |
| model_path="/path/to/model", |
| mem_fraction_static=0.85, |
| context_length=2048, |
| tp_size=1 |
| ) |
| |
| assert backend.model_path == "/path/to/model" |
| assert backend.mem_fraction_static == 0.85 |
| assert backend.context_length == 2048 |
| assert backend.tp_size == 1 |
| assert backend._loaded is False |
|
|
| def test_sglang_load_success(self): |
| """测试 SGLang 模型加载成功""" |
| with patch("sglang.Runtime") as mock_runtime: |
| mock_rt = MagicMock() |
| mock_runtime.return_value = mock_rt |
| |
| backend = SGLangBackend("/path/to/model") |
| backend.load() |
| |
| assert backend._loaded is True |
| assert backend._runtime is not None |
|
|
| def test_sglang_load_import_error(self): |
| """测试 SGLang 导入失败""" |
| backend = SGLangBackend("/path/to/model") |
| |
| with patch.dict("sys.modules", {"sglang": None}): |
| with pytest.raises(ImportError) as exc_info: |
| backend.load() |
| |
| assert "SGLang 未安装" in str(exc_info.value) |
|
|
| def test_sglang_generate_not_loaded(self): |
| """测试未加载时生成抛出异常""" |
| backend = SGLangBackend("/path/to/model") |
| request = InferenceRequest(prompt="测试") |
| |
| with pytest.raises(RuntimeError) as exc_info: |
| backend.generate(request) |
| |
| assert "SGLang 引擎未加载" in str(exc_info.value) |
|
|
| def test_sglang_generate_success(self): |
| """测试 SGLang 生成成功""" |
| backend = SGLangBackend("/path/to/model") |
| backend._loaded = True |
| backend._runtime = MagicMock() |
| |
| |
| backend._runtime.generate.return_value = [{ |
| "text": "生成的文本", |
| "token_ids": [1, 2, 3] |
| }] |
| |
| request = InferenceRequest(prompt="测试", max_tokens=50) |
| result = backend.generate(request) |
| |
| assert isinstance(result, InferenceResult) |
| assert result.text == "生成的文本" |
| assert result.token_ids == [1, 2, 3] |
|
|
| def test_sglang_generate_with_json_schema(self): |
| """测试 SGLang 带 JSON Schema 约束生成""" |
| backend = SGLangBackend("/path/to/model") |
| backend._loaded = True |
| backend._runtime = MagicMock() |
| |
| backend._runtime.generate.return_value = [{ |
| "text": '{"key": "value"}', |
| "token_ids": [1, 2, 3] |
| }] |
| |
| request = InferenceRequest( |
| prompt="测试", |
| json_schema={"type": "object"} |
| ) |
| result = backend.generate(request) |
| |
| assert result.metadata.get("constrained") is True |
|
|
| def test_sglang_shutdown(self): |
| """测试 SGLang 关闭""" |
| backend = SGLangBackend("/path/to/model") |
| backend._loaded = True |
| backend._runtime = MagicMock() |
| |
| with patch("torch.cuda.is_available", return_value=True), \ |
| patch("torch.cuda.empty_cache"): |
| backend.shutdown() |
| |
| assert backend._loaded is False |
| assert backend._runtime is None |
|
|
|
|
| class TestUnifiedInferenceEngine: |
| """统一推理引擎测试""" |
|
|
| def test_unified_engine_initialization_auto_backend(self): |
| """测试统一引擎自动选择后端""" |
| with patch("hos_optimizer.inference.detect_best_backend", return_value="vllm"), \ |
| patch("hos_optimizer.inference.VLLMBackend") as mock_backend_cls: |
| |
| mock_backend = MagicMock() |
| mock_backend_cls.return_value = mock_backend |
| |
| engine = UnifiedInferenceEngine( |
| model_path="/path/to/model", |
| auto_load=False |
| ) |
| |
| assert engine.backend_name == "vllm" |
| assert engine._backend is not None |
|
|
| def test_unified_engine_initialization_specific_backend(self): |
| """测试统一引擎指定后端""" |
| with patch("hos_optimizer.inference.LlamaCppBackend") as mock_backend_cls: |
| mock_backend = MagicMock() |
| mock_backend_cls.return_value = mock_backend |
| |
| engine = UnifiedInferenceEngine( |
| model_path="/path/to/model.gguf", |
| backend="llama-cpp", |
| auto_load=False |
| ) |
| |
| assert engine.backend_name == "llama_cpp" |
|
|
| def test_unified_engine_invalid_backend(self): |
| """测试统一引擎无效后端""" |
| with pytest.raises(ValueError) as exc_info: |
| UnifiedInferenceEngine( |
| model_path="/path/to/model", |
| backend="invalid_backend", |
| auto_load=False |
| ) |
| |
| assert "不支持的后端" in str(exc_info.value) |
|
|
| def test_unified_engine_generate(self): |
| """测试统一引擎生成""" |
| with patch("hos_optimizer.inference.VLLMBackend") as mock_backend_cls: |
| mock_backend = MagicMock() |
| mock_result = InferenceResult(text="结果", prompt_tokens=10, completion_tokens=20) |
| mock_backend.generate.return_value = mock_result |
| mock_backend_cls.return_value = mock_backend |
| |
| engine = UnifiedInferenceEngine( |
| model_path="/path/to/model", |
| backend="vllm", |
| auto_load=False |
| ) |
| |
| result = engine.generate("测试提示", max_tokens=100) |
| |
| assert result.text == "结果" |
| mock_backend.generate.assert_called_once() |
|
|
| def test_unified_engine_generate_batch(self): |
| """测试统一引擎批量生成""" |
| with patch("hos_optimizer.inference.VLLMBackend") as mock_backend_cls: |
| mock_backend = MagicMock() |
| mock_results = [ |
| InferenceResult(text="结果1", prompt_tokens=10, completion_tokens=20), |
| InferenceResult(text="结果2", prompt_tokens=15, completion_tokens=25) |
| ] |
| mock_backend.generate_batch.return_value = mock_results |
| mock_backend_cls.return_value = mock_backend |
| |
| engine = UnifiedInferenceEngine( |
| model_path="/path/to/model", |
| backend="vllm", |
| auto_load=False |
| ) |
| |
| results = engine.generate_batch(["提示1", "提示2"]) |
| |
| assert len(results) == 2 |
| mock_backend.generate_batch.assert_called_once() |
|
|
| def test_unified_engine_get_stats(self): |
| """测试统一引擎获取性能统计""" |
| with patch("hos_optimizer.inference.VLLMBackend") as mock_backend_cls: |
| mock_backend = MagicMock() |
| mock_stats = PerformanceStats(total_requests=5) |
| mock_backend.get_performance_stats.return_value = mock_stats |
| mock_backend_cls.return_value = mock_backend |
| |
| engine = UnifiedInferenceEngine( |
| model_path="/path/to/model", |
| backend="vllm", |
| auto_load=False |
| ) |
| |
| stats = engine.get_stats() |
| |
| assert stats.total_requests == 5 |
|
|
| def test_unified_engine_shutdown(self): |
| """测试统一引擎关闭""" |
| with patch("hos_optimizer.inference.VLLMBackend") as mock_backend_cls: |
| mock_backend = MagicMock() |
| mock_backend_cls.return_value = mock_backend |
| |
| engine = UnifiedInferenceEngine( |
| model_path="/path/to/model", |
| backend="vllm", |
| auto_load=False |
| ) |
| |
| engine.shutdown() |
| |
| mock_backend.shutdown.assert_called_once() |
|
|
|
|
| if __name__ == "__main__": |
| pytest.main([__file__, "-v"]) |
|
|
