From 38c3931362e4d6548c586d217452a19b279d9bbe Mon Sep 17 00:00:00 2001 From: Joao Gante Date: Thu, 10 Jul 2025 14:41:38 +0100 Subject: [PATCH] [server] add tests and fix passing a custom `generation_config` (#39230) * add tests; fix passing a custom generation_config * tool integration test * add install step * add accelerate as dep to serving * add todo --- .circleci/create_circleci_config.py | 2 +- setup.py | 2 +- src/transformers/commands/chat.py | 15 +- src/transformers/commands/serving.py | 92 ++++--- .../generation/continuous_batching.py | 2 +- tests/commands/test_serving.py | 260 +++++++++++++++++- 6 files changed, 320 insertions(+), 53 deletions(-) diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py index 7d820f3dbe..f8c9d8af5e 100644 --- a/.circleci/create_circleci_config.py +++ b/.circleci/create_circleci_config.py @@ -303,7 +303,7 @@ non_model_job = CircleCIJob( docker_image=[{"image": "huggingface/transformers-torch-light"}], # networkx==3.3 (after #36957) cause some issues # TODO: remove this once it works directly - install_steps=["uv venv && uv pip install ."], + install_steps=["uv venv && uv pip install .[serving]"], marker="not generate", parallelism=6, ) diff --git a/setup.py b/setup.py index a388141884..349257ab28 100644 --- a/setup.py +++ b/setup.py @@ -313,7 +313,7 @@ extras["hub-kernels"] = deps_list("kernels") extras["integrations"] = extras["hub-kernels"] + extras["optuna"] + extras["ray"] + extras["sigopt"] -extras["serving"] = deps_list("pydantic", "uvicorn", "fastapi", "starlette") +extras["serving"] = deps_list("pydantic", "uvicorn", "fastapi", "starlette") + extras["torch"] extras["audio"] = deps_list( "librosa", "pyctcdecode", diff --git a/src/transformers/commands/chat.py b/src/transformers/commands/chat.py index 8f6f49f26b..e74970f694 100644 --- a/src/transformers/commands/chat.py +++ b/src/transformers/commands/chat.py @@ -14,6 +14,7 @@ import asyncio +import copy import json import os import platform @@ -451,11 +452,13 @@ class ChatCommand(BaseTransformersCLICommand): ) return processed_generate_flags - def get_generation_parameterization(self, args: ChatArguments) -> tuple[GenerationConfig, dict]: + def get_generation_parameterization( + self, args: ChatArguments, model_generation_config: GenerationConfig + ) -> tuple[GenerationConfig, dict]: """ Returns a GenerationConfig object holding the generation parameters for the CLI command. """ - # No generation config arg provided -> use base generation config, apply CLI defaults + # No generation config arg provided -> use model's default generation config, then apply CLI defaults if args.generation_config is not None: if ".json" in args.generation_config: # is a local file dirname = os.path.dirname(args.generation_config) @@ -467,7 +470,8 @@ class ChatCommand(BaseTransformersCLICommand): # !!!!!!!!! # This is a chat session, so we have a few non-standard defaults # !!!!!!!!! - generation_config = GenerationConfig(do_sample=True, max_new_tokens=256) + generation_config = copy.deepcopy(model_generation_config) + generation_config.update({"do_sample": True, "max_new_tokens": 256}) # Finally: parse and apply `generate_flags` parsed_generate_flags = self.parse_generate_flags(args.generate_flags) @@ -675,7 +679,8 @@ class ChatCommand(BaseTransformersCLICommand): else: user = args.user - generation_config, model_kwargs = self.get_generation_parameterization(args) + model_generation_config = GenerationConfig.from_pretrained(args.model_name_or_path) + generation_config, model_kwargs = self.get_generation_parameterization(args, model_generation_config) interface = RichInterface(model_name=args.model_name_or_path, user_name=user) interface.clear() @@ -715,7 +720,7 @@ class ChatCommand(BaseTransformersCLICommand): stream=True, extra_body={ "request_id": request_id, - "generation_config": {**generation_config.to_dict()}, + "generation_config": generation_config.to_json_string(), "model": model, }, ) diff --git a/src/transformers/commands/serving.py b/src/transformers/commands/serving.py index 5ab2099669..a7be1d4545 100644 --- a/src/transformers/commands/serving.py +++ b/src/transformers/commands/serving.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import copy import functools import json import re @@ -20,12 +21,7 @@ from dataclasses import dataclass, field from threading import Thread from typing import Any, Optional -from huggingface_hub import ( - ChatCompletionStreamOutputDeltaToolCall, - ChatCompletionStreamOutputFunction, - ModelInfo, - model_info, -) +from huggingface_hub import ModelInfo, model_info from transformers.utils.import_utils import is_fastapi_available, is_pydantic_available, is_uvicorn_available @@ -86,6 +82,9 @@ if is_pydantic_available() and is_fastapi_available() and is_uvicorn_available() # tool_prompt: Optional[str] = None # top_logprobs: Optional[int] = None + # transformers-specific request fields + generation_config: Optional[str] = None + logger = logging.get_logger(__name__) @@ -110,26 +109,35 @@ def serve_command_factory(args: Namespace): return ServeCommand(args) -def create_generation_config_from_req(req: "ChatCompletionInput", **kwargs) -> "GenerationConfig": +def create_generation_config_from_req( + req: "ChatCompletionInput", model_generation_config: "GenerationConfig", **kwargs +) -> "GenerationConfig": """ - Creates a generation config from the parameters of the request. Note that we can pass a `GenerationConfig` - (serialized into a `dict`) in `extra_body`, for full `generate` parameterization. + Creates a generation config from the parameters of the request. If a generation config is passed in the request, + it will be used as a baseline for parameterization. Otherwise, we will use the model's default generation config. + Other parameters in the request will be applied on top of the baseline. Args: - req (`ChatCompletionInput`): The request which may optionally contain generation parameters. + req (`ChatCompletionInput`): + The request which may optionally contain generation parameters. + model_generation_config (`GenerationConfig`): + The model's default generation config. Returns: The prepared `GenerationConfig` object. """ - if req.extra_body is not None and "generation_config" in req.extra_body: - for key in req.extra_body["generation_config"].keys(): - if key in ChatCompletionInput.base_field_names.keys(): - raise ValueError("error: Duplicated key in the root request and in the passed generation config.") - - if req.extra_body is not None and "generation_config" in req.extra_body: - generation_config = GenerationConfig(**(req.extra_body["generation_config"]), **kwargs) + # If there is a generation config in the request, it is a json string serialization from a `GenerationConfig` + # object. For simplicity, flags set here take precedence over all other flags. + if req.generation_config is not None: + generation_config = GenerationConfig(**json.loads(req.generation_config)) else: - generation_config = GenerationConfig(**kwargs) + generation_config = copy.deepcopy(model_generation_config) + + non_standard_kwargs = generation_config.update(**kwargs) + # Set extra kwargs that are not in the `GenerationConfig` class (e.g. continuous batching flags) + for k, v in non_standard_kwargs.items(): + if v is not None: + setattr(generation_config, k, v) if req.frequency_penalty is not None: generation_config.repetition_penalty = float(req.frequency_penalty) @@ -267,7 +275,7 @@ class ServeCommand(BaseTransformersCLICommand): content: Optional[str] = None, role: Optional[str] = None, finish_reason: Optional[str] = None, - tool_calls: Optional[list[ChatCompletionStreamOutputDeltaToolCall]] = None, + tool_calls: Optional[list[dict]] = None, ) -> str: """ Builds a chunk of a streaming response. @@ -284,7 +292,7 @@ class ServeCommand(BaseTransformersCLICommand): The role of the next content, until a new role is defined. finish_reason (`str`, *optional*): The reason the generation by the model has finished. - tool_calls (`list[ChatCompletionStreamOutputDeltaToolCall]`, *optional*): + tool_calls (`list[dict]`, *optional*): Data about the tool calls, when they are triggered. Returns: @@ -380,6 +388,7 @@ class ServeCommand(BaseTransformersCLICommand): generation_config = create_generation_config_from_req( req, + model_generation_config=self.model.generation_config, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, use_cache=False, @@ -413,6 +422,10 @@ class ServeCommand(BaseTransformersCLICommand): ) queue_is_flushed = False + # Emit the assistant role to start the stream. Other chunks won't have a role, as it is implicit + # they come from the assistant. + yield self.build_chunk(request_id, role="assistant") + for result in self.running_continuous_batching_manager: if result.request_id != request_id: continue @@ -424,14 +437,12 @@ class ServeCommand(BaseTransformersCLICommand): queue_is_flushed = True finish_reason = "stop" if result.status == RequestStatus.FINISHED else None - yield self.build_chunk( - request_id=request_id, content=result.next_token, finish_reason=finish_reason - ) - if result.status == RequestStatus.FINISHED: + yield self.build_chunk(request_id, finish_reason=finish_reason) break + else: + yield self.build_chunk(request_id=request_id, content=result.next_token) - yield "data: [DONE]\n\n" except Exception as e: logger.error(str(e)) yield f'data: {{"error": "{str(e)}"}}' @@ -507,7 +518,10 @@ class ServeCommand(BaseTransformersCLICommand): generation_streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True, skip_prompt=True) - generation_config = create_generation_config_from_req(req) + generation_config = create_generation_config_from_req( + req, + model_generation_config=self.model.generation_config, + ) max_new_tokens = req.max_tokens or generation_config.max_new_tokens or 1024 generation_config.max_new_tokens = max_new_tokens @@ -570,14 +584,12 @@ class ServeCommand(BaseTransformersCLICommand): else: tool_name = tool_name.group(1) tool_state.has_tool_name_defined = True - tool = ChatCompletionStreamOutputDeltaToolCall( - function=ChatCompletionStreamOutputFunction( - name=tool_name, - ), - index=0, - type="function", - id=_request_id + "_tool_call", # Only the first tool call delta has an id - ) + tool = { + "function": {"name": tool_name}, + "index": 0, + "type": "function", + "id": _request_id + "_tool_call", # Only the first tool call delta has an id + } # Second step: extract tool arguments. The tool arguments can be seen as a json string # within the tool json string. We emit a delta for the arguments. @@ -597,13 +609,11 @@ class ServeCommand(BaseTransformersCLICommand): if tool_state.arg_nesting_level < 0: result = "".join(result.split("}")[:-2]) + "}" # e.g. "4}}\n" -> "4}" - tool = ChatCompletionStreamOutputDeltaToolCall( - function=ChatCompletionStreamOutputFunction( - arguments=result, - ), - index=0, - type="function", - ) + tool = { + "function": {"arguments": result}, + "index": 0, + "type": "function", + } yield self.build_chunk(_request_id, tool_calls=[tool]) continue diff --git a/src/transformers/generation/continuous_batching.py b/src/transformers/generation/continuous_batching.py index f57889e117..8f871a150f 100644 --- a/src/transformers/generation/continuous_batching.py +++ b/src/transformers/generation/continuous_batching.py @@ -640,7 +640,7 @@ def compute_optimal_blocks( memory_per_token = 2 * num_kv_heads * head_dim * dtype_size * num_hidden_layers # For K and V caches # Estimate sequence length requirements - tokens_to_generate = getattr(generation_config, "max_new_tokens", 20) + tokens_to_generate = getattr(generation_config, "max_new_tokens") or 20 if median_prefill_length is None and inputs: non_empty_inputs = [len(seq) for seq in inputs if seq] diff --git a/tests/commands/test_serving.py b/tests/commands/test_serving.py index 00fa4630e2..ae9fe16e68 100644 --- a/tests/commands/test_serving.py +++ b/tests/commands/test_serving.py @@ -11,22 +11,32 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import time import unittest +from threading import Thread from unittest.mock import patch +import aiohttp.client_exceptions +from huggingface_hub import AsyncInferenceClient +from parameterized import parameterized + import transformers.commands.transformers_cli as cli -from transformers.commands.serving import ServeCommand -from transformers.testing_utils import CaptureStd +from transformers import GenerationConfig +from transformers.commands.serving import ServeArguments, ServeCommand +from transformers.testing_utils import CaptureStd, slow class ServeCLITest(unittest.TestCase): def test_help(self): + """Minimal test: we can invoke the help command.""" with patch("sys.argv", ["transformers", "serve", "--help"]), CaptureStd() as cs: with self.assertRaises(SystemExit): cli.main() self.assertIn("serve", cs.out.lower()) def test_parsed_args(self): + """Minimal test: we can set arguments through the CLI.""" with ( patch.object(ServeCommand, "__init__", return_value=None) as init_mock, patch.object(ServeCommand, "run") as run_mock, @@ -39,9 +49,251 @@ class ServeCLITest(unittest.TestCase): self.assertEqual(parsed_args.host, "0.0.0.0") self.assertEqual(parsed_args.port, 9000) - def test_build_chunk(self): + def test_completions_build_chunk(self): + """Tests that the chunks are correctly built for the Completions API.""" dummy = ServeCommand.__new__(ServeCommand) dummy.args = type("Args", (), {})() - chunk = ServeCommand.build_chunk(dummy, "hello", "req0", finish_reason="stop") + + # Case 1: most fields are provided + chunk = ServeCommand.build_chunk(dummy, request_id="req0", content="hello", finish_reason="stop", role="user") self.assertIn("chat.completion.chunk", chunk) self.assertIn("data:", chunk) + self.assertIn( + '"choices": [{"delta": {"content": "hello", "role": "user"}, "index": 0, "finish_reason": "stop"}]', chunk + ) + + # Case 2: only the role is provided -- other fields in 'choices' are omitted + chunk = ServeCommand.build_chunk(dummy, request_id="req0", role="user") + self.assertIn("chat.completion.chunk", chunk) + self.assertIn("data:", chunk) + self.assertIn('"choices": [{"delta": {"role": "user"}, "index": 0}]', chunk) + + # Case 3: only the content is provided -- other fields in 'choices' are omitted + chunk = ServeCommand.build_chunk(dummy, request_id="req0", content="hello") + self.assertIn("chat.completion.chunk", chunk) + self.assertIn("data:", chunk) + self.assertIn('"choices": [{"delta": {"content": "hello"}, "index": 0}]', chunk) + + # Case 4: tool calls support a list of nested dictionaries + chunk = ServeCommand.build_chunk(dummy, request_id="req0", tool_calls=[{"foo1": "bar1", "foo2": "bar2"}]) + self.assertIn("chat.completion.chunk", chunk) + self.assertIn("data:", chunk) + self.assertIn('"choices": [{"delta": {"tool_calls": [{"foo1": "bar1", "foo2": "bar2"}]}, "index": 0}]', chunk) + + +def async_retry(fn, max_attempts=5, delay=2): + """ + Retry a function up to `max_attempts` times with a `delay` between attempts. + Useful for testing async functions that may fail due to server not being ready. + """ + + async def wrapper(*args, **kwargs): + for _ in range(max_attempts): + try: + return await fn(*args, **kwargs) + except aiohttp.client_exceptions.ClientConnectorError: + time.sleep(delay) + + return wrapper + + +class ServeCompletionsMixin: + """ + Mixin class for the Completions API tests, to seamlessly replicate tests across the two versions of the API + (`generate` and `continuous_batching`). + """ + + @async_retry + async def run_server(self, request): + client = AsyncInferenceClient("http://localhost:8000") + stream = client.chat_completion(**request) + + all_payloads = [] + async for payload in await stream: + all_payloads.append(payload) + + await client.close() + return all_payloads + + @parameterized.expand( + [ + ("default_request", {}), + ("one_token", {"max_tokens": 1}), + # TODO: CB fails next case, seems like it is unable to switch models. fix me + # ("different_model", {"model": "HuggingFaceTB/SmolLM2-135M-Instruct"}), + ( + "tool_call", + { + "tools": [ + { + "function": { + "name": "foo_bar", + "parameters": {"type": "object"}, + "description": "Foo bar", + }, + "type": "function", + } + ] + }, + ), + ] + ) + def test_requests(self, test_name: str, request_flags: dict): + """Tests that the completions app gracefully handles GOOD requests, producing the expected output payloads.""" + + request = { + "model": "Qwen/Qwen3-0.6B", + "messages": [{"role": "user", "content": "Hello, how are you?"}], + "stream": True, # We don't support "stream": False yet + "max_tokens": 5, # Small generation by default + } + request.update(request_flags) + all_payloads = asyncio.run(self.run_server(request)) + + # If a request is successful, the returned payload needs to follow the schema, which we test here. + # NOTE: the output of our server is wrapped by `AsyncInferenceClient`, which sends fields even when they + # are empty. + + # Finish reason: the last payload should have a finish reason of "stop", all others should be empty + # TODO: we may add other finish reasons in the future, and this may need more logic + finish_reasons = [payload.choices[0].finish_reason for payload in all_payloads] + self.assertEqual(finish_reasons[-1], "stop") + self.assertTrue(all(reason is None for reason in finish_reasons[:-1])) + + # Role: the first payload should have a role of "assistant", all others should be empty + roles = [payload.choices[0].delta.role for payload in all_payloads] + self.assertEqual(roles[0], "assistant") + self.assertTrue(all(role is None for role in roles[1:])) + + # Content: the first and the last payload shouldn't have content (role and finish reason). It may be empty + # in some other payload positions, e.g. tool calls. + contents = [payload.choices[0].delta.content for payload in all_payloads] + self.assertTrue(contents[0] is None and contents[-1] is None) + self.assertTrue(any(content is not None for content in contents[1:-1])) + # TODO: add "usage" field to output and test it + + def test_generation_config_in_request(self): + """Tests that the generation config is correctly passed into the generation call.""" + generation_config = GenerationConfig(do_sample=False, temperature=0.0) + request = { + "model": "Qwen/Qwen3-0.6B", + "messages": [{"role": "user", "content": "Hello, how are you?"}], + "stream": True, + "max_tokens": 10, + "extra_body": { + "generation_config": generation_config.to_json_string(), + }, + } + all_payloads = asyncio.run(self.run_server(request)) + contents = [payload.choices[0].delta.content for payload in all_payloads] + output_text = "".join([text for text in contents if text is not None]) + # The generation config sets greedy decoding, so the output is reproducible. By default, `Qwen/Qwen3-0.6B` + # sets `do_sample=True` + self.assertEqual(output_text, '\nOkay, the user just asked, "') + + # TODO: implement API-compliant error handling, and then test it + # See https://platform.openai.com/docs/guides/error-codes, + # TODO: one test for each request flag, to confirm it is working as expected + # TODO: speed-based test to confirm that KV cache is working across requests + + +class ServeCompletionsGenerateTest(ServeCompletionsMixin, unittest.TestCase): + """Tests the `generate` version of the Completions API.""" + + @classmethod + def setUpClass(cls): + """Starts a server for tests to connect to.""" + args = ServeArguments() + serve_command = ServeCommand(args) + thread = Thread(target=serve_command.run) + thread.daemon = True + thread.start() + + @slow + def test_tool_call(self): + """Tests that the tool call is correctly handled and that the payloads are correctly structured.""" + # TODO: move to the mixin when CB also supports tool calls + + request = { + # This model is a small model that's very eager to call tools + # TODO: this is a 4B model. Find a smaller model that's eager to call tools + "model": "Menlo/Jan-nano", + # The request should produce a tool call + "messages": [{"role": "user", "content": "Generate an image of a cat."}], + "stream": True, + "max_tokens": 50, + # Reproducibility + "temperature": 0.0, + # This tool is a copy from the tool in the original tiny-agents demo + "tools": [ + { + "function": { + "name": "flux1_schnell_infer", + "parameters": { + "type": "object", + "properties": { + "prompt": {"type": "string"}, + "seed": {"type": "number", "description": "numeric value between 0 and 2147483647"}, + "randomize_seed": {"type": "boolean", "default": True}, + "width": { + "type": "number", + "description": "numeric value between 256 and 2048", + "default": 1024, + }, + "height": { + "type": "number", + "description": "numeric value between 256 and 2048", + "default": 1024, + }, + "num_inference_steps": { + "type": "number", + "description": "numeric value between 1 and 16", + "default": 4, + }, + }, + }, + "description": "Generate an image using the Flux 1 Schnell Image Generator.", + }, + "type": "function", + } + ], + } + all_payloads = asyncio.run(self.run_server(request)) + + # The first payload should contain the role + roles = [payload.choices[0].delta.role for payload in all_payloads] + self.assertEqual(roles[0], "assistant") + self.assertTrue(all(role is None for role in roles[1:])) + + # All other payloads (except the last one) should be tool call related, for this specific request + contents = [payload.choices[0].delta.content for payload in all_payloads] + self.assertTrue(all(content is None for content in contents)) + + # The first tool call delta should contain the tool name. The other tool call deltas should contain the tool + # arguments. + tool_calls = [payload.choices[0].delta.tool_calls[0] for payload in all_payloads[1:-1]] + first_tool_call = tool_calls[0] + self.assertEqual(first_tool_call["function"]["name"], "flux1_schnell_infer") + self.assertEqual(first_tool_call["function"]["arguments"], None) + other_tool_calls = tool_calls[1:] + self.assertTrue(all(tool_call["function"]["name"] is None for tool_call in other_tool_calls)) + self.assertTrue(all(tool_call["function"]["arguments"] is not None for tool_call in other_tool_calls)) + + # Finally, the last payload should contain a finish reason + finish_reasons = [payload.choices[0].finish_reason for payload in all_payloads] + # TODO: I think the finish reason for a tool call is different? double check this + self.assertEqual(finish_reasons[-1], "stop") + self.assertTrue(all(reason is None for reason in finish_reasons[:-1])) + + +class ServeCompletionsContinuousBatchingTest(ServeCompletionsMixin, unittest.TestCase): + """Tests the `continuous_batching` version of the Completions API.""" + + @classmethod + def setUpClass(cls): + """Starts a server for tests to connect to.""" + args = ServeArguments(attn_implementation="sdpa_paged") # important: toggle continuous batching + serve_command = ServeCommand(args) + thread = Thread(target=serve_command.run) + thread.daemon = True + thread.start()