From e25239c2ecde8a3c07c9d2e624a1c13761551872 Mon Sep 17 00:00:00 2001 From: ChenWenjing <54166744+Shirley125@users.noreply.github.com> Date: Sat, 6 Jun 2026 00:05:07 +0800 Subject: [PATCH] [bugfix] fix realtime ci timeout error (#4187) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn> --- vllm_omni/core/sched/omni_ar_scheduler.py | 22 +++++++++---- vllm_omni/engine/orchestrator.py | 31 ++++++++++++++++++- .../entrypoints/openai/realtime_connection.py | 9 ++++-- 3 files changed, 53 insertions(+), 9 deletions(-) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index a6d5edd8f..67514994e 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -13,7 +13,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.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all -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 @@ -69,6 +69,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # alive. Stopped on the first scheduler step after extraction ack. self.pending_stop_after_extraction: set[str] = set() + self.finished_req_ids_dict = defaultdict(set) + # [Omni] Pre-parse KV transfer criteria self.kv_transfer_criteria = self._get_kv_transfer_criteria() @@ -556,16 +558,24 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # outputs in this step. engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()} + # FIXME: finished_req_ids_dict is unconditionally initialized as + # defaultdict(set) in __init__ (not gated by include_finished_set). + # This branch is therefore always eligible once any client_index is + # populated; revisit when wiring streaming-only / upstream semantics. finished_req_ids = self.finished_req_ids_dict if finished_req_ids: # Include ids of requests that finished since last outputs # were sent. for client_index, finished_set in finished_req_ids.items(): - # Set finished request set in EngineCoreOutputs for this client. - if (eco := engine_core_outputs.get(client_index)) is not None: - eco.finished_requests = finished_set - else: - engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set) + eco = engine_core_outputs.get(client_index) + if eco is None: + eco = EngineCoreOutputs() + engine_core_outputs[client_index] = eco + emitted = {o.request_id for o in eco.outputs} + for req_id in finished_set: + if req_id not in emitted: + eco.outputs.append(EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)) + eco.finished_requests = finished_set finished_req_ids.clear() if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 017362812..ea20c650b 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -697,6 +697,8 @@ class Orchestrator: "new_prompt_len_snapshot", None, ) + if req_state.streaming.enabled: + await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state) # OmniSchedulerMixin.make_stats() already throttles # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats # being non-None means this replica passed its own gate. @@ -829,6 +831,34 @@ class Orchestrator: if self.request_states.pop(request_id, None) is not None and self._running_counter is not None: self._running_counter.decrement() + async def _apply_raw_terminal_stage_finish( + self, + stage_id: int, + eco: Any, + req_state: OrchestratorRequestState, + ) -> None: + """Record session-level finish markers dropped by the streaming output processor. + + Streaming segment stops set ``is_segment_finished=True`` and are handled + via processed outputs. Session termination (e.g. ``finish_requests`` after + ``resumable=False``) emits a terminal ``finish_reason`` with + ``is_segment_finished=False``, but vLLM's output processor may remove the + request state before that EngineCoreOutput is processed. + + Only update ``finished_final_output_stage_ids`` here. Request cleanup stays + in ``_route_output`` so downstream async-chunk stages can still deliver + outputs after stage-0 session end. + """ + if getattr(eco, "finish_reason", None) is None: + return + if getattr(eco, "is_segment_finished", False): + return + + final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} + if stage_id not in final_output_stage_ids: + return + req_state.finished_final_output_stage_ids.add(stage_id) + def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" companion_request_ids = self._cfg_tracker.get_companion_request_ids(request_id) @@ -872,7 +902,6 @@ class Orchestrator: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) - if self.stage_pools[stage_id].final_output: await self.output_async_queue.put( OutputMessage( diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 1bc3df64e..5cf936e48 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -86,9 +86,14 @@ class RealtimeConnection(VllmRealtimeConnection): return [], 24000 sr = mm.get("sr") or mm.get("sample_rate") or mm.get("audio_sample_rate") or 24000 + if isinstance(sr, (list, tuple)) and sr: + sr = sr[-1] + if hasattr(sr, "item"): + sr = sr.item() + sample_rate_hz = int(sr) key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) if key is None: - return [], int(sr) + return [], sample_rate_hz raw_audio = mm.get(key) chunks: list[np.ndarray] = [] @@ -101,7 +106,7 @@ class RealtimeConnection(VllmRealtimeConnection): arr = self._tensor_to_numpy(raw_audio) if arr is not None and arr.size > 0: chunks.extend(self._raw_waveform_to_deltas(arr)) - return chunks, int(sr) + return chunks, sample_rate_hz @staticmethod def _pcm16_b64(audio_f32: np.ndarray) -> str: