@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 126 lines 3.57 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import openai
5 import pytest
6 import pytest_asyncio
7
8 from tests.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 monkeypatch_module():
16 from _pytest.monkeypatch import MonkeyPatch
17
18 mpatch = MonkeyPatch()
19 yield mpatch
20 mpatch.undo()
21
22
23 @pytest.fixture(scope="module", params=[True])
24 def server(request, monkeypatch_module):
25 args = [
26 "--dtype",
27 "bfloat16",
28 "--max-model-len",
29 "8192",
30 "--enforce-eager",
31 ]
32
33 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
34 yield remote_server
35
36
37 @pytest_asyncio.fixture
38 async def client(server):
39 async with server.get_async_client() as async_client:
40 yield async_client
41
42
43 @pytest.mark.asyncio
44 async def test_chat_completion_with_orca_header(server: RemoteOpenAIServer):
45 messages = [
46 {"role": "system", "content": "you are a helpful assistant"},
47 {"role": "user", "content": "what is 1+1?"},
48 ]
49
50 client = openai.OpenAI(
51 api_key="EMPTY",
52 base_url=f"http://localhost:{server.port}/v1",
53 default_headers={"endpoint-load-metrics-format": "TEXT"},
54 )
55
56 # 1. Use raw client to get response headers.
57 raw_client = client.with_raw_response
58
59 # 2. Make the API call using the raw_client
60 response_with_raw = raw_client.chat.completions.create(
61 model=MODEL_NAME,
62 messages=messages,
63 extra_headers={"endpoint-load-metrics-format": "TEXT"},
64 )
65
66 # 3. Access the raw httpx.Response object
67 raw_http_response = response_with_raw.http_response
68
69 # 4. Get the headers from the httpx.Response object
70 response_headers = raw_http_response.headers
71
72 assert "endpoint-load-metrics" in response_headers
73
74
75 @pytest.mark.asyncio
76 async def test_completion_with_orca_header(client: openai.AsyncOpenAI):
77 # 1. Use raw client to get response headers.
78 raw_client = client.with_raw_response
79
80 # 2. Make the API call using the raw_client
81 completion = await raw_client.completions.create(
82 model=MODEL_NAME,
83 prompt="Hello, my name is",
84 max_tokens=5,
85 extra_headers={"endpoint-load-metrics-format": "JSON"},
86 )
87
88 # 3. Access the raw httpx.Response object
89 raw_http_response = completion.http_response
90
91 # 4. Get the headers from the httpx.Response object
92 response_headers = raw_http_response.headers
93
94 assert "endpoint-load-metrics" in response_headers
95
96
97 @pytest.mark.asyncio
98 async def test_single_completion(client: openai.AsyncOpenAI):
99 completion = await client.completions.create(
100 model=MODEL_NAME,
101 prompt="Hello, my name is",
102 max_tokens=5,
103 extra_headers={"endpoint-load-metrics-format": "JSON"},
104 temperature=0.0,
105 )
106
107 assert completion.id is not None
108 assert completion.choices is not None and len(completion.choices) == 1
109
110 choice = completion.choices[0]
111 assert len(choice.text) >= 5
112 assert choice.finish_reason == "length"
113 # When using Qwen3-0.6B, prompt tokens=[9707, 11, 847, 829, 374]
114 assert completion.usage == openai.types.CompletionUsage(
115 completion_tokens=5, prompt_tokens=5, total_tokens=10
116 )
117
118 # test using token IDs
119 completion = await client.completions.create(
120 model=MODEL_NAME,
121 prompt=[0, 0, 0, 0, 0],
122 max_tokens=5,
123 temperature=0.0,
124 )
125 assert len(completion.choices[0].text) >= 1
126 assert completion.choices[0].prompt_logprobs is None