[CI] Add test for omni_llm.py.

[congw729]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

We focus on the OmniLLM class (vllm_omni/entrypoints/omni_llm.py), which coordinates multi-stage inference. Each stage runs in its own process, communicates via queues, and is driven by the orchestrator loop. The tests cover the following areas:

1. Initialization

   • Load stage configs automatically from either the model registry or a YAML file.
   • Create stage instances in parallel and attach their input/output queues.
   • Start stage workers (mocked in tests) without spawning real processes.
   • Wait for every stage to emit a stage_ready message before generation begins.

2. Generation

   • Validate sampling_params_list (non-null, matching number of stages).
   • Seed stage 0 by submitting queue tasks that contain prompts and sampling params.
   • Poll each stage output queue, decode IPC payloads, and persist results through set_engine_outputs.
   • Forward intermediate outputs using process_engine_inputs and submit when more stages exist.
   • Aggregate OmniRequestOutput items coming from stages flagged with final_output=True.

3. Error handling and shutdown

   • Log error messages emitted by stages without tearing down the pipeline immediately.
   • Ensure close() pushes None shutdown markers into every input queue and invokes stop_stage_worker() on all stages.

Test Plan

pytest tests/test_omni_llm.py

Test	Purpose	Relevant Code Path
test_initialize_stage_configs_called_when_none	Verifies OmniLLM.__init__ loads configs via load_stage_configs_from_model, attaches queues, and waits for all stage_ready messages.	omni_llm.py lines 67-82
test_generate_raises_on_length_mismatch	Confirms generate() raises ValueError when the number of sampling params does not match the number of stages.	omni_llm.py lines 296-301
test_generate_sampling_params_none_raises	Ensures generate() rejects sampling_params_list=None with a ValueError.	omni_llm.py lines 182-184
test_generate_pipeline_and_final_outputs	Exercises the full multi-stage pipeline (task submission, IPC decode, forwarding, final output aggregation).	omni_llm.py lines 226-400
test_generate_no_final_output_returns_empty	Checks that the pipeline returns an empty list when every stage has final_output=False.	omni_llm.py lines 293-300
test_wait_for_stages_ready_timeout	Validates _wait_for_stages_ready times out cleanly, logs missing stages, and exits.	omni_llm.py lines 404-464
test_generate_handles_error_messages	Ensures error messages from stages are logged and the loop keeps processing remaining results.	omni_llm.py lines 257-264
test_close_sends_shutdown_signal	Verifies close() pushes None to every input queue and calls stop_stage_worker for each stage.	omni_llm.py lines 134-147

Test Result

All tests are passed.

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Avoid 20s sleeps in tests by setting init_sleep_seconds

OmniLLM sleeps init_sleep_seconds after starting each stage (omni_llm.py:131-132) and the default is 20s. The tests instantiate OmniLLM without overriding that default here, so every run waits 20s per stage (40s for the two-stage cases) before exercising any logic, making the new suite take minutes and even the timeout test blocks for ~20s first. Please pass init_sleep_seconds=0 or patch time.sleep so the tests execute promptly.

Useful? React with 👍 / 👎.

[ywang96]

Closing as superseded by #93