@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 133 lines 4.01 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import openai # use the official client for correctness check
5 import pytest
6 import pytest_asyncio
7
8 from ...utils import RemoteOpenAIServer
9
10 # any model with a chat template should work here
11 MODEL_NAME = "Qwen/Qwen3-0.6B"
12
13
14 @pytest.fixture(scope="module")
15 def server():
16 args = [
17 # use half precision for speed and memory savings in CI environment
18 "--dtype",
19 "bfloat16",
20 "--max-model-len",
21 "8192",
22 "--enforce-eager",
23 "--max-num-seqs",
24 "128",
25 "--enable-chunked-prefill",
26 "--max-num-batched-tokens",
27 "1000",
28 ]
29
30 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
31 yield remote_server
32
33
34 @pytest_asyncio.fixture
35 async def client(server):
36 async with server.get_async_client() as async_client:
37 yield async_client
38
39
40 @pytest.mark.asyncio
41 async def test_completion_stream_options_and_logprobs_with_long_prompts(
42 client: openai.AsyncOpenAI,
43 ):
44 # Test stream with long prompt
45 prompt = "What is the capital of France?" * 400
46
47 stream = await client.completions.create(
48 model=MODEL_NAME,
49 prompt=prompt,
50 max_tokens=5,
51 temperature=0.0,
52 stream=True,
53 stream_options={
54 "include_usage": True,
55 "continuous_usage_stats": True,
56 },
57 logprobs=5,
58 )
59
60 tokens_received = 0
61 finished = False
62 async for chunk in stream:
63 assert chunk.usage.prompt_tokens >= 0
64 assert chunk.usage.completion_tokens >= 0
65 assert chunk.usage.total_tokens == (
66 chunk.usage.prompt_tokens + chunk.usage.completion_tokens
67 )
68 if not finished:
69 assert chunk.choices[0].text
70 # Count actual tokens from logprobs since multiple tokens
71 # can be batched into a single chunk
72 assert chunk.choices[0].logprobs and chunk.choices[0].logprobs.tokens
73 tokens_received += len(chunk.choices[0].logprobs.tokens)
74
75 if chunk.choices[0].finish_reason is not None:
76 finished = True
77
78 if finished:
79 assert chunk.usage.completion_tokens == tokens_received
80
81
82 @pytest.mark.asyncio
83 async def test_chat_completion_stream_options_and_logprobs_with_long_prompts(
84 client: openai.AsyncOpenAI,
85 ):
86 # Test stream with long prompt
87 messages = [
88 {"role": "system", "content": "You are a helpful assistant."},
89 {"role": "user", "content": "What is the capital of France?" * 400},
90 ]
91 stream = await client.chat.completions.create(
92 model=MODEL_NAME,
93 messages=messages,
94 max_tokens=5,
95 temperature=0.0,
96 stream=True,
97 stream_options={
98 "include_usage": True,
99 "continuous_usage_stats": True,
100 },
101 logprobs=True,
102 top_logprobs=5,
103 )
104
105 tokens_received = 0
106 empty_chunks_received = 0
107 finished = False
108 async for chunk in stream:
109 assert chunk.usage.prompt_tokens >= 0
110 assert chunk.usage.completion_tokens >= 0
111 assert chunk.usage.total_tokens == (
112 chunk.usage.prompt_tokens + chunk.usage.completion_tokens
113 )
114
115 if not finished:
116 if chunk.choices[0].delta.content == "":
117 # when there is no tokens generated
118 assert chunk.usage.completion_tokens == 0
119 assert chunk.choices[0].logprobs is None
120 empty_chunks_received += 1
121 else:
122 # Count actual tokens from logprobs since multiple tokens
123 # can be batched into a single chunk
124 assert chunk.choices[0].logprobs and chunk.choices[0].logprobs.content
125 tokens_received += len(chunk.choices[0].logprobs.content)
126
127 if chunk.choices[0].finish_reason is not None:
128 finished = True
129
130 if finished:
131 assert chunk.usage.completion_tokens == tokens_received
132
133 assert empty_chunks_received <= 1