[Feat] : Support Async chunk cleanup (#1087)

Signed-off-by: Sy03 <1370724210@qq.com>

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"

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),
+    }

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:

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:

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
     ########################################################################

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),
     }