From 0cc3481acf276304eb94ffac2fbce6b7f2760064 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 27 Feb 2026 14:39:33 +0800 Subject: [PATCH] [Feat] : Support Async chunk cleanup (#1087) Signed-off-by: Sy03 <1370724210@qq.com> --- .../test_chunk_transfer_adapter.py | 230 +++++++++++++++++- .../test_qwen3_tts_async_chunk.py | 72 ++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 5 + .../core/sched/omni_generation_scheduler.py | 12 +- .../chunk_transfer_adapter.py | 46 +++- .../stage_input_processors/qwen3_tts.py | 40 +-- 6 files changed, 382 insertions(+), 23 deletions(-) create mode 100644 tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index da2c051e8..f1b808912 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -54,7 +54,6 @@ def build_adapter(monkeypatch, mocker: MockerFixture): self._pending_save_reqs = deque() self._finished_save_reqs = set() self.stop_event = threading.Event() - self.lock = threading.Lock() monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init) monkeypatch.setattr( @@ -178,3 +177,232 @@ def test_postprocess_scheduler_output(build_adapter): assert cached_info["cached-ready"] == {"k": "v"} assert cached_info["missing"] is None assert adapter.requests_with_ready_chunks == {"leftover"} + + +# --------------------------------------------------------------- +# Cleanup tests +# --------------------------------------------------------------- + + +def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"): + """Fill every per-request structure so cleanup can be verified.""" + adapter.finished_requests.add(req_id) + adapter.get_req_chunk[req_id] = 3 + adapter.requests_with_ready_chunks.add(req_id) + adapter.request_ids_mapping[req_id] = ext_id + adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id)) + adapter._finished_load_reqs.add(req_id) + + adapter.put_req_chunk[ext_id] = 5 + adapter.request_payload[ext_id] = {"hidden": [1, 2]} + adapter.code_prompt_token_ids[ext_id] = [[10, 20]] + + +def test_cleanup_clears_all_state(build_adapter): + """After cleanup, no per-request key should remain in any dict/set.""" + adapter, _ = build_adapter(stage_id=1) + req_id, ext_id = "req-1", "ext-1" + _populate_adapter_state(adapter, req_id, ext_id) + + adapter.cleanup(req_id, ext_id) + + assert req_id not in adapter.finished_requests + assert req_id not in adapter.get_req_chunk + assert req_id not in adapter.requests_with_ready_chunks + assert req_id not in adapter.request_ids_mapping + assert all(getattr(r, "request_id", None) != req_id for r in adapter._pending_load_reqs) + assert req_id not in adapter._finished_load_reqs + + assert ext_id not in adapter.put_req_chunk + assert ext_id not in adapter.request_payload + assert ext_id not in adapter.code_prompt_token_ids + + +def test_cleanup_infers_external_id(build_adapter): + """When external_req_id is None, cleanup should look it up from the mapping.""" + adapter, _ = build_adapter(stage_id=1) + req_id, ext_id = "req-2", "ext-2" + _populate_adapter_state(adapter, req_id, ext_id) + + adapter.cleanup(req_id) + + assert ext_id not in adapter.put_req_chunk + assert ext_id not in adapter.request_payload + + +def test_cleanup_idempotent(build_adapter): + """Calling cleanup multiple times for the same (or nonexistent) request must not raise.""" + adapter, _ = build_adapter(stage_id=1) + + try: + adapter.cleanup("nonexistent") + adapter.cleanup("nonexistent") + except Exception as e: + pytest.fail(f"cleanup should be idempotent: {e}") + + req_id, ext_id = "req-3", "ext-3" + _populate_adapter_state(adapter, req_id, ext_id) + adapter.cleanup(req_id, ext_id) + + try: + adapter.cleanup(req_id, ext_id) + except Exception as e: + pytest.fail(f"second cleanup should be idempotent: {e}") + + +def test_cleanup_request_id_reuse_not_polluted(build_adapter): + """After cleanup, reusing the same request_id must not be treated as finished.""" + adapter, _ = build_adapter(stage_id=1) + req_id, ext_id = "req-reuse", "ext-reuse" + _populate_adapter_state(adapter, req_id, ext_id) + + adapter.cleanup(req_id, ext_id) + + assert req_id not in adapter.finished_requests + assert req_id not in adapter.get_req_chunk + + +def test_cleanup_preserves_pending_save(build_adapter): + """Cleanup must NOT remove _pending_save_reqs to avoid losing unsent chunks.""" + adapter, _ = build_adapter(stage_id=1) + req_id, ext_id = "req-4", "ext-4" + _populate_adapter_state(adapter, req_id, ext_id) + + pending_task = {"put_key": f"{ext_id}_1_0", "data": {"x": 1}} + adapter._pending_save_reqs.append(pending_task) + + adapter.cleanup(req_id, ext_id) + + assert len(adapter._pending_save_reqs) == 1 + + +def test_cleanup_only_affects_target_request(build_adapter): + """Cleanup for one request must not affect another request's state.""" + adapter, _ = build_adapter(stage_id=1) + _populate_adapter_state(adapter, "req-a", "ext-a") + _populate_adapter_state(adapter, "req-b", "ext-b") + + adapter.cleanup("req-a", "ext-a") + + assert "req-b" in adapter.finished_requests + assert "req-b" in adapter.get_req_chunk + assert "ext-b" in adapter.put_req_chunk + assert "ext-b" in adapter.request_payload + assert "ext-b" in adapter.code_prompt_token_ids + assert "req-b" in adapter.request_ids_mapping + + +def test_cleanup_after_poll_flow(build_adapter): + """Simulate full load_async -> poll -> finished -> cleanup cycle.""" + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + request = _req("req-flow", RequestStatus.WAITING, external_req_id="ext-flow") + + adapter.load_async(request) + + adapter.request_ids_mapping["req-flow"] = "ext-flow" + payload = {"hidden_states": torch.tensor([[1.0]]), "finished": True} + connector.get.return_value = (payload, 8) + adapter._poll_single_request(request) + + assert "req-flow" in adapter.finished_requests + assert adapter.get_req_chunk["req-flow"] == 1 + assert "req-flow" in adapter.request_ids_mapping + + adapter.cleanup("req-flow", "ext-flow") + + assert "req-flow" not in adapter.finished_requests + assert "req-flow" not in adapter.get_req_chunk + assert "req-flow" not in adapter.request_ids_mapping + assert "ext-flow" not in adapter.request_payload + + +# --------------------------------------------------------------- +# Scheduler trigger tests +# --------------------------------------------------------------- + + +class _HashableRequest(SimpleNamespace): + """SimpleNamespace that can be added to a set (needed by scheduler internals).""" + + def __hash__(self): + return hash(self.request_id) + + def __eq__(self, other): + return getattr(other, "request_id", None) == self.request_id + + +def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: MockerFixture): + """OmniGenerationScheduler must call adapter.cleanup when request finishes.""" + cleanup_calls = [] + + adapter_mock = mocker.MagicMock() + adapter_mock.finished_requests = {"req-s1"} + adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw)) + + from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + + scheduler = mocker.MagicMock() + scheduler.chunk_transfer_adapter = adapter_mock + scheduler.connector = None + scheduler.ec_connector = None + scheduler.perf_metrics = None + scheduler.log_stats = False + scheduler.recompute_kv_load_failures = False + scheduler.structured_output_manager = mocker.MagicMock() + scheduler.structured_output_manager.should_advance.return_value = False + scheduler.finished_req_ids_dict = {} + scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_event_publisher = mocker.MagicMock() + + request = _HashableRequest( + request_id="req-s1", + external_req_id="ext-s1", + status=RequestStatus.RUNNING, + is_finished=lambda: False, + num_computed_tokens=10, + num_prompt_tokens=10, + prompt_token_ids=list(range(10)), + num_output_placeholders=0, + sampling_params=None, + pooling_params=None, + stop_reason=None, + client_index=0, + take_events=lambda: [], + trace_headers=None, + num_cached_tokens=0, + num_external_computed_tokens=0, + num_nans_in_logits=0, + get_finished_reason=lambda: "stop", + ) + scheduler.requests = {"req-s1": request} + + scheduler._handle_stopped_request = mocker.MagicMock(return_value=True) + scheduler._free_request = mocker.MagicMock(return_value=None) + scheduler._get_routed_experts = mocker.MagicMock(return_value=None) + scheduler.running = [request] + scheduler.waiting = mocker.MagicMock() + scheduler.waiting.remove_requests = mocker.MagicMock() + scheduler.make_stats = mocker.MagicMock(return_value=None) + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req-s1": 10}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + model_runner_output = SimpleNamespace( + sampled_token_ids=None, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=None, + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={"req-s1": 0}, + ) + + OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output) + + assert len(cleanup_calls) == 1 + args, _ = cleanup_calls[0] + assert args[0] == "req-s1" + assert args[1] == "ext-s1" diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py new file mode 100644 index 000000000..e29b6af82 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from types import SimpleNamespace + +import torch + +from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_async_chunk + + +def _req(external_req_id: str, *, finished: bool): + return SimpleNamespace( + external_req_id=external_req_id, + is_finished=lambda: finished, + ) + + +def test_talker2code2wav_async_chunk_does_not_emit_empty_chunk_when_not_finished(): + transfer_manager = SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}), + ) + request = _req("rid-empty", finished=False) + + payload = talker2code2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output={"audio_codes": torch.zeros((0,))}, + request=request, + ) + + assert payload is None + + +def test_talker2code2wav_async_chunk_flushes_tail_when_finished_without_pooler_output(): + transfer_manager = SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}), + ) + request_id = "rid-tail" + transfer_manager.code_prompt_token_ids[request_id] = [[1, 2, 3, 4] for _ in range(24)] + request = _req(request_id, finished=True) + + payload = talker2code2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, # e.g. EOS step with no audio_codes + request=request, + ) + + assert payload is not None + assert payload["finished"].item() is True + # ctx_frames header + flat codes + assert len(payload["code_predictor_codes"]) == 1 + 4 * 24 + + +def test_talker2code2wav_async_chunk_emits_eof_marker_when_finished_with_no_frames(): + transfer_manager = SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}), + ) + request = _req("rid-eof", finished=True) + + payload = talker2code2wav_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload == { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7b0a1e54d..82894236e 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -380,6 +380,11 @@ class OmniARScheduler(VLLMScheduler): num_cached_tokens=request.num_cached_tokens, ) ) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.cleanup( + request.request_id, + getattr(request, "external_req_id", None), + ) # [Omni] Cleanup state for finished requests for req in stopped_running_reqs: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 1e7b6a215..75ffa8083 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -77,7 +77,6 @@ class OmniGenerationScheduler(VLLMScheduler): self.chunk_transfer_adapter is not None and request.request_id in self.chunk_transfer_adapter.finished_requests ): - request.status = RequestStatus.FINISHED_STOPPED # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED. if len(request.prompt_token_ids) <= num_computed_tokens: request.prompt_token_ids.append(0) @@ -132,7 +131,6 @@ class OmniGenerationScheduler(VLLMScheduler): # async_chunk: wait for the first upstream chunk (don't start with placeholders). if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0: if request.request_id in self.chunk_transfer_adapter.finished_requests: - request.status = RequestStatus.FINISHED_STOPPED request.prompt_token_ids.append(0) try: request._all_token_ids.append(0) # type: ignore[attr-defined] @@ -416,6 +414,11 @@ class OmniGenerationScheduler(VLLMScheduler): finished = self._handle_stopped_request(request) if finished: kv_transfer_params = self._free_request(request) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.cleanup( + request.request_id, + getattr(request, "external_req_id", None), + ) if status_before_stop == RequestStatus.WAITING_FOR_CHUNK: stopped_running_reqs.add(request) stopped_preempted_reqs.add(request) @@ -486,6 +489,11 @@ class OmniGenerationScheduler(VLLMScheduler): num_cached_tokens=request.num_cached_tokens, ) ) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.cleanup( + request.request_id, + getattr(request, "external_req_id", None), + ) # KV Connector: update state for finished KV Transfers. if kv_connector_output: 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 b46191539..a7f069beb 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 @@ -96,11 +96,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): request.additional_information = None self._pending_load_reqs.append(request) - def save_async(self, pooling_output: torch.Tensor | None = None, request: Request | None = None): + def save_async( + self, + pooling_output: torch.Tensor | None = None, + request: Request | None = None, + ): """Build and enqueue one chunk for asynchronous sending. - Payload extraction is executed in the caller thread via - ``custom_process_next_stage_input_func`` + Payload extraction happens in ``_send_single_request`` on the + background save_loop thread. Args: pooling_output: Partial pooling output dictionary @@ -194,7 +198,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): request_id = request.external_req_id chunk_id = self.put_req_chunk[request_id] connector_put_key = f"{request_id}_{stage_id}_{chunk_id}" - # Process payload in main thread to avoid race conditions on request state + # Process payload in save_loop thread payload_data = None if self.custom_process_next_stage_input_func: try: @@ -222,6 +226,40 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.put_req_chunk[request_id] += 1 logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") + ######################################################################## + # Cleanup + ######################################################################## + + def cleanup( + self, + request_id: str, + external_req_id: str | None = None, + ) -> None: + """Reclaim all per-request state after a request finishes. + + Idempotent: calling with an already-cleaned or unknown id is safe. + + Args: + request_id: Internal request id (receive / scheduler side key). + external_req_id: External request id (send / payload side key). + When *None*, looked up from ``request_ids_mapping``. + """ + if external_req_id is None: + external_req_id = self.request_ids_mapping.get(request_id, request_id) + + self.finished_requests.discard(request_id) + self.get_req_chunk.pop(request_id, None) + self.requests_with_ready_chunks.discard(request_id) + self.request_ids_mapping.pop(request_id, None) + + remaining = deque(r for r in self._pending_load_reqs if getattr(r, "request_id", None) != request_id) + self._pending_load_reqs = remaining + self._finished_load_reqs.discard(request_id) + + self.put_req_chunk.pop(external_req_id, None) + self.request_payload.pop(external_req_id, None) + self.code_prompt_token_ids.pop(external_req_id, None) + ######################################################################## # Schedule Helper ######################################################################## diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 3d9617bf3..72e17bf4f 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -21,14 +21,21 @@ def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None: def talker2code2wav_async_chunk( transfer_manager: Any, - pooling_output: dict[str, Any], + pooling_output: dict[str, Any] | None, request: Any, is_finished: bool = False, ) -> dict[str, Any] | None: - if not isinstance(pooling_output, dict): - return None - request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + if isinstance(pooling_output, dict): + frame = _extract_last_frame(pooling_output) + if frame is not None: + codec_codes = frame.cpu().tolist() + transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + elif not finished: + # Some steps may not produce pooling_output. Only flush on finish. + return None connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} @@ -40,23 +47,24 @@ def talker2code2wav_async_chunk( f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " f"codec_left_context_frames={left_context_size}" ) - - frame = _extract_last_frame(pooling_output) - if frame is not None: - codec_codes = frame.cpu().tolist() - transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) - length = len(transfer_manager.code_prompt_token_ids[request_id]) + + # Avoid emitting empty chunks during normal streaming. If the request is + # finished and nothing was produced, emit an EOF marker. + if length <= 0: + if finished: + return { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + return None + chunk_length = length % chunk_size - if chunk_length != 0 and not is_finished: + if chunk_length != 0 and not finished: return None context_length = chunk_length if chunk_length != 0 else chunk_size - - if length <= 0: - return None - end_index = min(length, left_context_size + context_length) ctx_frames = max(0, int(end_index - context_length)) window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] @@ -68,5 +76,5 @@ def talker2code2wav_async_chunk( # The model expects input_ids layout: [ctx_frames, *flat_codes]. return { "code_predictor_codes": [int(ctx_frames)] + code_predictor_codes, - "finished": torch.tensor(is_finished, dtype=torch.bool), + "finished": torch.tensor(finished, dtype=torch.bool), }