From a951dd0cb4eeb9a6c28b2a075e62c2144b1bc272 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Thu, 30 Jul 2026 23:17:24 +0800 Subject: [PATCH] [BugFix] Never dispatch diffusion with an incomplete CFG companion bundle (#5482) Signed-off-by: tzhouam Co-authored-by: tzhouam Co-authored-by: Claude Fable 5 --- tests/engine/test_cfg_companion_lifecycle.py | 334 +++++++++++++++++++ tests/engine/test_orchestrator.py | 3 + vllm_omni/engine/cfg_companion_tracker.py | 30 +- vllm_omni/engine/orchestrator.py | 87 ++++- 4 files changed, 437 insertions(+), 17 deletions(-) create mode 100644 tests/engine/test_cfg_companion_lifecycle.py diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py new file mode 100644 index 000000000..5cf2fd072 --- /dev/null +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deterministic tests for the CFG-companion output lifecycle. + +Regression suite for the companion-output race behind the nightly Bagel +failures: ``kv_ready`` used to mark a companion done before its processed +output was stashed, letting the parent dispatch to the diffusion stage with +0/N companion outputs (degraded CFG conditioning, and after an abort a TP +rank desync). The invariant under test: **dispatch to the diffusion stage +never happens with an incomplete companion bundle** — the parent is +re-deferred while outputs are pending and failed when the bundle can no +longer complete. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker +from vllm_omni.engine.messages import AbortRequestMessage, ErrorMessage +from vllm_omni.engine.orchestrator import Orchestrator + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ───────── tracker-level ordering semantics ───────── + + +def test_companion_first_then_parent_is_not_deferred(): + tracker = CfgCompanionTracker() + tracker.register_companion("p", "neg", "p__neg") + tracker.set_companion_output("p__neg", {"o": 1}) + assert tracker.on_companion_completed("p__neg") is None # parent not deferred yet + # Parent arrives afterwards: bundle already complete, no defer needed. + assert tracker.all_companions_done("p") + assert tracker.get_companion_outputs("p") == [{"o": 1}] + + +def test_parent_first_deferred_then_companion_flushes(): + tracker = CfgCompanionTracker() + tracker.register_companion("p", "neg", "p__neg") + tracker.defer_parent("p", {"parent": True}, stage_id=0) + tracker.set_companion_output("p__neg", {"o": 1}) + assert tracker.on_companion_completed("p__neg") == "p" + deferred = tracker.pop_pending_parent("p") + assert deferred is not None and deferred["engine_outputs"] == {"parent": True} + + +def test_duplicate_companion_completion_is_idempotent(): + tracker = CfgCompanionTracker() + tracker.register_companion("p", "neg", "p__neg") + tracker.defer_parent("p", {}, stage_id=0) + tracker.set_companion_output("p__neg", {"o": 1}) + assert tracker.on_companion_completed("p__neg") == "p" + assert tracker.on_companion_completed("p__neg") is None # duplicate: no re-flush + assert tracker.get_companion_outputs("p") == [{"o": 1}] + + +def test_bundle_survives_resubmission_until_cleanup(): + """get_companion_outputs is non-destructive: a streaming re-submission of + the parent bundles the complete set again (never a partial/empty one). + Only cleanup_parent releases the outputs.""" + tracker = CfgCompanionTracker() + tracker.register_companion("p", "neg", "p__neg") + tracker.set_companion_output("p__neg", {"o": 1}) + assert tracker.get_companion_outputs("p") == [{"o": 1}] + assert tracker.get_companion_outputs("p") == [{"o": 1}] # re-submit: full bundle again + tracker.cleanup_parent("p") + assert tracker.get_companion_outputs("p") == [] + + +# ───────── orchestrator-method-level lifecycle ───────── + + +class _FakePool: + """Minimal stage pool satisfying _forward_to_next_stage's preamble and cleanup.""" + + def __init__(self, stage_type: str = "diffusion") -> None: + self.stage_type = stage_type + self.stage_client = SimpleNamespace(requires_multimodal_data=False, custom_process_input_func=None) + self.aborted: list[list[str]] = [] + + async def abort_requests(self, request_ids): + self.aborted.append(list(request_ids)) + + def release_bindings(self, request_ids): + pass + + +def _make_orchestrator(num_stages: int = 2) -> Orchestrator: + orch = Orchestrator.__new__(Orchestrator) + orch._cfg_tracker = CfgCompanionTracker() + orch.request_states = {} + orch.stage_pools = [_FakePool("llm"), _FakePool("diffusion")][:num_stages] + orch.output_async_queue = asyncio.Queue() + orch.async_chunk = False + orch.duplex_control_plane = None + orch._pd_kv_params = {} + orch._pd_pair = None + orch._running_counter = None + orch._transfer_emitter = None + return orch + + +def _req_state(final_stage_id: int = 1) -> SimpleNamespace: + return SimpleNamespace( + final_stage_id=final_stage_id, + stage_submit_ts={}, + sampling_params_list=[SimpleNamespace(output_kind=None)] * (final_stage_id + 1), + streaming=SimpleNamespace(enabled=False, segment_finished=False), + running_counter_registered=False, + prompt={"prompt": "x"}, + pipeline_timings={}, + ) + + +@pytest.mark.asyncio +async def test_kv_ready_does_not_mark_companion_done(): + """kv_ready alone must not count a companion as done — only the processed + output (set_companion_output) may. Guards the original race.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch.request_states["p__neg"] = _req_state() + + raw = SimpleNamespace(request_id="p__neg", kv_transfer_params={"kv_ready": True}) + await orch._handle_kv_ready_raw_outputs(0, SimpleNamespace(outputs=[raw])) + + assert not orch._cfg_tracker.all_companions_done("p") + assert orch._cfg_tracker.pop_pending_parent("p") is None + + +@pytest.mark.asyncio +async def test_parent_kv_ready_defers_until_companion_outputs_stashed(): + """Parent kv_ready with an incomplete bundle defers; the companion's + processed output later flushes it (companion-first vs parent-first order + both funnel through the same defer/flush path).""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch.request_states["p"] = _req_state() + + raw = SimpleNamespace(request_id="p", kv_transfer_params={"kv_ready": True}) + await orch._handle_kv_ready_raw_outputs(0, SimpleNamespace(outputs=[raw])) + + deferred = orch._cfg_tracker.pop_pending_parent("p") + assert deferred is not None and deferred["stage_id"] == 0 + + +@pytest.mark.asyncio +async def test_forward_re_defers_when_bundle_incomplete(): + """The forward gate re-defers (non-destructively) instead of dispatching + with 0/N companion outputs.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.set_companion_output("p__neg", {"o": 1}) # stashed but NOT done-marked + state = _req_state() + orch.request_states["p"] = state + + await orch._forward_to_next_stage("p", 0, {"parent": True}, state) + + # Not dispatched: re-deferred with the stashed output left intact. + deferred = orch._cfg_tracker.pop_pending_parent("p") + assert deferred is not None + assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}] + assert orch.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_forward_fails_request_when_bundle_lost(): + """Companions done but outputs gone without a prior bundle: the request is + failed (terminal ErrorMessage + abort) instead of dispatching degraded CFG.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.set_companion_output("p__neg", {"o": 1}) + orch._cfg_tracker.on_companion_completed("p__neg") + orch._cfg_tracker._companion_outputs.pop("p__neg") # simulate lost output (unreachable normally) + state = _req_state() + orch.request_states["p"] = state + + await orch._forward_to_next_stage("p", 0, {"parent": True}, state) + + msg = orch.output_async_queue.get_nowait() + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "p" + assert "incomplete" in msg.error + assert "p" not in orch.request_states + assert any("p" in batch for pool in orch.stage_pools for batch in pool.aborted) + + +@pytest.mark.asyncio +async def test_companion_abort_fails_deferred_parent(): + """A companion aborted directly (parent not in the batch) can never + complete the bundle — the deferred parent is failed immediately.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0) + orch.request_states["p"] = _req_state() + + await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"])) + + msg = orch.output_async_queue.get_nowait() + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "p" + assert "p__neg" in msg.error + assert "p" not in orch.request_states + assert not orch._cfg_tracker.is_companion("p__neg") + # Idempotent: a repeated abort of the same companion is a no-op. + await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"])) + assert orch.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_companion_abort_after_completion_does_not_fail_parent(): + """A late/duplicate abort of a companion whose output is already stashed + must not cancel the (possibly already dispatched) parent.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.set_companion_output("p__neg", {"o": 1}) + orch._cfg_tracker.on_companion_completed("p__neg") + orch.request_states["p"] = _req_state() + + await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"])) + + assert orch.output_async_queue.empty() # no parent failure emitted + assert "p" in orch.request_states # parent untouched + assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}] # bundle intact + + +class _CapturedDispatchError(Exception): + """Sentinel raised by the capture hook to stop the forward after bundling.""" + + +@pytest.mark.asyncio +async def test_repeated_forward_bundles_full_set_again(): + """A second forward of the same parent (streaming terminal update) must + dispatch the complete [parent, *companions] bundle again — never [parent] + alone. Captured at the diffusion input hook, i.e. actual dispatch content.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.set_companion_output("p__neg", {"o": 1}) + orch._cfg_tracker.on_companion_completed("p__neg") + state = _req_state() + orch.request_states["p"] = state + + captured: list[list] = [] + + def _capture(source_outputs, prompt, requires_multimodal_data): + captured.append(list(source_outputs)) + raise _CapturedDispatchError + + orch.stage_pools[1].stage_client.custom_process_input_func = _capture + + for _ in range(2): + with pytest.raises(_CapturedDispatchError): + await orch._forward_to_next_stage("p", 0, {"parent": True}, state) + + assert captured == [[{"parent": True}, {"o": 1}], [{"parent": True}, {"o": 1}]] + + +@pytest.mark.asyncio +async def test_replica_loss_cleanup_releases_parent_tracker_state(): + """Replica-loss/membership teardown calls _cleanup_request_ids directly + (no cleanup_parent at the call site): the parent's tracker state must be + released and its companions folded into the same cleanup batch.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.set_companion_output("p__neg", {"o": 1}) + orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0) + orch.request_states["p"] = _req_state() + + await orch._cleanup_request_ids(["p"], abort=True) + + assert not orch._cfg_tracker.has_companions("p") + assert not orch._cfg_tracker.is_companion("p__neg") + assert orch._cfg_tracker.get_companion_outputs("p") == [] + assert orch._cfg_tracker.pop_pending_parent("p") is None + aborted = {rid for pool in orch.stage_pools for b in pool.aborted for rid in b} + assert {"p", "p__neg"} <= aborted + + +@pytest.mark.asyncio +async def test_replica_loss_of_companion_fails_deferred_parent(): + """A companion torn down by replica loss before its output arrived means + the parent can never complete its bundle — the cleanup fails the parent.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0) + orch.request_states["p"] = _req_state() + + await orch._cleanup_request_ids(["p__neg"]) + + msg = orch.output_async_queue.get_nowait() + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "p" + assert "p__neg" in msg.error + assert "p" not in orch.request_states + assert not orch._cfg_tracker.has_companions("p") + + +@pytest.mark.asyncio +async def test_finished_companion_cleanup_preserves_parent_bundle(): + """The normal companion-finished cleanup (_route_output cleans just the + companion id) must not disturb the parent's stashed bundle.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.set_companion_output("p__neg", {"o": 1}) + orch._cfg_tracker.on_companion_completed("p__neg") + orch.request_states["p"] = _req_state() + + await orch._cleanup_request_ids(["p__neg"]) + + assert orch.output_async_queue.empty() + assert "p" in orch.request_states + assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}] + + +@pytest.mark.asyncio +async def test_parent_abort_still_expands_to_companions(): + """Regression-guard the existing behavior: aborting the parent aborts its + companions and clears tracker state, without emitting an error.""" + orch = _make_orchestrator() + orch._cfg_tracker.register_companion("p", "neg", "p__neg") + orch._cfg_tracker.defer_parent("p", {"parent": True}, stage_id=0) + orch.request_states["p"] = _req_state() + + await orch._handle_abort(AbortRequestMessage(request_ids=["p"])) + + assert orch.output_async_queue.empty() + assert "p" not in orch.request_states + assert not orch._cfg_tracker.is_companion("p__neg") + aborted = {rid for pool in orch.stage_pools for batch in pool.aborted for rid in batch} + assert {"p", "p__neg"} <= aborted diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 5c7066e0c..c058ed907 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2282,12 +2282,15 @@ async def test_request_cleanup_failure_is_deferred_to_control_plane(): def finalize_closed_sessions(self, session_ids: list[str]) -> None: self.finalized.extend(session_ids) + from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker + orchestrator = object.__new__(Orchestrator) orchestrator.stage_pools = [_FailingPool()] orchestrator.duplex_control_plane = _Plane() orchestrator._pd_kv_params = {} orchestrator.request_states = {} orchestrator._running_counter = None + orchestrator._cfg_tracker = CfgCompanionTracker() with pytest.raises(RuntimeError, match="stage abort failed"): await orchestrator._cleanup_request_ids( diff --git a/vllm_omni/engine/cfg_companion_tracker.py b/vllm_omni/engine/cfg_companion_tracker.py index aeec9e848..904a608fb 100644 --- a/vllm_omni/engine/cfg_companion_tracker.py +++ b/vllm_omni/engine/cfg_companion_tracker.py @@ -57,16 +57,28 @@ class CfgCompanionTracker: """Stash companion engine output for the parent to bundle at forward time.""" self._companion_outputs[companion_id] = output - def pop_companion_outputs(self, parent_id: str) -> list[Any]: - """Pop companion outputs (role-registration order) for bundling into set_engine_outputs.""" + def get_companion_outputs(self, parent_id: str) -> list[Any]: + """Return companion outputs (role-registration order) for bundling. + + Non-destructive: outputs stay stashed until ``cleanup_parent`` so a + re-submission of the same parent (e.g. a streaming terminal update) + bundles the complete set again instead of an empty one. + """ role_map = self._companion_map.get(parent_id, {}) outputs = [] for cid in role_map.values(): - out = self._companion_outputs.pop(cid, None) + out = self._companion_outputs.get(cid) if out is not None: outputs.append(out) return outputs + def is_companion_done(self, companion_id: str) -> bool: + """True once this companion's processed output was stashed and marked done.""" + parent_id = self._companion_to_parent.get(companion_id) + if parent_id is None: + return False + return companion_id in self._done.get(parent_id, set()) + def on_companion_completed(self, companion_id: str) -> str | None: """Mark done. Returns parent_id only if parent is pending and all companions finished.""" parent_id = self._companion_to_parent.get(companion_id) @@ -83,10 +95,14 @@ class CfgCompanionTracker: return None def defer_parent(self, parent_id: str, engine_outputs: Any, stage_id: int) -> None: - """Hold parent result while waiting for companions to finish.""" - # TODO: Add timeout/error recovery when the orchestrator grows a - # companion-failure path. Today deferred parents are released only when - # companions finish or the external layer aborts the request. + """Hold parent result while waiting for companions to finish. + + Deferred parents are released when their companions' processed outputs + arrive, and failed by the orchestrator when a companion errors or is + aborted. TODO: a deadline for deferred parents and detection of a + companion replica dying silently (no error output, no abort) still + need a liveness owner in the orchestrator loop. + """ self._pending_parents[parent_id] = { "engine_outputs": engine_outputs, "stage_id": stage_id, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 53f30ea48..963054688 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -772,10 +772,10 @@ class Orchestrator: async def _handle_abort(self, msg: AbortRequestMessage) -> None: """Handle an abort message from the main thread.""" request_ids = msg.request_ids - await self._cleanup_request_ids( - self._cfg_tracker.abort_parents(request_ids), - abort=True, - ) + # _cleanup_request_ids is CFG-aware: it expands aborted parents to + # their companions and fails a deferred parent whose companion is + # aborted before its output arrived. + await self._cleanup_request_ids(list(request_ids), abort=True) logger.info("[Orchestrator] Aborted request(s) %s", request_ids) async def _handle_interaction(self, msg: InteractionMessage) -> None: @@ -1098,11 +1098,41 @@ class Orchestrator: abort: bool = False, close_duplex_sessions: bool = False, ) -> None: - """Release pool bindings and logical request state for the given ids.""" + """Release pool bindings and logical request state for the given ids. + + CFG-aware: cleaning a parent releases its tracker state and pulls its + companions into the batch; cleaning a companion whose parent is NOT in + the batch and whose output never arrived fails that parent (its bundle + can never complete). Every teardown path (stage error, abort, replica + loss, membership unregister) funnels through here, so tracker state + cannot outlive its requests. + """ if not request_ids: return cleanup_ids = list(dict.fromkeys(request_ids)) + batch = set(cleanup_ids) + orphaned_parents: dict[str, str] = {} + for rid in cleanup_ids: + pid = self._cfg_tracker.get_parent_id(rid) + if pid is not None and pid not in batch and not self._cfg_tracker.is_companion_done(rid): + orphaned_parents.setdefault(pid, rid) + for pid, cid in orphaned_parents.items(): + deferred = self._cfg_tracker.pop_pending_parent(pid) + await self.output_async_queue.put( + ErrorMessage( + request_id=pid, + stage_id=deferred["stage_id"] if deferred is not None else None, + error=f"CFG companion {cid} was aborted or lost before its outputs arrived", + ) + ) + batch.add(pid) + cleanup_ids.append(pid) + for rid in list(cleanup_ids): + for cid in self._cfg_tracker.cleanup_parent(rid): + if cid not in batch: + batch.add(cid) + cleanup_ids.append(cid) closing_session_ids: list[str] = [] if close_duplex_sessions and self.duplex_control_plane is not None: closed_sessions = self.duplex_control_plane.close_sessions_for_request_ids( @@ -1596,7 +1626,12 @@ class Orchestrator: if req_state is None: continue if self._cfg_tracker.is_companion(req_id): - await self._handle_cfg_companion_ready(req_id) + # kv_ready only says the companion's KV hit the connector; its + # processed output has not been stashed yet. Counting it as + # done here let the parent pass the all_companions_done gate + # and dispatch with 0/N companion outputs (degraded CFG). + # Done-marking now happens solely on the processed-output path + # in _route_output, after set_companion_output. continue if stage_id >= req_state.final_stage_id: continue @@ -1711,16 +1746,48 @@ class Orchestrator: _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": - companion_outputs = self._cfg_tracker.pop_companion_outputs(req_id) + # Gate: never dispatch with an incomplete CFG bundle. Checked + # non-destructively BEFORE popping — a pop-then-redefer would lose + # the partial outputs. all_companions_done is trustworthy here + # because done-marking happens only after set_companion_output. + if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): + self._cfg_tracker.defer_parent(req_id, output, src_stage_id) + logger.info( + "[Orchestrator] req=%s: CFG companion outputs not all stashed yet; re-deferring parent", + req_id, + ) + return + # Peek, don't pop: outputs stay stashed until cleanup_parent so a + # streaming re-submission bundles the complete set again rather + # than an empty one. + companion_outputs = self._cfg_tracker.get_companion_outputs(req_id) expected = len(self._cfg_tracker.get_companion_request_ids(req_id)) if expected > len(companion_outputs): - logger.warning( - "[Orchestrator] req=%s: only %d/%d CFG companion outputs arrived; " - "downstream CFG conditioning may degrade", + # Companions are done but outputs are missing — inconsistent + # tracker state (should be unreachable with peek semantics). + # Fail the request rather than degrade CFG conditioning. + logger.error( + "[Orchestrator] req=%s: only %d/%d CFG companion outputs available; " + "failing the request instead of dispatching degraded CFG", req_id, len(companion_outputs), expected, ) + await self.output_async_queue.put( + ErrorMessage( + request_id=req_id, + stage_id=src_stage_id, + error=( + f"CFG companion outputs incomplete ({len(companion_outputs)}/{expected}); " + "request aborted to avoid degraded CFG conditioning" + ), + ) + ) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + ) + return diffusion_source_outputs = [output, *companion_outputs] if next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter()