@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 304 lines 9.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import io
5 import json
6
7 import openai # use the official client for correctness check
8 import pybase64 as base64
9 import pytest
10 import pytest_asyncio
11 import torch
12
13 # downloading lora to test lora requests
14 from openai import BadRequestError
15 from transformers import AutoConfig
16
17 from tests.utils import RemoteOpenAIServer
18
19 # any model with a chat template should work here
20 MODEL_NAME = "facebook/opt-125m"
21 LORA_SERVING_MODEL_NAME = "opt125m-lora"
22
23 CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
24
25
26 @pytest.fixture(scope="module", params=["use-lora"])
27 def default_server_args(
28 request: pytest.FixtureRequest, opt125_lora_files: str
29 ) -> list[str]:
30 args = [
31 # use half precision for speed and memory savings in CI environment
32 "--dtype",
33 "bfloat16",
34 "--max-model-len",
35 "2048",
36 "--max-num-seqs",
37 "128",
38 "--enforce-eager",
39 # Prompt Embeds server args
40 "--enable-prompt-embeds",
41 ]
42
43 if request.param == "use-lora":
44 lora_module_1 = {
45 "name": LORA_SERVING_MODEL_NAME,
46 "path": opt125_lora_files,
47 "base_model_name": MODEL_NAME,
48 }
49
50 args.extend(
51 [
52 "--enable-lora",
53 "--lora-module",
54 json.dumps(lora_module_1),
55 "--max-lora-rank",
56 "64",
57 "--max-cpu-loras",
58 "2",
59 ]
60 )
61
62 return args
63
64
65 EXAMPLE_PROMPTS = [
66 "Hello, my name is",
67 "What is an LLM?",
68 ]
69
70
71 def _encode_embeds(embeds: torch.Tensor):
72 buffer = io.BytesIO()
73 torch.save(embeds, buffer)
74 return base64.b64encode(buffer.getvalue()).decode("utf-8")
75
76
77 @pytest.fixture(scope="module")
78 def example_prompt_embeds(hf_runner):
79 """Create example embeddings and return them as base64 encoded string."""
80 with hf_runner(MODEL_NAME) as hf_model:
81 example_embeddings = hf_model.get_prompt_embeddings(EXAMPLE_PROMPTS)
82
83 return [_encode_embeds(item) for item in example_embeddings]
84
85
86 @pytest.fixture(scope="module")
87 def server_with_prompt_embeds(default_server_args):
88 with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
89 yield remote_server
90
91
92 @pytest_asyncio.fixture
93 async def client_with_prompt_embeds(server_with_prompt_embeds):
94 async with server_with_prompt_embeds.get_async_client() as async_client:
95 yield async_client
96
97
98 @pytest.mark.asyncio
99 @pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
100 async def test_completions_with_prompt_embeds(
101 example_prompt_embeds,
102 client_with_prompt_embeds: openai.AsyncOpenAI,
103 model_name: str,
104 ):
105 encoded_embeds, encoded_embeds2 = example_prompt_embeds
106
107 # Test case: Single prompt embeds input
108 completion = await client_with_prompt_embeds.completions.create(
109 model=model_name,
110 prompt=None,
111 max_tokens=5,
112 temperature=0.0,
113 extra_body={"prompt_embeds": encoded_embeds},
114 )
115 assert len(completion.choices[0].text) >= 1
116 assert completion.choices[0].prompt_logprobs is None
117
118 # Test case: batch completion with prompt_embeds
119 completion = await client_with_prompt_embeds.completions.create(
120 model=model_name,
121 prompt=None,
122 max_tokens=5,
123 temperature=0.0,
124 extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]},
125 )
126 assert len(completion.choices) == 2
127 assert len(completion.choices[0].text) >= 1
128 assert len(completion.choices[1].text) >= 1
129
130 # Test case: streaming with prompt_embeds
131 single_completion = await client_with_prompt_embeds.completions.create(
132 model=model_name,
133 prompt=None,
134 max_tokens=5,
135 temperature=0.0,
136 extra_body={"prompt_embeds": encoded_embeds},
137 )
138 single_output = single_completion.choices[0].text
139
140 stream = await client_with_prompt_embeds.completions.create(
141 model=model_name,
142 prompt=None,
143 max_tokens=5,
144 temperature=0.0,
145 stream=True,
146 extra_body={"prompt_embeds": encoded_embeds},
147 )
148 chunks = []
149 finish_reason_count = 0
150 async for chunk in stream:
151 chunks.append(chunk.choices[0].text)
152 if chunk.choices[0].finish_reason is not None:
153 finish_reason_count += 1
154 assert finish_reason_count == 1
155 assert chunk.choices[0].finish_reason == "length"
156 assert chunk.choices[0].text
157 assert "".join(chunks) == single_output
158
159 # Test case: batch streaming with prompt_embeds
160 stream = await client_with_prompt_embeds.completions.create(
161 model=model_name,
162 prompt=None,
163 max_tokens=5,
164 temperature=0.0,
165 stream=True,
166 extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]},
167 )
168 chunks_stream_embeds: list[list[str]] = [[], []]
169 finish_reason_count = 0
170 async for chunk in stream:
171 chunks_stream_embeds[chunk.choices[0].index].append(chunk.choices[0].text)
172 if chunk.choices[0].finish_reason is not None:
173 finish_reason_count += 1
174 assert finish_reason_count == 2
175 assert chunk.choices[0].finish_reason == "length"
176 assert chunk.choices[0].text
177 assert len(chunks_stream_embeds[0]) > 0
178 assert len(chunks_stream_embeds[1]) > 0
179
180 # Test case: mixed text and prompt_embeds
181 completion_mixed = await client_with_prompt_embeds.completions.create(
182 model=model_name,
183 prompt="This is a prompt",
184 max_tokens=5,
185 temperature=0.0,
186 extra_body={"prompt_embeds": encoded_embeds},
187 )
188 assert len(completion.choices) == 2
189 completion_text_only = await client_with_prompt_embeds.completions.create(
190 model=model_name,
191 prompt="This is a prompt",
192 max_tokens=5,
193 temperature=0.0,
194 )
195 completion_embeds_only = await client_with_prompt_embeds.completions.create(
196 model=model_name,
197 prompt=None,
198 max_tokens=5,
199 temperature=0.0,
200 extra_body={"prompt_embeds": encoded_embeds},
201 )
202 # Embeddings responses should be handled first
203 assert completion_mixed.choices[0].text == completion_embeds_only.choices[0].text
204 assert completion_mixed.choices[1].text == completion_text_only.choices[0].text
205
206
207 @pytest.mark.asyncio
208 @pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
209 async def test_completions_errors_with_prompt_embeds(
210 client_with_prompt_embeds: openai.AsyncOpenAI, model_name: str
211 ):
212 # Test error case: invalid prompt_embeds
213 with pytest.raises(BadRequestError):
214 await client_with_prompt_embeds.completions.create(
215 prompt=None,
216 model=model_name,
217 max_tokens=5,
218 temperature=0.0,
219 extra_body={"prompt_embeds": "invalid_base64"},
220 )
221
222
223 @pytest.mark.asyncio
224 @pytest.mark.parametrize("logprobs_arg", [1, 0])
225 @pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
226 async def test_completions_with_logprobs_and_prompt_embeds(
227 example_prompt_embeds,
228 client_with_prompt_embeds: openai.AsyncOpenAI,
229 logprobs_arg: int,
230 model_name: str,
231 ):
232 encoded_embeds, encoded_embeds2 = example_prompt_embeds
233
234 # Test case: Logprobs using prompt_embeds
235 completion = await client_with_prompt_embeds.completions.create(
236 model=model_name,
237 prompt=None,
238 max_tokens=5,
239 temperature=0.0,
240 echo=False,
241 logprobs=logprobs_arg,
242 extra_body={"prompt_embeds": encoded_embeds},
243 )
244
245 logprobs = completion.choices[0].logprobs
246 assert logprobs is not None
247 assert len(logprobs.text_offset) == 5
248 assert len(logprobs.token_logprobs) == 5
249 assert len(logprobs.top_logprobs) == 5
250 for top_logprobs in logprobs.top_logprobs[1:]:
251 assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
252 assert len(logprobs.tokens) == 5
253
254 # Test case: Log probs with batch completion and prompt_embeds
255 completion = await client_with_prompt_embeds.completions.create(
256 model=model_name,
257 prompt=None,
258 max_tokens=5,
259 temperature=0.0,
260 echo=False,
261 logprobs=logprobs_arg,
262 extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]},
263 )
264
265 assert len(completion.choices) == 2
266 for choice in completion.choices:
267 logprobs = choice.logprobs
268 assert logprobs is not None
269 assert len(logprobs.text_offset) == 5
270 assert len(logprobs.token_logprobs) == 5
271 assert len(logprobs.top_logprobs) == 5
272 for top_logprobs in logprobs.top_logprobs[1:]:
273 assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
274 assert len(logprobs.tokens) == 5
275
276
277 @pytest.mark.asyncio
278 async def test_prompt_logprobs_raises_error(
279 example_prompt_embeds,
280 client_with_prompt_embeds: openai.AsyncOpenAI,
281 ):
282 encoded_embeds, _ = example_prompt_embeds
283
284 with pytest.raises(BadRequestError, match="not compatible"):
285 await client_with_prompt_embeds.completions.create(
286 model=MODEL_NAME,
287 prompt=None,
288 max_tokens=5,
289 temperature=0.0,
290 extra_body={"prompt_embeds": encoded_embeds, "prompt_logprobs": True},
291 )
292
293
294 @pytest.mark.asyncio
295 async def test_empty_prompt_embeds(
296 client_with_prompt_embeds: openai.AsyncOpenAI,
297 ) -> None:
298 await client_with_prompt_embeds.completions.create(
299 model=MODEL_NAME,
300 prompt="Hello",
301 max_tokens=5,
302 temperature=0.0,
303 extra_body={"prompt_embeds": []},
304 )