From 8484d85d1f78bf9e9002af879d0a47849957cadf Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Wed, 7 Oct 2026 15:04:01 -0600 Subject: [PATCH] [Bugfix] Fix Sender Address Propagation (#8599) Signed-off-by: Alex Brooks --- .../sched/test_omni_scheduling_coordinator.py | 21 +++++++++++++++++++ .../core/sched/omni_scheduling_coordinator.py | 10 ++++----- vllm_omni/request.py | 6 +++++- 3 files changed, 31 insertions(+), 6 deletions(-) diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 8a62cbbfd..34f208b6f 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -16,6 +16,7 @@ from types import SimpleNamespace import pytest import torch +from vllm import SamplingParams import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod from vllm_omni.core.sched.omni_scheduling_coordinator import ( @@ -23,6 +24,9 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import ( uses_native_mrv2_data_plane, ) from vllm_omni.core.sched.output import OmniChunkRecvHandle +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import _LoadEntry +from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens +from vllm_omni.request import OmniRequest pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -440,3 +444,20 @@ class TestTimeoutDetection(unittest.TestCase): self.assertNotIn("r1", coord.requests_with_ready_chunks) self.assertEqual([r.request_id for r in coord._waiting_for_chunk_running], ["r2"]) self.assertEqual([h.request_id for h in coord.pending_chunk_registrations], ["r2"]) + + +def test_sender_address_on_the_engine_request_reaches_both_receive_paths(): + """Ensure the sender address the orchestrator puts on an engine request reaches the receivers.""" + sender = {"host": "10.0.0.2", "zmq_port": 50071} + engine_request = build_engine_core_request_from_tokens( + "req-1", {"prompt_token_ids": [0]}, SamplingParams(max_tokens=1) + ) + engine_request.payload_sender_info = sender + request = OmniRequest.from_engine_core_request(engine_request, block_hasher=None) + # Build the coordinator and process the request with sender info + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.process_pending_full_payload_inputs(MockQueue([request]), stage_recv_req_ids=set()) + + # Ensure that the payload send info is accessible on both pending input registrations and a wrapped load entry + assert [handle.payload_sender_info for handle in coordinator.pending_input_registrations] == [sender] + assert _LoadEntry(request).source_metadata == {"source_host": "10.0.0.2", "source_port": 50071} diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index f3042b736..9ad4ec5a9 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -183,8 +183,8 @@ class OmniSchedulingCoordinator: self.pending_input_registrations.append( OmniChunkRecvHandle( request_id=request.request_id, - external_req_id=getattr(request, "external_req_id", None), - payload_sender_info=getattr(request, "payload_sender_info", None), + external_req_id=request.external_req_id, + payload_sender_info=request.payload_sender_info, ) ) elif request.status == RequestStatus.WAITING_FOR_INPUT: @@ -197,8 +197,8 @@ class OmniSchedulingCoordinator: self.pending_input_registrations.append( OmniChunkRecvHandle( request_id=request.request_id, - external_req_id=getattr(request, "external_req_id", None), - payload_sender_info=getattr(request, "payload_sender_info", None), + external_req_id=request.external_req_id, + payload_sender_info=request.payload_sender_info, ) ) if to_remove: @@ -456,7 +456,7 @@ class OmniSchedulingCoordinator: self.pending_chunk_registrations.append( OmniChunkRecvHandle( request_id=request.request_id, - external_req_id=getattr(request, "external_req_id", None), + external_req_id=request.external_req_id, ) ) request.status = RequestStatus.WAITING_FOR_CHUNK diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 643ab020a..3445f7009 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -3,7 +3,7 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -39,6 +39,7 @@ class OmniRequest(Request): external_req_id: str | None = None, additional_information: AdditionalInformationPayload | None = None, model_intermediate_buffer: dict | None = None, + payload_sender_info: dict[str, Any] | None = None, **kwargs, ): if prompt_embeds is not None: @@ -58,6 +59,8 @@ class OmniRequest(Request): self.additional_information: AdditionalInformationPayload | None = additional_information # Runner-owned runtime payload. self.model_intermediate_buffer: dict | None = model_intermediate_buffer + # Sender's connector address for this request's stage payload. + self.payload_sender_info: dict[str, Any] | None = payload_sender_info @staticmethod def _maybe_decode_prompt_embeds( @@ -105,6 +108,7 @@ class OmniRequest(Request): block_hasher=block_hasher, additional_information=request.additional_information, model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None), + payload_sender_info=request.payload_sender_info, resumable=request.resumable, session_id=request.session_id, kv_hints=request.kv_hints,