From 1ab88823d499d2e65355dc17a274e6109c04f8a3 Mon Sep 17 00:00:00 2001
From: Roger Wang <hey@rogerw.io>
Date: Fri, 28 Nov 2025 13:05:07 -0800
Subject: [PATCH] Simple test from PR#88 on Buildkite (#93)

Signed-off-by: WANG Cong <115451386+congw729@users.noreply.github.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: WANG Cong <115451386+congw729@users.noreply.github.com>
---
 .buildkite/pipeline.yml           |    9 +-
 .buildkite/scripts/simple_test.sh |   52 ++
 tests/test_omni_llm.py            | 1019 +++++++++++++++++++++++++++++
 3 files changed, 1075 insertions(+), 5 deletions(-)
 create mode 100755 .buildkite/scripts/simple_test.sh
 create mode 100644 tests/test_omni_llm.py

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index a67d038e1..7bc41512e 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -1,7 +1,6 @@
 steps:
-  - label: "Dummy pipeline"
-    if: "false"
-    agents:
-      queue: cpu_queue_premerge_us_east_1
+  - label: "Simple Unit Test"
     commands:
-      - "echo 'This pipeline is intentionally a no-op.'"
+      - ".buildkite/scripts/simple_test.sh"
+    agents:
+      queue: "cpu_queue_premerge"
diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh
new file mode 100755
index 000000000..8821a03ce
--- /dev/null
+++ b/.buildkite/scripts/simple_test.sh
@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
+UV_BIN="uv"
+PYTHON_BOOTSTRAP="${PYTHON:-python3}"
+VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
+TARGET_PY_VERSION="3.12"
+
+if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
+  PYTHON_BOOTSTRAP="python"
+fi
+
+cd "${REPO_ROOT}"
+
+# Ensure pip-installed scripts (including uv) are on PATH
+PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
+USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
+PATH="${PYTHON_SCRIPTS}:${PATH}"
+if [[ -n "${USER_BASE}" ]]; then
+  PATH="${USER_BASE}/bin:${PATH}"
+fi
+
+if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
+  "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
+  "${PYTHON_BOOTSTRAP}" -m pip install uv
+  hash -r
+fi
+
+TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
+if [[ -z "${TARGET_PY_PATH}" ]]; then
+  "${UV_BIN}" python install "${TARGET_PY_VERSION}"
+  TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
+fi
+
+if [[ -z "${TARGET_PY_PATH}" ]]; then
+  echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
+  exit 1
+fi
+
+if [[ ! -d "${VENV_DIR}" ]]; then
+  "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
+fi
+
+VENV_PYTHON="${VENV_DIR}/bin/python"
+[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
+
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.11.0
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
+"${VENV_PYTHON}" -m pytest tests/test_omni_llm.py
diff --git a/tests/test_omni_llm.py b/tests/test_omni_llm.py
new file mode 100644
index 000000000..12073576a
--- /dev/null
+++ b/tests/test_omni_llm.py
@@ -0,0 +1,1019 @@
+import warnings
+from queue import Empty, Queue
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+# Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
+warnings.filterwarnings(
+    "ignore",
+    message=r"builtin type SwigPy.*has no __module__ attribute",
+    category=DeprecationWarning,
+)
+
+
+class _FakeEngineArgs(dict):
+    """Fake engine args that can be used both as object attributes and as **kwargs."""
+
+    def __init__(self, args_dict: dict[str, Any]):
+        super().__init__(args_dict)
+        # Add required attributes if not present
+        if "model_stage" not in self:
+            self["model_stage"] = None
+        if "engine_output_type" not in self:
+            self["engine_output_type"] = None
+        # Also set as attributes for object-style access
+        for key, value in self.items():
+            setattr(self, key, value)
+
+
+class _FakeStageConfig:
+    """Fake stage config object that mimics the real stage config structure."""
+
+    def __init__(self, config_dict: dict[str, Any]):
+        # engine_args needs to work both as object (for OmniStage) and as dict (for **kwargs)
+        engine_args_dict = config_dict.get("engine_args", {})
+        self.engine_args = _FakeEngineArgs(engine_args_dict)
+        self.final_output = config_dict.get("final_output", False)
+        self.final_output_type = config_dict.get("final_output_type", None)
+        self.stage_id = config_dict.get("stage_id", 0)
+        # Store original dict for reference
+        self._config_dict = config_dict
+
+
+class _FakeQueue:
+    """Fake queue using standard library Queue to replace mp.Queue."""
+
+    def __init__(self, maxsize=0):
+        self._queue = Queue(maxsize=maxsize)
+
+    def put(self, item):
+        self._queue.put(item)
+
+    def put_nowait(self, item):
+        self._queue.put_nowait(item)
+
+    def get(self):
+        return self._queue.get()
+
+    def get_nowait(self):
+        return self._queue.get_nowait()
+
+    def empty(self):
+        return self._queue.empty()
+
+
+class _FakeStage:
+    """Lightweight Stage stub for multi-process pipeline version with queue support."""
+
+    def __init__(self, config):
+        # Handle both dict and object configs
+        if isinstance(config, dict):
+            config = _FakeStageConfig(config)
+        self.config = config
+        self.stage_config = config
+        self.engine = None
+        self.engine_outputs = None
+        # Set attributes that OmniStage expects
+        self.stage_id = getattr(config, "stage_id", 0)
+        self.engine_args = config.engine_args
+        self.model_stage = getattr(config.engine_args, "model_stage", None)
+        # Allow configuring final_output and final_output_type
+        self.final_output = config.final_output if hasattr(config, "final_output") else False
+        self.final_output_type = getattr(config, "final_output_type", None)
+        # Configurable processing logic, default returns placeholder
+        processed_input = getattr(config, "_config_dict", {}).get("processed_input", ["processed"])
+        self._processed_input = processed_input
+        # Queue references (set by attach_queues)
+        self._in_q = None
+        self._out_q = None
+        self._proc = None  # Mock process reference
+
+    def attach_queues(self, in_q, out_q):
+        """Attach input and output queues."""
+        self._in_q = in_q
+        self._out_q = out_q
+
+    def init_stage_worker(
+        self,
+        model: str,
+        *,
+        is_async: bool = False,
+        log_file=None,
+        shm_threshold_bytes: int = 65536,
+        ctx=None,
+        batch_timeout: int = 10,
+    ):
+        """Mock init_stage_worker: don't start real process, just send stage_ready message."""
+        # Create a mock process object
+        self._proc = MagicMock()
+        self._proc.start = MagicMock()
+        self._proc.join = MagicMock()
+        self._proc.is_alive = MagicMock(return_value=False)
+        self._proc.terminate = MagicMock()
+        # Send stage_ready message to output queue
+        if self._out_q is not None:
+            try:
+                self._out_q.put_nowait({"type": "stage_ready", "stage_id": self.stage_id})
+            except Exception:
+                pass
+
+    def stop_stage_worker(self):
+        """Mock stop_stage_worker: clean up queue references."""
+        if self._in_q is not None:
+            try:
+                self._in_q.put_nowait(None)
+            except Exception:
+                pass
+
+    def submit(self, payload: dict[str, Any]):
+        """Submit task to input queue."""
+        if self._in_q is not None:
+            self._in_q.put(payload)
+
+    def try_collect(self) -> Any:
+        """Non-blocking collect from output queue."""
+        if self._out_q is None:
+            return None
+        try:
+            return self._out_q.get_nowait()
+        except Empty:
+            return None
+
+    def set_engine_outputs(self, outputs):
+        """Set engine outputs for the stage."""
+        self.engine_outputs = outputs
+
+    def process_engine_inputs(self, stage_list, prompts):
+        """Process engine inputs: return preset processed result."""
+        return self._processed_input
+
+
+class _FakeEngine:
+    """Lightweight Engine stub: provides generate iterator output."""
+
+    def __init__(self, outputs: list[Any]):
+        self._outputs = outputs
+
+    def generate(self, prompts, sampling_params):
+        # Record the most recent prompts for outer assertions
+        self._last_prompts = prompts
+        # Simplified: return preset list at once, ensuring iterability
+        yield from self._outputs
+
+
+class _FakeStageLLM:
+    """Replace OmniStageLLM to avoid constructing real engine."""
+
+    def __init__(self, **kwargs):
+        # Allow injecting custom fake outputs, default returns single placeholder output
+        fake_outputs = kwargs.get("_fake_outputs", [[{"text": "ok"}]])
+        self._fake_engine = _FakeEngine(fake_outputs)
+
+    def generate(self, prompts, sampling_params, use_tqdm=False):
+        yield from self._fake_engine.generate(prompts, sampling_params)
+
+
+@pytest.fixture
+def fake_stage_config():
+    return {
+        # Don't include 'model' in engine_args since it's passed separately
+        "engine_args": {},
+        "final_output": True,
+        "final_output_type": "text",
+        # Second stage will use processed_input to verify the chain
+        "processed_input": ["processed-by-stage"],
+    }
+
+
+def _setup_engine_mocks(monkeypatch):
+    """Helper function to set up common engine mocks."""
+    fake_engine = MagicMock()
+    # Add necessary attributes to fake_engine
+    fake_engine.tokenizer = MagicMock()
+    fake_engine.log_stats = False
+    fake_engine.vllm_config = MagicMock()
+    fake_engine.vllm_config.model_config = MagicMock()
+    fake_engine.vllm_config.model_config.io_processor_plugin = None
+    fake_engine.get_supported_tasks = MagicMock(return_value=[])
+    fake_engine.model_config = MagicMock()
+    fake_engine.model_config.io_processor_plugin = None
+    # Add registry with resolve_model_cls method
+    fake_registry = MagicMock()
+    fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch"))
+    fake_engine.model_config.registry = fake_registry
+    fake_engine.vllm_config.model_config.registry = fake_registry
+
+    monkeypatch.setattr(
+        "vllm.v1.engine.llm_engine.LLMEngine.from_engine_args",
+        lambda **kw: fake_engine,
+        raising=False,
+    )
+
+    # Mock model_config.registry.resolve_model_cls to return a tuple
+    # Use a real class instead of MagicMock to avoid inspect.getsource issues
+    class FakeModelClass:
+        pass
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.utils.get_model_architecture",
+        lambda model_config: (FakeModelClass, "test_arch"),
+        raising=False,
+    )
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.utils._get_model_architecture",
+        lambda model_config: (FakeModelClass, "test_arch"),
+        raising=False,
+    )
+
+    # Mock try_create_mm_pooling_model_cls to return the class as-is
+    monkeypatch.setattr(
+        "vllm.model_executor.models.adapters.try_create_mm_pooling_model_cls",
+        lambda model_cls: model_cls,
+        raising=False,
+    )
+
+    # Mock _enable_processor_cache to return False
+    monkeypatch.setattr(
+        "vllm.multimodal.cache._enable_processor_cache",
+        lambda model_config, mm_registry: False,
+        raising=False,
+    )
+
+    # Mock get_io_processor to return None
+    monkeypatch.setattr(
+        "vllm.plugins.io_processors.get_io_processor",
+        lambda vllm_config, io_processor_plugin: None,
+        raising=False,
+    )
+
+
+def _setup_multiprocessing_mocks(monkeypatch):
+    """Helper function to set up multiprocessing mocks."""
+    import multiprocessing as mp
+
+    # Mock Process
+    fake_process_class = MagicMock()
+    fake_process_instance = MagicMock()
+    fake_process_instance.start = MagicMock()
+    fake_process_instance.join = MagicMock()
+    fake_process_instance.is_alive = MagicMock(return_value=False)
+    fake_process_instance.terminate = MagicMock()
+    fake_process_class.return_value = fake_process_instance
+
+    # Mock get_context to return a context with Queue that returns _FakeQueue
+    fake_ctx = MagicMock()
+    fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
+    fake_ctx.Process = fake_process_class
+
+    def _mock_get_context(method):
+        return fake_ctx
+
+    monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False)
+    monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)
+
+
+def _setup_ipc_mocks(monkeypatch):
+    """Helper function to set up IPC function mocks."""
+
+    # Mock _encode: simple serialization
+    def _fake_encode(obj, threshold, obj_key, shm_key):
+        return {obj_key: obj}
+
+    # Mock _load: extract object from result
+    def _fake_load(result, obj_key, shm_key):
+        return result.get(obj_key)
+
+    # Mock _set: calculate serialization size
+    def _fake_set(obj):
+        return str(obj).encode()
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni_llm._encode", _fake_encode, raising=False)
+    monkeypatch.setattr("vllm_omni.entrypoints.omni_llm._load", _fake_load, raising=False)
+    monkeypatch.setattr("vllm_omni.entrypoints.omni_llm._set", _fake_set, raising=False)
+
+
+def _setup_log_mocks(monkeypatch):
+    """Helper function to set up logging and stats mocks."""
+    # Mock init_stats_paths to return None files (no stats logging)
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_llm.init_stats_paths",
+        lambda enable_stats, log_file: (None, None),
+        raising=False,
+    )
+    # Mock configure_orchestrator_logger to do nothing
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_llm.configure_orchestrator_logger",
+        lambda logger, log_file: None,
+        raising=False,
+    )
+    # Mock OrchestratorMetrics to be a simple class that doesn't require file operations
+
+    class _FakeOrchestratorMetrics:
+        def __init__(self, num_stages, enable_stats, stats_file, overall_stats_file, wall_start_ts):
+            self.num_stages = num_stages
+            self.enable_stats = enable_stats
+            self.stage_first_ts = [None] * num_stages
+            self.stage_last_ts = [None] * num_stages
+            self.e2e_done = set()
+
+        def on_stage_metrics(self, stage_id, req_id, metrics):
+            pass
+
+        def on_finalize_request(self, stage_id, req_id, engine_outputs, start_ts):
+            self.e2e_done.add(req_id)
+
+        def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm):
+            pass
+
+        def build_and_log_summary(self, final_stage_id):
+            return "Fake summary"
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_llm.OrchestratorMetrics",
+        _FakeOrchestratorMetrics,
+        raising=False,
+    )
+
+
+@pytest.fixture(autouse=True)
+def mock_get_config(monkeypatch):
+    """Auto-mock get_config and related model loading functions to avoid model path validation."""
+    # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
+    # This prevents ImportError when async_omni_llm is imported (which happens via omni_stage)
+    import sys
+
+    fake_tokenizer = MagicMock()
+    fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
+    fake_tokenizer.decode = MagicMock(return_value="test")
+
+    # Mock init_tokenizer_from_configs (used in async_omni_llm)
+    def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
+        return fake_tokenizer
+
+    # Strategy 1: Mock in the original location (vllm.transformers_utils.tokenizer)
+    # This works if the module hasn't been imported yet
+    monkeypatch.setattr(
+        "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs",
+        _mock_init_tokenizer_from_configs,
+        raising=False,
+    )
+
+    # Strategy 2: If the module is already in sys.modules, patch it directly
+    tokenizer_module_path = "vllm.transformers_utils.tokenizer"
+    if tokenizer_module_path in sys.modules:
+        tokenizer_module = sys.modules[tokenizer_module_path]
+        setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+
+    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni_llm
+    # This is because async_omni_llm imports processor.py, which imports this function at module level
+    # Mock length_from_prompt_token_ids_or_embeds (used in processor.py)
+    def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
+        # Return a reasonable default length
+        if prompt_token_ids is not None:
+            if isinstance(prompt_token_ids, list):
+                return len(prompt_token_ids)
+            elif hasattr(prompt_token_ids, "shape"):
+                return prompt_token_ids.shape[-1] if len(prompt_token_ids.shape) > 0 else 1
+        if prompt_embeds is not None:
+            if hasattr(prompt_embeds, "shape"):
+                return prompt_embeds.shape[-2] if len(prompt_embeds.shape) > 1 else 1
+        return 10  # Default length
+
+    # Mock in vllm.utils
+    monkeypatch.setattr(
+        "vllm.utils.length_from_prompt_token_ids_or_embeds",
+        _mock_length_from_prompt_token_ids_or_embeds,
+        raising=False,
+    )
+    # Also mock in processor module if it's imported
+    monkeypatch.setattr(
+        "vllm_omni.engine.processor.length_from_prompt_token_ids_or_embeds",
+        _mock_length_from_prompt_token_ids_or_embeds,
+        raising=False,
+    )
+    # If processor module is already imported, patch it directly
+    processor_module_path = "vllm_omni.engine.processor"
+    if processor_module_path in sys.modules:
+        processor_module = sys.modules[processor_module_path]
+        setattr(
+            processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds
+        )
+
+    # Strategy 3: Now mock async_omni_llm AFTER length_from_prompt_token_ids_or_embeds is mocked
+    # This prevents ImportError when async_omni_llm imports processor.py
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.async_omni_llm.init_tokenizer_from_configs",
+        _mock_init_tokenizer_from_configs,
+        raising=False,
+    )
+
+    # Strategy 4: If async_omni_llm is already imported, patch it directly
+    async_omni_llm_path = "vllm_omni.entrypoints.async_omni_llm"
+    if async_omni_llm_path in sys.modules:
+        async_omni_llm_module = sys.modules[async_omni_llm_path]
+        setattr(async_omni_llm_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+
+    # Now mock get_config and other functions
+    fake_hf_config = MagicMock()
+    fake_hf_config.model_type = "qwen2_5_omni"
+
+    def _mock_get_config(model, **kwargs):
+        return fake_hf_config
+
+    monkeypatch.setattr(
+        "vllm.transformers_utils.config.get_config",
+        _mock_get_config,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.get_config",
+        _mock_get_config,
+        raising=False,
+    )
+
+    # Mock transformers' cached_file to avoid downloading model configs
+    def _mock_cached_file(path_or_repo_id, *args, **kwargs):
+        import os
+        import tempfile
+
+        fake_config_file = os.path.join(tempfile.gettempdir(), "fake_config.json")
+        if not os.path.exists(fake_config_file):
+            with open(fake_config_file, "w") as f:
+                f.write('{"model_type": "qwen2_5_omni"}')
+        return fake_config_file
+
+    monkeypatch.setattr(
+        "transformers.utils.hub.cached_file",
+        _mock_cached_file,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        "transformers.utils.hub.cached_files",
+        lambda path_or_repo_id, filenames, **kwargs: (
+            [_mock_cached_file(path_or_repo_id, filenames[0])] if filenames else None
+        ),
+        raising=False,
+    )
+
+
+def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
+    """Test that stage configs are auto-loaded when stage_configs_path is None."""
+
+    def _fake_loader(model: str):
+        return [
+            _FakeStageConfig(fake_stage_config),
+            _FakeStageConfig(fake_stage_config),
+        ]
+
+    # Remove modules from cache BEFORE setting mocks
+    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]
+
+    # Set up mocks
+    _setup_engine_mocks(monkeypatch)
+    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    # Mock remove_old_logs to avoid bug (called before stage_list is created)
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+        lambda log_file, num_stages: None,
+        raising=False,
+    )
+
+    # Mock load_stage_configs_from_model
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    # Replace OmniStage
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg: _FakeStage(cfg),
+        raising=False,
+    )
+
+    # Import the module after mocks are set
+    import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+    # Patch the imported function and class in the 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)
+    # Verify: auto-loaded stage_configs and stage_list have consistent count
+    assert isinstance(llm.stage_configs, list)
+    assert len(llm.stage_configs) == 2
+    assert len(llm.stage_list) == 2
+    # Verify: each Stage is _FakeStage instance
+    for st in llm.stage_list:
+        assert isinstance(st, _FakeStage)
+    # Verify: queues are attached
+    for st in llm.stage_list:
+        assert st._in_q is not None
+        assert st._out_q is not None
+    # Verify: all stages are ready
+    assert len(llm._stages_ready) == 2
+
+
+def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
+    """Test that generate raises ValueError when sampling_params_list length doesn't match."""
+
+    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=["hi"], sampling_params_list=[])
+
+
+def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
+    """Test multi-stage generation pipeline with queue polling."""
+    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["processed_input"] = ["processed-for-stage-1"]
+
+    def _fake_loader(model: str):
+        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+    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)
+
+    # Simulate worker behavior: manually put results into output queues
+    # Note: We put results before calling generate, which simulates worker processes
+    # that have already completed. The polling loop will collect them in stage order.
+    # Stage 0 output (will be collected first)
+    llm.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": 0,
+            "engine_outputs": [{"stage": 0, "text": "s0"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    # Stage 1 output (will be collected after stage 0 forwards to it)
+    # Note: In real flow, stage 1 result would appear after stage 0 forwards,
+    # but for testing we pre-populate it. The polling loop processes stages
+    # in order, so stage 0 result will be collected first, then forwarded,
+    # then stage 1 result will be collected.
+    llm.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": 0,
+            "engine_outputs": [{"stage": 1, "text": "s1"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    sampling_params_list = [object(), object()]
+    prompts = ["hi"]
+    outputs = llm.generate(prompts=prompts, sampling_params_list=sampling_params_list)
+
+    # Both stages have final_output=True, so should aggregate two OmniRequestOutput
+    assert len(outputs) == 2
+    # Verify stage outputs are set
+    assert llm.stage_list[0].engine_outputs == [{"stage": 0, "text": "s0"}]
+    assert llm.stage_list[1].engine_outputs == [{"stage": 1, "text": "s1"}]
+    # Verify stage 0 input queue received the task
+    assert not llm.stage_list[0]._in_q.empty()
+    # Verify stage 1 received forwarded task (process_engine_inputs was called)
+    assert llm.stage_list[1].process_engine_inputs([], []) is not None
+
+
+def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
+    """Test that generate returns empty list when all stages have final_output=False."""
+    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["final_output"] = False
+    stage_cfg1["final_output"] = False
+
+    def _fake_loader(model: str):
+        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+    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)
+
+    # Simulate worker behavior: put results into output queues
+    llm.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": 0,
+            "engine_outputs": [{"stage": 0}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+    llm.stage_list[1]._out_q.put_nowait(
+        {
+            "request_id": 0,
+            "engine_outputs": [{"stage": 1}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    outputs = llm.generate(prompts=["p"], sampling_params_list=[object(), object()])
+    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."""
+
+    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,
+    )
+
+    # Create a stage that doesn't send stage_ready message
+    class _FakeStageNoReady(_FakeStage):
+        def init_stage_worker(self, *args, **kwargs):
+            # Don't send stage_ready message
+            self._proc = MagicMock()
+            self._proc.start = MagicMock()
+            self._proc.join = MagicMock()
+            self._proc.is_alive = MagicMock(return_value=False)
+            self._proc.terminate = MagicMock()
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.OmniStage",
+        lambda cfg: _FakeStageNoReady(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: _FakeStageNoReady(cfg))
+
+    from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+    # Use very short timeout
+    with pytest.raises(SystemExit):
+        OmniLLM(model="any", init_timeout=0.01)
+
+
+def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
+    """Test that generate handles error messages from stages correctly."""
+
+    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)
+
+    # Put error message in output queue
+    llm.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": 0,
+            "error": "test error",
+        }
+    )
+    # Also put a valid result after error to allow the loop to complete
+    # (error handling continues the loop, so we need a valid result to finish)
+    llm.stage_list[0]._out_q.put_nowait(
+        {
+            "request_id": 0,
+            "engine_outputs": [{"stage": 0, "text": "result"}],
+            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+        }
+    )
+
+    # Generate should handle error gracefully (log but continue)
+    sampling_params_list = [object()]
+    outputs = llm.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
+    # Should return final output (error was logged but didn't stop processing)
+    assert isinstance(outputs, list)
+    # Since final_output=True, should have one output
+    assert len(outputs) == 1
+
+
+def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
+    """Test that close() sends shutdown signal to all input queues."""
+
+    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)
+
+    # Call close
+    llm.close()
+
+    # Verify shutdown signal (None) was sent to input queue
+    # Use get_nowait to avoid blocking (close() uses put_nowait, so should be safe)
+    try:
+        shutdown_signal = llm.stage_list[0]._in_q.get_nowait()
+        assert shutdown_signal is None
+    except Empty:
+        # If queue was already empty or only had stage_ready, that's also acceptable
+        # The important thing is that close() was called without error
+        pass
+
+    # Verify stop_stage_worker was called (process should be set)
+    assert llm.stage_list[0]._proc is not None