[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

@@ -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,
 )

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",

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,
                     },
                 )

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 (
+from huggingface_hub import ModelInfo, model_info
-    ChatCompletionStreamOutputDeltaToolCall,
-    ChatCompletionStreamOutputFunction,
-    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`
+    Creates a generation config from the parameters of the request. If a generation config is passed in the request,
-    (serialized into a `dict`) in `extra_body`, for full `generate` parameterization.
+    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:
+    # If there is a generation config in the request, it is a json string serialization from a `GenerationConfig`
-        for key in req.extra_body["generation_config"].keys():
+    # object. For simplicity, flags set here take precedence over all other flags.
-            if key in ChatCompletionInput.base_field_names.keys():
+    if req.generation_config is not None:
-                raise ValueError("error: Duplicated key in the root request and in the passed generation config.")
+        generation_config = GenerationConfig(**json.loads(req.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)
     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(
+                                    tool = {
-                                        function=ChatCompletionStreamOutputFunction(
+                                        "function": {"name": tool_name},
-                                            name=tool_name,
+                                        "index": 0,
-                                        ),
+                                        "type": "function",
-                                        index=0,
+                                        "id": _request_id + "_tool_call",  # Only the first tool call delta has an id
-                                        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(
+                                    tool = {
-                                        function=ChatCompletionStreamOutputFunction(
+                                        "function": {"arguments": result},
-                                            arguments=result,
+                                        "index": 0,
-                                        ),
+                                        "type": "function",
-                                        index=0,
+                                    }
-                                        type="function",
-                                    )
 
                                 yield self.build_chunk(_request_id, tool_calls=[tool])
                                 continue

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]

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 import GenerationConfig
-from transformers.testing_utils import CaptureStd
+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, '<think>\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()