|
2 | 2 | # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
3 | 3 | """Tests for content-format-driven chat template dispatch and placeholder handling.""" |
4 | 4 |
|
| 5 | +import threading |
| 6 | + |
5 | 7 | import pytest |
6 | 8 |
|
7 | 9 | from tensorrt_llm.inputs.content_format import ContentFormat |
|
15 | 17 | _build_openai_content, |
16 | 18 | _resolve_content_format, |
17 | 19 | add_multimodal_placeholders, |
| 20 | + async_apply_chat_template, |
18 | 21 | interleave_mm_placeholders, |
19 | 22 | ) |
20 | 23 |
|
@@ -324,3 +327,131 @@ def test_excess_existing_placeholders_preserved(self): |
324 | 327 | ) |
325 | 328 | assert result == text |
326 | 329 | assert result.count("<image>") == 3 |
| 330 | + |
| 331 | + |
| 332 | +class TestAsyncApplyChatTemplate: |
| 333 | + @pytest.mark.asyncio |
| 334 | + async def test_runs_in_worker_thread(self): |
| 335 | + event_loop_thread_id = threading.current_thread().ident |
| 336 | + |
| 337 | + class TrackingTokenizer: |
| 338 | + def __init__(self): |
| 339 | + self.worker_thread_id = None |
| 340 | + |
| 341 | + def apply_chat_template(self, **_): |
| 342 | + self.worker_thread_id = threading.current_thread().ident |
| 343 | + return "rendered" |
| 344 | + |
| 345 | + tokenizer = TrackingTokenizer() |
| 346 | + |
| 347 | + result = await async_apply_chat_template( |
| 348 | + model_type="test_string_model", |
| 349 | + tokenizer=tokenizer, |
| 350 | + processor=None, |
| 351 | + conversation=[ConversationMessage(role="user", content="hello", media=[])], |
| 352 | + add_generation_prompt=True, |
| 353 | + mm_placeholder_counts=[{}], |
| 354 | + chat_template="{{ messages }}", |
| 355 | + ) |
| 356 | + |
| 357 | + assert result == "rendered" |
| 358 | + assert tokenizer.worker_thread_id is not None |
| 359 | + assert tokenizer.worker_thread_id != event_loop_thread_id |
| 360 | + |
| 361 | + |
| 362 | +class TestServingChatTemplateGather: |
| 363 | + """Cover the asyncio.gather integration in the serving chat-template paths.""" |
| 364 | + |
| 365 | + @pytest.mark.asyncio |
| 366 | + async def test_resource_governor_convert_messages(self, monkeypatch): |
| 367 | + from unittest.mock import Mock |
| 368 | + |
| 369 | + import tensorrt_llm.serve.resource_governor as rg |
| 370 | + |
| 371 | + governor = object.__new__(rg.ResourceGovernor) |
| 372 | + governor.model_config = Mock() |
| 373 | + governor.tokenizer = Mock() |
| 374 | + governor.processor = None |
| 375 | + |
| 376 | + async def fake_mm_coroutine(): |
| 377 | + # parse_chat_messages_coroutines' coroutine yields |
| 378 | + # (mm_data, mm_embeddings). |
| 379 | + return ({"image": ["data"]}, None) |
| 380 | + |
| 381 | + monkeypatch.setattr( |
| 382 | + rg, |
| 383 | + "parse_chat_messages_coroutines", |
| 384 | + lambda messages, model_config, _: ([], fake_mm_coroutine(), [{}]), |
| 385 | + ) |
| 386 | + # Must resolve the top-level model type, matching the serving call |
| 387 | + # sites (not the raw model_config.model_type). |
| 388 | + monkeypatch.setattr(rg, "resolve_top_level_model_type", lambda cfg: "resolved-model-type") |
| 389 | + |
| 390 | + captured = {} |
| 391 | + |
| 392 | + async def fake_async_apply(**kwargs): |
| 393 | + captured.update(kwargs) |
| 394 | + return [1, 2, 3] |
| 395 | + |
| 396 | + monkeypatch.setattr(rg, "async_apply_chat_template", fake_async_apply) |
| 397 | + |
| 398 | + token_ids = await governor._convert_messages( |
| 399 | + messages=[{"role": "user", "content": "hi"}], |
| 400 | + tool_dicts=None, |
| 401 | + add_generation_prompt=True, |
| 402 | + documents=None, |
| 403 | + chat_template=None, |
| 404 | + chat_template_kwargs=None, |
| 405 | + ) |
| 406 | + |
| 407 | + # Returns only token_ids, not the (mm_data, mm_embeddings) tuple. |
| 408 | + assert token_ids == [1, 2, 3] |
| 409 | + # Uses the top-level resolver and forwards the real placeholder counts. |
| 410 | + assert captured["model_type"] == "resolved-model-type" |
| 411 | + assert captured["mm_placeholder_counts"] == [{}] |
| 412 | + |
| 413 | + @pytest.mark.asyncio |
| 414 | + async def test_responses_create_input_tokens_unpacks_mm_tuple(self, monkeypatch): |
| 415 | + """_create_input_tokens must return mm_data, not the whole gather tuple.""" |
| 416 | + from unittest.mock import Mock |
| 417 | + |
| 418 | + import tensorrt_llm.serve.responses_utils as ru |
| 419 | + |
| 420 | + async def fake_create_input_messages(request, prev_msgs): |
| 421 | + return [{"role": "user", "content": "hi"}] |
| 422 | + |
| 423 | + async def fake_mm_coroutine(): |
| 424 | + return ({"image": ["data"]}, {"image": ["embed"]}) |
| 425 | + |
| 426 | + monkeypatch.setattr(ru, "_create_input_messages", fake_create_input_messages) |
| 427 | + monkeypatch.setattr( |
| 428 | + ru, |
| 429 | + "parse_chat_messages_coroutines", |
| 430 | + lambda messages, model_config: ([], fake_mm_coroutine(), [{}]), |
| 431 | + ) |
| 432 | + monkeypatch.setattr(ru, "resolve_top_level_model_type", lambda cfg: "resolved-model-type") |
| 433 | + monkeypatch.setattr(ru, "_get_chat_completion_function_tools", lambda tools: []) |
| 434 | + |
| 435 | + async def fake_async_apply(**kwargs): |
| 436 | + return [1, 2, 3] |
| 437 | + |
| 438 | + monkeypatch.setattr(ru, "async_apply_chat_template", fake_async_apply) |
| 439 | + |
| 440 | + request = Mock() |
| 441 | + request.tools = None |
| 442 | + request.store = False |
| 443 | + |
| 444 | + token_ids, mm_data = await ru._create_input_tokens( |
| 445 | + request=request, |
| 446 | + prev_response=None, |
| 447 | + prev_msgs=None, |
| 448 | + conversation_store=None, |
| 449 | + enable_store=False, |
| 450 | + tokenizer=Mock(), |
| 451 | + model_config=Mock(), |
| 452 | + processor=None, |
| 453 | + ) |
| 454 | + |
| 455 | + assert token_ids == [1, 2, 3] |
| 456 | + # mm_data is the data dict, not the (mm_data, mm_embeddings) tuple. |
| 457 | + assert mm_data == {"image": ["data"]} |
0 commit comments