[BugFix] initial sampling_params_list when offline inference (#468)

Signed-off-by: Junhong <liujunhong11@huawei.com>
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>

tests/entrypoints/test_omni_llm.py

@@ -5,6 +5,9 @@ from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.stage_utils import _to_dict
 
 # Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
 warnings.filterwarnings(
@@ -91,6 +94,9 @@ class _FakeStage:
         self._out_q = None
         self._proc = None  # Mock process reference
 
+        default_sampling_params = getattr(config, "default_sampling_params", {})
+        self.default_sampling_params = SamplingParams(**_to_dict(default_sampling_params))
+
     def attach_queues(self, in_q, out_q):
         """Attach input and output queues."""
         self._in_q = in_q
@@ -777,61 +783,6 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     assert outputs == []
 
 
-def test_generate_sampling_params_none_raises(monkeypatch, fake_stage_config):
-    """Test that generate raises ValueError when sampling_params_list is None."""
-
-    def _fake_loader(model: str):
-        return [_FakeStageConfig(fake_stage_config)]
-
-    import sys
-
-    for module_name in [
-        "vllm_omni.entrypoints.utils",
-        "vllm_omni.entrypoints.omni_llm",
-        "vllm_omni.entrypoints.omni_stage",
-    ]:
-        if module_name in sys.modules:
-            del sys.modules[module_name]
-
-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
-    _setup_ipc_mocks(monkeypatch)
-    _setup_log_mocks(monkeypatch)
-
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.omni_llm.remove_old_logs",
-        lambda log_file, num_stages: None,
-        raising=False,
-    )
-
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg: _FakeStage(cfg),
-        raising=False,
-    )
-
-    import vllm_omni.entrypoints.omni_llm as omni_llm_module
-
-    monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
-
-    from vllm_omni.entrypoints.omni_llm import OmniLLM
-
-    llm = OmniLLM(model="any", init_timeout=1)
-    with pytest.raises(ValueError):
-        llm.generate(prompts=["p"], sampling_params_list=None)
-
-
 def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     """Test that _wait_for_stages_ready handles timeout correctly."""
 

vllm_omni/entrypoints/omni_llm.py

@@ -114,6 +114,9 @@ class OmniLLM:
             self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
         )
 
+        self.stage_list: list[OmniStage] = []
+        self.default_sampling_params_list: list[SamplingParams] = []
+
         self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
 
     def _initialize_stages(
@@ -124,6 +127,7 @@ class OmniLLM:
         init_timeout: int,
     ) -> None:
         self.stage_list: list[OmniStage] = []
+        self.default_sampling_params_list: list[SamplingParams] = []
 
         # Build OmniStage instances in parallel, preserve original order
         def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
@@ -137,6 +141,7 @@ class OmniLLM:
                 results.append(fut.result())
         results.sort(key=lambda x: x[0])
         self.stage_list = [st for _, st in results]
+        self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
         self.output_modalities = [st.final_output_type for st in self.stage_list]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
@@ -243,6 +248,8 @@ class OmniLLM:
             ValueError: If sampling_params_list is None or has incorrect length.
         """
         try:
+            if sampling_params_list is None:
+                sampling_params_list = self.default_sampling_params_list
             return self._run_generation(prompts, sampling_params_list)
         except Exception as e:
             logger.exception("[Orchestrator] Failed to run generation: %s", e)