From b3c7ffcab82c2439726f8cb213800f6f38c023d3 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 22 May 2026 05:43:33 -0500 Subject: [PATCH] [Misc] Replace assert with proper exceptions for security and validation in pooling (#43286) Signed-off-by: Taneem Ibrahim Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Nick Hill --- .../layers/test_pooler_activations.py | 2 +- .../layers/pooler/activations.py | 12 +++++---- vllm/pooling_params.py | 14 ++++++---- vllm/v1/pool/metadata.py | 26 ++++++++++++++----- 4 files changed, 37 insertions(+), 17 deletions(-) diff --git a/tests/model_executor/layers/test_pooler_activations.py b/tests/model_executor/layers/test_pooler_activations.py index 587156a741..697207a68b 100644 --- a/tests/model_executor/layers/test_pooler_activations.py +++ b/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) diff --git a/vllm/model_executor/layers/pooler/activations.py b/vllm/model_executor/layers/pooler/activations.py index 94a77382fb..5fd05ef682 100644 --- a/vllm/model_executor/layers/pooler/activations.py +++ b/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 diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index b4cf23abac..3cfe9b427b 100644 --- a/vllm/pooling_params.py +++ b/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}" + ) diff --git a/vllm/v1/pool/metadata.py b/vllm/v1/pool/metadata.py index ac98f6fe0d..2c317873af 100644 --- a/vllm/v1/pool/metadata.py +++ b/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: