@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 175 lines 5.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import asyncio
4 import os
5
6 import openai # use the official client for correctness check
7 import pytest
8 import pytest_asyncio
9
10 from tests.utils import RemoteOpenAIServer
11 from tests.v1.utils import check_request_balancing
12
13 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
14
15 DP_SIZE = os.getenv("DP_SIZE", "1")
16
17
18 @pytest.fixture(scope="module")
19 def default_server_args():
20 return [
21 # use half precision for speed and memory savings in CI environment
22 "--dtype",
23 "bfloat16",
24 "--max-model-len",
25 "2048",
26 "--max-num-seqs",
27 "128",
28 "--enforce-eager",
29 "--api-server-count",
30 "4",
31 "--data_parallel_size",
32 DP_SIZE,
33 ]
34
35
36 @pytest.fixture(scope="module")
37 def server(default_server_args):
38 with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
39 yield remote_server
40
41
42 @pytest_asyncio.fixture
43 async def client(server):
44 async with server.get_async_client() as async_client:
45 yield async_client
46
47
48 @pytest.mark.asyncio
49 @pytest.mark.parametrize(
50 "model_name",
51 [MODEL_NAME],
52 )
53 async def test_single_completion(
54 client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model_name: str
55 ) -> None:
56 async def make_request():
57 completion = await client.completions.create(
58 model=model_name, prompt="Hello, my name is", max_tokens=10, temperature=1.0
59 )
60
61 assert completion.id is not None
62 assert completion.choices is not None and len(completion.choices) == 1
63
64 choice = completion.choices[0]
65 # The exact number of tokens can vary slightly with temperature=1.0,
66 # so we check for a reasonable minimum length.
67 assert len(choice.text) >= 1
68 # Finish reason might not always be 'length' if the model finishes early
69 # or due to other reasons, especially with high temperature.
70 # So, we'll accept 'length' or 'stop'.
71 assert choice.finish_reason in ("length", "stop")
72
73 # Token counts can also vary, so we check they are positive.
74 assert completion.usage.completion_tokens > 0
75 assert completion.usage.prompt_tokens > 0
76 assert completion.usage.total_tokens > 0
77 return completion
78
79 # Test single request
80 result = await make_request()
81 assert result is not None
82
83 await asyncio.sleep(0.5)
84
85 # Send two bursts of requests
86 num_requests = 100
87 tasks = [make_request() for _ in range(num_requests)]
88 results = await asyncio.gather(*tasks)
89 assert len(results) == num_requests
90 assert all(completion is not None for completion in results)
91
92 await asyncio.sleep(0.5)
93
94 tasks = [make_request() for _ in range(num_requests)]
95 results = await asyncio.gather(*tasks)
96 assert len(results) == num_requests
97 assert all(completion is not None for completion in results)
98
99 # Check request balancing via Prometheus metrics if DP_SIZE > 1
100 check_request_balancing(server, int(DP_SIZE))
101
102
103 @pytest.mark.asyncio
104 @pytest.mark.parametrize(
105 "model_name",
106 [MODEL_NAME],
107 )
108 async def test_completion_streaming(
109 client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model_name: str
110 ) -> None:
111 prompt = "What is an LLM?"
112
113 async def make_streaming_request():
114 # Perform a non-streaming request to get the expected full output
115 single_completion = await client.completions.create(
116 model=model_name,
117 prompt=prompt,
118 max_tokens=5,
119 temperature=0.0,
120 )
121 single_output = single_completion.choices[0].text
122
123 # Perform the streaming request
124 stream = await client.completions.create(
125 model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
126 )
127 chunks: list[str] = []
128 finish_reason_count = 0
129 last_chunk = None
130 async for chunk in stream:
131 chunks.append(chunk.choices[0].text)
132 if chunk.choices[0].finish_reason is not None:
133 finish_reason_count += 1
134 last_chunk = chunk # Keep track of the last chunk
135
136 # finish reason should only return in the last block for OpenAI API
137 assert finish_reason_count == 1, "Finish reason should appear exactly once."
138 assert last_chunk is not None, "Stream should have yielded at least one chunk."
139 assert last_chunk.choices[0].finish_reason == "length", (
140 "Finish reason should be 'length'."
141 )
142 # Check that the combined text matches the non-streamed version.
143 assert "".join(chunks) == single_output, (
144 "Streamed output should match non-streamed output."
145 )
146 return True # Indicate success for this request
147
148 # Test single request
149 result = await make_streaming_request()
150 assert result is not None
151
152 await asyncio.sleep(0.5)
153
154 # Send two bursts of requests
155 num_requests = 100
156 tasks = [make_streaming_request() for _ in range(num_requests)]
157 results = await asyncio.gather(*tasks)
158
159 assert len(results) == num_requests, (
160 f"Expected {num_requests} results, got {len(results)}"
161 )
162 assert all(results), "Not all streaming requests completed successfully."
163
164 await asyncio.sleep(0.5)
165
166 tasks = [make_streaming_request() for _ in range(num_requests)]
167 results = await asyncio.gather(*tasks)
168
169 assert len(results) == num_requests, (
170 f"Expected {num_requests} results, got {len(results)}"
171 )
172 assert all(results), "Not all streaming requests completed successfully."
173
174 # Check request balancing via Prometheus metrics if DP_SIZE > 1
175 check_request_balancing(server, int(DP_SIZE))