feature: support X-Data-Parallel-Rank header to specific dp-rank. (#19832)

Signed-off-by: doujiang24 <doujiang24@gmail.com>

python/sglang/srt/entrypoints/openai/serving_base.py

@@ -268,3 +268,34 @@ class OpenAIServingBase(ABC):
         if raw_request is None:
             return None
         return raw_request.headers.get("x-smg-routing-key")
+
+    def extract_routed_dp_rank_from_header(
+        self, raw_request: Request, body_routed_dp_rank: Optional[int] = None
+    ) -> Optional[int]:
+        """Extract routed_dp_rank from HTTP header, with higher priority than routed_dp_rank in body.
+
+        Header name: X-Data-Parallel-Rank (case-insensitive in HTTP/1.1/2)
+        """
+        if raw_request is None:
+            return body_routed_dp_rank
+
+        header_value = raw_request.headers.get("x-data-parallel-rank")
+        if header_value is not None:
+            try:
+                header_dp_rank = int(header_value)
+                if (
+                    body_routed_dp_rank is not None
+                    and header_dp_rank != body_routed_dp_rank
+                ):
+                    logger.debug(
+                        f"X-Data-Parallel-Rank header ({header_dp_rank}) overrides "
+                        f"body routed_dp_rank ({body_routed_dp_rank})"
+                    )
+                return header_dp_rank
+            except ValueError:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Invalid X-Data-Parallel-Rank header: must be an integer, got '{header_value}'",
+                )
+
+        return body_routed_dp_rank

python/sglang/srt/entrypoints/openai/serving_chat.py

@@ -276,6 +276,11 @@ class OpenAIServingChat(OpenAIServingBase):
         # Extract custom labels from raw request headers
         custom_labels = self.extract_custom_labels(raw_request)
 
+        # Extract routed_dp_rank from header (has higher priority than body)
+        effective_routed_dp_rank = self.extract_routed_dp_rank_from_header(
+            raw_request, request.routed_dp_rank
+        )
+
         # Resolve LoRA adapter from model parameter or explicit lora_path
         lora_path = self._resolve_lora_path(request.model, request.lora_path)
         img_max_dynamic_patch, vid_max_dynamic_patch = _extract_max_dynamic_patch(
@@ -297,7 +302,7 @@ class OpenAIServingChat(OpenAIServingBase):
             bootstrap_host=request.bootstrap_host,
             bootstrap_port=request.bootstrap_port,
             bootstrap_room=request.bootstrap_room,
-            routed_dp_rank=request.routed_dp_rank,
+            routed_dp_rank=effective_routed_dp_rank,
             disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
             return_hidden_states=request.return_hidden_states,
             return_routed_experts=request.return_routed_experts,

python/sglang/srt/entrypoints/openai/serving_completions.py

@@ -97,6 +97,11 @@ class OpenAIServingCompletion(OpenAIServingBase):
         # Extract custom labels from raw request headers
         custom_labels = self.extract_custom_labels(raw_request)
 
+        # Extract routed_dp_rank from header (has higher priority than body)
+        effective_routed_dp_rank = self.extract_routed_dp_rank_from_header(
+            raw_request, request.routed_dp_rank
+        )
+
         # Resolve LoRA adapter from model parameter or explicit lora_path
         lora_path = self._resolve_lora_path(request.model, request.lora_path)
 
@@ -112,7 +117,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
             bootstrap_host=request.bootstrap_host,
             bootstrap_port=request.bootstrap_port,
             bootstrap_room=request.bootstrap_room,
-            routed_dp_rank=request.routed_dp_rank,
+            routed_dp_rank=effective_routed_dp_rank,
             disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
             return_hidden_states=request.return_hidden_states,
             return_routed_experts=request.return_routed_experts,