A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI/Build] Fix flaky entrypoints test (#25663)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Sep 25, 2025 at 20:19 UTC
0bcc3a160d95d778cd53c96984d2b2b50fa68342
1 file changed
+41
-23
tests/entrypoints/openai/test_completion_with_prompt_embeds.py
+41
-23
@@ -15,7 +15,7 @@ from transformers import AutoConfig
15
from ...utils import RemoteOpenAIServer
16
17
# any model with a chat template should work here
18
-MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
18
+MODEL_NAME = "facebook/opt-125m"
19
20
CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
21
@@ -27,7 +27,7 @@ def default_server_args() -> list[str]:
27
"--dtype",
28
"bfloat16",
29
"--max-model-len",
30
- "8192",
30
+ "2048",
31
"--max-num-seqs",
32
"128",
33
"--enforce-eager",
@@ -36,6 +36,27 @@ def default_server_args() -> list[str]:
36
]
37
38
39
+EXAMPLE_PROMPTS = [
40
+ "Hello, my name is",
41
+ "What is an LLM?",
42
+]
43
+
44
+
45
+def _encode_embeds(embeds: torch.Tensor):
46
+ buffer = io.BytesIO()
47
+ torch.save(embeds, buffer)
48
+ return base64.b64encode(buffer.getvalue()).decode('utf-8')
49
+
50
+
51
+@pytest.fixture(scope="module")
52
+def example_prompt_embeds(hf_runner):
53
+ """Create example embeddings and return them as base64 encoded string."""
54
+ with hf_runner(MODEL_NAME) as hf_model:
55
+ example_embeddings = hf_model.get_prompt_embeddings(EXAMPLE_PROMPTS)
56
+
57
+ return [_encode_embeds(item) for item in example_embeddings]
58
+
59
+
60
@pytest.fixture(scope="module",
61
params=["", "--disable-frontend-multiprocessing"])
62
def server_with_prompt_embeds(default_server_args, request):
@@ -52,21 +73,16 @@ async def client_with_prompt_embeds(server_with_prompt_embeds):
73
yield async_client
74
75
55
-def create_dummy_embeds(num_tokens: int = 5) -> str:
56
- """Create dummy embeddings and return them as base64 encoded string."""
57
- dummy_embeds = torch.randn(num_tokens, CONFIG.hidden_size)
58
- buffer = io.BytesIO()
59
- torch.save(dummy_embeds, buffer)
60
- return base64.b64encode(buffer.getvalue()).decode('utf-8')
61
-
62
-
63
-@pytest.mark.skip("This test is skipped because it is flaky.")
76
@pytest.mark.asyncio
77
@pytest.mark.parametrize("model_name", [MODEL_NAME])
78
async def test_completions_with_prompt_embeds(
67
- client_with_prompt_embeds: openai.AsyncOpenAI, model_name: str):
79
+ example_prompt_embeds,
80
+ client_with_prompt_embeds: openai.AsyncOpenAI,
81
+ model_name: str,
82
+):
83
+ encoded_embeds, encoded_embeds2 = example_prompt_embeds
84
+
85
# Test case: Single prompt embeds input
69
- encoded_embeds = create_dummy_embeds()
86
completion = await client_with_prompt_embeds.completions.create(
87
model=model_name,
88
prompt="", # Add empty prompt as required parameter
@@ -77,7 +93,6 @@ async def test_completions_with_prompt_embeds(
93
assert completion.choices[0].prompt_logprobs is None
94
95
# Test case: batch completion with prompt_embeds
80
- encoded_embeds2 = create_dummy_embeds()
96
completion = await client_with_prompt_embeds.completions.create(
97
model=model_name,
98
prompt="", # Add empty prompt as required parameter
@@ -89,7 +104,6 @@ async def test_completions_with_prompt_embeds(
104
assert len(completion.choices[1].text) >= 1
105
106
# Test case: streaming with prompt_embeds
92
- encoded_embeds = create_dummy_embeds()
107
single_completion = await client_with_prompt_embeds.completions.create(
108
model=model_name,
109
prompt="", # Add empty prompt as required parameter
@@ -117,7 +131,6 @@ async def test_completions_with_prompt_embeds(
131
assert "".join(chunks) == single_output
132
133
# Test case: batch streaming with prompt_embeds
120
- encoded_embeds2 = create_dummy_embeds()
134
stream = await client_with_prompt_embeds.completions.create(
135
model=model_name,
136
prompt="", # Add empty prompt as required parameter
@@ -139,7 +152,6 @@ async def test_completions_with_prompt_embeds(
152
assert len(chunks_stream_embeds[1]) > 0
153
154
# Test case: mixed text and prompt_embeds
142
- encoded_embeds = create_dummy_embeds()
155
completion_mixed = await client_with_prompt_embeds.completions.create(
156
model=model_name,
157
prompt="This is a prompt",
@@ -184,10 +196,14 @@ async def test_completions_errors_with_prompt_embeds(
196
@pytest.mark.parametrize("logprobs_arg", [1, 0])
197
@pytest.mark.parametrize("model_name", [MODEL_NAME])
198
async def test_completions_with_logprobs_and_prompt_embeds(
187
- client_with_prompt_embeds: openai.AsyncOpenAI, logprobs_arg: int,
188
- model_name: str):
199
+ example_prompt_embeds,
200
+ client_with_prompt_embeds: openai.AsyncOpenAI,
201
+ logprobs_arg: int,
202
+ model_name: str,
203
+):
204
+ encoded_embeds, encoded_embeds2 = example_prompt_embeds
205
+
206
# Test case: Logprobs using prompt_embeds
190
- encoded_embeds = create_dummy_embeds()
207
completion = await client_with_prompt_embeds.completions.create(
208
model=model_name,
209
prompt="", # Add empty prompt as required parameter
@@ -207,7 +223,6 @@ async def test_completions_with_logprobs_and_prompt_embeds(
223
assert len(logprobs.tokens) == 5
224
225
# Test case: Log probs with batch completion and prompt_embeds
210
- encoded_embeds2 = create_dummy_embeds()
226
completion = await client_with_prompt_embeds.completions.create(
227
model=model_name,
228
prompt="", # Add empty prompt as required parameter
@@ -232,9 +247,12 @@ async def test_completions_with_logprobs_and_prompt_embeds(
247
248
@pytest.mark.asyncio
249
async def test_prompt_logprobs_raises_error(
235
- client_with_prompt_embeds: openai.AsyncOpenAI):
250
+ example_prompt_embeds,
251
+ client_with_prompt_embeds: openai.AsyncOpenAI,
252
+):
253
+ encoded_embeds, _ = example_prompt_embeds
254
+
255
with pytest.raises(BadRequestError, match="not compatible"):
237
- encoded_embeds = create_dummy_embeds()
256
await client_with_prompt_embeds.completions.create(
257
model=MODEL_NAME,
258
prompt="",