fix: mark None-default parameters as Optional in vLLMExporter

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,20 +104,20 @@ 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,
         cpu_offload_gb: float = 0,
         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.

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"
+                )