riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 33 Packages Projects Releases Wiki Activity

[server] add tests and fix passing a custom generation_config (#39230)

Browse Source 
* add tests; fix passing a custom generation_config

* tool integration test

* add install step

* add accelerate as dep to serving

* add todo
This commit is contained in:
Joao Gante 2025-07-10 14:41:38 +01:00 committed by GitHub
parent 6b09c8eab0
commit 38c3931362
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 320 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.circleci/create_circleci_config.py
View File

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

2
setup.py
View File

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

15
src/transformers/commands/chat.py
View File

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

92
src/transformers/commands/serving.py
View File

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

2
src/transformers/generation/continuous_batching.py
View File

@ -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]

260
tests/commands/test_serving.py
View File

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