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

Co-authored-by: Chang Su <chang.s.su@oracle.com>
This commit is contained in:
Simo Lin
2025-12-29 08:34:05 -08:00
committed by GitHub
parent 8e08207c18
commit 162d1cf9be
9 changed files with 167 additions and 14 deletions

View File

@@ -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),
"architectures": getattr(model_config.hf_config, "architectures", None),
"model_type": getattr(hf_config, "model_type", 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

View File

@@ -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

File diff suppressed because one or more lines are too long

View File

@@ -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__ = ()