Merge branch 'main' into fix/cfg-companion-output-race

Nightly Bagel failures (builds 2892/2918) traced to a companion-output
race: _handle_kv_ready_raw_outputs marked a companion done on kv_ready
alone, before its processed output was stashed, so the parent could pass
the all_companions_done gate and _forward_to_next_stage dispatched with
0/N companion outputs - silently degraded CFG conditioning, and after an
abort a TP rank desync (mismatched all_reduce sizes, NCCL watchdog kill)
in the Bagel text-embedding path.

Lifecycle now enforces one invariant - diffusion dispatch never happens
with an incomplete companion bundle:
- a companion counts as done only after set_companion_output stashes its
  processed output (_route_output is the sole done-marking site);
- the forward gate checks completeness non-destructively and re-defers
  the parent instead of popping partial state; a bundle that is done but
  gone without ever being consumed fails the request instead of
  dispatching degraded CFG (was_bundled distinguishes streaming
  re-submissions);
- a companion aborted directly (parent not in the batch) fails its
  deferred parent immediately via the existing cleanup plumbing.

Deferred to a follow-up (needs a liveness owner in the orchestrator
loop): a deadline on deferred parents and detection of a companion
replica dying silently.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@gmail.com>

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

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(

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,

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()