From 51eefbfebc3e85dbe2ae2b5681c217617cef1e77 Mon Sep 17 00:00:00 2001 From: Junhong Liu Date: Fri, 26 Dec 2025 13:54:12 +0800 Subject: [PATCH] [BugFix] initial sampling_params_list when offline inference (#468) Signed-off-by: Junhong Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Signed-off-by: Junhong Liu --- tests/entrypoints/test_omni_llm.py | 61 +++--------------------------- vllm_omni/entrypoints/omni_llm.py | 7 ++++ 2 files changed, 13 insertions(+), 55 deletions(-) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index cef5dce43..c6dd04241 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/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.""" diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 36cd60a5a..d2dcba697 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/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)