[Core] Refactor CFG companion tracker and use in Orchestrator

[yinpeiqi]

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

Purpose

The CfgCompaionTracker, originally used in the AsyncOmni and Omni class (before #1908 refactor), and now is deperated.

Before this change, the CFG companion lifecycle was managed directly inside Orchestrator. That made the orchestration layer responsible for too many low-level details:

• parent -> companion mapping
• companion completion tracking
• deferred parent forwarding
• cleanup and abort handling
• diffusion-stage CFG KV request ID attachment

This had a few problems:

• Orchestrator exposed too much CFG-specific state and bookkeeping logic.
• The CFG flow was spread across multiple methods, which made the control flow harder to follow.
• State ownership was unclear: companion-related data lived in Orchestrator, while a separate CfgCompanionTracker already existed but did not match the current runtime path.
• The abstraction boundary was weak, which made future CFG changes riskier and harder to reason about.

What changed

This refactor moves CFG companion state ownership out of Orchestrator and into CfgCompanionTracker.

Orchestrator now initializes a single tracker instance and delegates CFG-specific state management to it. The tracker is responsible for:

• registering parent/companion relationships
• checking whether a request is a companion
• tracking companion completion
• deferring parent forwarding until companions are ready
• attaching cfg_kv_request_ids for diffusion requests
• cleaning up companion state
• expanding aborts from parent requests to companion requests

At the same time, old tracker logic that no longer matches the current orchestrator flow was removed, including unused prompt-expansion, timeout, and failure-propagation paths.

cc: @fake0fan @princepride @natureofnature @gcanlin @hsliuustc0106

Test Plan

Test Result

---

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. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (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)

[hsliuustc0106]

do we have any perf/acc regression example tests or just run the CI？

[yinpeiqi]

do we have any perf/acc regression example tests or just run the CI？

Just need to run the CI for BAGEL, this PR do not change the execution logic.

[natureofnature]

do we have any perf/acc regression example tests or just run the CI？

Just need to run the CI for BAGEL, this PR do not change the execution logic.

You can refer to

VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_CLEAN_GPU_MEMORY=1 VLLM_IMAGE_FETCH_TIMEOUT=60  pytest tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/offline_inference/test_bagel_text2img.py tests/e2e/online_serving/test_bagel_online.py tests/e2e/online_serving/test_bagel_expansion.py -v -m "advanced_model" --run-level advanced_model

for both online and offline L3 test @yinpeiqi

[fake0fan]

At first glance, it feels very strange to put this file under the entrypoint.

[hsliuustc0106]

agree, this too feature specific

[yinpeiqi]

Moved to engine folder.

[yinpeiqi]

This PR could be merge? @gcanlin @hsliuustc0106

[gcanlin]

LGTM. @princepride Could you double-check it?

[lishunyang12]

Good refactor. Moving CFG companion state out of Orchestrator and into a dedicated CfgCompanionTracker class makes the responsibilities clearer and the orchestrator loop easier to follow. The key rename from "output" to "engine_outputs" is consistently applied across defer_parent, pop_pending_parent, and the consuming code in _handle_cfg_companion_ready. Tests cover the main paths well.

A few observations:

1. abort_parents does not clean up companion-only aborts
If abort_parents receives a companion ID (not a parent), it passes it through without removing it from _companion_ids, _companion_to_parent, or _done. This leaves stale tracking state. In practice this is probably fine because _handle_abort is always called with parent IDs from the external layer, but it is worth a comment or a defensive cleanup to avoid confusion if the method is reused elsewhere.

2. on_companion_completed assumes _done[parent_id] exists
Line self._done[parent_id].add(companion_id) will KeyError if the parent was never registered. This is safe today because register_companion always initializes _done via register_parent, but a guard or assertion (assert parent_id in self._done) would make the invariant explicit and give a better error message if the contract is violated.

3. Removed _upgrade_to_omni_request call in async_omni_engine.py
The deletion of request = _upgrade_to_omni_request(request, companion_prompt) at line 1075 looks intentional but is not mentioned in the PR description. Is this line dead code after #1908, or was it still doing something? Worth confirming this does not affect the companion request payload.

4. Removal of timeout / failure-propagation paths
The old tracker had check_timeouts(), on_companion_error(), is_parent_failed(), and consume_parent_failure(). The PR description says these are unused in the current orchestrator flow. That seems correct from the diff -- the orchestrator never called them. Just confirming this is intentional: if a companion hangs indefinitely the parent will remain deferred forever. If that is an acceptable risk today, a TODO comment noting the lack of timeout would be helpful for future readers.

5. Minor nit: typo in PR title
"compaion" should be "companion".

Overall this is clean and well-scoped. LGTM.

Replacing with inline comments

[yinpeiqi]

Good refactor. Moving CFG companion state out of Orchestrator and into a dedicated CfgCompanionTracker class makes the responsibilities clearer and the orchestrator loop easier to follow. The key rename from "output" to "engine_outputs" is consistently applied across defer_parent, pop_pending_parent, and the consuming code in _handle_cfg_companion_ready. Tests cover the main paths well.

A few observations:

1. abort_parents does not clean up companion-only aborts If abort_parents receives a companion ID (not a parent), it passes it through without removing it from _companion_ids, _companion_to_parent, or _done. This leaves stale tracking state. In practice this is probably fine because _handle_abort is always called with parent IDs from the external layer, but it is worth a comment or a defensive cleanup to avoid confusion if the method is reused elsewhere.

2. on_companion_completed assumes _done[parent_id] exists Line self._done[parent_id].add(companion_id) will KeyError if the parent was never registered. This is safe today because register_companion always initializes _done via register_parent, but a guard or assertion (assert parent_id in self._done) would make the invariant explicit and give a better error message if the contract is violated.

3. Removed _upgrade_to_omni_request call in async_omni_engine.py The deletion of request = _upgrade_to_omni_request(request, companion_prompt) at line 1075 looks intentional but is not mentioned in the PR description. Is this line dead code after #1908, or was it still doing something? Worth confirming this does not affect the companion request payload.

4. Removal of timeout / failure-propagation paths The old tracker had check_timeouts(), on_companion_error(), is_parent_failed(), and consume_parent_failure(). The PR description says these are unused in the current orchestrator flow. That seems correct from the diff -- the orchestrator never called them. Just confirming this is intentional: if a companion hangs indefinitely the parent will remain deferred forever. If that is an acceptable risk today, a TODO comment noting the lack of timeout would be helpful for future readers.

5. Minor nit: typo in PR title "compaion" should be "companion".

Overall this is clean and well-scoped. LGTM.

Fixed according to the comments. For 3, _upgrade_to_omni_request is used only for the TTS model with input additional_information. Not affect to BAGEL, so i directly remove this one here.