@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 223 lines 7.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5
6 import pytest
7
8 from vllm import PoolingParams
9 from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
10 from vllm.entrypoints.pooling.typing import OfflineEncodeInputsContext
11 from vllm.exceptions import VLLMValidationError
12 from vllm.renderers import TokenizeParams
13
14
15 @pytest.fixture
16 def processor() -> PoolingIOProcessor:
17 return object.__new__(PoolingIOProcessor)
18
19
20 def test_rejects_untrusted_request_chat_template(processor: PoolingIOProcessor):
21 with pytest.raises(VLLMValidationError) as exc_info:
22 processor._validate_chat_template("template", None, False)
23
24 assert str(exc_info.value) == (
25 "Chat template is passed with request, but "
26 "--trust-request-chat-template is not set. "
27 "Refused request with untrusted chat template."
28 )
29 assert exc_info.value.parameter is None
30 assert exc_info.value.value is None
31
32
33 def test_rejects_mismatched_pooling_params(processor: PoolingIOProcessor):
34 with pytest.raises(VLLMValidationError) as exc_info:
35 processor._params_to_seq([PoolingParams()], num_requests=2)
36
37 assert str(exc_info.value) == (
38 "The lengths of prompts (2) and params (1) must be the same."
39 )
40 assert exc_info.value.parameter is None
41 assert exc_info.value.value is None
42
43
44 def test_rejects_mismatched_lora_requests(processor: PoolingIOProcessor):
45 with pytest.raises(VLLMValidationError) as exc_info:
46 processor._lora_request_to_seq([None], num_requests=2)
47
48 assert str(exc_info.value) == (
49 "The lengths of prompts (2) and lora_request (1) must be the same."
50 )
51 assert exc_info.value.parameter is None
52 assert exc_info.value.value is None
53
54
55 def test_rejects_conflicting_pooling_task(processor: PoolingIOProcessor):
56 processor.model_config = SimpleNamespace(is_encoder_decoder=False)
57 processor.renderer = SimpleNamespace(
58 default_cmpl_tok_params=TokenizeParams(max_total_tokens=None)
59 )
60 ctx = OfflineEncodeInputsContext(
61 pooling_task="embed",
62 tokenization_kwargs=None,
63 lora_request=None,
64 priorities=None,
65 prompts=[[1]],
66 pooling_params=PoolingParams(task="classify"),
67 )
68
69 with pytest.raises(VLLMValidationError) as exc_info:
70 processor.get_request_factory_offline(ctx)
71
72 assert str(exc_info.value) == (
73 "You cannot overwrite param.task='classify' with pooling_task='embed'!"
74 )
75 assert exc_info.value.parameter is None
76 assert exc_info.value.value is None
77
78
79 def test_qwen3_reranker_warns_without_chat_template(monkeypatch):
80 from unittest.mock import MagicMock
81
82 from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
83
84 monkeypatch.setattr(
85 "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock()
86 )
87 monkeypatch.setattr(
88 "vllm.model_executor.models.interfaces.supports_score_template",
89 lambda *_: False,
90 )
91 monkeypatch.setattr(
92 "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer",
93 lambda *_: False,
94 )
95
96 mock_logger = MagicMock()
97 monkeypatch.setattr(
98 "vllm.entrypoints.pooling.scoring.io_processor.logger",
99 mock_logger,
100 )
101
102 model_config = MagicMock()
103 model_config.architecture = "Qwen3ForSequenceClassification"
104 model_config.model = "Qwen/Qwen3-Reranker-0.6B"
105 model_config.hf_config.is_original_qwen3_reranker = True
106 model_config.use_sep_token = False
107 model_config.is_multimodal_model = False
108 model_config.max_model_len = 8192
109
110 vllm_config = MagicMock(model_config=model_config)
111 tokenizer = MagicMock()
112 tokenizer.pad_token_id = 0
113 renderer = MagicMock()
114 renderer.get_tokenizer.return_value = tokenizer
115 chat_template_config = MagicMock(chat_template=None)
116
117 CrossEncoderIOProcessor(
118 vllm_config=vllm_config,
119 renderer=renderer,
120 chat_template_config=chat_template_config,
121 )
122
123 mock_logger.warning.assert_called_once()
124 warning_call = mock_logger.warning.call_args
125 assert "Qwen/Qwen3-Reranker-0.6B" in warning_call[0][1]
126 assert "qwen3_reranker.jinja" in warning_call[0][2]
127
128
129 def test_qwen3_vl_reranker_warns_with_vl_template(monkeypatch):
130 from unittest.mock import MagicMock
131
132 from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
133
134 monkeypatch.setattr(
135 "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock()
136 )
137 monkeypatch.setattr(
138 "vllm.model_executor.models.interfaces.supports_score_template",
139 lambda *_: False,
140 )
141 monkeypatch.setattr(
142 "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer",
143 lambda *_: False,
144 )
145
146 mock_logger = MagicMock()
147 monkeypatch.setattr(
148 "vllm.entrypoints.pooling.scoring.io_processor.logger",
149 mock_logger,
150 )
151
152 model_config = MagicMock()
153 model_config.architecture = "Qwen3VLForSequenceClassification"
154 model_config.model = "Qwen/Qwen3-VL-Reranker-2B"
155 model_config.hf_config.is_original_qwen3_reranker = True
156 model_config.use_sep_token = False
157 model_config.is_multimodal_model = True
158 model_config.max_model_len = 8192
159
160 vllm_config = MagicMock(model_config=model_config)
161 tokenizer = MagicMock()
162 tokenizer.pad_token_id = 0
163 renderer = MagicMock()
164 renderer.get_tokenizer.return_value = tokenizer
165 chat_template_config = MagicMock(chat_template=None)
166
167 CrossEncoderIOProcessor(
168 vllm_config=vllm_config,
169 renderer=renderer,
170 chat_template_config=chat_template_config,
171 )
172
173 mock_logger.warning.assert_called_once()
174 warning_call = mock_logger.warning.call_args
175 assert "Qwen/Qwen3-VL-Reranker-2B" in warning_call[0][1]
176 assert "qwen3_vl_reranker.jinja" in warning_call[0][2]
177
178
179 def test_qwen3_reranker_no_warning_when_template_provided(monkeypatch):
180 from unittest.mock import MagicMock
181
182 from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
183
184 monkeypatch.setattr(
185 "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock()
186 )
187 monkeypatch.setattr(
188 "vllm.model_executor.models.interfaces.supports_score_template",
189 lambda *_: False,
190 )
191 monkeypatch.setattr(
192 "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer",
193 lambda *_: False,
194 )
195
196 mock_logger = MagicMock()
197 monkeypatch.setattr(
198 "vllm.entrypoints.pooling.scoring.io_processor.logger",
199 mock_logger,
200 )
201
202 model_config = MagicMock()
203 model_config.architecture = "Qwen3ForSequenceClassification"
204 model_config.model = "Qwen/Qwen3-Reranker-0.6B"
205 model_config.hf_config.is_original_qwen3_reranker = True
206 model_config.use_sep_token = False
207 model_config.is_multimodal_model = False
208 model_config.max_model_len = 8192
209
210 vllm_config = MagicMock(model_config=model_config)
211 tokenizer = MagicMock()
212 tokenizer.pad_token_id = 0
213 renderer = MagicMock()
214 renderer.get_tokenizer.return_value = tokenizer
215 chat_template_config = MagicMock(chat_template="custom template")
216
217 CrossEncoderIOProcessor(
218 vllm_config=vllm_config,
219 renderer=renderer,
220 chat_template_config=chat_template_config,
221 )
222
223 mock_logger.warning.assert_not_called()