test_serving_chat.py 16.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
from __future__ import annotations

6
import asyncio
7
from contextlib import suppress
8
from dataclasses import dataclass, field
9
from typing import TYPE_CHECKING, Any, Optional
10
from unittest.mock import AsyncMock, MagicMock
11

12
import pytest
13
import pytest_asyncio
14

15
from vllm.config.multimodal import MultiModalConfig
16
from vllm.entrypoints.openai.protocol import ChatCompletionRequest
17
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
18
19
from vllm.entrypoints.openai.serving_models import (BaseModelPath,
                                                    OpenAIServingModels)
20
from vllm.transformers_utils.tokenizer import get_tokenizer
21
from vllm.v1.engine.async_llm import AsyncLLM
22

23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
from ...utils import RemoteOpenAIServer

if TYPE_CHECKING:
    from openai import OpenAI

GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b"


@pytest.fixture(scope="module")
def monkeypatch_module():
    from _pytest.monkeypatch import MonkeyPatch
    mpatch = MonkeyPatch()
    yield mpatch
    mpatch.undo()


39
40
41
42
43
44
45
@pytest.fixture(scope="module",
                params=[True, False],
                ids=["with_tool_parser", "without_tool_parser"])
def with_tool_parser(request) -> bool:
    return request.param


46
@pytest.fixture(scope="module")
47
48
49
50
51
52
53
54
55
56
57
58
59
def default_server_args(with_tool_parser: bool):
    args = [
        # use half precision for speed and memory savings in CI environment
        "--enforce-eager",
        "--max-model-len",
        "4096",
        "--reasoning-parser",
        "openai_gptoss",
        "--gpu-memory-utilization",
        "0.8",
    ]
    if with_tool_parser:
        args.extend([
60
61
62
            "--tool-call-parser",
            "openai",
            "--enable-auto-tool-choice",
63
64
65
66
67
68
69
70
        ])
    return args


@pytest.fixture(scope="module")
def gptoss_server(monkeypatch_module: pytest.MonkeyPatch,
                  default_server_args: list[str]):
    with monkeypatch_module.context() as m:
71
        m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN")
72
73
        with RemoteOpenAIServer(GPT_OSS_MODEL_NAME,
                                default_server_args) as remote_server:
74
75
76
77
78
79
80
81
82
83
            yield remote_server


@pytest_asyncio.fixture
async def gptoss_client(gptoss_server):
    async with gptoss_server.get_async_client() as async_client:
        yield async_client


@pytest.mark.asyncio
84
85
async def test_gpt_oss_chat_tool_call_streaming(gptoss_client: OpenAI,
                                                with_tool_parser: bool):
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
    tools = [{
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get the current weather in a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string"
                    },
                    "state": {
                        "type": "string"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                    },
                },
                "required": ["city", "state", "unit"],
            },
        },
    }]

    messages = [
        {
            "role": "user",
            "content": "What is the weather in Dallas, TX?"
        },
    ]

    stream = await gptoss_client.chat.completions.create(
118
119
120
121
        model=GPT_OSS_MODEL_NAME,
        messages=messages,
        tools=tools if with_tool_parser else None,
        stream=True)
122
123
124

    name = None
    args_buf = ""
125
    content_buf = ""
126
127
128
129
130
131
132
133
    async for chunk in stream:
        delta = chunk.choices[0].delta
        if delta.tool_calls:
            tc = delta.tool_calls[0]
            if tc.function and tc.function.name:
                name = tc.function.name
            if tc.function and tc.function.arguments:
                args_buf += tc.function.arguments
134
135
136
137
138
139
140
141
142
        if getattr(delta, "content", None):
            content_buf += delta.content
    if with_tool_parser:
        assert name is not None
        assert len(args_buf) > 0
    else:
        assert name is None
        assert len(args_buf) == 0
        assert len(content_buf) > 0
143
144
145


@pytest.mark.asyncio
146
147
148
149
async def test_gpt_oss_multi_turn_chat(gptoss_client: OpenAI,
                                       with_tool_parser: bool):
    if not with_tool_parser:
        pytest.skip("skip non-tool for multi-turn tests")
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
    tools = [{
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get the current weather in a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string"
                    },
                    "state": {
                        "type": "string"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                    },
                },
                "required": ["city", "state", "unit"],
            },
        },
    }]

    messages = [
        {
            "role": "system",
            "content": "you are a helpful assistant"
        },
        {
            "role": "user",
181
            "content": "What is the weather in Dallas, TX with celsius?"
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
        },
    ]

    first = await gptoss_client.chat.completions.create(
        model=GPT_OSS_MODEL_NAME,
        messages=messages,
        tools=tools,
        temperature=0.0,
    )
    first_msg = first.choices[0].message
    assert first_msg.tool_calls is not None and len(first_msg.tool_calls) > 0
    tc = first_msg.tool_calls[0]
    assert tc.function is not None and tc.function.name == "get_current_weather"
    args1 = tc.function.arguments
    assert args1 is not None and len(args1) > 0
197
    assert not first_msg.content
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212

    messages.append({"role": "assistant", "content": args1})
    messages.append({
        "role": "user",
        "content": "Now convert to celsius and return JSON only"
    })

    second = await gptoss_client.chat.completions.create(
        model=GPT_OSS_MODEL_NAME,
        messages=messages,
        tools=tools,
        temperature=0.0,
    )
    second_msg = second.choices[0].message
    assert (second_msg.content is not None and len(second_msg.content) > 0) or \
213
        (second_msg.tool_calls is not None and len(second_msg.tool_calls) > 0)
214
215


216
MODEL_NAME = "openai-community/gpt2"
217
MODEL_NAME_SHORT = "gpt2"
218
CHAT_TEMPLATE = "Dummy chat template for testing {}"
219
220
221
222
BASE_MODEL_PATHS = [
    BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
    BaseModelPath(name=MODEL_NAME_SHORT, model_path=MODEL_NAME_SHORT)
]
223
224


225
226
227
228
229
@dataclass
class MockHFConfig:
    model_type: str = "any"


230
231
@dataclass
class MockModelConfig:
232
    task = "generate"
233
    runner_type = "generate"
234
235
236
237
238
    tokenizer = MODEL_NAME
    trust_remote_code = False
    tokenizer_mode = "auto"
    max_model_len = 100
    tokenizer_revision = None
239
    multimodal_config = MultiModalConfig()
240
    hf_config = MockHFConfig()
241
    logits_processor_pattern = None
242
    diff_sampling_param: Optional[dict] = None
243
    allowed_local_media_path: str = ""
244
    allowed_media_domains: Optional[list[str]] = None
245
    encoder_config = None
246
    generation_config: str = "auto"
247
    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
248
    skip_tokenizer_init = False
249
250
251

    def get_diff_sampling_param(self):
        return self.diff_sampling_param or {}
252
253


254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def _build_serving_chat(engine: AsyncLLM,
                        model_config: MockModelConfig) -> OpenAIServingChat:
    models = OpenAIServingModels(engine_client=engine,
                                 base_model_paths=BASE_MODEL_PATHS,
                                 model_config=model_config)
    serving_chat = OpenAIServingChat(engine,
                                     model_config,
                                     models,
                                     response_role="assistant",
                                     chat_template=CHAT_TEMPLATE,
                                     chat_template_content_format="auto",
                                     request_logger=None)

    async def _fake_process_inputs(request_id, engine_prompt, sampling_params,
                                   *, lora_request, trace_headers, priority):
        return dict(engine_prompt), {}

    serving_chat._process_inputs = AsyncMock(side_effect=_fake_process_inputs)
    return serving_chat


275
276
277
278
@dataclass
class MockEngine:

    async def get_model_config(self):
279
        return MockModelConfig()
280
281
282


async def _async_serving_chat_init():
283
284
285
    engine = MockEngine()
    model_config = await engine.get_model_config()

286
    models = OpenAIServingModels(engine, model_config, BASE_MODEL_PATHS)
287
288
    serving_completion = OpenAIServingChat(engine,
                                           model_config,
289
                                           models,
290
                                           response_role="assistant",
291
                                           chat_template=CHAT_TEMPLATE,
292
                                           chat_template_content_format="auto",
293
                                           request_logger=None)
294
295
296
297
298
    return serving_completion


def test_async_serving_chat_init():
    serving_completion = asyncio.run(_async_serving_chat_init())
299
    assert serving_completion.chat_template == CHAT_TEMPLATE
300
301


302
303
@pytest.mark.asyncio
async def test_serving_chat_returns_correct_model_name():
304
    mock_engine = MagicMock(spec=AsyncLLM)
305
306
307
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False

308
    serving_chat = _build_serving_chat(mock_engine, MockModelConfig())
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
    messages = [{"role": "user", "content": "what is 1+1?"}]

    async def return_model_name(*args):
        return args[3]

    serving_chat.chat_completion_full_generator = return_model_name

    # Test that full name is returned when short name is requested
    req = ChatCompletionRequest(model=MODEL_NAME_SHORT, messages=messages)
    assert await serving_chat.create_chat_completion(req) == MODEL_NAME

    # Test that full name is returned when empty string is specified
    req = ChatCompletionRequest(model="", messages=messages)
    assert await serving_chat.create_chat_completion(req) == MODEL_NAME

    # Test that full name is returned when no model is specified
    req = ChatCompletionRequest(messages=messages)
    assert await serving_chat.create_chat_completion(req) == MODEL_NAME


329
330
@pytest.mark.asyncio
async def test_serving_chat_should_set_correct_max_tokens():
331
    mock_engine = MagicMock(spec=AsyncLLM)
332
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
333
    mock_engine.errored = False
334

335
    serving_chat = _build_serving_chat(mock_engine, MockModelConfig())
336

337
338
339
340
341
342
343
344
345
    req = ChatCompletionRequest(
        model=MODEL_NAME,
        messages=[{
            "role": "user",
            "content": "what is 1+1?"
        }],
    )

    with suppress(Exception):
346
        await serving_chat.create_chat_completion(req)
347
348
349
350
351

    assert mock_engine.generate.call_args.args[1].max_tokens == 93

    req.max_tokens = 10
    with suppress(Exception):
352
        await serving_chat.create_chat_completion(req)
353
354

    assert mock_engine.generate.call_args.args[1].max_tokens == 10
355

356
357
358
359
360
361
362
363
    # Setting server's max_tokens in the generation_config.json
    # lower than context_window - prompt_tokens
    mock_model_config = MockModelConfig()
    mock_model_config.diff_sampling_param = {
        "max_tokens": 10  # Setting server-side max_tokens limit
    }

    # Reinitialize the engine with new settings
364
    mock_engine = MagicMock(spec=AsyncLLM)
365
366
367
368
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False

    # Initialize the serving chat
369
    serving_chat = _build_serving_chat(mock_engine, mock_model_config)
370
371
372
373
374
375
376
377
378
379
380

    # Test Case 1: No max_tokens specified in request
    req = ChatCompletionRequest(
        model=MODEL_NAME,
        messages=[{
            "role": "user",
            "content": "what is 1+1?"
        }],
    )

    with suppress(Exception):
381
        await serving_chat.create_chat_completion(req)
382
383
384
385
386
387
388

    assert mock_engine.generate.call_args.args[1].max_tokens == 10

    # Test Case 2: Request's max_tokens set higher than server accepts
    req.max_tokens = 15

    with suppress(Exception):
389
        await serving_chat.create_chat_completion(req)
390
391
392
393
394
395
396

    assert mock_engine.generate.call_args.args[1].max_tokens == 10

    # Test Case 3: Request's max_tokens set lower than server accepts
    req.max_tokens = 5

    with suppress(Exception):
397
        await serving_chat.create_chat_completion(req)
398
399
400
401
402
403
404
405
406
407
408

    assert mock_engine.generate.call_args.args[1].max_tokens == 5

    # Setting server's max_tokens in the generation_config.json
    # higher than context_window - prompt_tokens
    mock_model_config = MockModelConfig()
    mock_model_config.diff_sampling_param = {
        "max_tokens": 200  # Setting server-side max_tokens limit
    }

    # Reinitialize the engine with new settings
409
    mock_engine = MagicMock(spec=AsyncLLM)
410
411
412
413
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False

    # Initialize the serving chat
414
    serving_chat = _build_serving_chat(mock_engine, mock_model_config)
415
416
417
418
419
420
421
422
423
424
425

    # Test case 1: No max_tokens specified, defaults to context_window
    req = ChatCompletionRequest(
        model=MODEL_NAME,
        messages=[{
            "role": "user",
            "content": "what is 1+1?"
        }],
    )

    with suppress(Exception):
426
        await serving_chat.create_chat_completion(req)
427
428
429
430
431
432
433

    assert mock_engine.generate.call_args.args[1].max_tokens == 93

    # Test Case 2: Request's max_tokens set higher than server accepts
    req.max_tokens = 100

    with suppress(Exception):
434
        await serving_chat.create_chat_completion(req)
435
436
437
438
439
440
441

    assert mock_engine.generate.call_args.args[1].max_tokens == 93

    # Test Case 3: Request's max_tokens set lower than server accepts
    req.max_tokens = 5

    with suppress(Exception):
442
        await serving_chat.create_chat_completion(req)
443
444
445

    assert mock_engine.generate.call_args.args[1].max_tokens == 5

446

447
448
@pytest.mark.asyncio
async def test_serving_chat_could_load_correct_generation_config():
449
450
451
452
453
454
455

    mock_model_config = MockModelConfig()
    mock_model_config.diff_sampling_param = {
        "temperature": 0.5,
        "repetition_penalty": 1.05
    }

456
    mock_engine = MagicMock(spec=AsyncLLM)
457
458
459
460
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False

    # Initialize the serving chat
461
    serving_chat = _build_serving_chat(mock_engine, mock_model_config)
462

463
464
465
466
467
468
469
470
471
    req = ChatCompletionRequest(
        model=MODEL_NAME,
        messages=[{
            "role": "user",
            "content": "what is 1+1?"
        }],
    )

    with suppress(Exception):
472
        await serving_chat.create_chat_completion(req)
473
474
475
476
477
478
479
480

    assert mock_engine.generate.call_args.args[1].temperature == 0.5
    assert mock_engine.generate.call_args.args[1].repetition_penalty == 1.05

    # Test the param when user set it
    req.temperature = 0.1

    with suppress(Exception):
481
        await serving_chat.create_chat_completion(req)
482
483
484
485
486
487
488
489

    assert mock_engine.generate.call_args.args[1].temperature == 0.1
    assert mock_engine.generate.call_args.args[1].repetition_penalty == 1.05

    # Test When temperature==0.0
    req.temperature = 0.0

    with suppress(Exception):
490
        await serving_chat.create_chat_completion(req)
491
492
493

    assert mock_engine.generate.call_args.args[1].temperature == 0.0
    assert mock_engine.generate.call_args.args[1].repetition_penalty == 1.05
494
495


496
@pytest.mark.parametrize("model_type", ["gpt_oss", "any"])
497
@pytest.mark.asyncio
498
async def test_serving_chat_did_set_correct_cache_salt(model_type):
499
    mock_model_config = MockModelConfig()
500
    mock_model_config.hf_config.model_type = model_type
501

502
    mock_engine = MagicMock(spec=AsyncLLM)
503
504
505
    mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
    mock_engine.errored = False

506
    serving_chat = _build_serving_chat(mock_engine, mock_model_config)
507
508
509
510
511
512
513
514
515
516

    # Test cache_salt
    req = ChatCompletionRequest(
        model=MODEL_NAME,
        messages=[{
            "role": "user",
            "content": "what is 1+1?"
        }],
    )

517
    # By default, cache_salt in the engine prompt is not set
518
    with suppress(Exception):
519
        await serving_chat.create_chat_completion(req)
520
521
    engine_prompt = serving_chat._process_inputs.await_args_list[0].args[1]
    assert "cache_salt" not in engine_prompt
522
523
524
525

    # Test with certain cache_salt
    req.cache_salt = "test_salt"
    with suppress(Exception):
526
        await serving_chat.create_chat_completion(req)
527
528
    engine_prompt = serving_chat._process_inputs.await_args_list[1].args[1]
    assert engine_prompt.get("cache_salt") == "test_salt"