From ec1f5ee97eed4423cb8f231bd55390af788331f0 Mon Sep 17 00:00:00 2001 From: jschen069 <3563624058@qq.com> Date: Wed, 12 Aug 2026 20:51:48 +0800 Subject: [PATCH 1/4] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E5=88=9D=E7=89=88?= =?UTF-8?q?=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../stable_perf_metric_calculator.py | 115 +++++++----------- .../test_stable_perf_metric_calculator.py | 11 +- 2 files changed, 50 insertions(+), 76 deletions(-) diff --git a/ais_bench/benchmark/calculators/stable_perf_metric_calculator.py b/ais_bench/benchmark/calculators/stable_perf_metric_calculator.py index c8cc1ce0..8597f320 100644 --- a/ais_bench/benchmark/calculators/stable_perf_metric_calculator.py +++ b/ais_bench/benchmark/calculators/stable_perf_metric_calculator.py @@ -10,18 +10,17 @@ from ais_bench.benchmark.utils.logging.error_codes import CALC_CODES from ais_bench.benchmark.utils.logging.exceptions import AISBenchDataContentError -WAVE_OFFSET = 0.02 -INTERVAL_OFFSET = 0.001 @PERF_METRIC_CALCULATORS.register_module() class StablePerfMetricCalculator(BasePerfMetricCalculator): """ Performance metric calculator for stable stage analysis. - This calculator focuses on analyzing the stable phase of benchmark execution, - where the system operates at maximum concurrency with minimal fluctuations. - It identifies and analyzes the stable period to provide more accurate - performance metrics. + This calculator identifies the stable phase as the interval from the first + time concurrency reaches max_concurrency to the last time it is at that level. + All requests starting within this interval are included regardless of + concurrency fluctuations, providing a robust steady-state performance measurement + without requiring manual parameter tuning. Args: stats_list (list, optional): List of statistics to calculate @@ -63,6 +62,11 @@ def _get_requests_id(self, perf_details: dict) -> list: """ Identify requests that belong to the stable stage. + The stable stage is defined as the interval from the first moment + concurrency reaches max_concurrency to the last moment concurrency + is at max_concurrency. All requests with start_time within this + closed interval are included, regardless of concurrency fluctuations. + Args: perf_details (dict): Performance details dictionary @@ -70,11 +74,12 @@ def _get_requests_id(self, perf_details: dict) -> list: list: List of request IDs in the stable stage Raises: - RuntimeError: If no stable stage can be identified + AISBenchDataContentError: If no stable stage can be identified """ # Calculate the minimum start time as the baseline min_start_time = min(perf_details["start_time"]) - time_point_concurrency = [0] * 2 * len(perf_details["id"]) + + # Build sorted list of start/end events for all requests request_time_sections = [] for id in range(len(perf_details["id"])): request_time_sections.append( @@ -92,84 +97,52 @@ def _get_requests_id(self, perf_details: dict) -> list: } ) sorted_time_sections = sorted(request_time_sections, key=lambda x: x["time"]) - id_lists = [] + self.logger.info("Starting stable stage calculation...") - requested = 0 + first_max_time = None + last_max_time = None + concurrency = 0 + progress_bar = tqdm( total=len(sorted_time_sections), desc="Calculating stable stage", unit=" req", ) - for i, section in enumerate(sorted_time_sections): + + for section in sorted_time_sections: if section["attr"] == "start": - time_point_concurrency[i] = time_point_concurrency[i - 1] + 1 - requested += 1 + concurrency += 1 else: - time_point_concurrency[i] = time_point_concurrency[i - 1] - 1 - if ( - section["attr"] == "start" - and time_point_concurrency[i] == self.max_concurrency - ): - id_lists.append(section["id"]) - if len(id_lists) == 2: - self.stage_section[0] = section["time"] # total start time - elif ( - section["attr"] == "start" - and time_point_concurrency[i] - >= int(self.max_concurrency * (1 - WAVE_OFFSET)) - and len(id_lists) > 2 - ): - id_lists.append(section["id"]) - elif requested == len(perf_details["id"]) and section["attr"] == "end": - self.stage_section[1] = section["time"] - progress_bar.update(len(sorted_time_sections) - progress_bar.n) - break - elif ( - len(id_lists) > 1 - and section["attr"] == "end" - and time_point_concurrency[i] - < int(self.max_concurrency * (1 - WAVE_OFFSET)) - ): - # Check if there's a start event within INTERVAL_OFFSET that recovers concurrency - # Skip consecutive end events and look ahead for start events - should_exit_stable = True - current_time = section["time"] - current_concurrency = time_point_concurrency[i] - - # Look ahead to find the next start event within INTERVAL_OFFSET - for j in range(i + 1, len(sorted_time_sections)): - next_section = sorted_time_sections[j] - time_interval = next_section["time"] - current_time - - # If time interval exceeds INTERVAL_OFFSET, exit stable stage - if time_interval > INTERVAL_OFFSET: - break - - # Update concurrency based on the event type - if next_section["attr"] == "end": - current_concurrency -= 1 - else: # start event - current_concurrency += 1 - # If concurrency recovers to threshold, don't exit stable stage - if current_concurrency >= int( - self.max_concurrency * (1 - WAVE_OFFSET) - ): - should_exit_stable = False - break - - if should_exit_stable: - self.stage_section[1] = section["time"] - progress_bar.update(len(sorted_time_sections) - progress_bar.n) - break + concurrency -= 1 + + if concurrency == self.max_concurrency: + if first_max_time is None: + first_max_time = section["time"] + last_max_time = section["time"] + progress_bar.update(1) progress_bar.close() - if len(id_lists) > 0: - id_lists.pop(0) # ignore first request that reached max concurrency + + if first_max_time is None: + raise AISBenchDataContentError( + CALC_CODES.CAN_NOT_FIND_STABLE_STAGE, + "Can not find a stable stage from performance results! Please check the conccurency plot.", + ) + + self.stage_section = [first_max_time, last_max_time] + + # Collect all requests whose start_time falls within the stable interval + id_lists = [ + id for id in range(len(perf_details["id"])) + if first_max_time <= perf_details["start_time"][id] <= last_max_time + ] + if len(id_lists) == 0: raise AISBenchDataContentError( CALC_CODES.CAN_NOT_FIND_STABLE_STAGE, "Can not find a stable stage from performance results! Please check the conccurency plot.", ) + # Convert to relative time based on minimum start time relative_start_time = self.stage_section[0] - min_start_time relative_end_time = self.stage_section[1] - min_start_time diff --git a/tests/UT/calculators/test_stable_perf_metric_calculator.py b/tests/UT/calculators/test_stable_perf_metric_calculator.py index cf8c3542..426875d2 100644 --- a/tests/UT/calculators/test_stable_perf_metric_calculator.py +++ b/tests/UT/calculators/test_stable_perf_metric_calculator.py @@ -300,11 +300,12 @@ def test_convert_result_integration(self): def test_edge_case_empty_stable_stage(self): # 测试边缘情况 - 稳定阶段只有少量请求 + # 新算法：单个请求达到 max_concurrency=1 时也是有效的稳定阶段 calculator = StablePerfMetricCalculator() calculator.max_concurrency = 1 calculator.logger = self.mock_logger calculator.stage_section = [0, 0] # 初始化必要的属性 - + # 准备测试数据 - 只有一个请求 perf_details = { "id": [0], @@ -312,10 +313,10 @@ def test_edge_case_empty_stable_stage(self): "end_time": [2.0], "success": [True] } - - # 验证抛出异常 - with self.assertRaises(AISBenchDataContentError): - calculator._get_requests_id(perf_details) + + result = calculator._get_requests_id(perf_details) + self.assertEqual(result, [0]) + self.assertEqual(calculator.stage_section, [1.0, 1.0]) if __name__ == "__main__": From 2c0635d8cf2e352ab05ced18ab98317034f2f3e2 Mon Sep 17 00:00:00 2001 From: jschen069 <3563624058@qq.com> Date: Wed, 19 Aug 2026 17:00:56 +0800 Subject: [PATCH 2/4] docs: design vllm itl timing alignment --- .../2026-08-19-vllm-itl-alignment-design.md | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md diff --git a/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md b/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md new file mode 100644 index 00000000..6b5ff933 --- /dev/null +++ b/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md @@ -0,0 +1,76 @@ +# Align aisbench ITL Timing with vLLM Completions + +## Goal + +Make aisbench's VLLM streaming timing semantics match `vllm bench serve --backend vllm`: only nonempty `choices` responses count as output timing events, while usage-only and `[DONE]` messages do not affect ITL, TPOT, or E2EL. + +## Scope + +- Apply token-event filtering to `VLLMCustomAPI` and `VLLMCustomAPIChat`. +- Preserve existing streaming timing behavior for TGI, Triton, MindIE, and VITA adapters. +- Continue parsing usage-only responses so prompt and completion token counts remain available. +- Remove the current uncommitted `[ITL-DIAG]` instrumentation from `base_api.py`. +- Add regression tests before changing production behavior. +- Do not change ITL aggregation, percentiles, CSV columns, or the meaning of the reported `N` field in this change. + +## Design + +### Protocol-aware timing predicate + +Add an overridable method to `BaseAPIModel`: + +```python +def should_record_stream_time_point(self, data: dict) -> bool: + return True +``` + +The base implementation preserves all current non-vLLM adapters. Both vLLM adapters override it with: + +```python +def should_record_stream_time_point(self, data: dict) -> bool: + return bool(data.get("choices")) +``` + +### Stream data flow + +For each nonempty, non-comment, non-`[DONE]` stream message: + +1. Decode and parse the JSON payload. +2. Call `should_record_stream_time_point(data)`. +3. Record a time point only when the predicate returns `True`. +4. Always call `parse_stream_response(data, output)` so usage-only responses still populate token counts. + +The request-start timestamp remains unchanged. For 1024 choice-bearing output chunks, the final timing array contains the start plus 1024 output events, producing 1023 ITLs. + +## Compatibility + +- This matches the vLLM v0.26 Completions client, which timestamps inside its nonempty-`choices` branch. +- The chat adapter uses the same filtering so the currently configured aisbench Chat endpoint can be compared using the same timing semantics. +- Other adapters inherit the default predicate and retain unconditional JSON-chunk timing. +- JSON parse failures retain the existing error path and do not create timing points. + +## Testing + +Add focused async tests that feed a stream containing two choice-bearing chunks, one usage-only chunk, and `[DONE]`. + +Required assertions: + +- The VLLM adapter records three time points: request start plus two choice events. +- The usage-only chunk does not add a timing point. +- Completion and prompt token counts are populated from usage. +- The derived ITL count is one. +- A representative non-vLLM adapter continues recording each parsed data chunk. + +The tests must fail against the current implementation before production code changes are made, then pass after the minimal implementation. + +## Error Handling + +No new error category is introduced. Existing HTTP, JSON parsing, retry, and empty-response behavior remains unchanged. + +## Acceptance Criteria + +- No sub-millisecond usage-tail ITL is produced by a VLLM usage-only response. +- VLLM usage token accounting remains correct. +- Non-vLLM streaming timing behavior does not change. +- Targeted tests and the relevant existing API-model test suite pass. +- `base_api.py` parses successfully after the temporary diagnostics are removed. From 1dae9af6901e86c927e0d025f252af4a888d376d Mon Sep 17 00:00:00 2001 From: jschen069 <3563624058@qq.com> Date: Wed, 19 Aug 2026 17:04:07 +0800 Subject: [PATCH 3/4] docs: translate vllm itl alignment design --- .../2026-08-19-vllm-itl-alignment-design.md | 86 +++++++++---------- 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md b/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md index 6b5ff933..8b84126a 100644 --- a/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md +++ b/docs/superpowers/specs/2026-08-19-vllm-itl-alignment-design.md @@ -1,76 +1,76 @@ -# Align aisbench ITL Timing with vLLM Completions +# aisbench ITL 计时与 vLLM Completions 对齐设计 -## Goal +## 目标 -Make aisbench's VLLM streaming timing semantics match `vllm bench serve --backend vllm`: only nonempty `choices` responses count as output timing events, while usage-only and `[DONE]` messages do not affect ITL, TPOT, or E2EL. +使 aisbench 的 VLLM 流式计时语义与 `vllm bench serve --backend vllm` 保持一致：仅将 `choices` 非空的响应视为输出计时事件；usage-only 响应和 `[DONE]` 消息均不影响 ITL、TPOT 或 E2EL。 -## Scope +## 范围 -- Apply token-event filtering to `VLLMCustomAPI` and `VLLMCustomAPIChat`. -- Preserve existing streaming timing behavior for TGI, Triton, MindIE, and VITA adapters. -- Continue parsing usage-only responses so prompt and completion token counts remain available. -- Remove the current uncommitted `[ITL-DIAG]` instrumentation from `base_api.py`. -- Add regression tests before changing production behavior. -- Do not change ITL aggregation, percentiles, CSV columns, or the meaning of the reported `N` field in this change. +- 对 `VLLMCustomAPI` 和 `VLLMCustomAPIChat` 应用 token 事件过滤。 +- 保持 TGI、Triton、MindIE 和 VITA 适配器现有的流式计时行为不变。 +- 继续解析 usage-only 响应，确保 prompt token 和 completion token 数量仍能正确获取。 +- 删除 `base_api.py` 中当前未提交的 `[ITL-DIAG]` 临时诊断代码。 +- 修改生产代码前先添加回归测试。 +- 本次变更不修改 ITL 聚合方式、百分位计算、CSV 列或报告中 `N` 字段的含义。 -## Design +## 设计 -### Protocol-aware timing predicate +### 协议感知的计时判断方法 -Add an overridable method to `BaseAPIModel`: +在 `BaseAPIModel` 中增加一个可覆盖的方法： ```python def should_record_stream_time_point(self, data: dict) -> bool: return True ``` -The base implementation preserves all current non-vLLM adapters. Both vLLM adapters override it with: +基类默认返回 `True`，从而保持所有非 VLLM 适配器的现有行为。两个 VLLM 适配器均覆盖该方法： ```python def should_record_stream_time_point(self, data: dict) -> bool: return bool(data.get("choices")) ``` -### Stream data flow +### 流式数据处理流程 -For each nonempty, non-comment, non-`[DONE]` stream message: +对于每一条非空、非注释且不等于 `[DONE]` 的流式消息： -1. Decode and parse the JSON payload. -2. Call `should_record_stream_time_point(data)`. -3. Record a time point only when the predicate returns `True`. -4. Always call `parse_stream_response(data, output)` so usage-only responses still populate token counts. +1. 解码并解析 JSON 数据。 +2. 调用 `should_record_stream_time_point(data)`。 +3. 仅当该方法返回 `True` 时记录时间点。 +4. 始终调用 `parse_stream_response(data, output)`，确保 usage-only 响应仍可写入 token 数量。 -The request-start timestamp remains unchanged. For 1024 choice-bearing output chunks, the final timing array contains the start plus 1024 output events, producing 1023 ITLs. +请求开始时间点的记录方式保持不变。如果响应中包含 1024 个携带 `choices` 的输出 chunk，最终计时数组将包含一个请求开始时间点和 1024 个输出事件，因此生成 1023 个 ITL。 -## Compatibility +## 兼容性 -- This matches the vLLM v0.26 Completions client, which timestamps inside its nonempty-`choices` branch. -- The chat adapter uses the same filtering so the currently configured aisbench Chat endpoint can be compared using the same timing semantics. -- Other adapters inherit the default predicate and retain unconditional JSON-chunk timing. -- JSON parse failures retain the existing error path and do not create timing points. +- 该设计与 vLLM v0.26 Completions 客户端一致：仅在 `choices` 非空分支内记录时间点。 +- Chat 适配器采用相同过滤规则，使当前配置的 aisbench Chat 接口也可以按照同一套计时语义进行比较。 +- 其他适配器继承基类默认判断，继续为每个已解析的 JSON chunk 记录时间点。 +- JSON 解析失败时继续沿用现有错误处理路径，并且不记录时间点。 -## Testing +## 测试设计 -Add focused async tests that feed a stream containing two choice-bearing chunks, one usage-only chunk, and `[DONE]`. +新增针对性的异步测试，模拟包含以下内容的响应流：两个携带 `choices` 的 chunk、一个 usage-only chunk，以及 `[DONE]`。 -Required assertions: +必须验证： -- The VLLM adapter records three time points: request start plus two choice events. -- The usage-only chunk does not add a timing point. -- Completion and prompt token counts are populated from usage. -- The derived ITL count is one. -- A representative non-vLLM adapter continues recording each parsed data chunk. +- VLLM 适配器共记录三个时间点：一个请求开始时间点和两个 choice 事件时间点。 +- usage-only chunk 不增加时间点。 +- prompt token 和 completion token 数量可从 usage 响应中正确获取。 +- 最终只产生一个 ITL。 +- 选取一个有代表性的非 VLLM 适配器，验证其仍为每个已解析的数据 chunk 记录时间点。 -The tests must fail against the current implementation before production code changes are made, then pass after the minimal implementation. +在修改生产代码之前，测试必须在当前实现上因 usage-only chunk 被错误计时而失败。完成最小实现后，测试必须通过。 -## Error Handling +## 错误处理 -No new error category is introduced. Existing HTTP, JSON parsing, retry, and empty-response behavior remains unchanged. +本次变更不引入新的错误类型。现有 HTTP 错误、JSON 解析错误、重试和空响应处理逻辑均保持不变。 -## Acceptance Criteria +## 验收标准 -- No sub-millisecond usage-tail ITL is produced by a VLLM usage-only response. -- VLLM usage token accounting remains correct. -- Non-vLLM streaming timing behavior does not change. -- Targeted tests and the relevant existing API-model test suite pass. -- `base_api.py` parses successfully after the temporary diagnostics are removed. +- VLLM usage-only 响应不再产生亚毫秒级的尾部 ITL。 +- VLLM usage token 统计保持正确。 +- 非 VLLM 流式计时行为不发生变化。 +- 针对性测试及相关现有 API 模型测试全部通过。 +- 删除临时诊断代码后，`base_api.py` 能够通过 Python 语法解析。 From cac4aee15924ee5f04861af0448057f565d2cfe3 Mon Sep 17 00:00:00 2001 From: jschen069 <3563624058@qq.com> Date: Wed, 19 Aug 2026 18:21:17 +0800 Subject: [PATCH 4/4] =?UTF-8?q?=E6=B7=BB=E5=8A=A0itl=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../benchmark/models/api_models/base_api.py | 7 +++- .../models/api_models/vllm_custom_api.py | 3 ++ .../models/api_models/vllm_custom_api_chat.py | 3 ++ tests/UT/models/api_models/test_base_api.py | 29 ++++++++++++++ .../models/api_models/test_vllm_custom_api.py | 38 +++++++++++++++++-- .../api_models/test_vllm_custom_api_chat.py | 38 +++++++++++++++++-- 6 files changed, 111 insertions(+), 7 deletions(-) diff --git a/ais_bench/benchmark/models/api_models/base_api.py b/ais_bench/benchmark/models/api_models/base_api.py index db5eef7f..29145fde 100644 --- a/ais_bench/benchmark/models/api_models/base_api.py +++ b/ais_bench/benchmark/models/api_models/base_api.py @@ -226,6 +226,10 @@ async def parse_stream_response(self, data, output): f"{self.__class__.__name__} should be implemented if stream is True", ) + def should_record_stream_time_point(self, data: dict) -> bool: + """Return whether a parsed stream response represents an output event.""" + return True + async def generate( self, input_data: PromptType, @@ -295,7 +299,6 @@ async def stream_infer(self, request_body: dict, output: Output): chunk = chunk.removeprefix("data:").strip() if chunk == "[DONE]": break - await output.record_time_point() try: data = json.loads(chunk) except json.JSONDecodeError as e: @@ -305,6 +308,8 @@ async def stream_infer(self, request_body: dict, output: Output): MODEL_CODES.PARSE_TEXT_RSP_INVALID_FORMAT, f"Unexpected response format. Please check 'error_info' in ***_failed.jsonl for more information.", ) + if self.should_record_stream_time_point(data): + await output.record_time_point() await self.parse_stream_response(data, output) output.success = True else: diff --git a/ais_bench/benchmark/models/api_models/vllm_custom_api.py b/ais_bench/benchmark/models/api_models/vllm_custom_api.py index 3e4f5083..b1dec711 100644 --- a/ais_bench/benchmark/models/api_models/vllm_custom_api.py +++ b/ais_bench/benchmark/models/api_models/vllm_custom_api.py @@ -148,6 +148,9 @@ async def parse_text_response(self, api_response: dict, output: Output): await self._parse_usage(api_response, output) self.logger.debug(f"Output content: {output.content}") + def should_record_stream_time_point(self, data: dict) -> bool: + return bool(data.get("choices")) + async def parse_stream_response(self, api_response: dict, output: Output): generated_text = "" if len(api_response.get("choices", [])) > 0: diff --git a/ais_bench/benchmark/models/api_models/vllm_custom_api_chat.py b/ais_bench/benchmark/models/api_models/vllm_custom_api_chat.py index 03413794..cd9f6de2 100644 --- a/ais_bench/benchmark/models/api_models/vllm_custom_api_chat.py +++ b/ais_bench/benchmark/models/api_models/vllm_custom_api_chat.py @@ -186,6 +186,9 @@ async def _parse_usage(self, json_content: dict, output: Output): output.input_tokens = json_content["usage"].get("prompt_tokens", 0) output.output_tokens = json_content["usage"].get("completion_tokens", 0) + def should_record_stream_time_point(self, data: dict) -> bool: + return bool(data.get("choices")) + async def parse_stream_response(self, json_content, output): for item in json_content.get("choices", []): if item["delta"].get("content"): diff --git a/tests/UT/models/api_models/test_base_api.py b/tests/UT/models/api_models/test_base_api.py index 83b05d0d..cb80403a 100644 --- a/tests/UT/models/api_models/test_base_api.py +++ b/tests/UT/models/api_models/test_base_api.py @@ -284,6 +284,35 @@ async def run(): asyncio.run(run()) + def test_stream_infer_default_records_each_json_chunk(self): + async def mock_content(): + yield b'data: {"chunk": "hello"}\n\n' + yield b'data: {"chunk": " world"}\n\n' + yield b"data: [DONE]\n\n" + + class MockPostContext: + async def __aenter__(self): + response = mock.MagicMock() + response.status = 200 + response.content = mock_content() + return response + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + model = self.model_class(**self.default_kwargs) + model.stream = True + model.session = mock.MagicMock() + model.session.post.return_value = MockPostContext() + output = Output(perf_mode=True) + output.text = "" + + asyncio.run(model.stream_infer({"prompt": "test"}, output)) + + self.assertTrue(output.success) + self.assertEqual(output.text, "hello world") + self.assertEqual(len(output.time_points), 3) + @mock.patch("aiohttp.ClientSession.post") async def test_stream_infer_json_decode_error(self, mock_post): async def mock_content(): diff --git a/tests/UT/models/api_models/test_vllm_custom_api.py b/tests/UT/models/api_models/test_vllm_custom_api.py index 9f0e6238..f24a7f7e 100644 --- a/tests/UT/models/api_models/test_vllm_custom_api.py +++ b/tests/UT/models/api_models/test_vllm_custom_api.py @@ -360,6 +360,40 @@ def test_calc_ppl_with_none(self): expected_ppl = -(-0.5 - 0.7) / 2 self.assertAlmostEqual(ppl, expected_ppl, places=5) + def test_stream_infer_excludes_usage_only_chunk_from_timing(self): + async def mock_content(): + yield b'data: {"choices": [{"text": "Hello "}]}\n\n' + yield b'data: {"choices": [{"text": "world"}]}\n\n' + yield ( + b'data: {"choices": [], "usage": ' + b'{"prompt_tokens": 7, "completion_tokens": 2}}\n\n' + ) + yield b"data: [DONE]\n\n" + + class MockPostContext: + async def __aenter__(self): + response = MagicMock() + response.status = 200 + response.content = mock_content() + return response + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + kwargs = self.default_kwargs.copy() + kwargs["stream"] = True + model = VLLMCustomAPI(**kwargs) + model.session = MagicMock() + model.session.post.return_value = MockPostContext() + output = Output(perf_mode=True) + + asyncio.run(model.stream_infer({"prompt": "test"}, output)) + + self.assertEqual(output.content, "Hello world") + self.assertEqual(output.input_tokens, 7) + self.assertEqual(output.output_tokens, 2) + self.assertEqual(len(output.time_points), 3) + class TestVLLMCustomAPILora(unittest.TestCase): """针对 Multi-LoRA 兼容性扩展的 UT（completions API）。""" @@ -485,7 +519,5 @@ def test_request_body_no_lora_config_keeps_base_model(self): self.assertEqual(body["model"], "base-model-name") self.assertNotIn("adapter_id", body) - - if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/UT/models/api_models/test_vllm_custom_api_chat.py b/tests/UT/models/api_models/test_vllm_custom_api_chat.py index 93c3115a..b0183251 100644 --- a/tests/UT/models/api_models/test_vllm_custom_api_chat.py +++ b/tests/UT/models/api_models/test_vllm_custom_api_chat.py @@ -350,6 +350,40 @@ def test_calc_ppl_with_none(self): expected_ppl = -(-0.5 - 0.7) / 2 self.assertAlmostEqual(ppl, expected_ppl, places=5) + def test_stream_infer_excludes_usage_only_chunk_from_timing(self): + async def mock_content(): + yield b'data: {"choices": [{"delta": {"content": "Hello "}}]}\n\n' + yield b'data: {"choices": [{"delta": {"content": "world"}}]}\n\n' + yield ( + b'data: {"choices": [], "usage": ' + b'{"prompt_tokens": 7, "completion_tokens": 2}}\n\n' + ) + yield b"data: [DONE]\n\n" + + class MockPostContext: + async def __aenter__(self): + response = MagicMock() + response.status = 200 + response.content = mock_content() + return response + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + kwargs = self.default_kwargs.copy() + kwargs["stream"] = True + model = VLLMCustomAPIChat(**kwargs) + model.session = MagicMock() + model.session.post.return_value = MockPostContext() + output = RequestOutput(perf_mode=True) + + asyncio.run(model.stream_infer({"messages": []}, output)) + + self.assertEqual(output.content, "Hello world") + self.assertEqual(output.input_tokens, 7) + self.assertEqual(output.output_tokens, 2) + self.assertEqual(len(output.time_points), 3) + class TestVLLMCustomAPIChatLora(unittest.TestCase): """针对 Multi-LoRA 兼容性扩展的 UT（base model 内嵌，无新增子类）。""" @@ -533,7 +567,5 @@ def test_request_body_lora_hit_with_promptlist(self): self.assertEqual(body["model"], "LoraB") self.assertEqual(len(body["messages"]), 1) - - if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main()