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
[PD] add test for chat completions endpoint (#21925)
Signed-off-by: Abirdcfly <fp544037857@gmail.com>
Abirdcfly committed
Aug 4, 2025 at 10:57 UTC
0d7db16a92afd9fc005ed0fba73356845586f5e7
2 files changed
+29
-14
tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py
+27
-14
@@ -51,20 +51,31 @@ def check_vllm_server(url: str, timeout=5, retries=3) -> bool:
51
return False
52
53
54
-def run_simple_prompt(base_url: str, model_name: str,
55
- input_prompt: str) -> str:
54
+def run_simple_prompt(base_url: str, model_name: str, input_prompt: str,
55
+ use_chat_endpoint: bool) -> str:
56
client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
57
- completion = client.completions.create(model=model_name,
58
- prompt=input_prompt,
59
- max_tokens=MAX_OUTPUT_LEN,
60
- temperature=0.0,
61
- seed=42)
57
+ if use_chat_endpoint:
58
+ completion = client.chat.completions.create(
59
+ model=model_name,
60
+ messages=[{
61
+ "role": "user",
62
+ "content": [{
63
+ "type": "text",
64
+ "text": input_prompt
65
+ }]
66
+ }],
67
+ max_completion_tokens=MAX_OUTPUT_LEN,
68
+ temperature=0.0,
69
+ seed=42)
70
+ return completion.choices[0].message.content
71
+ else:
72
+ completion = client.completions.create(model=model_name,
73
+ prompt=input_prompt,
74
+ max_tokens=MAX_OUTPUT_LEN,
75
+ temperature=0.0,
76
+ seed=42)
77
63
- # print("-" * 50)
64
- # print(f"Completion results for {model_name}:")
65
- # print(completion)
66
- # print("-" * 50)
67
- return completion.choices[0].text
78
+ return completion.choices[0].text
79
80
81
def main():
@@ -125,10 +136,12 @@ def main():
136
f"vllm server: {args.service_url} is not ready yet!")
137
138
output_strs = dict()
128
- for prompt in SAMPLE_PROMPTS:
139
+ for i, prompt in enumerate(SAMPLE_PROMPTS):
140
+ use_chat_endpoint = (i % 2 == 1)
141
output_str = run_simple_prompt(base_url=service_url,
142
model_name=args.model_name,
131
- input_prompt=prompt)
143
+ input_prompt=prompt,
144
+ use_chat_endpoint=use_chat_endpoint)
145
print(f"Prompt: {prompt}, output: {output_str}")
146
output_strs[prompt] = output_str
147
tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
+2
@@ -162,6 +162,8 @@ async def send_request_to_service(client_info: dict, endpoint: str,
162
}
163
req_data["stream"] = False
164
req_data["max_tokens"] = 1
165
+ if "max_completion_tokens" in req_data:
166
+ req_data["max_completion_tokens"] = 1
167
if "stream_options" in req_data:
168
del req_data["stream_options"]
169
headers = {