@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 229 lines 8.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from argparse import ArgumentError
5 from unittest.mock import patch
6
7 import pytest
8
9 from vllm.config import ModelConfig
10 from vllm.engine.arg_utils import EngineArgs
11 from vllm.usage.usage_lib import UsageContext
12 from vllm.utils.argparse_utils import FlexibleArgumentParser
13 from vllm.utils.hashing import _xxhash
14
15
16 def test_prefix_caching_from_cli():
17 parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
18 args = parser.parse_args([])
19 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
20 assert vllm_config.cache_config.enable_prefix_caching, (
21 "V1 turns on prefix caching by default."
22 )
23 assert vllm_config.cache_config.prefix_cache_retention_interval == 0
24
25 # Turn it off possible with flag.
26 args = parser.parse_args(["--no-enable-prefix-caching"])
27 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
28 assert not vllm_config.cache_config.enable_prefix_caching
29
30 # Turn it on with flag.
31 args = parser.parse_args(["--enable-prefix-caching"])
32 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
33 assert vllm_config.cache_config.enable_prefix_caching
34
35 # default hash algorithm is "builtin"
36 assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
37
38 # set hash algorithm to sha256_cbor
39 args = parser.parse_args(["--prefix-caching-hash-algo", "sha256_cbor"])
40 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
41 assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256_cbor"
42
43 # set hash algorithm to sha256
44 args = parser.parse_args(["--prefix-caching-hash-algo", "sha256"])
45 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
46 assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
47
48 # an invalid hash algorithm raises an error
49 parser.exit_on_error = False
50 with pytest.raises(ArgumentError):
51 args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
52
53 args = parser.parse_args(["--prefix-cache-retention-interval", "64"])
54 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
55 assert vllm_config.cache_config.prefix_cache_retention_interval == 64
56
57
58 @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
59 def test_prefix_caching_xxhash_from_cli():
60 parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
61
62 # set hash algorithm to xxhash (pickle)
63 args = parser.parse_args(["--prefix-caching-hash-algo", "xxhash"])
64 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
65 assert vllm_config.cache_config.prefix_caching_hash_algo == "xxhash"
66
67 # set hash algorithm to xxhash_cbor
68 args = parser.parse_args(["--prefix-caching-hash-algo", "xxhash_cbor"])
69 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
70 assert vllm_config.cache_config.prefix_caching_hash_algo == "xxhash_cbor"
71
72
73 def test_mm_prefix_lm_raises_batched_tokens_floor():
74 """Verify that prefix-LM multimodal models auto-raise
75 max_num_batched_tokens to fit at least one multimodal item.
76
77 Regression test for https://github.com/vllm-project/vllm/issues/42687
78 """
79 from unittest.mock import patch
80
81 # Simulate a prefix-LM multimodal model whose largest modality
82 # (video) requires 2496 tokens — more than the 2048 default.
83 fake_mm_min = (2496, "video")
84
85 engine_args = EngineArgs(
86 model="facebook/opt-125m",
87 max_model_len=2048,
88 enforce_eager=True,
89 )
90
91 with (
92 patch.object(
93 type(engine_args),
94 "_get_min_mm_batched_tokens",
95 staticmethod(lambda _mc: fake_mm_min),
96 ),
97 patch(
98 "vllm.config.ModelConfig.is_multimodal_model",
99 new_callable=lambda: property(lambda self: True),
100 ),
101 patch(
102 "vllm.config.ModelConfig.is_mm_prefix_lm",
103 new_callable=lambda: property(lambda self: True),
104 ),
105 ):
106 vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
107
108 assert vllm_config.scheduler_config.max_num_batched_tokens >= 2496
109
110
111 def test_custom_histogram_buckets_from_cli():
112 parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
113
114 # Unset: byte-identical default behavior.
115 args = parser.parse_args([])
116 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
117 assert vllm_config.observability_config.custom_histogram_buckets is None
118
119 # A JSON mapping round-trips into the config verbatim.
120 args = parser.parse_args(
121 [
122 "--custom-histogram-buckets",
123 '{"request_latency": [0.01, 0.05, 0.1, 0.5]}',
124 ]
125 )
126 vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
127 assert vllm_config.observability_config.custom_histogram_buckets == {
128 "request_latency": [0.01, 0.05, 0.1, 0.5]
129 }
130
131 # Malformed JSON is rejected at argument-parsing time.
132 parser.exit_on_error = False
133 with pytest.raises(ArgumentError):
134 parser.parse_args(["--custom-histogram-buckets", "{not json"])
135
136 # An unknown family key parses as JSON but fails config validation.
137 args = parser.parse_args(["--custom-histogram-buckets", '{"bogus": [1.0, 2.0]}'])
138 with pytest.raises(ValueError, match="unknown bucket family"):
139 EngineArgs.from_cli_args(args=args).create_engine_config()
140
141
142 def test_data_parallel_start_rank_zero_infers_hybrid_lb():
143 """An explicit --data-parallel-start-rank 0 must be treated the same as
144 any other explicit start rank when inferring hybrid LB mode, not as
145 "unset" (regression test for a truthiness-vs-`is not None` bug).
146 """
147 engine_args = EngineArgs(
148 model="facebook/opt-125m",
149 data_parallel_size=4,
150 data_parallel_size_local=2,
151 data_parallel_start_rank=0,
152 )
153 vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
154
155 assert vllm_config.parallel_config.data_parallel_hybrid_lb is True
156 assert vllm_config.parallel_config.data_parallel_rank == 0
157
158
159 def test_external_lb_preserves_explicit_rank_when_dp_exceeds_nodes():
160 engine_args = EngineArgs(
161 model="facebook/opt-125m",
162 data_parallel_size=4,
163 data_parallel_rank=3,
164 data_parallel_external_lb=True,
165 nnodes=2,
166 node_rank=1,
167 )
168
169 with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
170 vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
171
172 assert vllm_config.parallel_config.data_parallel_rank == 3
173
174
175 @pytest.mark.parametrize(
176 ("data_parallel_size", "nnodes", "tensor_parallel_size"),
177 [(4, 2, 1), (2, 3, 3)],
178 )
179 def test_external_lb_requires_explicit_rank_when_nodes_are_not_evenly_partitioned(
180 data_parallel_size, nnodes, tensor_parallel_size
181 ):
182 engine_args = EngineArgs(
183 model="facebook/opt-125m",
184 data_parallel_size=data_parallel_size,
185 data_parallel_external_lb=True,
186 tensor_parallel_size=tensor_parallel_size,
187 nnodes=nnodes,
188 node_rank=1,
189 )
190
191 with (
192 patch.object(ModelConfig, "is_moe", new=property(lambda self: True)),
193 pytest.raises(
194 ValueError,
195 match="Set a unique `--data-parallel-rank`",
196 ),
197 ):
198 engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
199
200
201 def test_external_lb_infers_rank_when_dp_does_not_exceed_nodes():
202 engine_args = EngineArgs(
203 model="facebook/opt-125m",
204 data_parallel_size=2,
205 data_parallel_external_lb=True,
206 nnodes=2,
207 node_rank=1,
208 )
209
210 with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
211 vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
212
213 assert vllm_config.parallel_config.data_parallel_rank == 1
214
215
216 def test_external_lb_infers_rank_for_multinode_replicas():
217 engine_args = EngineArgs(
218 model="facebook/opt-125m",
219 data_parallel_size=2,
220 data_parallel_external_lb=True,
221 tensor_parallel_size=12,
222 nnodes=4,
223 node_rank=1,
224 )
225
226 with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
227 vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
228
229 assert vllm_config.parallel_config.data_parallel_rank == 0