minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model-gateway] add JWT/OIDC authentication for control plane APIs (#15850)

Browse Source
This commit is contained in:
Simo Lin
2025-12-26 18:30:00 -08:00
committed by GitHub
parent 67caea6fe4
commit 4edee6954a
17 changed files with 3738 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/bindings/python/sglang_router/router.py
+68
View File
@@ -5,8 +5,12 @@ from sglang_router.sglang_router_rs import (
BackendType,
HistoryBackendType,
PolicyType,
PyApiKeyEntry,
PyControlPlaneAuthConfig,
PyJwtConfig,
PyOracleConfig,
PyPostgresConfig,
PyRole,
)
from sglang_router.sglang_router_rs import Router as _Router
@@ -61,6 +65,60 @@ def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType:
raise ValueError(f"Unknown history backend: {backend_str}")
def role_from_str(role_str: str) -> PyRole:
"""Convert role string to PyRole enum."""
if role_str.lower() == "admin":
return PyRole.Admin
return PyRole.User
def build_control_plane_auth_config(
args_dict: dict,
) -> Optional[PyControlPlaneAuthConfig]:
"""Build control plane auth config from args dict."""
api_keys = args_dict.get("control_plane_api_keys", [])
jwt_issuer = args_dict.get("jwt_issuer")
jwt_audience = args_dict.get("jwt_audience")
audit_enabled = args_dict.get("control_plane_audit_enabled", False)
# Check if any auth is configured
has_api_keys = bool(api_keys)
has_jwt = jwt_issuer is not None and jwt_audience is not None
if not has_api_keys and not has_jwt:
return None
# Build API key entries
py_api_keys = []
for key_tuple in api_keys:
# Tuple format: (id, name, key, role)
key_id, name, key, role = key_tuple
py_api_keys.append(
PyApiKeyEntry(
id=key_id,
name=name,
key=key,
role=role_from_str(role),
)
)
# Build JWT config if present
jwt_config = None
if has_jwt:
jwt_config = PyJwtConfig(
issuer=jwt_issuer,
audience=jwt_audience,
jwks_uri=args_dict.get("jwt_jwks_uri"),
role_mapping=args_dict.get("jwt_role_mapping", {}),
)
return PyControlPlaneAuthConfig(
jwt=jwt_config,
api_keys=py_api_keys,
audit_enabled=audit_enabled,
)
class Router:
"""
A high-performance router for distributing requests across worker nodes.
@@ -202,6 +260,9 @@ class Router:
)
args_dict["postgres_config"] = postgres_config
# Build control plane auth config
args_dict["control_plane_auth"] = build_control_plane_auth_config(args_dict)
# Remove fields that shouldn't be passed to Rust Router constructor
fields_to_remove = [
"mini_lb",
@@ -215,6 +276,13 @@ class Router:
"oracle_pool_timeout_secs",
"postgres_db_url",
"postgres_pool_max",
# Control plane auth fields (converted to control_plane_auth)
"control_plane_api_keys",
"control_plane_audit_enabled",
"jwt_issuer",
"jwt_audience",
"jwt_jwks_uri",
"jwt_role_mapping",
]
for field in fields_to_remove:
args_dict.pop(field, None)
sgl-model-gateway/bindings/python/sglang_router/router_args.py
+109
View File
@@ -126,6 +126,15 @@ class RouterArgs:
# Trace
enable_trace: bool = False
otlp_traces_endpoint: str = "localhost:4317"
# Control plane authentication
# API keys for control plane auth (list of tuples: id, name, key, role)
control_plane_api_keys: List[tuple] = dataclasses.field(default_factory=list)
control_plane_audit_enabled: bool = False
# JWT/OIDC configuration for control plane auth
jwt_issuer: Optional[str] = None
jwt_audience: Optional[str] = None
jwt_jwks_uri: Optional[str] = None
jwt_role_mapping: Dict[str, str] = dataclasses.field(default_factory=dict)
@staticmethod
def add_cli_args(
@@ -686,6 +695,47 @@ class RouterArgs:
default="localhost:4317",
help="Config opentelemetry collector endpoint if --enable-trace is set. format: <ip>:<port>",
)
# Control plane authentication
parser.add_argument(
f"--{prefix}control-plane-api-keys",
type=str,
nargs="*",
default=[],
help="API keys for control plane authentication. Format: 'id:name:role:key' where role is 'admin' or 'user'. "
"Example: --control-plane-api-keys 'key1:Service Account:admin:secret123' 'key2:Read Only:user:secret456'",
)
parser.add_argument(
f"--{prefix}control-plane-audit-enabled",
action="store_true",
default=False,
help="Enable audit logging for control plane operations",
)
parser.add_argument(
f"--{prefix}jwt-issuer",
type=str,
default=None,
help="OIDC issuer URL for JWT authentication (e.g., https://login.microsoftonline.com/{tenant}/v2.0)",
)
parser.add_argument(
f"--{prefix}jwt-audience",
type=str,
default=None,
help="Expected audience claim for JWT tokens (usually the client ID or API identifier)",
)
parser.add_argument(
f"--{prefix}jwt-jwks-uri",
type=str,
default=None,
help="Explicit JWKS URI. If not provided, discovered from issuer via .well-known/openid-configuration",
)
parser.add_argument(
f"--{prefix}jwt-role-mapping",
type=str,
nargs="*",
default=[],
help="Mapping from IDP role/group names to gateway roles. Format: 'idp_role=gateway_role'. "
"Example: --jwt-role-mapping 'Gateway.Admin=admin' 'Gateway.User=user'",
)
@classmethod
def from_cli_args(
@@ -741,6 +791,16 @@ class RouterArgs:
# Mooncake-specific annotation
args_dict["bootstrap_port_annotation"] = "sglang.ai/bootstrap-port"
# Parse control plane API keys
args_dict["control_plane_api_keys"] = cls._parse_control_plane_api_keys(
cli_args_dict.get(f"{prefix}control_plane_api_keys", [])
)
# Parse JWT role mapping
args_dict["jwt_role_mapping"] = cls._parse_jwt_role_mapping(
cli_args_dict.get(f"{prefix}jwt_role_mapping", [])
)
return cls(**args_dict)
def _validate_router_args(self):
@@ -830,3 +890,52 @@ class RouterArgs:
# decode_list is a list of single-element lists due to nargs=1
return [url[0] for url in decode_list]
@staticmethod
def _parse_control_plane_api_keys(api_keys_list):
"""Parse control plane API keys from --control-plane-api-keys arguments.
Format: id:name:role:key
Example: --control-plane-api-keys 'key1:Service Account:admin:secret123'
"""
if not api_keys_list:
return []
parsed_keys = []
for key_str in api_keys_list:
parts = key_str.split(":", 3) # Split into at most 4 parts
if len(parts) != 4:
raise ValueError(
f"Invalid API key format: '{key_str}'. Expected 'id:name:role:key'"
)
key_id, name, role, key = parts
role_lower = role.lower()
if role_lower not in ("admin", "user"):
raise ValueError(f"Invalid role: '{role}'. Must be 'admin' or 'user'")
parsed_keys.append((key_id, name, key, role_lower))
return parsed_keys
@staticmethod
def _parse_jwt_role_mapping(role_mapping_list):
"""Parse JWT role mapping from --jwt-role-mapping arguments.
Format: idp_role=gateway_role
Example: --jwt-role-mapping 'Gateway.Admin=admin' 'Gateway.User=user'
"""
if not role_mapping_list:
return {}
mapping = {}
for mapping_str in role_mapping_list:
if "=" not in mapping_str:
raise ValueError(
f"Invalid role mapping format: '{mapping_str}'. Expected 'idp_role=gateway_role'"
)
idp_role, gateway_role = mapping_str.split("=", 1)
gateway_role_lower = gateway_role.lower()
if gateway_role_lower not in ("admin", "user"):
raise ValueError(
f"Invalid gateway role: '{gateway_role}'. Must be 'admin' or 'user'"
)
mapping[idp_role] = gateway_role_lower
return mapping
sgl-model-gateway/bindings/python/src/lib.rs
+158
View File
@@ -31,6 +31,152 @@ pub enum HistoryBackendType {
Postgres,
}
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug, Default)]
pub enum PyRole {
Admin,
#[default]
User,
}
impl PyRole {
pub fn to_auth_role(&self) -> auth::Role {
match self {
PyRole::Admin => auth::Role::Admin,
PyRole::User => auth::Role::User,
}
}
}
#[pyclass]
#[derive(Clone, Debug, PartialEq)]
pub struct PyApiKeyEntry {
#[pyo3(get, set)]
pub id: String,
#[pyo3(get, set)]
pub name: String,
#[pyo3(get, set)]
pub key: String,
#[pyo3(get, set)]
pub role: PyRole,
}
#[pymethods]
impl PyApiKeyEntry {
#[new]
#[pyo3(signature = (id, name, key, role = PyRole::User))]
fn new(id: String, name: String, key: String, role: PyRole) -> Self {
PyApiKeyEntry { id, name, key, role }
}
}
impl PyApiKeyEntry {
pub fn to_auth_api_key_entry(&self) -> auth::ApiKeyEntry {
auth::ApiKeyEntry::new(&self.id, &self.name, &self.key, self.role.to_auth_role())
}
}
#[pyclass]
#[derive(Clone, Debug, PartialEq)]
pub struct PyJwtConfig {
#[pyo3(get, set)]
pub issuer: String,
#[pyo3(get, set)]
pub audience: String,
#[pyo3(get, set)]
pub jwks_uri: Option<String>,
#[pyo3(get, set)]
pub role_mapping: HashMap<String, String>,
}
#[pymethods]
impl PyJwtConfig {
#[new]
#[pyo3(signature = (
issuer,
audience,
jwks_uri = None,
role_mapping = HashMap::new(),
))]
fn new(
issuer: String,
audience: String,
jwks_uri: Option<String>,
role_mapping: HashMap<String, String>,
) -> Self {
PyJwtConfig {
issuer,
audience,
jwks_uri,
role_mapping,
}
}
}
impl PyJwtConfig {
pub fn to_auth_jwt_config(&self) -> auth::JwtConfig {
let mut config = auth::JwtConfig::new(&self.issuer, &self.audience);
// Conditionally set JWKS URI
if let Some(ref uri) = self.jwks_uri {
config = config.with_jwks_uri(uri);
}
// Add role mappings
for (idp_role, gateway_role) in &self.role_mapping {
let role = match gateway_role.to_lowercase().as_str() {
"admin" => auth::Role::Admin,
_ => auth::Role::User,
};
config = config.with_role_mapping(idp_role, role);
}
config
}
}
#[pyclass]
#[derive(Clone, Debug, Default, PartialEq)]
pub struct PyControlPlaneAuthConfig {
#[pyo3(get, set)]
pub jwt: Option<PyJwtConfig>,
#[pyo3(get, set)]
pub api_keys: Vec<PyApiKeyEntry>,
#[pyo3(get, set)]
pub audit_enabled: bool,
}
#[pymethods]
impl PyControlPlaneAuthConfig {
#[new]
#[pyo3(signature = (
jwt = None,
api_keys = vec![],
audit_enabled = true,
))]
fn new(
jwt: Option<PyJwtConfig>,
api_keys: Vec<PyApiKeyEntry>,
audit_enabled: bool,
) -> Self {
PyControlPlaneAuthConfig {
jwt,
api_keys,
audit_enabled,
}
}
}
impl PyControlPlaneAuthConfig {
pub fn to_auth_control_plane_config(&self) -> auth::ControlPlaneAuthConfig {
auth::ControlPlaneAuthConfig {
jwt: self.jwt.as_ref().map(|j| j.to_auth_jwt_config()),
api_keys: self.api_keys.iter().map(|k| k.to_auth_api_key_entry()).collect(),
audit_enabled: self.audit_enabled,
}
}
}
#[pyclass]
#[derive(Clone, PartialEq)]
pub struct PyOracleConfig {
@@ -232,6 +378,7 @@ struct Router {
server_key_path: Option<String>,
enable_trace: bool,
otlp_traces_endpoint: String,
control_plane_auth: Option<PyControlPlaneAuthConfig>,
}
impl Router {
@@ -502,6 +649,7 @@ impl Router {
server_key_path = None,
enable_trace = false,
otlp_traces_endpoint = String::from("localhost:4317"),
control_plane_auth = None,
))]
#[allow(clippy::too_many_arguments)]
fn new(
@@ -583,6 +731,7 @@ impl Router {
server_key_path: Option<String>,
enable_trace: bool,
otlp_traces_endpoint: String,
control_plane_auth: Option<PyControlPlaneAuthConfig>,
) -> PyResult<Self> {
let mut all_urls = worker_urls.clone();
@@ -678,6 +827,7 @@ impl Router {
server_key_path,
enable_trace,
otlp_traces_endpoint,
control_plane_auth,
})
}
@@ -736,6 +886,10 @@ impl Router {
request_timeout_secs: self.request_timeout_secs,
request_id_headers: self.request_id_headers.clone(),
shutdown_grace_period_secs: self.shutdown_grace_period_secs,
control_plane_auth: self
.control_plane_auth
.as_ref()
.map(|c| c.to_auth_control_plane_config()),
})
.await
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
@@ -772,6 +926,10 @@ fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PolicyType>()?;
m.add_class::<BackendType>()?;
m.add_class::<HistoryBackendType>()?;
m.add_class::<PyRole>()?;
m.add_class::<PyApiKeyEntry>()?;
m.add_class::<PyJwtConfig>()?;
m.add_class::<PyControlPlaneAuthConfig>()?;
m.add_class::<PyOracleConfig>()?;
m.add_class::<PyPostgresConfig>()?;
m.add_class::<Router>()?;