@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 179 lines 6.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from argparse import Namespace
5
6 import pytest
7
8 from vllm.entrypoints.generate.base.protocol import StreamOptions
9 from vllm.entrypoints.serve.utils import api_utils
10 from vllm.entrypoints.serve.utils.api_utils import (
11 get_max_tokens,
12 redact_sensitive_args,
13 should_include_usage,
14 )
15
16
17 @pytest.mark.parametrize(
18 ("stream_options", "expected"),
19 [
20 (None, (True, True)),
21 (StreamOptions(include_usage=False), (True, True)),
22 (
23 StreamOptions(include_usage=False, continuous_usage_stats=False),
24 (True, True),
25 ),
26 (
27 StreamOptions(include_usage=True, continuous_usage_stats=False),
28 (True, True),
29 ),
30 ],
31 )
32 def test_should_include_usage_force_enables_continuous_usage(stream_options, expected):
33 assert should_include_usage(stream_options, True) == expected
34
35
36 class TestGetMaxTokens:
37 """Tests for get_max_tokens() to ensure generation_config's max_tokens
38 acts as a default when from model author, and as a ceiling when
39 explicitly set by the user."""
40
41 def test_default_sampling_params_used_when_no_request_max_tokens(self):
42 """When user doesn't specify max_tokens, generation_config default
43 should apply."""
44 result = get_max_tokens(
45 max_model_len=24000,
46 max_tokens=None,
47 input_length=100,
48 default_sampling_params={"max_tokens": 2048},
49 )
50 assert result == 2048
51
52 def test_request_max_tokens_not_capped_by_default_sampling_params(self):
53 """When user specifies max_tokens in request, model author's
54 generation_config max_tokens must NOT cap it (fixes #34005)."""
55 result = get_max_tokens(
56 max_model_len=24000,
57 max_tokens=5000,
58 input_length=100,
59 default_sampling_params={"max_tokens": 2048},
60 )
61 assert result == 5000
62
63 def test_override_max_tokens_caps_request(self):
64 """When user explicitly sets max_tokens, it acts as a ceiling."""
65 result = get_max_tokens(
66 max_model_len=24000,
67 max_tokens=5000,
68 input_length=100,
69 default_sampling_params={"max_tokens": 2048},
70 override_max_tokens=2048,
71 )
72 assert result == 2048
73
74 def test_override_max_tokens_used_as_default(self):
75 """When no request max_tokens, override still applies as default."""
76 result = get_max_tokens(
77 max_model_len=24000,
78 max_tokens=None,
79 input_length=100,
80 default_sampling_params={"max_tokens": 2048},
81 override_max_tokens=2048,
82 )
83 assert result == 2048
84
85 def test_max_model_len_still_caps_output(self):
86 """max_model_len - input_length is always the hard ceiling."""
87 result = get_max_tokens(
88 max_model_len=3000,
89 max_tokens=5000,
90 input_length=100,
91 default_sampling_params={"max_tokens": 2048},
92 )
93 assert result == 2900 # 3000 - 100
94
95 def test_request_max_tokens_smaller_than_default(self):
96 """When user explicitly requests fewer tokens than gen_config default,
97 that should be respected."""
98 result = get_max_tokens(
99 max_model_len=24000,
100 max_tokens=512,
101 input_length=100,
102 default_sampling_params={"max_tokens": 2048},
103 )
104 assert result == 512
105
106 def test_input_length_exceeds_max_model_len(self):
107 with pytest.raises(
108 ValueError,
109 match="Input length .* exceeds model's maximum context length .*",
110 ):
111 get_max_tokens(
112 max_model_len=100,
113 max_tokens=50,
114 input_length=150,
115 default_sampling_params={"max_tokens": 2048},
116 )
117
118
119 class TestRedactSensitiveArgs:
120 API_KEY = "sk-test-secret-12345"
121
122 def test_redact_replaces_sensitive_values_only(self):
123 args = {
124 "api_key": self.API_KEY,
125 "hf_token": "hf_secret",
126 "watermark_config": {"algorithm": "gumbel", "key": 42},
127 "other": "visible",
128 }
129 redacted = redact_sensitive_args(args)
130 assert redacted == {
131 "api_key": "***",
132 "hf_token": "***",
133 "watermark_config": "***",
134 "other": "visible",
135 }
136 # original dict must not be mutated
137 assert args == {
138 "api_key": self.API_KEY,
139 "hf_token": "hf_secret",
140 "watermark_config": {"algorithm": "gumbel", "key": 42},
141 "other": "visible",
142 }
143
144 def test_no_sensitive_fields_returns_original(self):
145 args = {"model_tag": "org/model", "other": "visible"}
146 assert redact_sensitive_args(args) is args
147
148 def test_api_key_not_in_log(self, monkeypatch, caplog):
149 non_default = {
150 "model_tag": "org/model",
151 "default_chat_template_kwargs": {"enable_thinking": False},
152 "api_key": self.API_KEY,
153 "enable_auto_tool_choice": True,
154 "tool_call_parser": "qwen3_coder",
155 }
156 monkeypatch.setattr(api_utils, "get_non_default_args", lambda args: non_default)
157 with caplog.at_level("INFO", logger="vllm.entrypoints.serve.utils.api_utils"):
158 api_utils.log_non_default_args(args=Namespace())
159 message = caplog.text
160 assert self.API_KEY not in message
161 assert "'api_key': '***'" in message
162 # non-sensitive args are still logged
163 assert "org/model" in message
164 assert "qwen3_coder" in message
165
166 def test_hf_token_not_in_log(self, monkeypatch, caplog):
167 hf_token = "hf_TESTTOKENVALUE123"
168 non_default = {
169 "model_tag": "org/model",
170 "hf_token": hf_token,
171 "tool_call_parser": "qwen3_coder",
172 }
173 monkeypatch.setattr(api_utils, "get_non_default_args", lambda args: non_default)
174 with caplog.at_level("INFO", logger="vllm.entrypoints.serve.utils.api_utils"):
175 api_utils.log_non_default_args(args=Namespace())
176 message = caplog.text
177 assert hf_token not in message
178 assert "'hf_token': '***'" in message
179 assert "org/model" in message