From 637a7bf1df5781017fde1ad418bcfec989f76079 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Fri, 24 Jul 2026 10:59:45 -0500 Subject: [PATCH] fix: mark None-default parameters as Optional in vLLMExporter Several arguments in vLLMExporter.export and vLLMExporter.forward use None as a default but were annotated as non-optional str/int. This causes type-checker warnings and violates the documented signatures. Annotate them as Optional. Adds a static regression test that asserts every parameter with a None default in these methods is Optional. Signed-off-by: Andrew White --- nemo_export/vllm_exporter.py | 16 ++++++++-------- tests/unit_tests/export/test_vllm_exporter.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/nemo_export/vllm_exporter.py b/nemo_export/vllm_exporter.py index 6fd15ae77..2d0edef73 100644 --- a/nemo_export/vllm_exporter.py +++ b/nemo_export/vllm_exporter.py @@ -16,7 +16,7 @@ import logging import tempfile from pathlib import Path -from typing import Any, Dict, List, Literal +from typing import Any, Dict, List, Literal, Optional import numpy as np @@ -104,12 +104,12 @@ def __init__(self): def export( self, model_path_id: str, - tokenizer: str = None, + tokenizer: Optional[str] = None, trust_remote_code: bool = False, enable_lora: bool = False, tensor_parallel_size: int = 1, dtype: str = "auto", - quantization: str = None, + quantization: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: float = 4, @@ -117,7 +117,7 @@ def export( enforce_eager: bool = False, task: Literal["auto", "generate", "embedding"] = "auto", model_format: Literal["hf", "megatron_bridge"] = "megatron_bridge", - hf_model_id: str = None, + hf_model_id: Optional[str] = None, ): """ Exports a Hugging Face or Megatron-Bridge checkpoint to vLLM and initializes the engine. @@ -630,10 +630,10 @@ def forward( top_k: int = 1, top_p: float = 0.1, temperature: float = 1.0, - n_log_probs: int = None, - n_prompt_log_probs: int = None, - seed: int = None, - lora_model_name: str = None, + n_log_probs: Optional[int] = None, + n_prompt_log_probs: Optional[int] = None, + seed: Optional[int] = None, + lora_model_name: Optional[str] = None, ): """ Generate text completions for a list of input prompts using the vLLM model. diff --git a/tests/unit_tests/export/test_vllm_exporter.py b/tests/unit_tests/export/test_vllm_exporter.py index 7de8f3f46..0721ed331 100644 --- a/tests/unit_tests/export/test_vllm_exporter.py +++ b/tests/unit_tests/export/test_vllm_exporter.py @@ -13,6 +13,8 @@ # limitations under the License. +import inspect +import typing from unittest.mock import MagicMock, patch import numpy as np @@ -860,3 +862,20 @@ def test_export_megatron_bridge_with_all_vllm_params(exporter, mock_llm): assert call_kwargs["cpu_offload_gb"] == 2 assert call_kwargs["enforce_eager"] is False assert call_kwargs["runner"] == "generate" + + +def test_none_default_args_are_marked_optional(): + """Parameters that default to None must be annotated as Optional.""" + try: + from nemo_export.vllm_exporter import vLLMExporter + except ImportError: + pytest.skip("vllm_exporter dependencies not available") + + for method_name in ("export", "forward"): + sig = inspect.signature(getattr(vLLMExporter, method_name)) + for param in sig.parameters.values(): + if param.default is None and param.annotation is not inspect.Parameter.empty: + origin = getattr(param.annotation, "__origin__", None) + assert origin is typing.Union, ( + f"{method_name}.{param.name} defaults to None but is not Optional" + )