[BugFix] Modify the method of obtaining external_request_id (#961)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

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,

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,

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,

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,

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]

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 = {

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

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,