diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index a45a52273..de75ed202 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -90,28 +90,36 @@ pub struct PyJwtConfig { pub jwks_uri: Option, #[pyo3(get, set)] pub role_mapping: HashMap, + #[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, role_mapping: HashMap, + 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, + // 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, + pool_idle_timeout_secs: u64, + connect_timeout_secs: u64, + pool_max_idle_per_host: usize, + tcp_keepalive_secs: u64, + enable_wasm: bool, ) -> PyResult { 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, }) } diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/router.py b/sgl-model-gateway/bindings/python/src/sglang_router/router.py index 7d37da6f5..fa051b576 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/router.py +++ b/sgl-model-gateway/bindings/python/src/sglang_router/router.py @@ -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: diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py index aa0353a08..d02a75992 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py +++ b/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py @@ -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", []) diff --git a/sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py b/sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py new file mode 100644 index 000000000..feea8329f --- /dev/null +++ b/sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py @@ -0,0 +1,726 @@ +""" +Tests that exercise the PyO3 boundary directly (no mocking of `_Router`). + +These guard against drift between the Python `RouterArgs` dataclass, the +`Router.from_args` mapping, and the Rust `Router::new` signature in lib.rs. +The tests construct `_Router` without dispatching to remote workers, so they +run quickly and don't require GPU/network. +""" + +import pytest +from sglang_router.router import ( + Router, + backend_from_str, + build_control_plane_auth_config, + history_backend_from_str, + policy_from_str, + role_from_str, +) +from sglang_router.router_args import RouterArgs +from sglang_router.sglang_router_rs import ( + BackendType, + HistoryBackendType, + PolicyType, + PyApiKeyEntry, + PyControlPlaneAuthConfig, + PyJwtConfig, + PyOracleConfig, + PyPostgresConfig, + PyRedisConfig, + PyRole, +) +from sglang_router.sglang_router_rs import Router as _Router + + +class TestEnumConversions: + """All Python ↔ Rust enum conversion helpers cover every variant.""" + + def test_policy_from_str_covers_all_variants(self): + # Mirrors the PolicyType enum in lib.rs. Adding a variant on the Rust + # side without updating policy_from_str / _POLICY_CHOICES will fail here. + cases = { + "random": PolicyType.Random, + "round_robin": PolicyType.RoundRobin, + "cache_aware": PolicyType.CacheAware, + "power_of_two": PolicyType.PowerOfTwo, + "bucket": PolicyType.Bucket, + "manual": PolicyType.Manual, + "consistent_hashing": PolicyType.ConsistentHashing, + "prefix_hash": PolicyType.PrefixHash, + } + for s, expected in cases.items(): + assert policy_from_str(s) == expected + + def test_policy_from_str_none(self): + assert policy_from_str(None) is None + + def test_backend_from_str(self): + assert backend_from_str("sglang") == BackendType.Sglang + assert backend_from_str("openai") == BackendType.Openai + assert backend_from_str("SGLANG") == BackendType.Sglang + assert backend_from_str(None) == BackendType.Sglang + assert backend_from_str(BackendType.Openai) == BackendType.Openai + with pytest.raises(ValueError, match="Unknown backend"): + backend_from_str("vllm") + + def test_history_backend_from_str(self): + assert history_backend_from_str("memory") == HistoryBackendType.Memory + assert history_backend_from_str("none") == getattr(HistoryBackendType, "None") + assert history_backend_from_str("oracle") == HistoryBackendType.Oracle + assert history_backend_from_str("postgres") == HistoryBackendType.Postgres + assert history_backend_from_str("redis") == HistoryBackendType.Redis + assert history_backend_from_str(None) == HistoryBackendType.Memory + assert ( + history_backend_from_str(HistoryBackendType.Redis) + == HistoryBackendType.Redis + ) + with pytest.raises(ValueError, match="Unknown history backend"): + history_backend_from_str("dynamodb") + + def test_role_from_str(self): + assert role_from_str("admin") == PyRole.Admin + assert role_from_str("ADMIN") == PyRole.Admin + assert role_from_str("user") == PyRole.User + # Unknown roles fall through to User + assert role_from_str("unknown") == PyRole.User + + +class TestPyOracleConfig: + """PyOracleConfig PyO3 validation.""" + + def test_defaults(self): + cfg = PyOracleConfig() + assert cfg.pool_min == 1 + assert cfg.pool_max == 16 + assert cfg.pool_timeout_secs == 30 + assert cfg.username is None + assert cfg.password is None + assert cfg.connect_descriptor is None + assert cfg.wallet_path is None + + def test_invalid_pool_min_zero(self): + with pytest.raises(ValueError, match="pool_min must be at least 1"): + PyOracleConfig(pool_min=0) + + def test_invalid_pool_max_below_min(self): + with pytest.raises(ValueError, match="pool_max must be >= pool_min"): + PyOracleConfig(pool_min=5, pool_max=2) + + def test_full_config(self): + cfg = PyOracleConfig( + password="secret", + username="orcl", + connect_descriptor="dsn", + wallet_path="/path/to/wallet", + pool_min=2, + pool_max=20, + pool_timeout_secs=45, + ) + assert cfg.username == "orcl" + assert cfg.pool_min == 2 + assert cfg.pool_max == 20 + assert cfg.pool_timeout_secs == 45 + + +class TestPyPostgresConfig: + def test_defaults(self): + cfg = PyPostgresConfig() + assert cfg.db_url is None + assert cfg.pool_max == 16 + + def test_with_values(self): + cfg = PyPostgresConfig(db_url="postgres://localhost/db", pool_max=32) + assert cfg.db_url == "postgres://localhost/db" + assert cfg.pool_max == 32 + + +class TestPyRedisConfig: + def test_defaults(self): + cfg = PyRedisConfig(url="redis://localhost:6379") + assert cfg.url == "redis://localhost:6379" + assert cfg.pool_max == 16 + assert cfg.retention_days == 30 + + def test_persistent_retention(self): + cfg = PyRedisConfig(url="redis://localhost", retention_days=None) + assert cfg.retention_days is None + + +class TestPyApiKeyEntry: + def test_default_role_is_user(self): + entry = PyApiKeyEntry(id="k1", name="svc", key="secret") + assert entry.id == "k1" + assert entry.name == "svc" + assert entry.key == "secret" + assert entry.role == PyRole.User + + def test_admin_role(self): + entry = PyApiKeyEntry(id="k1", name="svc", key="secret", role=PyRole.Admin) + assert entry.role == PyRole.Admin + + +class TestPyJwtConfig: + def test_defaults_have_role_claim(self): + # role_claim defaults to "roles" matching the smg-auth crate; without + # this surfaced through PyO3, OIDC role mapping silently breaks. + cfg = PyJwtConfig(issuer="https://issuer", audience="api") + assert cfg.issuer == "https://issuer" + assert cfg.audience == "api" + assert cfg.role_claim == "roles" + assert cfg.role_mapping == {} + assert cfg.jwks_uri is None + + def test_custom_role_claim(self): + cfg = PyJwtConfig( + issuer="https://issuer", + audience="api", + role_claim="groups", + role_mapping={"AdminGroup": "admin"}, + ) + assert cfg.role_claim == "groups" + assert cfg.role_mapping == {"AdminGroup": "admin"} + + +class TestPyControlPlaneAuthConfig: + def test_default_audit_enabled(self): + # PyO3 default mirrors the smg-auth crate: ControlPlaneAuthConfig + # constructed without arguments has audit_enabled = true. + cfg = PyControlPlaneAuthConfig() + assert cfg.audit_enabled is True + assert cfg.api_keys == [] + assert cfg.jwt is None + + def test_with_jwt_and_keys(self): + # `PyJwtConfig` doesn't implement Python __eq__, so compare by field. + jwt = PyJwtConfig(issuer="i", audience="a") + keys = [PyApiKeyEntry(id="k", name="n", key="s", role=PyRole.Admin)] + cfg = PyControlPlaneAuthConfig(jwt=jwt, api_keys=keys, audit_enabled=False) + assert cfg.audit_enabled is False + assert cfg.jwt is not None + assert cfg.jwt.issuer == "i" + assert cfg.jwt.audience == "a" + assert len(cfg.api_keys) == 1 + assert cfg.api_keys[0].id == "k" + + +class TestBuildControlPlaneAuthConfig: + def test_returns_none_when_no_auth(self): + assert build_control_plane_auth_config({}) is None + + def test_returns_none_when_only_audit_set(self): + # Audit-only without keys/JWT shouldn't materialize a config object. + assert ( + build_control_plane_auth_config({"control_plane_audit_enabled": True}) + is None + ) + + def test_audit_default_when_unspecified(self): + # The Python wrapper has historically defaulted audit_enabled to False + # when the user doesn't pass control_plane_audit_enabled. Lock that in + # so a future change can't silently flip it. + cfg = build_control_plane_auth_config( + { + "control_plane_api_keys": [("id1", "Svc", "secret", "admin")], + } + ) + assert cfg is not None + assert cfg.audit_enabled is False + assert len(cfg.api_keys) == 1 + assert cfg.api_keys[0].role == PyRole.Admin + + def test_jwt_role_claim_threaded_through(self): + # jwt_role_claim must reach PyJwtConfig — without this the helper + # silently drops the user's claim name. + cfg = build_control_plane_auth_config( + { + "jwt_issuer": "https://issuer", + "jwt_audience": "api", + "jwt_role_claim": "groups", + "jwt_role_mapping": {"Admins": "admin"}, + } + ) + assert cfg is not None + assert cfg.jwt is not None + assert cfg.jwt.role_claim == "groups" + assert cfg.jwt.role_mapping == {"Admins": "admin"} + + def test_jwt_default_role_claim(self): + cfg = build_control_plane_auth_config( + {"jwt_issuer": "https://issuer", "jwt_audience": "api"} + ) + assert cfg is not None and cfg.jwt is not None + assert cfg.jwt.role_claim == "roles" + + def test_warns_when_jwt_incomplete(self, caplog): + # If the user sets jwt_role_claim/jwks_uri/role_mapping but forgets + # issuer/audience, the helper drops them silently. Emit a warning so + # users notice their JWT auth isn't actually enabled. + with caplog.at_level("WARNING", logger="sglang_router.router"): + cfg = build_control_plane_auth_config( + {"jwt_role_claim": "groups", "jwt_role_mapping": {"X": "admin"}} + ) + assert cfg is None + assert any( + "jwt_issuer/jwt_audience missing" in record.message + for record in caplog.records + ) + + +class TestParseControlPlaneApiKeys: + def test_valid(self): + result = RouterArgs._parse_control_plane_api_keys( + ["k1:Service Account:admin:secret123", "k2:Read Only:user:secret456"] + ) + assert result == [ + ("k1", "Service Account", "secret123", "admin"), + ("k2", "Read Only", "secret456", "user"), + ] + + def test_invalid_format(self): + with pytest.raises(ValueError, match="Invalid API key format"): + RouterArgs._parse_control_plane_api_keys(["just-a-key"]) + + def test_invalid_role(self): + with pytest.raises(ValueError, match="Invalid role"): + RouterArgs._parse_control_plane_api_keys(["id:name:superuser:secret"]) + + def test_key_with_colons_preserved(self): + # The split limit of 4 means the key portion can itself contain colons. + result = RouterArgs._parse_control_plane_api_keys( + ["id:name:user:sk-abc:def:ghi"] + ) + assert result == [("id", "name", "sk-abc:def:ghi", "user")] + + def test_empty(self): + assert RouterArgs._parse_control_plane_api_keys([]) == [] + assert RouterArgs._parse_control_plane_api_keys(None) == [] + + +class TestParseJwtRoleMapping: + def test_valid(self): + result = RouterArgs._parse_jwt_role_mapping( + ["Gateway.Admin=admin", "Gateway.User=user"] + ) + assert result == {"Gateway.Admin": "admin", "Gateway.User": "user"} + + def test_invalid_format(self): + with pytest.raises(ValueError, match="Invalid role mapping format"): + RouterArgs._parse_jwt_role_mapping(["no-equals"]) + + def test_invalid_role(self): + with pytest.raises(ValueError, match="Invalid gateway role"): + RouterArgs._parse_jwt_role_mapping(["X=superuser"]) + + def test_empty(self): + assert RouterArgs._parse_jwt_role_mapping([]) == {} + + +class TestRouterFromArgsKitchenSink: + """End-to-end tests of `Router.from_args(RouterArgs(...))`. + + These instantiate a real PyO3 `_Router` (no mocking) so any drift between + the Python dataclass fields and the Rust constructor signature surfaces here. + """ + + def test_minimal_regular_mode(self): + args = RouterArgs( + host="127.0.0.1", + port=30000, + worker_urls=["http://w1:8000"], + policy="round_robin", + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + + def test_pd_mode(self): + args = RouterArgs( + pd_disaggregation=True, + prefill_urls=[("http://prefill1:8000", 9000)], + decode_urls=["http://decode1:8001"], + policy="cache_aware", + prefill_policy="power_of_two", + decode_policy="round_robin", + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + + def test_all_policies_construct(self): + # Ensures every PolicyType the binding accepts is reachable through + # RouterArgs without exploding (e.g. unknown assignment_mode panics). + for policy in ( + "random", + "round_robin", + "cache_aware", + "power_of_two", + "bucket", + "manual", + "consistent_hashing", + "prefix_hash", + ): + args = RouterArgs( + worker_urls=["http://w1:8000"], + policy=policy, + pd_disaggregation=True, + prefill_urls=[("http://prefill1:8000", None)], + decode_urls=["http://decode1:8001"], + prefill_policy=policy, + decode_policy=policy, + ) + Router.from_args(args) + + def test_kitchen_sink_passes_every_field(self): + # Touches every dataclass field that maps directly to a parameter of + # Rust's Router::new (see lib.rs). PD-mode fields are exercised in + # test_pd_mode/test_all_policies_construct, and history-backend + # sub-configs are exercised in test_{oracle,postgres,redis}_history_backend. + # If you add a field on the Rust side without wiring it through + # RouterArgs/from_args, _Router(**args_dict) raises + # TypeError("got an unexpected keyword argument ..."). + args = RouterArgs( + worker_urls=["http://w1:8000", "http://w2:8000"], + host="127.0.0.1", + port=30001, + policy="cache_aware", + worker_startup_timeout_secs=60, + worker_startup_check_interval=5, + cache_threshold=0.5, + balance_abs_threshold=32, + balance_rel_threshold=1.2, + eviction_interval_secs=30, + max_tree_size=2**20, + max_idle_secs=600, + assignment_mode="min_load", + max_payload_size=1024 * 1024, + bucket_adjust_interval_secs=10, + dp_aware=True, + enable_igw=False, + api_key="key123", + log_dir="/tmp/router-logs", + log_level="debug", + json_log=True, + service_discovery=False, + selector={"app": "worker"}, + service_discovery_port=8080, + service_discovery_namespace="default", + prefill_selector={"role": "prefill"}, + decode_selector={"role": "decode"}, + bootstrap_port_annotation="custom.io/bootstrap-port", + prometheus_port=29000, + prometheus_host="127.0.0.1", + prometheus_duration_buckets=[0.1, 0.5, 1.0], + request_id_headers=["x-trace-id"], + request_timeout_secs=600, + shutdown_grace_period_secs=30, + max_concurrent_requests=128, + queue_size=50, + queue_timeout_secs=30, + rate_limit_tokens_per_second=64, + cors_allowed_origins=["http://localhost:3000"], + retry_max_retries=2, + retry_initial_backoff_ms=10, + retry_max_backoff_ms=1000, + retry_backoff_multiplier=2.0, + retry_jitter_factor=0.3, + disable_retries=False, + cb_failure_threshold=5, + cb_success_threshold=2, + cb_timeout_duration_secs=30, + cb_window_duration_secs=60, + disable_circuit_breaker=False, + health_failure_threshold=2, + health_success_threshold=1, + health_check_timeout_secs=3, + health_check_interval_secs=15, + health_check_endpoint="/healthz", + disable_health_check=False, + model_path="meta-llama/Llama-3-8B", + tokenizer_path=None, + chat_template=None, + tokenizer_cache_enable_l0=True, + tokenizer_cache_l0_max_entries=1000, + tokenizer_cache_enable_l1=True, + tokenizer_cache_l1_max_memory=1024 * 1024, + reasoning_parser="qwen3", + tool_call_parser=None, + mcp_config_path=None, + backend="sglang", + history_backend="memory", + client_cert_path=None, + client_key_path=None, + ca_cert_paths=[], + server_cert_path=None, + server_key_path=None, + enable_trace=True, + otlp_traces_endpoint="otel-collector:4317", + control_plane_api_keys=[("k1", "svc", "secret", "admin")], + control_plane_audit_enabled=False, + jwt_issuer="https://issuer", + jwt_audience="api", + jwt_jwks_uri="https://issuer/.well-known/jwks.json", + jwt_role_claim="groups", + jwt_role_mapping={"Admins": "admin"}, + pool_idle_timeout_secs=20, + connect_timeout_secs=5, + pool_max_idle_per_host=100, + tcp_keepalive_secs=15, + enable_wasm=True, + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + # Confirm the new fields actually carry the right value into Rust — + # an isinstance check alone wouldn't catch a typo'd builder call like + # `.pool_idle_timeout_secs(self.connect_timeout_secs)`. + assert router._router.pool_idle_timeout_secs == 20 + assert router._router.connect_timeout_secs == 5 + assert router._router.pool_max_idle_per_host == 100 + assert router._router.tcp_keepalive_secs == 15 + assert router._router.enable_wasm is True + + def test_oracle_history_backend(self): + args = RouterArgs( + worker_urls=["http://w1:8000"], + history_backend="oracle", + oracle_username="user", + oracle_password="pw", + oracle_connect_descriptor="dsn", + oracle_pool_min=2, + oracle_pool_max=8, + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + + def test_postgres_history_backend(self): + args = RouterArgs( + worker_urls=["http://w1:8000"], + history_backend="postgres", + postgres_db_url="postgres://localhost/db", + postgres_pool_max=8, + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + + def test_redis_history_backend(self): + args = RouterArgs( + worker_urls=["http://w1:8000"], + history_backend="redis", + redis_url="redis://localhost:6379", + redis_pool_max=8, + redis_retention_days=7, + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + + def test_redis_persistent_retention(self): + # redis_retention_days < 0 means persistent (None on the Rust side). + args = RouterArgs( + worker_urls=["http://w1:8000"], + history_backend="redis", + redis_url="redis://localhost:6379", + redis_retention_days=-1, + ) + Router.from_args(args) + + +class TestBootstrapPortAnnotation: + """Regression: the wrapper must not silently override user-supplied values.""" + + def test_user_value_preserved_through_from_cli_args(self): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args( + [ + "--bootstrap-port-annotation", + "custom.io/bootstrap-port", + ] + ) + assert args.bootstrap_port_annotation == "custom.io/bootstrap-port" + + def test_default_value(self): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args([]) + assert args.bootstrap_port_annotation == "sglang.ai/bootstrap-port" + + +class TestNewBindingFields: + """Round-trip checks for fields whose CLI flag, dataclass attribute, and + PyO3 constructor parameter were historically out of sync.""" + + def test_jwt_role_claim_default(self): + args = RouterArgs() + assert args.jwt_role_claim == "roles" + + def test_audit_enabled_default_off(self): + # The Python wrapper defaults audit_enabled to False even though the + # Rust standalone binary defaults it to True (main.rs:614, + # disable_audit_logging = false → audit_enabled = true). The + # divergence is intentional: changing the wrapper default is a + # behavior change that needs an explicit migration. + args = RouterArgs() + assert args.control_plane_audit_enabled is False + + def test_audit_enabled_via_cli_flag(self): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args(["--control-plane-audit-enabled"]) + assert args.control_plane_audit_enabled is True + + def test_default_audit_via_cli_is_off(self): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args([]) + assert args.control_plane_audit_enabled is False + + def test_http_pool_defaults(self): + args = RouterArgs() + assert args.pool_idle_timeout_secs == 50 + assert args.connect_timeout_secs == 10 + assert args.pool_max_idle_per_host == 500 + assert args.tcp_keepalive_secs == 30 + + def test_enable_wasm_default(self): + args = RouterArgs() + assert args.enable_wasm is False + + def test_http_pool_via_cli(self): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args( + [ + "--pool-idle-timeout-secs", + "120", + "--connect-timeout-secs", + "20", + "--pool-max-idle-per-host", + "200", + "--tcp-keepalive-secs", + "45", + "--enable-wasm", + ] + ) + assert args.pool_idle_timeout_secs == 120 + assert args.connect_timeout_secs == 20 + assert args.pool_max_idle_per_host == 200 + assert args.tcp_keepalive_secs == 45 + assert args.enable_wasm is True + + def test_jwt_role_claim_via_cli(self): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args( + [ + "--jwt-issuer", + "https://issuer", + "--jwt-audience", + "api", + "--jwt-role-claim", + "groups", + ] + ) + assert args.jwt_role_claim == "groups" + + def test_jwt_role_claim_end_to_end(self): + # Full pipeline: CLI parser → RouterArgs → Router.from_args → + # `_Router(**args_dict)`. Pins the invariant that `jwt_role_claim` is + # consumed by `build_control_plane_auth_config` AND popped from + # args_dict before reaching the Rust constructor (otherwise _Router + # would raise TypeError on the unknown kwarg). + from sglang_router.launch_router import parse_router_args + + args = parse_router_args( + [ + "--jwt-issuer", + "https://issuer", + "--jwt-audience", + "api", + "--jwt-role-claim", + "groups", + "--jwt-role-mapping", + "Admins=admin", + ] + ) + router = Router.from_args(args) + assert isinstance(router._router, _Router) + + +class TestPolicyChoiceListConsistency: + """Every policy in the binding's PolicyType must be a CLI choice on every + policy flag. Catches drift if someone hard-codes a list at one of the three + argparse `choices=` sites instead of using `_POLICY_CHOICES`.""" + + @pytest.mark.parametrize( + "policy", + [ + "random", + "round_robin", + "cache_aware", + "power_of_two", + "bucket", + "manual", + "consistent_hashing", + "prefix_hash", + ], + ) + def test_main_policy_accepts(self, policy): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args(["--policy", policy]) + assert args.policy == policy + + @pytest.mark.parametrize( + "policy", + [ + "random", + "round_robin", + "cache_aware", + "power_of_two", + "bucket", + "manual", + "consistent_hashing", + "prefix_hash", + ], + ) + def test_prefill_policy_accepts(self, policy): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args( + [ + "--pd-disaggregation", + "--prefill", + "http://p:8000", + "--decode", + "http://d:8001", + "--prefill-policy", + policy, + ] + ) + assert args.prefill_policy == policy + + @pytest.mark.parametrize( + "policy", + [ + "random", + "round_robin", + "cache_aware", + "power_of_two", + "bucket", + "manual", + "consistent_hashing", + "prefix_hash", + ], + ) + def test_decode_policy_accepts(self, policy): + from sglang_router.launch_router import parse_router_args + + args = parse_router_args( + [ + "--pd-disaggregation", + "--prefill", + "http://p:8000", + "--decode", + "http://d:8001", + "--decode-policy", + policy, + ] + ) + assert args.decode_policy == policy