[Misc] Replace assert with proper exceptions for security and validation in pooling (#43286)

Signed-off-by: Taneem Ibrahim <[email protected]>
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: Nick Hill <[email protected]>

tests/model_executor/layers/test_pooler_activations.py

@@ -212,7 +212,7 @@ class TestGetActFn:
             problem_type="",
             sentence_transformers={"activation_fn": "os.system"},
         )
-        with pytest.raises(AssertionError, match="restricted"):
+        with pytest.raises(ValueError, match="restricted"):
             get_act_fn(cfg)
 
 

vllm/model_executor/layers/pooler/activations.py

@@ -49,10 +49,11 @@ def get_act_fn(
         function_name = config.sbert_ce_default_activation_function
 
     if function_name is not None:
-        assert function_name.startswith("torch.nn.modules."), (
-            "Loading of activation functions is restricted to "
-            "torch.nn.modules for security reasons"
-        )
+        if not function_name.startswith("torch.nn.modules."):
+            raise ValueError(
+                "Loading of activation functions is restricted to "
+                "torch.nn.modules for security reasons"
+            )
         fn = resolve_obj_by_qualname(function_name)()
         return PoolerActivation.wraps(fn)
 
@@ -67,7 +68,8 @@ def resolve_classifier_act_fn(
     if act_fn is None:
         return get_act_fn(model_config.hf_config, static_num_labels)
 
-    assert callable(act_fn)
+    if not callable(act_fn):
+        raise TypeError(f"Expected a callable activation function, got {type(act_fn)}")
     return act_fn
 
 

vllm/pooling_params.py

@@ -110,7 +110,8 @@ class PoolingParams(
         if pooler_config is None:
             return
 
-        assert self.task is not None, "task must be set"
+        if self.task is None:
+            raise ValueError("task must be set before merging parameters")
         valid_parameters = self.valid_parameters[self.task]
 
         for k in valid_parameters:
@@ -189,7 +190,8 @@ class PoolingParams(
             raise ValueError(f"Unknown pooling task: {self.task!r}")
 
     def _verify_valid_parameters(self):
-        assert self.task is not None, "task must be set"
+        if self.task is None:
+            raise ValueError("task must be set before verifying parameters")
         valid_parameters = self.valid_parameters[self.task]
         invalid_parameters = []
         for k in self.all_parameters:
@@ -221,6 +223,8 @@ class PoolingParams(
         )
 
     def __post_init__(self) -> None:
-        assert self.output_kind == RequestOutputKind.FINAL_ONLY, (
-            "For pooling output_kind has to be FINAL_ONLY"
-        )
+        if self.output_kind != RequestOutputKind.FINAL_ONLY:
+            raise ValueError(
+                "For pooling output_kind has to be FINAL_ONLY, "
+                f"got {self.output_kind!r}"
+            )

vllm/v1/pool/metadata.py

@@ -64,7 +64,11 @@ class PoolingMetadata:
             for pooling_param in pooling_params
             if (task := pooling_param.task) is not None
         ]
-        assert len(pooling_params) == len(tasks)
+        if len(pooling_params) != len(tasks):
+            raise ValueError(
+                "Every pooling param must have a task set, but got "
+                f"{len(tasks)} tasks for {len(pooling_params)} pooling params"
+            )
 
         self.tasks = tasks
 
@@ -88,9 +92,11 @@ class PoolingMetadata:
         self,
         prompt_token_ids: torch.Tensor | None,
     ) -> list[torch.Tensor]:
-        assert prompt_token_ids is not None, (
-            "Please set `requires_token_ids=True` in `get_pooling_updates`"
-        )
+        if prompt_token_ids is None:
+            raise ValueError(
+                "prompt_token_ids is required but was not set. "
+                "Please set `requires_token_ids=True` in `get_pooling_updates`"
+            )
         return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
 
     def get_prompt_token_ids(self) -> list[torch.Tensor]:
@@ -101,7 +107,11 @@ class PoolingMetadata:
 
     def get_pooling_cursor(self) -> PoolingCursor:
         pooling_cursor = self.pooling_cursor
-        assert pooling_cursor is not None, "Should call `build_pooling_cursor` first"
+        if pooling_cursor is None:
+            raise RuntimeError(
+                "pooling_cursor has not been initialized. "
+                "Call `build_pooling_cursor` before accessing it"
+            )
 
         return pooling_cursor
 
@@ -115,7 +125,11 @@ class PoolingMetadata:
         n_seq = len(num_scheduled_tokens_np)
         prompt_lens = self.prompt_lens
 
-        assert len(prompt_lens) == n_seq
+        if len(prompt_lens) != n_seq:
+            raise ValueError(
+                f"prompt_lens length ({len(prompt_lens)}) does not match "
+                f"the number of sequences ({n_seq})"
+            )
 
         num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)
         if query_start_loc_gpu is None: