diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index c72cc91d0..8fb6278c1 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -591,6 +591,25 @@ def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapt assert adapter._pending_save_reqs[-1]["request"] is next_segment +def test_send_single_request_drops_queued_old_segment_after_boundary(build_adapter): + adapter, connector = build_adapter(stage_id=1) + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct()) + request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old") + request.resumable = True + request._omni_segment_generation = 0 + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False) + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + + old_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(old_task) + connector.put.assert_not_called() + + boundary_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(boundary_task) + connector.put.assert_called_once() + + def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch): adapter, _ = build_adapter(stage_id=1) request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished") 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 7c9267162..7668fa5a0 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 @@ -370,6 +370,21 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): stage_id = self.connector.stage_id next_stage_id = stage_id + 1 external_req_id = request.external_req_id + task_generation = task.get("segment_generation") + expected_generation = self._segment_generation.get(external_req_id, task_generation) + if ( + isinstance(task_generation, int) + and isinstance(expected_generation, int) + and task_generation < expected_generation + and not is_segment_finished + ): + logger.warning( + "Skip queued late chunk for request %s, segment_generation=%s, expected=%s", + external_req_id, + task_generation, + expected_generation, + ) + return chunk_id = self.put_req_chunk[external_req_id] connector_put_key = f"{external_req_id}_{stage_id}_{chunk_id}" # Process payload in save_loop thread