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

[sgl-model-gateway] Close PyO3 binding gaps and add regression tests (#24719)

Browse Source 
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
2026-05-13 17:43:01 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 7618ad7075
commit af7511e0e8
4 changed files with 864 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/bindings/python/src/lib.rs
+43
View File
@@ -90,28 +90,36 @@ pub struct PyJwtConfig {
pub jwks_uri: Option<String>,
#[pyo3(get, set)]
pub role_mapping: HashMap<String, String>,
#[pyo3(get, set)]
pub role_claim: String,
}
#[pymethods]
impl PyJwtConfig {
#[new]
// `role_claim` is appended at the end with a default so existing positional
// callers — `PyJwtConfig(issuer, audience, jwks_uri, role_mapping)` — keep
// working unchanged.
#[pyo3(signature = (
issuer,
audience,
jwks_uri = None,
role_mapping = HashMap::new(),
role_claim = String::from("roles"),
))]
fn new(
issuer: String,
audience: String,
jwks_uri: Option<String>,
role_mapping: HashMap<String, String>,
role_claim: String,
) -> Self {
PyJwtConfig {
issuer,
audience,
jwks_uri,
role_mapping,
role_claim,
}
}
}
@@ -119,6 +127,7 @@ impl PyJwtConfig {
impl PyJwtConfig {
pub fn to_auth_jwt_config(&self) -> auth::JwtConfig {
let mut config = auth::JwtConfig::new(&self.issuer, &self.audience);
config.role_claim = self.role_claim.clone();
// Conditionally set JWKS URI
if let Some(ref uri) = self.jwks_uri {
@@ -421,6 +430,20 @@ struct Router {
enable_trace: bool,
otlp_traces_endpoint: String,
control_plane_auth: Option<PyControlPlaneAuthConfig>,
// The following five fields expose `#[pyo3(get)]` so tests can verify the
// Python kwargs landed in the right slot. Without getters, a typo'd builder
// call (e.g. `.pool_idle_timeout_secs(self.connect_timeout_secs)`) is
// undetectable from Python.
#[pyo3(get)]
pool_idle_timeout_secs: u64,
#[pyo3(get)]
connect_timeout_secs: u64,
#[pyo3(get)]
pool_max_idle_per_host: usize,
#[pyo3(get)]
tcp_keepalive_secs: u64,
#[pyo3(get)]
enable_wasm: bool,
}
impl Router {
@@ -623,6 +646,11 @@ impl Router {
.retries(!self.disable_retries)
.circuit_breaker(!self.disable_circuit_breaker)
.igw(self.enable_igw)
.pool_idle_timeout_secs(self.pool_idle_timeout_secs)
.connect_timeout_secs(self.connect_timeout_secs)
.pool_max_idle_per_host(self.pool_max_idle_per_host)
.tcp_keepalive_secs(self.tcp_keepalive_secs)
.enable_wasm(self.enable_wasm)
.maybe_client_cert_and_key(
self.client_cert_path.as_ref(),
self.client_key_path.as_ref(),
@@ -724,6 +752,11 @@ impl Router {
enable_trace = false,
otlp_traces_endpoint = String::from("localhost:4317"),
control_plane_auth = None,
pool_idle_timeout_secs = 50,
connect_timeout_secs = 10,
pool_max_idle_per_host = 500,
tcp_keepalive_secs = 30,
enable_wasm = false,
))]
#[allow(clippy::too_many_arguments)]
fn new(
@@ -811,6 +844,11 @@ impl Router {
enable_trace: bool,
otlp_traces_endpoint: String,
control_plane_auth: Option<PyControlPlaneAuthConfig>,
pool_idle_timeout_secs: u64,
connect_timeout_secs: u64,
pool_max_idle_per_host: usize,
tcp_keepalive_secs: u64,
enable_wasm: bool,
) -> PyResult<Self> {
let mut all_urls = worker_urls.clone();
@@ -912,6 +950,11 @@ impl Router {
enable_trace,
otlp_traces_endpoint,
control_plane_auth,
pool_idle_timeout_secs,
connect_timeout_secs,
pool_max_idle_per_host,
tcp_keepalive_secs,
enable_wasm,
})
}
sgl-model-gateway/bindings/python/src/sglang_router/router.py
+21
View File
@@ -1,6 +1,9 @@
import logging
from typing import Optional
from sglang_router.router_args import RouterArgs
logger = logging.getLogger(__name__)
from sglang_router.sglang_router_rs import (
BackendType,
HistoryBackendType,
@@ -90,6 +93,22 @@ def build_control_plane_auth_config(
has_api_keys = bool(api_keys)
has_jwt = jwt_issuer is not None and jwt_audience is not None
# Warn when JWT-related fields are set but the required pair is missing —
# otherwise the silent drop here disables JWT auth without explanation.
jwt_partial = bool(
args_dict.get("jwt_jwks_uri")
or args_dict.get("jwt_role_mapping")
or (
args_dict.get("jwt_role_claim") is not None
and args_dict.get("jwt_role_claim") != "roles"
)
)
if jwt_partial and not has_jwt:
logger.warning(
"JWT-related fields set but jwt_issuer/jwt_audience missing; "
"JWT auth will NOT be enabled."
)
if not has_api_keys and not has_jwt:
return None
@@ -114,6 +133,7 @@ def build_control_plane_auth_config(
issuer=jwt_issuer,
audience=jwt_audience,
jwks_uri=args_dict.get("jwt_jwks_uri"),
role_claim=args_dict.get("jwt_role_claim", "roles"),
role_mapping=args_dict.get("jwt_role_mapping", {}),
)
@@ -305,6 +325,7 @@ class Router:
"jwt_issuer",
"jwt_audience",
"jwt_jwks_uri",
"jwt_role_claim",
"jwt_role_mapping",
]
for field in fields_to_remove:
sgl-model-gateway/bindings/python/src/sglang_router/router_args.py
+74 -31
View File
@@ -17,6 +17,22 @@ except ModuleNotFoundError:
logger = logging.getLogger(__name__)
# Single source of truth for routing-policy CLI choices. Keep this in sync with
# `policy_from_str` in router.py and the `PolicyType` enum exposed by the Rust
# binding (sglang_router_rs). The Rust standalone binary (src/main.rs) accepts a
# subset of these — extending its `value_parser` and `parse_policy` to match is
# tracked separately.
_POLICY_CHOICES = (
"random",
"round_robin",
"cache_aware",
"power_of_two",
"bucket",
"manual",
"consistent_hashing",
"prefix_hash",
)
@dataclasses.dataclass
class RouterArgs:
@@ -151,7 +167,15 @@ class RouterArgs:
jwt_issuer: Optional[str] = None
jwt_audience: Optional[str] = None
jwt_jwks_uri: Optional[str] = None
jwt_role_claim: str = "roles"
jwt_role_mapping: Dict[str, str] = dataclasses.field(default_factory=dict)
# HTTP client connection pool tuning for upstream worker requests
pool_idle_timeout_secs: int = 50
connect_timeout_secs: int = 10
pool_max_idle_per_host: int = 500
tcp_keepalive_secs: int = 30
# Enable WebAssembly support
enable_wasm: bool = False
@staticmethod
def add_cli_args(
@@ -189,6 +213,9 @@ class RouterArgs:
request_group = parser.add_argument_group(
"Request Handling", "Request timeout and ID configuration"
)
http_client_group = parser.add_argument_group(
"HTTP Client", "Tuning for upstream HTTP client connection pooling"
)
rate_limit_group = parser.add_argument_group(
"Rate Limiting", "Concurrent request and queue limits"
)
@@ -257,46 +284,21 @@ class RouterArgs:
f"--{prefix}policy",
type=str,
default=RouterArgs.policy,
choices=[
"random",
"round_robin",
"cache_aware",
"power_of_two",
"manual",
"consistent_hashing",
"prefix_hash",
],
choices=_POLICY_CHOICES,
help="Load balancing policy to use. In PD mode, this is used for both prefill and decode unless overridden",
)
routing_group.add_argument(
f"--{prefix}prefill-policy",
type=str,
default=None,
choices=[
"random",
"round_robin",
"cache_aware",
"power_of_two",
"manual",
"bucket",
"consistent_hashing",
"prefix_hash",
],
choices=_POLICY_CHOICES,
help="Specific policy for prefill nodes in PD mode. If not specified, uses the main policy",
)
routing_group.add_argument(
f"--{prefix}decode-policy",
type=str,
default=None,
choices=[
"random",
"round_robin",
"cache_aware",
"power_of_two",
"manual",
"consistent_hashing",
"prefix_hash",
],
choices=_POLICY_CHOICES,
help="Specific policy for decode nodes in PD mode. If not specified, uses the main policy",
)
routing_group.add_argument(
@@ -467,6 +469,12 @@ class RouterArgs:
default={},
help="Label selector for decode server pods in PD mode (format: key1=value1 key2=value2)",
)
k8s_group.add_argument(
f"--{prefix}bootstrap-port-annotation",
type=str,
default=RouterArgs.bootstrap_port_annotation,
help="Kubernetes annotation key for bootstrap port (PD mode)",
)
# Prometheus configuration
prometheus_group.add_argument(
f"--{prefix}prometheus-port",
@@ -514,6 +522,32 @@ class RouterArgs:
help="CORS allowed origins (e.g., http://localhost:3000 https://example.com)",
)
# HTTP client connection pool tuning
http_client_group.add_argument(
f"--{prefix}pool-idle-timeout-secs",
type=int,
default=RouterArgs.pool_idle_timeout_secs,
help="Idle timeout in seconds for pooled upstream HTTP connections",
)
http_client_group.add_argument(
f"--{prefix}connect-timeout-secs",
type=int,
default=RouterArgs.connect_timeout_secs,
help="Timeout in seconds for new upstream HTTP connections",
)
http_client_group.add_argument(
f"--{prefix}pool-max-idle-per-host",
type=int,
default=RouterArgs.pool_max_idle_per_host,
help="Maximum idle upstream HTTP connections to keep per host",
)
http_client_group.add_argument(
f"--{prefix}tcp-keepalive-secs",
type=int,
default=RouterArgs.tcp_keepalive_secs,
help="TCP keepalive idle time in seconds for upstream HTTP connections",
)
# Rate limiting configuration
rate_limit_group.add_argument(
f"--{prefix}max-concurrent-requests",
@@ -726,6 +760,12 @@ class RouterArgs:
choices=["memory", "none", "oracle", "postgres", "redis"],
help="History storage backend for conversations and responses (default: memory)",
)
backend_group.add_argument(
f"--{prefix}enable-wasm",
action="store_true",
default=RouterArgs.enable_wasm,
help="Enable WebAssembly support",
)
# Oracle configuration
oracle_group.add_argument(
@@ -900,6 +940,12 @@ class RouterArgs:
default=None,
help="Explicit JWKS URI. If not provided, discovered from issuer via .well-known/openid-configuration",
)
auth_group.add_argument(
f"--{prefix}jwt-role-claim",
type=str,
default=RouterArgs.jwt_role_claim,
help="JWT claim name containing the role (default: 'roles')",
)
auth_group.add_argument(
f"--{prefix}jwt-role-mapping",
type=str,
@@ -960,9 +1006,6 @@ class RouterArgs:
cli_args_dict.get(f"{prefix}decode_selector", None)
)
# 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", [])