A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 01:45 UTC
[Misc] Replace assert with proper exceptions for security and validation in pooling (#43286)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Taneem Ibrahim committed
May 22, 2026 at 05:43 UTC
b3c7ffcab82c2439726f8cb213800f6f38c023d3
4 files changed
+37
-17
tests/model_executor/layers/test_pooler_activations.py
+1
-1
@@ -212,7 +212,7 @@ class TestGetActFn:
212
problem_type="",
213
sentence_transformers={"activation_fn": "os.system"},
214
)
215
- with pytest.raises(AssertionError, match="restricted"):
215
+ with pytest.raises(ValueError, match="restricted"):
216
get_act_fn(cfg)
217
218
vllm/model_executor/layers/pooler/activations.py
+7
-5
@@ -49,10 +49,11 @@ def get_act_fn(
49
function_name = config.sbert_ce_default_activation_function
50
51
if function_name is not None:
52
- assert function_name.startswith("torch.nn.modules."), (
53
- "Loading of activation functions is restricted to "
54
- "torch.nn.modules for security reasons"
55
- )
52
+ if not function_name.startswith("torch.nn.modules."):
53
+ raise ValueError(
54
+ "Loading of activation functions is restricted to "
55
+ "torch.nn.modules for security reasons"
56
+ )
57
fn = resolve_obj_by_qualname(function_name)()
58
return PoolerActivation.wraps(fn)
59
@@ -67,7 +68,8 @@ def resolve_classifier_act_fn(
68
if act_fn is None:
69
return get_act_fn(model_config.hf_config, static_num_labels)
70
70
- assert callable(act_fn)
71
+ if not callable(act_fn):
72
+ raise TypeError(f"Expected a callable activation function, got {type(act_fn)}")
73
return act_fn
74
75
vllm/pooling_params.py
+9
-5
@@ -110,7 +110,8 @@ class PoolingParams(
110
if pooler_config is None:
111
return
112
113
- assert self.task is not None, "task must be set"
113
+ if self.task is None:
114
+ raise ValueError("task must be set before merging parameters")
115
valid_parameters = self.valid_parameters[self.task]
116
117
for k in valid_parameters:
@@ -189,7 +190,8 @@ class PoolingParams(
190
raise ValueError(f"Unknown pooling task: {self.task!r}")
191
192
def _verify_valid_parameters(self):
192
- assert self.task is not None, "task must be set"
193
+ if self.task is None:
194
+ raise ValueError("task must be set before verifying parameters")
195
valid_parameters = self.valid_parameters[self.task]
196
invalid_parameters = []
197
for k in self.all_parameters:
@@ -221,6 +223,8 @@ class PoolingParams(
223
)
224
225
def __post_init__(self) -> None:
224
- assert self.output_kind == RequestOutputKind.FINAL_ONLY, (
225
- "For pooling output_kind has to be FINAL_ONLY"
226
- )
226
+ if self.output_kind != RequestOutputKind.FINAL_ONLY:
227
+ raise ValueError(
228
+ "For pooling output_kind has to be FINAL_ONLY, "
229
+ f"got {self.output_kind!r}"
230
+ )
vllm/v1/pool/metadata.py
+20
-6
@@ -64,7 +64,11 @@ class PoolingMetadata:
64
for pooling_param in pooling_params
65
if (task := pooling_param.task) is not None
66
]
67
- assert len(pooling_params) == len(tasks)
67
+ if len(pooling_params) != len(tasks):
68
+ raise ValueError(
69
+ "Every pooling param must have a task set, but got "
70
+ f"{len(tasks)} tasks for {len(pooling_params)} pooling params"
71
+ )
72
73
self.tasks = tasks
74
@@ -88,9 +92,11 @@ class PoolingMetadata:
92
self,
93
prompt_token_ids: torch.Tensor | None,
94
) -> list[torch.Tensor]:
91
- assert prompt_token_ids is not None, (
92
- "Please set `requires_token_ids=True` in `get_pooling_updates`"
93
- )
95
+ if prompt_token_ids is None:
96
+ raise ValueError(
97
+ "prompt_token_ids is required but was not set. "
98
+ "Please set `requires_token_ids=True` in `get_pooling_updates`"
99
+ )
100
return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
101
102
def get_prompt_token_ids(self) -> list[torch.Tensor]:
@@ -101,7 +107,11 @@ class PoolingMetadata:
107
108
def get_pooling_cursor(self) -> PoolingCursor:
109
pooling_cursor = self.pooling_cursor
104
- assert pooling_cursor is not None, "Should call `build_pooling_cursor` first"
110
+ if pooling_cursor is None:
111
+ raise RuntimeError(
112
+ "pooling_cursor has not been initialized. "
113
+ "Call `build_pooling_cursor` before accessing it"
114
+ )
115
116
return pooling_cursor
117
@@ -115,7 +125,11 @@ class PoolingMetadata:
125
n_seq = len(num_scheduled_tokens_np)
126
prompt_lens = self.prompt_lens
127
118
- assert len(prompt_lens) == n_seq
128
+ if len(prompt_lens) != n_seq:
129
+ raise ValueError(
130
+ f"prompt_lens length ({len(prompt_lens)}) does not match "
131
+ f"the number of sequences ({n_seq})"
132
+ )
133
134
num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)
135
if query_start_loc_gpu is None: