From fa639b889e50d28a3cbe88b20b2e94cdc56bdf65 Mon Sep 17 00:00:00 2001 From: Tianyao Wu <54675599+twu3202@users.noreply.github.com> Date: Wed, 16 Sep 2026 20:08:25 +0800 Subject: [PATCH] [Bugfix] Fail a duplex session that outgrows max_model_len instead of killing the EngineCore (#7277) Signed-off-by: Tianyao Wu Co-authored-by: Claude Co-authored-by: Gao Han Co-authored-by: amy-why-3459 --- tests/clients/test_duplex_client.py | 27 ++ .../sched/test_omni_ar_scheduler_logprobs.py | 3 + .../sched/test_omni_ar_scheduler_streaming.py | 303 ++++++++++++++++++ tests/engine/test_orchestrator.py | 59 ++++ vllm_omni/clients/duplex.py | 7 +- vllm_omni/core/sched/omni_ar_scheduler.py | 140 +++++++- vllm_omni/engine/orchestrator.py | 46 ++- 7 files changed, 582 insertions(+), 3 deletions(-) diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index 43b7547cb..75ef06e3f 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -570,6 +570,33 @@ async def test_error_event_surfaces_on_response_path(make_rig): rig.feed(SESSION_CLOSED) +@pytest.mark.asyncio +async def test_close_after_error_reports_the_server_reason_to_a_sender(): + # A client streaming input sits in send(), not in responses(), so the + # error text has to survive into the close marker or the session looks + # like it ended for no reason. + sock = FakeSocket() + sock.feed(SESSION_CREATED) + client, _ = make_client(sock) + async with client: + sock.feed( + { + "type": "error", + "error": { + "code": "runtime_data_plane_stream_failed", + "message": "context_length_exceeded: streaming session prompt would grow to 41207 tokens", + }, + "server_event_seq": 2, + } + ) + sock.feed(SESSION_CLOSED) + await _drain(lambda: client._closed.is_set()) + with pytest.raises(DuplexSessionClosedError) as excinfo: + await client.send({"type": "input_audio_buffer.append", "audio": ""}) + assert "runtime_data_plane_stream_failed" in str(excinfo.value) + assert "context_length_exceeded" in str(excinfo.value) + + @pytest.mark.asyncio async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): from vllm_omni.clients.duplex import AudioDelta, ResponseHandle diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 604904dd0..810c506c7 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -149,6 +149,9 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace: for name in _MIXIN_UPDATE_HELPERS: setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler)) scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler) + scheduler._emit_streaming_context_overflow_outputs = MethodType( + OmniARScheduler._emit_streaming_context_overflow_outputs, scheduler + ) scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None scheduler.make_stats = lambda *args, **kwargs: None return scheduler diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index bca7d095c..154644fb2 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -5,6 +5,7 @@ from __future__ import annotations +from collections import defaultdict from types import SimpleNamespace from unittest.mock import MagicMock @@ -17,6 +18,8 @@ import vllm_omni # noqa: F401 - import for side effects (patch vLLM) from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.engine import FinishReason +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( @@ -958,3 +961,303 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked( assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests sched.skipped_waiting.remove_requests.assert_not_called() sched._enqueue_waiting_request.assert_not_called() + + +def _make_live_session_scheduler(*, max_model_len: int) -> OmniARScheduler: + """A scheduler whose finish path is real: the queues, ``finish_requests``, + ``_free_request`` and the finished-request bookkeeping that wakes an idle + engine all run. Only block freeing and the encoder cache are mocked.""" + sched = _make_admission_scheduler(max_model_len=max_model_len) + del sched._free_request # the admission helper mocks it; run the real one + sched._omits_kv_transfer_cache = {} + sched.connector = None + sched.perf_metrics = None + sched.recompute_kv_load_failures = False + sched.finished_req_ids = set() + sched.finished_req_ids_dict = defaultdict(set) + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.waiting_for_transfer_free = set() + sched.pending_stop_after_extraction = set() + sched.requests_needing_kv_transfer = {} + sched._kv_wait_start_ts = {} + sched.kv_cache_manager = SimpleNamespace(take_events=lambda: None, estimate_cached_tokens=lambda _request: 0) + return sched + + +def _run_idle_step(sched: OmniARScheduler): + """One ``update_from_output`` with nothing scheduled: the step an idle + engine runs once ``has_finished_requests()`` reports a freed session.""" + scheduler_output = SimpleNamespace( + num_scheduled_tokens={}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + model_runner_output = SimpleNamespace( + sampled_token_ids=[], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=None, + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={}, + routed_experts=None, + ) + return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + +def _park_session(sched: OmniARScheduler, session: Request) -> None: + """Put a session where a 1 fps video stream leaves it between appends: + prompt partly computed, waiting for streaming input, in admission.""" + sched.requests[session.request_id] = session + session.append_output_token_ids([7, 8, 9]) + session.num_computed_tokens = 6 + session.num_output_placeholders = 0 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.skipped_waiting.add_request(session) + sched.num_waiting_for_streaming_input = 1 + + +def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_session() -> None: + """A native duplex session grows its prompt on every append; once the + extension cannot fit the model the worker crashes copying the prompt. + The scheduler must drop the update and fail only this session, at once, + through the real finish path: a parked session does not make the engine + schedule, so what wakes the engine is the freed request showing up in + ``has_finished_requests()``, and the step that follows must carry one + explicit ERROR rather than the synthesized ABORT.""" + sched = _make_live_session_scheduler(max_model_len=8) + session = _make_request() + _park_session(sched, session) + + sched._update_request_as_session(session, _make_update([10, 20, 30])) + + assert session.prompt_token_ids == [1, 2, 3] + assert session.num_prompt_tokens == 3 + assert session.status == RequestStatus.FINISHED_ERROR + assert session.request_id not in sched.requests + assert len(sched.waiting) == 0 + assert len(sched.skipped_waiting) == 0 + assert sched.num_waiting_for_streaming_input == 0 + sched._free_request_blocks.assert_called_once_with(session) + assert sched.finished_req_ids == {session.request_id} + assert sched.has_finished_requests() + client_index, reason = sched._streaming_context_overflow[session.request_id] + assert client_index == session.client_index + assert reason.startswith("context_length_exceeded: ") + assert "9 tokens" in reason and "max_model_len 8" in reason + + engine_core_outputs = _run_idle_step(sched) + + (output,) = engine_core_outputs[session.client_index].outputs + assert output.request_id == session.request_id + assert output.finish_reason == FinishReason.ERROR + assert output.stop_reason == reason + assert output.new_token_ids == [] + assert engine_core_outputs[session.client_index].finished_requests == {session.request_id} + assert sched._streaming_context_overflow == {} + # Nothing is left to report once the error went out. + assert _run_idle_step(sched) == {} + + +def test_stage0_streaming_update_that_fills_max_model_len_exactly_finishes_the_session() -> None: + """An exactly full prompt leaves no room for the token every duplex step + samples: upstream's running budget drops to -1, which ``schedule()`` does + not catch, so the exact fill is an overflow too.""" + sched = _make_scheduler(stage_id=0) + sched.max_model_len = 8 + sched.finish_requests = MagicMock() + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + session.append_output_token_ids([7, 8, 9]) + session.num_computed_tokens = 6 + session.num_output_placeholders = 0 + + sched._update_request_as_session(session, _make_update([10, 20])) + + assert session.prompt_token_ids == [1, 2, 3] + assert session.num_prompt_tokens == 3 + sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR) + _, reason = sched._streaming_context_overflow[session.request_id] + assert "8 tokens" in reason and "max_model_len 8" in reason + + +def test_stage0_streaming_update_that_leaves_room_to_sample_is_applied() -> None: + sched = _make_scheduler(stage_id=0) + sched.max_model_len = 9 + sched.finish_requests = MagicMock() + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + session.append_output_token_ids([7, 8, 9]) + session.num_computed_tokens = 6 + session.num_output_placeholders = 0 + + sched._update_request_as_session(session, _make_update([10, 20])) + + assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20] + assert session.num_prompt_tokens == 8 + assert session.status == RequestStatus.WAITING + sched.finish_requests.assert_not_called() + assert not getattr(sched, "_streaming_context_overflow", {}) + + +def test_stage0_streaming_update_keeps_room_for_every_token_sampled_per_step() -> None: + """With speculative decoding a step samples more than one token, and the + prompt must leave room for all of them.""" + sched = _make_scheduler(stage_id=0) + sched.max_model_len = 9 + sched.num_sampled_tokens_per_step = 2 + sched.finish_requests = MagicMock() + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + session.append_output_token_ids([7, 8, 9]) + session.num_computed_tokens = 6 + session.num_output_placeholders = 0 + + sched._update_request_as_session(session, _make_update([10, 20])) + + assert session.prompt_token_ids == [1, 2, 3] + sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR) + assert "8 tokens" in sched._streaming_context_overflow[session.request_id][1] + + +def test_stage0_streaming_update_overflow_counts_an_uncomputed_prompt() -> None: + """A preempted session keeps its prompt even though num_computed_tokens + fell behind it; the projection must use the longer of the two.""" + sched = _make_scheduler(stage_id=0) + sched.max_model_len = 4 + sched.finish_requests = MagicMock() + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + session.num_computed_tokens = 0 + + sched._update_request_as_session(session, _make_update([10, 20])) + + assert session.prompt_token_ids == [1, 2, 3] + sched.finish_requests.assert_called_once() + assert "5 tokens" in sched._streaming_context_overflow[session.request_id][1] + + +def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler: + """A scheduler with real admission queues, so nothing about the terminal + handling below is mocked away except block freeing.""" + sched = _make_scheduler(stage_id=0) + sched.max_model_len = max_model_len + sched.policy = SchedulingPolicy.FCFS + sched.waiting = create_request_queue(sched.policy) + sched.skipped_waiting = create_request_queue(sched.policy) + sched.running = [] + sched.requests = {} + sched._free_request = MagicMock() # needs a KV manager; not what this covers + return sched + + +def _make_queued_stop(session: Request, update: StreamingUpdate) -> None: + """Put the request in the state upstream hands to _handle_stopped_request: + an append that arrived during generation is queued, and the stop status is + still on the request.""" + from collections import deque + + session.resumable = True + session.streaming_queue = deque([update]) + session.status = RequestStatus.FINISHED_STOPPED + + +def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -> None: + sched = _make_admission_scheduler(max_model_len=4) + session = _make_request() + sched.requests[session.request_id] = session + session.num_computed_tokens = 3 + _make_queued_stop(session, _make_update([10, 20])) + + finished = sched._handle_stopped_request(session) + + assert finished is True + assert len(sched.waiting) == 0 + assert len(sched.skipped_waiting) == 0 + assert session.status == RequestStatus.FINISHED_ERROR + assert session.resumable is False + assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1] + # update_from_output frees every request reported as finished; freeing here + # as well deleted it from self.requests twice (KeyError in _free_blocks). + sched._free_request.assert_not_called() + + +def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None: + sched = _make_admission_scheduler(max_model_len=64) + session = _make_request() + sched.requests[session.request_id] = session + session.num_computed_tokens = 3 + _make_queued_stop(session, _make_update([10, 20])) + + finished = sched._handle_stopped_request(session) + + assert finished is False + assert len(sched.waiting) == 1 + assert session.status == RequestStatus.WAITING + assert not getattr(sched, "_streaming_context_overflow", {}) + sched._free_request.assert_not_called() + + +def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None: + """The overflow may already be on record when the stop is handled (an + earlier update of the same session). Upstream still re-enqueues the + request, so the override must take it back out either way instead of + leaving a FINISHED_ERROR request for admission to trip over.""" + sched = _make_admission_scheduler(max_model_len=64) + session = _make_request() + sched.requests[session.request_id] = session + session.num_computed_tokens = 3 + sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")} + _make_queued_stop(session, _make_update([10, 20])) # this update itself fits + + finished = sched._handle_stopped_request(session) + + assert finished is True + assert len(sched.waiting) == 0 + assert len(sched.skipped_waiting) == 0 + assert session.status == RequestStatus.FINISHED_ERROR + assert session.resumable is False + sched._free_request.assert_not_called() + + +def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanced() -> None: + """Same as above with nothing queued: upstream parks the request as + waiting for streaming input and counts it; the dequeue must uncount it.""" + from collections import deque + + sched = _make_admission_scheduler(max_model_len=64) + session = _make_request() + sched.requests[session.request_id] = session + session.num_computed_tokens = 3 + sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")} + session.resumable = True + session.streaming_queue = deque() + session.status = RequestStatus.FINISHED_STOPPED + + finished = sched._handle_stopped_request(session) + + assert finished is True + assert len(sched.skipped_waiting) == 0 + assert sched.num_waiting_for_streaming_input == 0 + assert session.status == RequestStatus.FINISHED_ERROR + sched._free_request.assert_not_called() + + +def test_context_overflow_emits_an_error_output_with_the_reason() -> None: + sched = _make_scheduler(stage_id=0) + sched._streaming_context_overflow = {"req-a": (2, "context_length_exceeded: too long")} + outputs: dict[int, list] = {} + + sched._emit_streaming_context_overflow_outputs(outputs) + + (output,) = outputs[2] + assert output.request_id == "req-a" + assert output.finish_reason == FinishReason.ERROR + assert output.stop_reason == "context_length_exceeded: too long" + assert output.new_token_ids == [] + assert sched._streaming_context_overflow == {} + + sched._emit_streaming_context_overflow_outputs(outputs) + assert len(outputs[2]) == 1 diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index beddd2106..3b6ee7e7b 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -21,6 +21,7 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats +from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.messages import ( AbortRequestMessage, AbortResultMessage, @@ -1891,3 +1892,61 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort(): assert first.physical_abort_calls == 1 assert second.physical_abort_calls == 2 assert second.bound == set() + + +@pytest.mark.asyncio +async def test_duplex_session_request_error_finish_is_delivered_as_request_error() -> None: + """A session-owned request the scheduler finished with FinishReason.ERROR + (e.g. its prompt could not grow past max_model_len) must reach the + session's consumer as a request-scoped error, not vanish: its terminal + outputs are not processed like other requests, and the output processor + may already have dropped the request state.""" + output_queue: asyncio.Queue = asyncio.Queue() + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=output_queue, + rpc_async_queue=asyncio.Queue(), + stage_pools=[], + ) + duplex_state = OrchestratorRequestState(request_id="duplex-req", session_owned=True) + plain_state = OrchestratorRequestState(request_id="plain-req") + reason = "context_length_exceeded: streaming session prompt would grow to 8220 tokens, above max_model_len 8192" + + await orchestrator._report_duplex_session_request_error( + 0, + 0, + OmniEngineCoreOutput( + request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.ERROR, stop_reason=reason + ), + duplex_state, + ) + error = output_queue.get_nowait() + assert isinstance(error, ErrorMessage) + assert error.request_id == "duplex-req" + assert error.stage_id == 0 + assert error.fatal is False + assert error.error == reason + + # An ordinary segment stop, a segment-finished error and a request that is + # not session-owned produce nothing. + await orchestrator._report_duplex_session_request_error( + 0, + 0, + OmniEngineCoreOutput(request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.STOP), + duplex_state, + ) + await orchestrator._report_duplex_session_request_error( + 0, + 0, + OmniEngineCoreOutput( + request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.ERROR, is_segment_finished=True + ), + duplex_state, + ) + await orchestrator._report_duplex_session_request_error( + 0, + 0, + OmniEngineCoreOutput(request_id="plain-req", new_token_ids=[], finish_reason=FinishReason.ERROR), + plain_state, + ) + assert output_queue.empty() diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index 52266b64e..17358906f 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -645,6 +645,7 @@ class DuplexClientBase(ABC): self._closing = False self._closed = asyncio.Event() self._closed_marker: _ClosedMarker | None = None + self._last_error_reason: str | None = None # -- transport hooks ------------------------------------------------------ @@ -1061,12 +1062,16 @@ class DuplexClientBase(ABC): # on the response path too so a consumer waiting in responses() # is not left waiting forever (see the responses() docstring). _put_drop_oldest(self._response_queue, _ErrorMarker(event)) + # Keep the text for the close that usually follows: a client that + # is streaming input is inside send(), not draining responses(), + # and would otherwise only ever see "closed". + self._last_error_reason = f"{event.code}: {event.message}" if event.code else event.message for queue in list(self._subscribers): _put_drop_oldest(queue, event) if isinstance(event, SessionClosed): - self._finalize("closed", expected=True) + self._finalize(self._last_error_reason or "closed", expected=True) elif isinstance(event, SessionExpired): reason = data.get("reason") self._finalize(f"expired: {reason}" if reason else "expired", expected=False) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 329e8ad81..3d866fda7 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -16,7 +16,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler -from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason from vllm.v1.metrics.perf import PerfStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate @@ -116,6 +116,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._init_omni_io_scheduling_state() # Snapshot prompt length for each streaming input update self._new_prompt_len_snapshot: dict[str, int] = {} + # Streaming sessions finished because their next prompt extension + # would exceed max_model_len: request_id -> (client_index, reason). + # Drained into an explicit FinishReason.ERROR output on the next + # update_from_output so the client learns why the session ended. + self._streaming_context_overflow: dict[str, tuple[int, str]] = {} def _get_confirmed_num_computed_tokens(self, request: Request) -> int: """num_computed_tokens minus async placeholders (KV actually on GPU).""" @@ -704,6 +709,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): failed_kv_load_req_ids, outputs, ) + self._emit_streaming_context_overflow_outputs(outputs) if self.chunk_transfer_adapter is not None: for request in failed_requests: self.chunk_transfer_adapter.cleanup_receiver(request.request_id) @@ -883,11 +889,143 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._release_replaced_streaming_prompt_cache(session) self._replace_streaming_session(session, update) return + if self._streaming_update_overflows(session, update): + return session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 super()._update_request_as_session(session, update) if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer + # Prefix of the stop_reason carried by the FinishReason.ERROR output, so the + # serving side can map it to a stable error code. + STREAMING_CONTEXT_OVERFLOW_STOP_REASON = "context_length_exceeded" + + def _streaming_update_overflows(self, session: Request, update: StreamingUpdate) -> bool: + """Finish a streaming session whose next extension cannot fit the model. + + Upstream ``_update_request_as_session`` appends the update to the + session prompt without checking ``max_model_len``. The worker's input + batch then fails to copy the prompt (``could not broadcast input array + from shape (N,) into shape (max_model_len,)``) and the EngineCore dies, + taking every session on the replica with it. A native duplex session + grows by tens to hundreds of tokens per second of input, so long + sessions reach this point in normal use. + + A prompt that fills the model exactly is over the line as well: the + session samples at least one listen/speak token after every append, + and upstream's running-request budget + ``max_model_len - num_computed_tokens - num_sampled_tokens_per_step`` + then goes negative, which the ``num_new_tokens == 0`` guard in + ``schedule()`` does not catch (``allocate_slots`` dies, or the worker + asserts ``max_model_len + 1`` sampled positions). So the extended + prompt must leave room for the tokens sampled in one step. + + The update is dropped and only this request is finished, right here: + a parked session does not make the engine schedule, so deferring the + finish to the next ``schedule()`` would leave the client waiting. The + reason is emitted with the terminal output (see + :meth:`_emit_streaming_context_overflow_outputs`). + """ + max_model_len = getattr(self, "max_model_len", None) + if max_model_len is None: + model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) + max_model_len = getattr(model_config, "max_model_len", None) + if not max_model_len: + return False + new_tokens = len(update.prompt_token_ids or ()) + # The extended prompt is the current prompt plus the computed output + # tokens upstream keeps, then the update: num_computed_tokens covers + # both when the prompt was fully computed. + projected = max(int(session.num_prompt_tokens), int(session.num_computed_tokens)) + new_tokens + # Room for the tokens one step samples on top of the prompt (1 without + # speculative decoding). __new__-built test schedulers carry no + # num_sampled_tokens_per_step. + sample_room = max(1, int(getattr(self, "num_sampled_tokens_per_step", 1) or 1)) + if projected + sample_room <= int(max_model_len): + return False + reason = ( + f"{self.STREAMING_CONTEXT_OVERFLOW_STOP_REASON}: streaming session prompt would grow to " + f"{projected} tokens, leaving no room to sample within max_model_len {int(max_model_len)}" + ) + logger.error( + "[Omni] %s: %s; finishing the request instead of extending it", + session.request_id, + reason, + ) + overflow = getattr(self, "_streaming_context_overflow", None) + if overflow is None: + overflow = self._streaming_context_overflow = {} + overflow[session.request_id] = (int(getattr(session, "client_index", 0) or 0), reason) + if session.is_finished(): + # Reached from ``_handle_stopped_request`` with a queued update. + # ``update_from_output`` frees every request that call reports as + # finished, so finishing the session here too would free it twice. + # ``_handle_stopped_request`` below takes it out of admission and + # leaves the single free to the caller. + return True + self.finish_requests((session.request_id,), RequestStatus.FINISHED_ERROR) + return True + + def _handle_stopped_request(self, request: Request) -> bool: + """Do not resume a session whose queued update overflowed the model. + + Upstream pops one queued ``StreamingUpdate``, applies it through + ``_update_request_as_session`` and then re-enqueues the request + unconditionally. When that update overflows, the request must not go + back into the waiting queue: it is terminal, and admission raises + ``RuntimeError: Invalid request status`` on anything that is neither + WAITING nor PREEMPTED, which would kill the EngineCore this guard + exists to keep alive. The same holds for a session whose overflow was + recorded before this call and that upstream still reports as resumed. + """ + finished = super()._handle_stopped_request(request) + if finished: + return True + overflow = getattr(self, "_streaming_context_overflow", None) + if not overflow or request.request_id not in overflow: + return False + # Whether the overflow was recorded by this call's queued update or + # earlier makes no difference: a session in the overflow map is + # terminal, and upstream has just put it back into admission. + self._finish_overflowed_streaming_session(request) + return True + + def _finish_overflowed_streaming_session(self, request: Request) -> None: + """Take a terminal session back out of admission. + + Queues and status only. ``update_from_output`` frees every request + ``_handle_stopped_request`` reports as finished, so freeing here as + well deletes it from ``self.requests`` twice (``KeyError`` in + ``_free_blocks``) and skips the caller's input-coordinator cleanup. + """ + self.waiting.remove_requests((request,)) + self.skipped_waiting.remove_requests((request,)) + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + + def _emit_streaming_context_overflow_outputs(self, outputs: dict[int, list[EngineCoreOutput]]) -> None: + """Turn recorded context overflows into explicit error outputs. + + Without this the finished session would only get the synthesized + ``FinishReason.ABORT`` output, which a client cannot tell apart from + its own cancel. + """ + overflow = getattr(self, "_streaming_context_overflow", None) + if not overflow: + return + for request_id, (client_index, reason) in list(overflow.items()): + outputs.setdefault(client_index, []).append( + OmniEngineCoreOutput( + request_id=request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + stop_reason=reason, + ) + ) + overflow.clear() + def _free_request( self, request: Request, delay_free_blocks: bool = False ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index aced2c1b3..aea5c8519 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -30,7 +30,7 @@ from vllm.logger import init_logger from vllm.outputs import CompletionOutput, RequestOutput from vllm.pooling_params import PoolingParams from vllm.sampling_params import RequestOutputKind, SamplingParams -from vllm.v1.engine import EngineCoreOutputs +from vllm.v1.engine import EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats @@ -777,6 +777,7 @@ class OrchestratorBase: ) if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state): raw_terminal_request_ids.add(req_state.request_id) + await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state) iteration_stats = IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None processed = await pool.process_llm_raw_outputs( replica_id, @@ -1503,6 +1504,49 @@ class OrchestratorBase: req_state.finished_final_output_stage_ids.add(stage_id) return True + async def _report_duplex_session_request_error( + self, + stage_id: int, + replica_id: int | None, + eco: Any, + req_state: OrchestratorRequestState, + ) -> None: + """Turn a scheduler-side error finish of a duplex session request into a request error. + + A duplex session request is resumable, so its terminal outputs are + deliberately not turned into processed outputs + (``_finish_raw_terminal_requests`` skips duplex sessions), and vLLM's + output processor may already have dropped the request state by the + time the raw ``FinishReason.ERROR`` output arrives. Without this the + session never learns that its request is gone: the next append + re-creates the request from scratch and the session silently loses + its context. Deliver the error to the request's consumer instead so + the serving layer can fail and close the session. + """ + if not req_state.session_owned: + return + if getattr(eco, "finish_reason", None) != FinishReason.ERROR: + return + if getattr(eco, "is_segment_finished", False): + return + stop_reason = getattr(eco, "stop_reason", None) + error = stop_reason if isinstance(stop_reason, str) and stop_reason else "duplex session request failed" + logger.error( + "[Orchestrator] Duplex session request %s failed at stage-%s replica-%s: %s", + req_state.request_id, + stage_id, + replica_id, + error, + ) + await self.output_async_queue.put( + ErrorMessage( + error=error, + fatal=False, + request_id=req_state.request_id, + stage_id=stage_id, + ) + ) + async def _finish_raw_terminal_requests( self, stage_id: int,