Tiny pass routing key to scheduler processes (#16839)

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

@@ -282,3 +282,8 @@ class OpenAIServingBase(ABC):
                 if label in self.allowed_custom_labels
             }
         return custom_labels
+
+    def extract_routing_key(self, raw_request):
+        if raw_request is None:
+            return None
+        return raw_request.headers.get("x-smg-routing-key")

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

@@ -246,6 +246,7 @@ class OpenAIServingChat(OpenAIServingBase):
             extra_key=self._compute_extra_key(request),
             require_reasoning=self._get_reasoning_from_request(request),
             priority=request.priority,
+            routing_key=self.extract_routing_key(raw_request),
             custom_labels=custom_labels,
             custom_logit_processor=request.custom_logit_processor,
             image_max_dynamic_patch=img_max_dynamic_patch,

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

@@ -121,6 +121,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
             rid=request.rid,
             extra_key=self._compute_extra_key(request),
             priority=request.priority,
+            routing_key=self.extract_routing_key(raw_request),
             custom_labels=custom_labels,
             custom_logit_processor=request.custom_logit_processor,
         )

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

@@ -130,6 +130,7 @@ class OpenAIServingEmbedding(OpenAIServingBase):
             **prompt_kwargs,
             rid=request.rid,
             priority=request.priority,
+            routing_key=self.extract_routing_key(raw_request),
             dimensions=request.dimensions,
         )