diff --git a/ais_bench/benchmark/models/api_models/base_api.py b/ais_bench/benchmark/models/api_models/base_api.py index 4fc85737..28bf87b5 100644 --- a/ais_bench/benchmark/models/api_models/base_api.py +++ b/ais_bench/benchmark/models/api_models/base_api.py @@ -249,6 +249,21 @@ def _logprobs_enabled(self) -> bool: return val > 0 return False + @staticmethod + def _update_usage_from_details(api_response: dict, output: Output) -> None: + """Fill token counters from a ``details`` block reported by the service. + + TGI / Triton / MindIE expose token usage as ``details.generated_tokens`` + (cumulative for streaming) instead of the OpenAI ``usage`` object. Without + this, every performance report for those backends shows zero throughput. + """ + details = api_response.get("details") + if not isinstance(details, dict): + return + generated_tokens = details.get("generated_tokens") + if isinstance(generated_tokens, int) and not isinstance(generated_tokens, bool): + output.output_tokens = generated_tokens + async def parse_text_response(self, data, output): raise AISBenchNotImplementedError( MODEL_CODES.PARSE_TEXT_RSP_NOT_IMPLEMENTED, diff --git a/ais_bench/benchmark/models/api_models/mindie_stream_api.py b/ais_bench/benchmark/models/api_models/mindie_stream_api.py index a86b838b..6cf9356c 100644 --- a/ais_bench/benchmark/models/api_models/mindie_stream_api.py +++ b/ais_bench/benchmark/models/api_models/mindie_stream_api.py @@ -122,4 +122,6 @@ async def get_request_body( async def parse_stream_response(self, api_response: dict, output: Output): generated_text = api_response.get("generated_text", "") - output.content = generated_text \ No newline at end of file + output.content = generated_text + # The token count is only reported on the final chunk. + self._update_usage_from_details(api_response, output) \ No newline at end of file diff --git a/ais_bench/benchmark/models/api_models/tgi_api.py b/ais_bench/benchmark/models/api_models/tgi_api.py index 63e6bfc6..0506daff 100644 --- a/ais_bench/benchmark/models/api_models/tgi_api.py +++ b/ais_bench/benchmark/models/api_models/tgi_api.py @@ -125,11 +125,14 @@ async def get_request_body( async def parse_text_response(self, api_response:dict, output:Output): generated_text = api_response.get("generated_text", "") output.content = generated_text + self._update_usage_from_details(api_response, output) async def parse_stream_response(self, api_response:dict, output:Output): generated_text = api_response.get("generated_text", "") # TGI will return the generated text in last chunk output.content = generated_text + # The token count is only reported on the final chunk. + self._update_usage_from_details(api_response, output) class TGICustomAPIStream(TGICustomAPI): diff --git a/ais_bench/benchmark/models/api_models/triton_api.py b/ais_bench/benchmark/models/api_models/triton_api.py index e663296a..6e71566b 100644 --- a/ais_bench/benchmark/models/api_models/triton_api.py +++ b/ais_bench/benchmark/models/api_models/triton_api.py @@ -92,10 +92,13 @@ async def get_request_body(self, input_data: PromptType, max_out_len: int, outpu async def parse_text_response(self, api_response:dict, output:Output): generated_text = api_response.get("text_output", "") output.content = generated_text + self._update_usage_from_details(api_response, output) async def parse_stream_response(self, api_response:dict, output:Output): generated_text = api_response.get("text_output", "") output.content += generated_text + # ``details.generated_tokens`` is cumulative across chunks. + self._update_usage_from_details(api_response, output) class TritonCustomAPIStream(TritonCustomAPI): diff --git a/ais_bench/benchmark/summarizers/default_perf.py b/ais_bench/benchmark/summarizers/default_perf.py index 80c596de..1076b37e 100644 --- a/ais_bench/benchmark/summarizers/default_perf.py +++ b/ais_bench/benchmark/summarizers/default_perf.py @@ -155,10 +155,13 @@ def recursive_update(detail_data): pass # Prefer service-returned "prompt_tokens". elif is_mm_prompt(perf_data["input"]): perf_data["input_tokens"] = 0 # multi-modal input does not support input_tokens - elif "input_tokens" not in perf_data or perf_data.get("input_tokens") is None: - perf_data["input_tokens"] = len(tokenizer.encode(perf_data["input"])) # input_tokens is not provided, calculate it + else: + # `Output` initializes the counters to 0, so a missing value is + # indistinguishable from "not reported" by the key alone. + perf_data["input_tokens"] = len(tokenizer.encode(perf_data["input"])) - if "output_tokens" not in perf_data or perf_data.get("output_tokens") is None: + output_tokens = perf_data.get("output_tokens") + if (output_tokens is None or output_tokens <= 0) and isinstance(perf_data.get("prediction"), str): perf_data["output_tokens"] = len(tokenizer.encode(perf_data["prediction"])) perf_data.pop("input") perf_data.pop("prediction") diff --git a/tests/UT/models/api_models/test_mindie_stream_api.py b/tests/UT/models/api_models/test_mindie_stream_api.py index a69e9f87..68207dc1 100644 --- a/tests/UT/models/api_models/test_mindie_stream_api.py +++ b/tests/UT/models/api_models/test_mindie_stream_api.py @@ -236,6 +236,33 @@ def test_is_api_property(self): model = MindieStreamApi(**self.default_kwargs) self.assertTrue(model.is_api) + def _run_parse(self, parse_name, api_response): + async def run(): + model = MindieStreamApi(**self.default_kwargs) + output = Output() + await getattr(model, parse_name)(api_response, output) + return output + return asyncio.run(run()) + + def test_parse_stream_response_records_output_tokens_from_final_chunk(self): + """MindIE reports token usage under details.generated_tokens on the final chunk""" + output = self._run_parse("parse_stream_response", {"token": {"text": "Hel"}}) + self.assertEqual(output.output_tokens, 0) + + output = self._run_parse( + "parse_stream_response", + {"generated_text": "Hello world!", "details": {"generated_tokens": 12}}, + ) + self.assertEqual(output.content, "Hello world!") + self.assertEqual(output.output_tokens, 12) + + def test_parse_stream_response_without_details_keeps_token_count(self): + """A response without a usage block must not crash or invent a count""" + output = self._run_parse("parse_stream_response", {"generated_text": "hi"}) + + self.assertEqual(output.content, "hi") + self.assertEqual(output.output_tokens, 0) + class TestMindieStreamApiLora(unittest.TestCase): """针对 Multi-LoRA 兼容性扩展的 UT（MindIE 通过 parameters.adapter_id 路由 LoRA）。""" diff --git a/tests/UT/models/api_models/test_tgi_api.py b/tests/UT/models/api_models/test_tgi_api.py index 60fe20aa..014031fa 100644 --- a/tests/UT/models/api_models/test_tgi_api.py +++ b/tests/UT/models/api_models/test_tgi_api.py @@ -245,6 +245,43 @@ async def test_generate_integration_stream_true(self, mock_session_class): await model.generate("test prompt", 100, output) mock_stream_infer.assert_called_once() + def _run_parse(self, parse_name, api_response): + async def run(): + model = TGICustomAPI(**self.default_kwargs) + output = Output() + await getattr(model, parse_name)(api_response, output) + return output + return asyncio.run(run()) + + def test_parse_text_response_records_output_tokens(self): + """TGI reports token usage under details.generated_tokens""" + output = self._run_parse( + "parse_text_response", + {"generated_text": "Hello world!", "details": {"generated_tokens": 12}}, + ) + + self.assertEqual(output.content, "Hello world!") + self.assertEqual(output.output_tokens, 12) + + def test_parse_stream_response_records_output_tokens_from_final_chunk(self): + """Streaming: intermediate chunks carry no details, the final one does""" + output = self._run_parse("parse_stream_response", {"token": {"text": "Hel"}}) + self.assertEqual(output.output_tokens, 0) + + output = self._run_parse( + "parse_stream_response", + {"generated_text": "Hello world!", "details": {"generated_tokens": 12}}, + ) + self.assertEqual(output.content, "Hello world!") + self.assertEqual(output.output_tokens, 12) + + def test_parse_response_without_details_keeps_token_count(self): + """A response without a usage block must not crash or invent a count""" + output = self._run_parse("parse_text_response", {"generated_text": "hi"}) + + self.assertEqual(output.content, "hi") + self.assertEqual(output.output_tokens, 0) + class TestTGICustomAPILora(unittest.TestCase): """针对 Multi-LoRA 兼容性扩展的 UT（TGI 通过 parameters.adapter_id 路由 LoRA）。""" diff --git a/tests/UT/models/api_models/test_triton_api.py b/tests/UT/models/api_models/test_triton_api.py index 4c8bbf76..34bb3108 100644 --- a/tests/UT/models/api_models/test_triton_api.py +++ b/tests/UT/models/api_models/test_triton_api.py @@ -294,5 +294,55 @@ def test_get_service_model_path_call(self, mock_get_model_path): mock_get_model_path.assert_called_once() + def _run_parse(self, parse_name, model, api_response): + async def run(): + output = Output() + await getattr(model, parse_name)(api_response, output) + return output + return asyncio.run(run()) + + def test_parse_text_response_records_output_tokens(self): + """Triton reports token usage under details.generated_tokens""" + with mock.patch.object(base_api.BaseAPIModel, '_get_service_model_path', return_value="m"): + model = TritonCustomAPI(path="test", model_name="m") + output = self._run_parse( + "parse_text_response", + model, + {"text_output": "Hello world!", "details": {"generated_tokens": 12}}, + ) + + self.assertEqual(output.content, "Hello world!") + self.assertEqual(output.output_tokens, 12) + + def test_parse_stream_response_uses_cumulative_output_tokens(self): + """Streaming: details.generated_tokens is cumulative, so the last chunk wins""" + with mock.patch.object(base_api.BaseAPIModel, '_get_service_model_path', return_value="m"): + model = TritonCustomAPI(path="test", model_name="m") + + async def run(): + output = Output() + await model.parse_stream_response( + {"text_output": "Hel", "details": {"generated_tokens": 3}}, output + ) + assert output.output_tokens == 3 + await model.parse_stream_response( + {"text_output": "lo", "details": {"generated_tokens": 5}}, output + ) + return output + + output = asyncio.run(run()) + self.assertEqual(output.content, "Hello") + self.assertEqual(output.output_tokens, 5) + + def test_parse_response_without_details_keeps_token_count(self): + """A response without a usage block must not crash or invent a count""" + with mock.patch.object(base_api.BaseAPIModel, '_get_service_model_path', return_value="m"): + model = TritonCustomAPI(path="test", model_name="m") + output = self._run_parse("parse_text_response", model, {"text_output": "hi"}) + + self.assertEqual(output.content, "hi") + self.assertEqual(output.output_tokens, 0) + + if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/tests/UT/summarizers/test_default_perf.py b/tests/UT/summarizers/test_default_perf.py index c8f67016..472b9de8 100644 --- a/tests/UT/summarizers/test_default_perf.py +++ b/tests/UT/summarizers/test_default_perf.py @@ -131,6 +131,9 @@ def test_calc_perf_data_success(self, mock_load_tokenizer, mock_aistokenizer, mo self.assertEqual(result["end_time"], 3) self.assertEqual(result["latency"], 3) self.assertEqual(result["ttft"], 1) + # The service reported output_tokens=0 (i.e. nothing reported); the + # summarizer must fall back to the tokenizer instead of keeping 0. + self.assertEqual(result["output_tokens"], len("test output")) mock_conn.close.assert_called_once() @patch('ais_bench.benchmark.summarizers.default_perf.build_model_from_cfg')