[model-gateway] Add classification model support infrastructure (#16061)

Co-authored-by: Chang Su <chang.s.su@oracle.com>

python/sglang/srt/entrypoints/grpc_server.py

@@ -5,6 +5,7 @@ Uses GrpcRequestManager for orchestration without tokenization.
 
 import asyncio
 import dataclasses
+import json
 import logging
 import os
 import signal
@@ -334,6 +335,9 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
             pad_token_id=self.model_info["pad_token_id"],
             bos_token_id=self.model_info["bos_token_id"],
             max_req_input_len=self.model_info["max_req_input_len"],
+            # Classification model support
+            id2label_json=self.model_info.get("id2label_json") or "",
+            num_labels=self.model_info.get("num_labels") or 0,
         )
 
     async def GetServerInfo(
@@ -743,6 +747,22 @@ async def serve_grpc(
 
     # Update model info from scheduler info and model config
     if model_info is None:
+        # Extract classification labels from HuggingFace config (if available)
+        # Match logic in serving_classify.py::_get_id2label_mapping
+        hf_config = model_config.hf_config
+        id2label = getattr(hf_config, "id2label", None)
+        num_labels = getattr(hf_config, "num_labels", 0) or 0
+
+        # If no id2label but num_labels exists, create default mapping
+        if not id2label and num_labels:
+            id2label = {i: f"LABEL_{i}" for i in range(num_labels)}
+        elif id2label and not num_labels:
+            num_labels = len(id2label)
+
+        # Convert to JSON string for proto transport
+        # id2label is a dict like {0: "negative", 1: "positive"}
+        id2label_json = json.dumps(id2label) if id2label else ""
+
         model_info = {
             "model_name": server_args.model_path,
             "max_context_length": scheduler_info.get(
@@ -750,12 +770,15 @@ async def serve_grpc(
             ),
             "vocab_size": scheduler_info.get("vocab_size", 128256),
             "supports_vision": scheduler_info.get("supports_vision", False),
-            "model_type": getattr(model_config.hf_config, "model_type", None),
+            "model_type": getattr(hf_config, "model_type", None),
-            "architectures": getattr(model_config.hf_config, "architectures", None),
+            "architectures": getattr(hf_config, "architectures", None),
             "max_req_input_len": scheduler_info.get("max_req_input_len", 8192),
             "eos_token_ids": scheduler_info.get("eos_token_ids", []),
             "pad_token_id": scheduler_info.get("pad_token_id", 0),
             "bos_token_id": scheduler_info.get("bos_token_id", 1),
+            # Classification model support
+            "id2label_json": id2label_json,
+            "num_labels": num_labels or 0,
         }
 
     # Create request manager with the correct port args

python/sglang/srt/grpc/sglang_scheduler.proto

@@ -428,6 +428,12 @@ message GetModelInfoResponse {
   int32 bos_token_id = 13;
   int32 max_req_input_len = 14;
   repeated string architectures = 15;
+
+  // Classification model support (from HuggingFace config.json)
+  // id2label maps class indices to label names, e.g., {"0": "negative", "1": "positive"}
+  string id2label_json = 16;
+  // Number of classification labels (0 if not a classifier)
+  int32 num_labels = 17;
 }
 
 // Get server information

python/sglang/srt/grpc/sglang_scheduler_pb2.pyi

@@ -432,7 +432,7 @@ class GetModelInfoRequest(_message.Message):
     def __init__(self) -> None: ...
 
 class GetModelInfoResponse(_message.Message):
-    __slots__ = ("model_path", "tokenizer_path", "is_generation", "preferred_sampling_params", "weight_version", "served_model_name", "max_context_length", "vocab_size", "supports_vision", "model_type", "eos_token_ids", "pad_token_id", "bos_token_id", "max_req_input_len", "architectures")
+    __slots__ = ("model_path", "tokenizer_path", "is_generation", "preferred_sampling_params", "weight_version", "served_model_name", "max_context_length", "vocab_size", "supports_vision", "model_type", "eos_token_ids", "pad_token_id", "bos_token_id", "max_req_input_len", "architectures", "id2label_json", "num_labels")
     MODEL_PATH_FIELD_NUMBER: _ClassVar[int]
     TOKENIZER_PATH_FIELD_NUMBER: _ClassVar[int]
     IS_GENERATION_FIELD_NUMBER: _ClassVar[int]
@@ -448,6 +448,8 @@ class GetModelInfoResponse(_message.Message):
     BOS_TOKEN_ID_FIELD_NUMBER: _ClassVar[int]
     MAX_REQ_INPUT_LEN_FIELD_NUMBER: _ClassVar[int]
     ARCHITECTURES_FIELD_NUMBER: _ClassVar[int]
+    ID2LABEL_JSON_FIELD_NUMBER: _ClassVar[int]
+    NUM_LABELS_FIELD_NUMBER: _ClassVar[int]
     model_path: str
     tokenizer_path: str
     is_generation: bool
@@ -463,7 +465,9 @@ class GetModelInfoResponse(_message.Message):
     bos_token_id: int
     max_req_input_len: int
     architectures: _containers.RepeatedScalarFieldContainer[str]
-    def __init__(self, model_path: _Optional[str] = ..., tokenizer_path: _Optional[str] = ..., is_generation: bool = ..., preferred_sampling_params: _Optional[str] = ..., weight_version: _Optional[str] = ..., served_model_name: _Optional[str] = ..., max_context_length: _Optional[int] = ..., vocab_size: _Optional[int] = ..., supports_vision: bool = ..., model_type: _Optional[str] = ..., eos_token_ids: _Optional[_Iterable[int]] = ..., pad_token_id: _Optional[int] = ..., bos_token_id: _Optional[int] = ..., max_req_input_len: _Optional[int] = ..., architectures: _Optional[_Iterable[str]] = ...) -> None: ...
+    id2label_json: str
+    num_labels: int
+    def __init__(self, model_path: _Optional[str] = ..., tokenizer_path: _Optional[str] = ..., is_generation: bool = ..., preferred_sampling_params: _Optional[str] = ..., weight_version: _Optional[str] = ..., served_model_name: _Optional[str] = ..., max_context_length: _Optional[int] = ..., vocab_size: _Optional[int] = ..., supports_vision: bool = ..., model_type: _Optional[str] = ..., eos_token_ids: _Optional[_Iterable[int]] = ..., pad_token_id: _Optional[int] = ..., bos_token_id: _Optional[int] = ..., max_req_input_len: _Optional[int] = ..., architectures: _Optional[_Iterable[str]] = ..., id2label_json: _Optional[str] = ..., num_labels: _Optional[int] = ...) -> None: ...
 
 class GetServerInfoRequest(_message.Message):
     __slots__ = ()

sgl-model-gateway/src/core/model_card.rs

@@ -7,6 +7,8 @@
 //!
 //! Inspired by Dynamo's ModelDeploymentCard but simplified for router needs.
 
+use std::collections::HashMap;
+
 use serde::{Deserialize, Serialize};
 
 use super::{
@@ -168,6 +170,21 @@ pub struct ModelCard {
     /// User-defined metadata (for fields not covered above)
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub metadata: Option<serde_json::Value>,
+
+    // === Classification Support ===
+    /// Classification label mapping (class index -> label name).
+    /// Empty if not a classification model.
+    /// Example: {0: "negative", 1: "positive"}
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub id2label: HashMap<u32, String>,
+
+    /// Number of classification labels (0 if not a classifier).
+    #[serde(default, skip_serializing_if = "is_zero")]
+    pub num_labels: u32,
+}
+
+fn is_zero(n: &u32) -> bool {
+    *n == 0
 }
 
 fn default_model_type() -> ModelType {
@@ -193,6 +210,8 @@ impl ModelCard {
             reasoning_parser: None,
             tool_parser: None,
             metadata: None,
+            id2label: HashMap::new(),
+            num_labels: 0,
         }
     }
 
@@ -276,6 +295,19 @@ impl ModelCard {
         self
     }
 
+    /// Set the id2label mapping for classification models
+    pub fn with_id2label(mut self, id2label: HashMap<u32, String>) -> Self {
+        self.num_labels = id2label.len() as u32;
+        self.id2label = id2label;
+        self
+    }
+
+    /// Set num_labels directly (alternative to with_id2label)
+    pub fn with_num_labels(mut self, num_labels: u32) -> Self {
+        self.num_labels = num_labels;
+        self
+    }
+
     // === Query methods ===
 
     /// Check if this model matches the given ID (including aliases)
@@ -334,6 +366,20 @@ impl ModelCard {
     pub fn supports_reasoning(&self) -> bool {
         self.model_type.supports_reasoning()
     }
+
+    /// Check if this is a classification model
+    #[inline]
+    pub fn is_classifier(&self) -> bool {
+        self.num_labels > 0
+    }
+
+    /// Get label for a class index, with fallback to generic label (LABEL_N)
+    pub fn get_label(&self, class_idx: u32) -> String {
+        self.id2label
+            .get(&class_idx)
+            .cloned()
+            .unwrap_or_else(|| format!("LABEL_{}", class_idx))
+    }
 }
 
 impl Default for ModelCard {

sgl-model-gateway/src/core/steps/worker/local/create_worker.rs

@@ -175,6 +175,40 @@ fn build_model_card(
         }
     }
 
+    // Parse classification model id2label mapping
+    // The proto field is id2label_json: JSON string like {"0": "negative", "1": "positive"}
+    if let Some(id2label_json) = labels.get("id2label_json") {
+        if !id2label_json.is_empty() {
+            // Parse JSON: keys are string indices, values are label names
+            if let Ok(string_map) = serde_json::from_str::<HashMap<String, String>>(id2label_json) {
+                // Convert string keys ("0", "1") to u32 keys (0, 1)
+                let id2label: HashMap<u32, String> = string_map
+                    .into_iter()
+                    .filter_map(|(k, v)| k.parse::<u32>().ok().map(|idx| (idx, v)))
+                    .collect();
+
+                if !id2label.is_empty() {
+                    card = card.with_id2label(id2label);
+                    debug!("Parsed id2label with {} classes", card.num_labels);
+                }
+            }
+        }
+    }
+    // Fallback: if num_labels is set but id2label wasn't parsed, create default labels
+    // Match logic in serving_classify.py::_get_id2label_mapping
+    else if let Some(num_labels_str) = labels.get("num_labels") {
+        if let Ok(num_labels) = num_labels_str.parse::<u32>() {
+            if num_labels > 0 {
+                // Create default mapping: {0: "LABEL_0", 1: "LABEL_1", ...}
+                let id2label: HashMap<u32, String> = (0..num_labels)
+                    .map(|i| (i, format!("LABEL_{}", i)))
+                    .collect();
+                card = card.with_id2label(id2label);
+                debug!("Created default id2label with {} classes", num_labels);
+            }
+        }
+    }
+
     card
 }
 

sgl-model-gateway/src/core/worker.rs

@@ -242,6 +242,40 @@ pub trait Worker: Send + Sync + fmt::Debug {
         self.metadata().provider_for_model(model_id)
     }
 
+    /// Check if a model is a classifier (has id2label mapping).
+    fn is_classifier(&self, model_id: &str) -> bool {
+        self.metadata()
+            .find_model(model_id)
+            .map(|m| m.is_classifier())
+            .unwrap_or(false)
+    }
+
+    /// Get the id2label mapping for a classification model.
+    /// Returns None if model is not a classifier or not found.
+    fn id2label(&self, model_id: &str) -> Option<&std::collections::HashMap<u32, String>> {
+        self.metadata()
+            .find_model(model_id)
+            .filter(|m| m.is_classifier())
+            .map(|m| &m.id2label)
+    }
+
+    /// Get the number of classification labels for a model.
+    fn num_labels(&self, model_id: &str) -> u32 {
+        self.metadata()
+            .find_model(model_id)
+            .map(|m| m.num_labels)
+            .unwrap_or(0)
+    }
+
+    /// Get label for a class index from a classification model.
+    /// Returns generic label (LABEL_N) if model not found or index not in mapping.
+    fn get_label(&self, model_id: &str, class_idx: u32) -> String {
+        self.metadata()
+            .find_model(model_id)
+            .map(|m| m.get_label(class_idx))
+            .unwrap_or_else(|| format!("LABEL_{}", class_idx))
+    }
+
     /// Check if this worker supports a specific model.
     /// If models list is empty, worker accepts any model.
     fn supports_model(&self, model_id: &str) -> bool {

sgl-model-gateway/src/proto/sglang_scheduler.proto

@@ -428,6 +428,12 @@ message GetModelInfoResponse {
   int32 bos_token_id = 13;
   int32 max_req_input_len = 14;
   repeated string architectures = 15;
+
+  // Classification model support (from HuggingFace config.json)
+  // id2label maps class indices to label names, e.g., {"0": "negative", "1": "positive"}
+  string id2label_json = 16;
+  // Number of classification labels (0 if not a classifier)
+  int32 num_labels = 17;
 }
 
 // Get server information

sgl-model-gateway/src/routers/grpc/client.rs

@@ -105,7 +105,7 @@ impl GrpcClient {
         match self {
             Self::Sglang(client) => {
                 let info = client.get_model_info().await?;
-                Ok(ModelInfo::Sglang(info))
+                Ok(ModelInfo::Sglang(Box::new(info)))
             }
             Self::Vllm(client) => {
                 let info = client.get_model_info().await?;
@@ -151,7 +151,7 @@ impl GrpcClient {
 
 /// Unified ModelInfo wrapper
 pub enum ModelInfo {
-    Sglang(crate::grpc_client::sglang_proto::GetModelInfoResponse),
+    Sglang(Box<crate::grpc_client::sglang_proto::GetModelInfoResponse>),
     Vllm(crate::grpc_client::vllm_proto::GetModelInfoResponse),
 }