From d2b3a24af32e0645226c7b82a1c50038f00d8740 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Tue, 27 Jan 2026 20:56:02 +0800 Subject: [PATCH] [BugFix] Modify the method of obtaining external_request_id (#961) Signed-off-by: amy-why-3459 --- tests/entrypoints/test_omni_new_request_data.py | 2 ++ vllm_omni/core/sched/omni_ar_scheduler.py | 1 + vllm_omni/core/sched/omni_generation_scheduler.py | 1 + vllm_omni/core/sched/output.py | 3 +++ vllm_omni/distributed/omni_connectors/adapter.py | 9 +++++---- .../omni_connectors/connectors/shm_connector.py | 1 + vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py | 2 -- vllm_omni/request.py | 6 ++++++ 8 files changed, 19 insertions(+), 6 deletions(-) diff --git a/tests/entrypoints/test_omni_new_request_data.py b/tests/entrypoints/test_omni_new_request_data.py index 81992fb39..776509d5b 100644 --- a/tests/entrypoints/test_omni_new_request_data.py +++ b/tests/entrypoints/test_omni_new_request_data.py @@ -13,6 +13,7 @@ def test_omni_new_request_data_copies_payloads(): } request = SimpleNamespace( request_id="req-1", + external_req_id="ext-1", prompt_token_ids=[101, 102], mm_features=None, sampling_params=None, @@ -33,6 +34,7 @@ def test_omni_new_request_data_copies_payloads(): def test_omni_new_request_data_allows_missing_payloads(): request = SimpleNamespace( request_id="req-2", + external_req_id="ext-2", prompt_token_ids=[201, 202], mm_features=None, sampling_params=None, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 03eb5ff21..647eb4867 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -166,6 +166,7 @@ class OmniARScheduler(VLLMScheduler): # Build omni entry preserving all base fields omni_nr = OmniNewRequestData( req_id=nr.req_id, + external_req_id=(getattr(request, "external_req_id", None) if request else None), prompt_token_ids=nr.prompt_token_ids, mm_features=nr.mm_features, sampling_params=nr.sampling_params, diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 0164029d2..5974adb52 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -210,6 +210,7 @@ class OmniGenerationScheduler(VLLMScheduler): # Build omni entry preserving all base fields omni_nr = OmniNewRequestData( req_id=nr.req_id, + external_req_id=(getattr(request, "external_req_id", None) if request else None), prompt_token_ids=nr.prompt_token_ids, mm_features=nr.mm_features, sampling_params=nr.sampling_params, diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 3176ec14a..f4acb034d 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -21,6 +21,8 @@ class OmniNewRequestData(NewRequestData): # Optional serialized prompt embeddings prompt_embeds: PromptEmbedsPayload | None = None + # Optional external request ID for tracking + external_req_id: str | None = None # Optional serialized additional information additional_information: AdditionalInformationPayload | None = None @@ -42,6 +44,7 @@ class OmniNewRequestData(NewRequestData): """ return cls( req_id=request.request_id, + external_req_id=request.external_req_id, prompt_token_ids=request.prompt_token_ids, mm_features=request.mm_features, sampling_params=request.sampling_params, diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index da3abf4a9..a634595ac 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -199,7 +199,8 @@ def get_chunk( target_stage_id = stage_id - 1 # Handle new requests for new_req_data in scheduler_output.scheduled_new_reqs: - req_id = new_req_data.req_id[0:25] + connector.request_ids_mapping[new_req_data.req_id] = new_req_data.external_req_id + req_id = new_req_data.external_req_id chunk_id = connector.get_requests[req_id] connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) @@ -214,7 +215,7 @@ def get_chunk( cached_reqs.additional_information = {} for i, cached_req_id in enumerate(cached_reqs.req_ids): - req_id = cached_req_id[0:25] + req_id = connector.request_ids_mapping.get(cached_req_id, cached_req_id) if req_id in connector.finished_requests: continue chunk_id = connector.get_requests[req_id] @@ -265,7 +266,7 @@ def get_chunk_for_generation( """ stage_id = connector.stage_id target_stage_id = stage_id - 1 - request_id = request.request_id[0:25] + request_id = request.external_req_id if request_id in connector.finished_requests: return @@ -306,7 +307,7 @@ def put_chunk( """ stage_id = connector.stage_id next_stage_id = stage_id + 1 - request_id = request.request_id[0:25] + request_id = request.external_req_id prompt_token_ids = request.prompt_token_ids connector.request_prompt_token_ids[request_id] = prompt_token_ids chunk_id = connector.put_requests[request_id] diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index a9587e21c..bd6ec8551 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -29,6 +29,7 @@ class SharedMemoryConnector(OmniConnectorBase): self.request_payload = {} self.request_prompt_token_ids: dict[str, list[int]] = defaultdict(list) self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) + self.request_ids_mapping: dict[str, str] = {} # Default threshold matches legacy behavior (64KB) self.threshold = int(config.get("shm_threshold_bytes", 65536)) self._metrics = { diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 0dd590710..a42bd8e76 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -180,8 +180,6 @@ class Qwen3OmniMoeForConditionalGeneration( self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None ) - self.chunk_segment_info = {} # request_id -> chunk_segment_info - # ==================== Device utilities ==================== @staticmethod diff --git a/vllm_omni/request.py b/vllm_omni/request.py index ca5bd437d..864d9dce4 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -26,6 +26,8 @@ class OmniRequest(Request): def __init__( self, prompt_embeds: PromptEmbedsPayload | None = None, + # Optional external request ID for tracking + external_req_id: str | None = None, additional_information: AdditionalInformationPayload | None = None, *args, **kwargs, @@ -33,6 +35,8 @@ class OmniRequest(Request): super().__init__(*args, **kwargs) # Serialized prompt embeddings payload (optional) self.prompt_embeds: PromptEmbedsPayload | None = prompt_embeds + # Optional external request ID for tracking + self.external_req_id: str | None = external_req_id # Serialized additional information payload (optional) self.additional_information: AdditionalInformationPayload | None = additional_information @@ -54,6 +58,8 @@ class OmniRequest(Request): """ return cls( request_id=request.request_id, + # Optional external request ID for tracking + external_req_id=request.external_req_id, client_index=request.client_index, prompt_token_ids=request.prompt_token_ids, prompt_embeds=request.prompt_embeds,