From 9ef9fc442a8a85c8e6b42a46412ceffaec7e0999 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:10:37 +0800 Subject: [PATCH] [Bugfix] Send a downstream terminal chunk when a parked stage ends (#6889) Signed-off-by: psv666 <2693925048@qq.com> --- .../test_chunk_transfer_adapter.py | 350 ++++++++++++++++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 62 ++++ .../chunk_transfer_adapter.py | 120 +++++- 3 files changed, 519 insertions(+), 13 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 40d29e969..a80ded5f0 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -2526,6 +2526,356 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): assert order == ["adapter", "super"] +def _parked_sender_scheduler(adapter, session): + """A minimal OmniARScheduler holding one parked streaming session.""" + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + scheduler.requests = {session.request_id: session} + scheduler.running = [] + scheduler.waiting = DummyWaitingQueue() + scheduler.skipped_waiting = DummyWaitingQueue([session]) + scheduler.num_waiting_for_streaming_input = 1 + scheduler.finished_req_ids_dict = {} + + def _free_request(self, request, delay_free_blocks=False): + del delay_free_blocks + self.requests.pop(request.request_id, None) + return None, None + + scheduler._free_request = MethodType(_free_request, scheduler) + return scheduler + + +@pytest.mark.parametrize("send_first_chunk", [False, True]) +def test_parked_streaming_final_update_emits_downstream_terminal(build_adapter, send_first_chunk): + """#6670: a final update on a parked sender must terminate the receiver. + + Upstream turns "final update while parked in WAITING_FOR_STREAMING_REQ" + into a silent local abort. The downstream stage only learns a stream ended + from a chunk carrying ``finished``, so the abort used to leave it in + WAITING_FOR_CHUNK until the 600s input deadline. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + session = _req( + "req-parked-final", + RequestStatus.WAITING_FOR_STREAMING_REQ, + external_req_id="ext-parked-final", + ) + session.resumable = True + session._omni_segment_generation = 0 + # A segment stop queues its boundary before the final update arrives. + # The downstream request is already prewarmed even if the sender has not + # dequeued this first task, so it still needs a terminal in both cases. + adapter.save_async(None, session, is_segment_finished=True) + assert session.external_req_id not in adapter.put_req_chunk + if send_first_chunk: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.reset_mock() + + scheduler = _parked_sender_scheduler(adapter, session) + + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + OmniARScheduler.add_request(scheduler, final_update) + + # Local teardown happened... + assert session.request_id not in scheduler.requests + # ...but only after a terminal chunk was queued for the next stage. + assert session.status == RequestStatus.FINISHED_STOPPED + assert any(task["is_finished"] for task in adapter._pending_save_reqs) + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + connector.put.assert_called_once() + assert connector.put.call_args.kwargs["put_key"] == f"ext-parked-final_1_{int(send_first_chunk)}" + payload = connector.put.call_args.kwargs["data"] + assert bool(payload.meta.finished.item()) is True + assert bool(payload.meta.is_segment_finished.item()) is False + + +def test_queued_terminal_chunk_survives_sender_cleanup(build_adapter): + """#6670: ``cleanup_sender`` must not discard an already-queued terminal.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal") + request.resumable = False + + adapter.save_async(None, request, is_segment_finished=False) + assert len(adapter._pending_save_reqs) == 1 + + # The scheduler tears the request down before the save_loop drains. + adapter.cleanup_sender("ext-terminal") + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_called_once() + assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True + # The terminal put reclaims the generation it kept alive. + assert "ext-terminal" not in adapter._sender_tokens + + +def test_non_terminal_chunk_is_still_dropped_by_sender_cleanup(build_adapter): + """The #6670 fence is terminal-only: ordinary queued chunks still drop.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-mid", RequestStatus.RUNNING, external_req_id="ext-mid") + request.resumable = True + + adapter.save_async(None, request, is_segment_finished=False) + adapter.cleanup_sender("ext-mid") + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_not_called() + + +def test_final_update_on_non_sender_stage_keeps_upstream_abort(build_adapter): + """A final stage never put a chunk, so it must keep the upstream path.""" + adapter, _connector = build_adapter(stage_id=2, model_mode="generation") + session = _req( + "req-final-stage", + RequestStatus.WAITING_FOR_STREAMING_REQ, + external_req_id="ext-final-stage", + ) + session.resumable = True + scheduler = _parked_sender_scheduler(adapter, session) + + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + + with patch.object(VLLMScheduler, "add_request") as base_add: + OmniARScheduler.add_request(scheduler, final_update) + + base_add.assert_called_once() + assert not adapter._pending_save_reqs + + +def _queue_terminal( + adapter, + external_req_id: str, + req_id: str = "req-terminal", + segment_generation: int = 0, +) -> None: + """Queue one request-terminal chunk on the adapter's save queue. + + ``segment_generation`` must match the adapter's dedup watermark, which a + preceding segment stop advances; otherwise ``save_async`` drops the chunk + as a late duplicate. + """ + request = _req(req_id, RequestStatus.FINISHED_STOPPED, external_req_id=external_req_id) + request.resumable = False + request._omni_segment_generation = segment_generation + adapter.save_async(None, request, is_segment_finished=False) + + +def _assert_terminal_put(connector) -> None: + connector.put.assert_called_once() + assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True + + +def test_terminal_chunk_survives_cleanup_after_dequeue(build_adapter): + """#6670: the fence must hold across the connector handoff, not just the queue. + + ``_finish_parked_streaming_session`` queues the terminal and then calls + ``finish_requests``, so ``cleanup_sender`` can land *after* the save_loop + has already dequeued the terminal. Releasing the fence at dequeue left + that interleaving cancelling the chunk before ``put`` -- the same lost + terminal, and the same downstream WAITING_FOR_CHUNK hang. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-late-cleanup") + task = adapter._pending_save_reqs.popleft() + + real_send = adapter._send_single_request_for_generation + + def send_with_cleanup_in_flight(inner_task, sender_token=None): + # The scheduler thread reaches cleanup_sender here: the save_loop owns + # the terminal but has not handed it to the connector yet. + adapter.cleanup_sender("ext-late-cleanup") + return real_send(inner_task, sender_token) + + adapter._send_single_request_for_generation = send_with_cleanup_in_flight + adapter._send_single_request(task) + + _assert_terminal_put(connector) + # The generation the fence kept alive is reclaimed, not leaked. + assert "ext-late-cleanup" not in adapter._sender_tokens + assert "ext-late-cleanup" not in adapter.put_req_chunk + + +def test_terminal_chunk_survives_cleanup_racing_an_inflight_sibling(build_adapter): + """#6670: an in-flight ordinary chunk must not reclaim a fenced generation. + + ``_send_single_request``'s ``finally`` also tears a cancelled generation + down. If the stop and the abort both land while an earlier chunk is on the + wire, that teardown would unregister the token and the queued terminal + would then fail the identity check. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + running = _req("req-sibling", RequestStatus.RUNNING, external_req_id="ext-sibling") + running.resumable = True + adapter.save_async(None, running, is_segment_finished=True) + ordinary_task = adapter._pending_save_reqs.popleft() + + real_send = adapter._send_single_request_for_generation + + def send_then_stop_and_abort(inner_task, sender_token=None): + result = real_send(inner_task, sender_token) + # Terminal stop and abort both land while this chunk is on the wire. + # The segment stop above already armed the next segment's watermark. + _queue_terminal(adapter, "ext-sibling", req_id="req-sibling", segment_generation=1) + adapter.cleanup_sender("ext-sibling") + return result + + adapter._send_single_request_for_generation = send_then_stop_and_abort + adapter._send_single_request(ordinary_task) + adapter._send_single_request_for_generation = real_send + + assert len(adapter._pending_save_reqs) == 1 + connector.put.reset_mock() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + _assert_terminal_put(connector) + assert "ext-sibling" not in adapter._sender_tokens + + +@pytest.mark.parametrize("finish_inside_put", [False, True]) +def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_inside_put): + """A cleanup during put must not reuse a key already consumed downstream.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + receiver, receiver_connector = build_adapter(stage_id=2, model_mode="generation") + session = _req("req-sequence", RequestStatus.WAITING_FOR_STREAMING_REQ, external_req_id="ext-sequence") + session.resumable = True + session._omni_segment_generation = 0 + scheduler = _parked_sender_scheduler(adapter, session) + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + downstream = _req("req-downstream", RequestStatus.WAITING, external_req_id=session.external_req_id) + receiver.request_ids_mapping[downstream.request_id] = session.external_req_id + payloads = {} + + def receive_next(): + entry = _dequeue_load_entry(receiver, downstream) + return receiver._poll_single_request(entry) + + def put_then_finish(**kwargs): + payload = kwargs["data"] + payloads[kwargs["put_key"]] = { + "meta": { + "finished": bool(payload.meta.finished.item()), + "is_segment_finished": bool(payload.meta.is_segment_finished.item()), + } + } + if not payload.meta.finished.item(): + # Consume the ordinary boundary while put is still in flight. + assert receive_next() + if finish_inside_put: + scheduler.add_request(final_update) + return True, 1, {} + + def get_by_key(_from_stage, _to_stage, key): + payload = payloads.pop(key, None) + return (payload, 1) if payload is not None else None + + connector.put.side_effect = put_then_finish + receiver_connector.get.side_effect = get_by_key + adapter.save_async(None, session, is_segment_finished=True) + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + if not finish_inside_put: + scheduler.add_request(final_update) + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert [call.kwargs["put_key"] for call in connector.put.call_args_list] == ["ext-sequence_1_0", "ext-sequence_1_1"] + assert receive_next() + assert downstream.request_id in receiver.upstream_exhausted_requests + assert not payloads + assert session.external_req_id not in adapter._sender_tokens + + +def test_late_successful_put_cannot_advance_replacement_generation(build_adapter): + """The post-put guard must retain identity checks while relaxing cancellation.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-old", RequestStatus.RUNNING, external_req_id="ext-replaced") + request.resumable = True + adapter.save_async(None, request, is_segment_finished=True) + task = adapter._pending_save_reqs.popleft() + old_token = task["sender_token"] + replacement_token = type(old_token)() + + def replace_during_put(**kwargs): + # Model a completion that no longer owns its external id. The old + # terminal fence must never authorize writes to the successor's state. + with adapter._sender_state_lock: + old_token.cancelled = True + old_token.terminal_pending = True + adapter._sender_tokens[request.external_req_id] = replacement_token + adapter.put_req_chunk[request.external_req_id] = 4 + adapter.ramp_chunk_count[request.external_req_id] = 2 + return True, 1, {} + + connector.put.side_effect = replace_during_put + adapter._send_single_request(task) + assert adapter._sender_tokens[request.external_req_id] is replacement_token + assert adapter.put_req_chunk[request.external_req_id] == 4 + assert adapter.ramp_chunk_count[request.external_req_id] == 2 + + +def test_new_request_cannot_join_a_generation_with_a_queued_terminal(build_adapter): + """#6670: the fence must not readmit a request reusing the external id. + + ``save_async`` rejects chunks whose generation is ``cancelled``; that is + what stops a new request from inheriting a retiring generation's token and + chunk counters. Deferring the reclaim must not also defer that retirement. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-shared", req_id="req-old") + adapter.cleanup_sender("ext-shared") + + newcomer = _req("req-new", RequestStatus.RUNNING, external_req_id="ext-shared") + newcomer.resumable = True + adapter.save_async(None, newcomer, is_segment_finished=True) + + # Rejected loudly, so the scheduler fails the newcomer now instead of its + # chunks being silently discarded at send time. + assert len(adapter._pending_save_reqs) == 1 + assert adapter.collect_failed_send_request_ids() == {"req-new": "previous sender generation is still draining"} + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + _assert_terminal_put(connector) + assert "ext-shared" not in adapter._sender_tokens + + +def test_failed_terminal_put_runs_the_deferred_cleanup(build_adapter): + """#6670: a silent connector drop must not strand the deferred cleanup. + + ``connector.put`` returning False skips the success path that would have + called ``cleanup``. The reclaim ``cleanup_sender`` handed over is owed on + that path too, or the whole per-request sender state leaks for good. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-put-failed") + adapter.cleanup_sender("ext-put-failed") + connector.put.return_value = (False, 0, {}) + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert "ext-put-failed" not in adapter._sender_tokens + assert "ext-put-failed" not in adapter.put_req_chunk + + +def test_raising_terminal_put_runs_the_deferred_cleanup(build_adapter): + """#6670: the same reclaim is owed when the connector raises.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-put-raised") + adapter.cleanup_sender("ext-put-raised") + connector.put.side_effect = RuntimeError("connector down") + + # ``save_loop`` catches this and records a send failure. + with pytest.raises(RuntimeError): + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert "ext-put-raised" not in adapter._sender_tokens + assert "ext-put-raised" not in adapter.put_req_chunk + + def test_wire_round_trip_struct_to_dict_contract(): """Pin the wire contract: encoding ``OmniPayloadStruct`` and decoding it yields a dict equivalent to ``to_dict(struct)``. diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 9123b794d..f33359835 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -688,6 +688,68 @@ class OmniSchedulerMixin: engine_core_outputs[0] = output = EngineCoreOutputs() output.scheduler_stats = stats + def add_request(self, request: Request) -> None: + """Route a terminal streaming update that lands on a parked session. + + Upstream ``Scheduler.add_request`` turns "final update (``resumable`` + is False) arrives while the session sits in + ``WAITING_FOR_STREAMING_REQ``" into a local + ``finish_requests(FINISHED_ABORTED)``. That is fine for a stage whose + output the orchestrator forwards, but an async-chunk *sender* also + owes its downstream stage a terminal chunk: the receiver learns that a + stream ended only from a connector payload carrying ``finished``, and + an abort emits none. The downstream stage then parks in + ``WAITING_FOR_CHUNK`` until ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` (600s by + default), long after the client's own deadline -- the realtime + async-chunk hang in vllm-project/vllm-omni#6670. + + Whether this races is timing, not configuration: the stage parks only + when its own stop beats both the upstream terminal chunk and the + orchestrator's terminal update, a window of tens of milliseconds at + the end of the last segment. Emit the terminal chunk here so the + downstream stage terminates on the payload path either way. + """ + existing = self.requests.get(request.request_id) + if existing is not None and existing.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + adapter = None if getattr(request, "resumable", False) else self._adapter_owing_terminal(existing) + if adapter is not None: + self._finish_parked_streaming_session(existing, adapter) + return + super().add_request(request) + + def _adapter_owing_terminal(self, request: Request) -> OmniChunkTransferAdapter | None: + """The adapter that owes *request*'s downstream stage a terminal chunk. + + Sender ownership starts when the first chunk is queued, before the + background thread initializes its chunk counter. A prewarmed receiver + needs a terminal even if that first chunk has not been sent yet. + Final stages never queue sends, so they never match. + """ + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is None: + return None + external_req_id = getattr(request, "external_req_id", None) or request.request_id + return adapter if adapter.has_active_sender(external_req_id) else None + + def _finish_parked_streaming_session(self, request: Request, adapter: OmniChunkTransferAdapter) -> None: + """End a parked async-chunk session with a downstream terminal chunk.""" + request.resumable = False + # A segment stop already advanced the adapter's dedup watermark to + # ``generation + 1`` for the segment that never arrived. Claim that + # generation, exactly as ``_update_request_as_session`` would have, + # or ``save_async`` drops this chunk as a late duplicate. + request._omni_segment_generation = int(getattr(request, "_omni_segment_generation", 0) or 0) + 1 + # ``force_request_finished``: the terminal must be enqueued *before* + # ``finish_requests``, which skips requests that already report + # ``is_finished()``. + adapter.save_async( + None, + request, + is_segment_finished=False, + force_request_finished=True, + ) + self.finish_requests(request.request_id, RequestStatus.FINISHED_STOPPED) + def finish_requests( self, request_ids: str | Iterable[str] | None, diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index d71d55bf1..163363931 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -32,6 +32,16 @@ class _SenderGeneration: def __init__(self) -> None: self.in_flight = False self.cancelled = False + # A request-terminal chunk is queued but not yet on the wire. The + # downstream stage has no other way to learn the stream ended -- it + # polls the connector and parks in WAITING_FOR_CHUNK until a chunk + # carrying ``finished`` arrives -- so cleanup must not discard it. + # See vllm-project/vllm-omni#6670. + self.terminal_pending = False + # ``cleanup_sender`` asked to reclaim this generation while the fence + # was up. The terminal task owes that cleanup once it is done with the + # connector -- on every exit path, including a failed or raising put. + self.cleanup_deferred = False class _LoadEntry: @@ -321,6 +331,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._processor_accepts_step_tokens[processor] = accepts return accepts + def has_active_sender(self, external_req_id: str) -> bool: + """Whether this request owns a sender, including work still queued. + + ``save_async`` registers the generation before the background thread + can send its first chunk. Chunk counters are initialized later and + cannot determine whether a prewarmed receiver needs a terminal. + """ + with self._sender_state_lock: + sender_token = self._sender_tokens.get(external_req_id) + return sender_token is not None and not sender_token.cancelled + def save_async( self, multimodal_output: dict[str, Any] | None = None, @@ -329,6 +350,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): new_token_ids: Iterable[int] | None = None, confirmed_num_computed_tokens: int | None = None, segment_generation: int | None = None, + force_request_finished: bool = False, ): """Build and enqueue one chunk for asynchronous sending. @@ -352,8 +374,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): before a streaming transition can mutate ``request`` segment_generation: generation captured before a resumable stop can apply a queued update to the mutable request + force_request_finished: send this chunk as the request terminal + even though ``request`` is not in a finished status yet. The + scheduler needs this when it ends a *parked* streaming session: + the terminal has to be enqueued before ``finish_requests``, + because ``finish_requests`` skips requests that already report + ``is_finished()``. See ``_finish_parked_streaming_session``. """ - is_finished = request.is_finished() and not request.resumable + is_finished = force_request_finished or (request.is_finished() and not request.resumable) if not hasattr(self, "_segment_generation"): self._segment_generation = defaultdict(int) external_req_id = request.external_req_id @@ -414,6 +442,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._sender_tokens[external_req_id] = sender_token task["sender_token"] = sender_token self._pending_save_reqs.append(task) + if is_finished: + # Fence this generation against cleanup until the terminal + # chunk has been handed to the connector (#6670). + sender_token.terminal_pending = True if is_segment_finished: # The queued FIFO item now owns the old segment. Start the next # segment's deduplication watermark before the worker sends it. @@ -624,32 +656,71 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if sender_token is None: self._send_single_request_for_generation(task) return + is_terminal_task = bool(task.get("is_finished")) with self._sender_state_lock: is_current = self._sender_tokens.get(external_req_id) is sender_token - if is_current and not sender_token.cancelled: + # ``cancelled`` does not veto a terminal chunk. It is the + # downstream stage's only end-of-stream signal, so dropping it + # leaves that stage parked in WAITING_FOR_CHUNK until the input + # deadline -- the hang this fence exists to prevent (#6670). + # Losing ownership of the external id still vetoes: the state the + # chunk describes then belongs to a different generation. + if is_current and (is_terminal_task or not sender_token.cancelled): sender_token.in_flight = True else: is_current = False if not is_current: logger.debug("Discarding stale queued chunk for aborted request %s", external_req_id) + if is_terminal_task: + self._release_terminal_fence(external_req_id, sender_token) return try: self._send_single_request_for_generation(task, sender_token) finally: + if is_terminal_task: + # Drop the fence and run the cleanup it deferred, whether the + # put succeeded, reported failure, or raised. Until this runs, + # the block below must not reclaim the generation either. + self._release_terminal_fence(external_req_id, sender_token) with self._sender_state_lock: if self._sender_tokens.get(external_req_id) is sender_token: sender_token.in_flight = False - if sender_token.cancelled: + if sender_token.cancelled and not sender_token.terminal_pending: self._sender_tokens.pop(external_req_id, None) self._clear_sender_state_locked(external_req_id) + def _release_terminal_fence(self, external_req_id: str, sender_token: _SenderGeneration) -> None: + """Drop the terminal fence and run the cleanup it deferred. + + Called on every exit path of a terminal task. ``cleanup_sender`` + cannot reclaim a generation whose terminal chunk is still queued, so + it records the request on the token instead; without this the only + path that ever replayed it was a *successful* put, and a connector + failure leaked the whole per-request sender state (#6670). + """ + with self._sender_state_lock: + sender_token.terminal_pending = False + deferred = sender_token.cleanup_deferred + sender_token.cleanup_deferred = False + # A generation that no longer owns the external id was already + # reclaimed by whoever replaced it; cleaning now would clear the + # successor's state. + still_current = self._sender_tokens.get(external_req_id) is sender_token + if deferred and still_current: + self.cleanup_sender(external_req_id) + def _sender_generation_is_active( self, external_req_id: str, sender_token: _SenderGeneration, + *, + is_terminal: bool = False, ) -> bool: with self._sender_state_lock: - return self._sender_tokens.get(external_req_id) is sender_token and not sender_token.cancelled + if self._sender_tokens.get(external_req_id) is not sender_token: + return False + # Terminal chunks outlive cancellation -- see _send_single_request. + return is_terminal or not sender_token.cancelled def _send_single_request_for_generation( self, @@ -712,7 +783,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if payload_data.meta.is_segment_finished is None: payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool) - if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): + if sender_token is not None and not self._sender_generation_is_active( + external_req_id, sender_token, is_terminal=is_finished + ): logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id) return @@ -723,16 +796,25 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): data=payload_data, ) - if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): - # cleanup_sender() may cancel this generation while put() is - # blocked. Do not let that stale completion recreate state for a - # request whose cleanup is already in progress. - logger.debug("Ignoring completed put for cancelled request %s", external_req_id) - return + with self._sender_state_lock: + if sender_token is not None: + still_current = self._sender_tokens.get(external_req_id) is sender_token + retiring_without_terminal = sender_token.cancelled and not ( + is_finished or sender_token.terminal_pending + ) + if not still_current or retiring_without_terminal: + logger.debug("Ignoring completed put for retired request %s", external_req_id) + return + if success: + # A receiver may already have consumed this key while put was + # in flight. Even if cleanup retired this generation, its + # pending terminal must use the next key. Check ownership and + # advance counters under the same lock so a stale completion + # cannot modify a successor's state. + self.put_req_chunk[external_req_id] += 1 + self.ramp_chunk_count[external_req_id] += 1 if success: - self.put_req_chunk[external_req_id] += 1 - self.ramp_chunk_count[external_req_id] += 1 logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") # Sender uses struct attr access here; the receive path in # `_load_one_request` / `_update_request_payload` reads dict keys. @@ -862,6 +944,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if sender_token is None: self._clear_sender_state_locked(external_req_id) return + if sender_token.terminal_pending: + # A request-terminal chunk is still queued. Reclaiming the + # state now would drop it, and the downstream stage would sit + # in WAITING_FOR_CHUNK until VLLM_OMNI_INPUT_WAIT_TIMEOUT_S + # with no way to know the stream ended (#6670). Retire the + # generation anyway -- ``save_async`` rejects new chunks on + # ``cancelled``, which is what keeps a request reusing this + # external id from joining a generation that is going away -- + # and hand the reclaim to the terminal task. + sender_token.cancelled = True + sender_token.cleanup_deferred = True + return sender_token.cancelled = True if sender_token.in_flight: # The sender's finally block is the first point where no