[Bugfix] Reject removed pooling parameters (#48984)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jul 19, 2026 at 07:18 UTC
e6d1310b2ac52e56e266450693af9de12a871a51
9 files changed
+108
-17
tests/test_pooling_params.py
+38
@@ -1,12 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
+from typing import Any
5
6
import pytest
7
+from pydantic import TypeAdapter, ValidationError
8
9
from tests.models.utils import EmbedModelInfo
10
from vllm import PoolingParams
11
from vllm.config import ModelConfig, PoolerConfig
12
+from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
13
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
14
+from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest
15
+from vllm.exceptions import VLLMValidationError
16
17
EMBEDDING_MODELS = [
18
EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
@@ -27,6 +33,38 @@ class MockModelConfig:
33
pooler_config: PoolerConfig
34
35
36
+@pytest.mark.parametrize(
37
+ ("parameter", "value", "message"),
38
+ [
39
+ (
40
+ "normalize",
41
+ False,
42
+ "Parameter `normalize` was removed; use `use_activation` instead.",
43
+ ),
44
+ ("task", "score", "`score` task was removed; use `classify` instead."),
45
+ (
46
+ "task",
47
+ "encode",
48
+ "`encode` task was removed; use `token_embed` or `token_classify` instead.",
49
+ ),
50
+ ],
51
+)
52
+def test_removed_pooling_parameters(parameter: str, value: Any, message: str):
53
+ data = {"input": "hello", parameter: value}
54
+ for request_type in (EmbeddingRequest, ClassificationRequest, PoolingRequest):
55
+ with pytest.raises(ValidationError, match=message) as exc_info:
56
+ TypeAdapter(request_type).validate_python(data)
57
+ assert len(exc_info.value.errors()) == 1
58
+
59
+ with pytest.raises(ValidationError, match=message) as exc_info:
60
+ TypeAdapter(PoolerConfig).validate_python({parameter: value})
61
+ assert len(exc_info.value.errors()) == 1
62
+
63
+ if parameter == "task":
64
+ with pytest.raises(VLLMValidationError, match=message):
65
+ PoolingParams(task=value)
66
+
67
+
68
def test_embed():
69
task = "embed"
70
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
vllm/config/pooler.py
+17
-1
@@ -3,9 +3,12 @@
3
4
from typing import Any, Literal, get_args
5
6
+from pydantic import model_validator
7
+from pydantic_core import ArgsKwargs
8
+
9
from vllm.config.utils import config
10
from vllm.logger import init_logger
8
-from vllm.tasks import PoolingTask
11
+from vllm.tasks import PoolingTask, check_removed_pooling_task
12
from vllm.utils.hashing import safe_hash
13
14
logger = init_logger(__name__)
@@ -112,6 +115,19 @@ class PoolerConfig:
115
`math-shepherd-mistral-7b-prm` model.
116
"""
117
118
+ @model_validator(mode="before")
119
+ @classmethod
120
+ def reject_removed_parameters(cls, data):
121
+ values = data.kwargs if isinstance(data, ArgsKwargs) else data
122
+ if not isinstance(values, dict):
123
+ return data
124
+ if "normalize" in values:
125
+ raise ValueError(
126
+ "Parameter `normalize` was removed; use `use_activation` instead."
127
+ )
128
+ check_removed_pooling_task(values.get("task"))
129
+ return data
130
+
131
def __post_init__(self) -> None:
132
if self.logit_sigma is not None and self.logit_sigma == 0:
133
raise ValueError("logit_sigma cannot be 0 (division by zero)")
vllm/entrypoints/pooling/base/protocol.py
+13
@@ -14,10 +14,23 @@ from vllm.entrypoints.chat_utils import (
14
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel
15
from vllm.exceptions import VLLMValidationError
16
from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
17
+from vllm.tasks import check_removed_pooling_task
18
from vllm.utils import random_uuid
19
from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
20
21
22
+def reject_removed_pooling_parameters(data):
23
+ if not isinstance(data, dict):
24
+ return data
25
+ if "normalize" in data:
26
+ raise VLLMValidationError(
27
+ "Parameter `normalize` was removed; use `use_activation` instead.",
28
+ parameter="normalize",
29
+ )
30
+ check_removed_pooling_task(data.get("task"))
31
+ return data
32
+
33
+
34
class PoolingBasicRequestMixin(OpenAIBaseModel):
35
# --8<-- [start:pooling-common-params]
36
model: str | None = None
vllm/entrypoints/pooling/classify/protocol.py
+7
-5
@@ -2,9 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import time
5
-from typing import TypeAlias
5
+from typing import Annotated, TypeAlias
6
7
-from pydantic import Field
7
+from pydantic import BeforeValidator, Field
8
9
from vllm import PoolingParams
10
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
@@ -17,6 +17,7 @@ from ..base.protocol import (
17
CompletionRequestMixin,
18
FixedMaxLenTokenizeParamsMixin,
19
PoolingBasicRequestMixin,
20
+ reject_removed_pooling_parameters,
21
)
22
23
logger = init_logger(__name__)
@@ -48,9 +49,10 @@ class ClassificationChatRequest(
49
)
50
51
51
-ClassificationRequest: TypeAlias = (
52
- ClassificationCompletionRequest | ClassificationChatRequest
53
-)
52
+ClassificationRequest: TypeAlias = Annotated[
53
+ ClassificationCompletionRequest | ClassificationChatRequest,
54
+ BeforeValidator(reject_removed_pooling_parameters),
55
+]
56
57
58
class ClassificationData(OpenAIBaseModel):
vllm/entrypoints/pooling/embed/protocol.py
+6
-4
@@ -13,7 +13,7 @@ from collections.abc import Sequence
13
from typing import Annotated, Any, Literal, TypeAlias
14
15
import pybase64 as base64
16
-from pydantic import BaseModel, Field, model_validator
16
+from pydantic import BaseModel, BeforeValidator, Field, model_validator
17
18
from vllm import PoolingParams
19
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
@@ -27,6 +27,7 @@ from ..base.protocol import (
27
EmbeddingTokenizeParamsMixin,
28
EmbedRequestMixin,
29
PoolingBasicRequestMixin,
30
+ reject_removed_pooling_parameters,
31
)
32
33
@@ -154,13 +155,14 @@ class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
155
return normalized
156
157
157
-EmbeddingRequest: TypeAlias = (
158
+EmbeddingRequest: TypeAlias = Annotated[
159
EmbeddingCompletionRequest
160
| EmbeddingChatRequest
161
| EmbeddingBatchChatRequest
162
| EmbeddingChatInputRequest
162
- | EmbeddingBatchChatInputRequest
163
-)
163
+ | EmbeddingBatchChatInputRequest,
164
+ BeforeValidator(reject_removed_pooling_parameters),
165
+]
166
167
168
# ---------------------------------------------------------------------------
vllm/entrypoints/pooling/pooling/protocol.py
+7
-5
@@ -1,9 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
-from typing import Generic, TypeAlias, TypeVar
4
+from typing import Annotated, Generic, TypeAlias, TypeVar
5
6
-from pydantic import Field
6
+from pydantic import BeforeValidator, Field
7
8
from vllm import PoolingParams
9
from vllm.config import ModelConfig
@@ -20,6 +20,7 @@ from ..base.protocol import (
20
EncodingRequestMixin,
21
FixedMaxLenTokenizeParamsMixin,
22
PoolingBasicRequestMixin,
23
+ reject_removed_pooling_parameters,
24
)
25
26
@@ -92,9 +93,10 @@ class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
93
"""
94
95
95
-PoolingRequest: TypeAlias = (
96
- PoolingCompletionRequest | PoolingChatRequest | IOProcessorRequest
97
-)
96
+PoolingRequest: TypeAlias = Annotated[
97
+ PoolingCompletionRequest | PoolingChatRequest | IOProcessorRequest,
98
+ BeforeValidator(reject_removed_pooling_parameters),
99
+]
100
101
102
class PoolingResponseData(OpenAIBaseModel):
vllm/entrypoints/pooling/pooling/serving.py
+3
-1
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import cast
4
+
5
from fastapi.responses import JSONResponse, Response, StreamingResponse
6
from typing_extensions import assert_never
7
@@ -52,7 +54,7 @@ class ServingPooling(PoolingBaseServing):
54
self.json_response_cls = get_json_response_cls()
55
56
def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor:
55
- assert isinstance(request, PoolingRequest)
57
+ request = cast(PoolingRequest, request)
58
pooling_task = self._verify_pooling_task(request)
59
return self.io_processors[pooling_task]
60
vllm/pooling_params.py
+2
-1
@@ -9,7 +9,7 @@ import msgspec
9
from vllm.config import ModelConfig, PoolerConfig
10
from vllm.logger import init_logger
11
from vllm.sampling_params import RequestOutputKind
12
-from vllm.tasks import PoolingTask
12
+from vllm.tasks import PoolingTask, check_removed_pooling_task
13
14
logger = init_logger(__name__)
15
@@ -229,6 +229,7 @@ class PoolingParams(
229
)
230
231
def __post_init__(self) -> None:
232
+ check_removed_pooling_task(self.task)
233
if self.output_kind != RequestOutputKind.FINAL_ONLY:
234
raise ValueError(
235
"For pooling output_kind has to be FINAL_ONLY, "
vllm/tasks.py
+15
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import Literal, get_args
4
5
+from vllm.exceptions import VLLMValidationError
6
+
7
GenerationTask = Literal["generate", "transcription", "realtime"]
8
GENERATION_TASKS: tuple[GenerationTask, ...] = get_args(GenerationTask)
9
@@ -15,6 +17,19 @@ PoolingTask = Literal[
17
]
18
POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
19
20
+_REMOVED_POOLING_TASK_MESSAGES = {
21
+ "score": "`score` task was removed; use `classify` instead.",
22
+ "encode": (
23
+ "`encode` task was removed; use `token_embed` or `token_classify` instead."
24
+ ),
25
+}
26
+
27
+
28
+def check_removed_pooling_task(task: object) -> None:
29
+ if isinstance(task, str) and (message := _REMOVED_POOLING_TASK_MESSAGES.get(task)):
30
+ raise VLLMValidationError(message, parameter="task")
31
+
32
+
33
ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
34
SCORE_TYPE_MAP: dict[PoolingTask, ScoreType] = {
35
"embed": "bi-encoder",