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

[model-gateway] move python to binding folder (#13295)

Browse Source
This commit is contained in:
Simo Lin
2025-11-15 00:32:21 -08:00
committed by GitHub
parent 56fc483073
commit b732ffa404
28 changed files with 928 additions and 819 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
.github/workflows/nightly-release-router.yml vendored
View File

@@ -66,16 +66,16 @@ jobs:
- name: Modify version for nightly release
run: |
# Get current version from pyproject.toml
CURRENT_VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])" 2>/dev/null || python -c "import tomli; print(tomli.load(open('pyproject.toml', 'rb'))['project']['version'])")
CURRENT_VERSION=$(python -c "import tomllib; print(tomllib.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])" 2>/dev/null || python -c "import tomli; print(tomli.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])")
# Create nightly version with date: e.g., 0.2.1.dev20250128
NIGHTLY_VERSION="${CURRENT_VERSION}.dev$(date +%Y%m%d)"
echo "Nightly version: $NIGHTLY_VERSION"
# Update pyproject.toml with nightly version (temporary, not committed)
sed -i.bak "s/version = \"${CURRENT_VERSION}\"/version = \"${NIGHTLY_VERSION}\"/" pyproject.toml
sed -i.bak "s/version = \"${CURRENT_VERSION}\"/version = \"${NIGHTLY_VERSION}\"/" bindings/python/pyproject.toml
# Verify the change
cat pyproject.toml | grep "^version"
cat bindings/python/pyproject.toml | grep "^version"
shell: bash
- name: Install twine and tomli
@@ -92,6 +92,7 @@ jobs:
- name: Build wheels
uses: PyO3/maturin-action@v1
with:
working-directory: bindings/python
target: ${{ matrix.target }}
manylinux: ${{ matrix.manylinux || 'auto' }}
args: --release --out dist --features vendored-openssl --interpreter ${{ matrix.interpreter || '3.9 3.10 3.11 3.12 3.13 3.14' }}
@@ -151,20 +152,21 @@ jobs:
- name: Modify version for nightly release
run: |
# Get current version from pyproject.toml
CURRENT_VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])" 2>/dev/null || python -c "import tomli; print(tomli.load(open('pyproject.toml', 'rb'))['project']['version'])")
CURRENT_VERSION=$(python -c "import tomllib; print(tomllib.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])" 2>/dev/null || python -c "import tomli; print(tomli.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])")
# Create nightly version with date: e.g., 0.2.1.dev20250128
NIGHTLY_VERSION="${CURRENT_VERSION}.dev$(date +%Y%m%d)"
echo "Nightly version: $NIGHTLY_VERSION"
# Update pyproject.toml with nightly version (temporary, not committed)
sed -i "s/version = \"${CURRENT_VERSION}\"/version = \"${NIGHTLY_VERSION}\"/" pyproject.toml
sed -i "s/version = \"${CURRENT_VERSION}\"/version = \"${NIGHTLY_VERSION}\"/" bindings/python/pyproject.toml
# Verify the change
cat pyproject.toml | grep "^version"
cat bindings/python/pyproject.toml | grep "^version"
- name: Build SDist
uses: PyO3/maturin-action@v1
with:
working-directory: bindings/python
command: sdist
args: --out dist
rust-toolchain: stable

2
.github/workflows/pr-test-pd-router.yml vendored
View File

@@ -147,7 +147,7 @@ jobs:
run: |
source "$HOME/.cargo/env"
echo "Building sgl-router..."
cd sgl-router
cd sgl-router/bindings/python
pip install maturin
maturin build --release --out dist --features vendored-openssl
pip install --force-reinstall dist/*.whl

19
.github/workflows/pr-test-rust.yml vendored
View File

@@ -57,16 +57,17 @@ jobs:
- name: Test maturin build
uses: PyO3/maturin-action@v1
with:
working-directory: bindings/python
args: --release --out dist --features vendored-openssl
rust-toolchain: stable
sccache: true
- name: List built wheel
run: ls -lh dist/
run: ls -lh bindings/python/dist/
- name: Test wheel install
run: |
pip install dist/*.whl
pip install bindings/python/dist/*.whl
python -c "import sglang_router; print('Python package: OK')"
python -c "from sglang_router.sglang_router_rs import Router; print('Rust extension: OK')"
python -m sglang_router.launch_router --help > /dev/null && echo "Entry point: OK"
@@ -165,18 +166,18 @@ jobs:
run: |
source "$HOME/.cargo/env"
export RUSTC_WRAPPER=sccache
cd sgl-router
cd sgl-router/bindings/python
python3 -m pip install --upgrade pip maturin
pip uninstall -y sglang-router
maturin build --profile ci --features vendored-openssl
pip install target/wheels/*.whl
maturin build --profile ci --features vendored-openssl --out dist
pip install dist/*.whl
- name: Run Python unit tests
run: |
cd sgl-router
source "$HOME/.cargo/env"
python3 -m pip install pytest pytest-cov pytest-xdist
pytest -q py_test/unit --cov=sglang_router --cov-report=term-missing --cov-fail-under=80
pytest -q py_test/unit --cov=sglang_router --cov-config=bindings/python/.coveragerc --cov-report=term-missing --cov-fail-under=80
- name: Run Python integration tests
run: |
@@ -274,11 +275,11 @@ jobs:
run: |
source "$HOME/.cargo/env"
export RUSTC_WRAPPER=sccache
cd sgl-router
cd sgl-router/bindings/python
python3 -m pip install --upgrade pip maturin
pip uninstall -y sglang-router
maturin build --profile ci --features vendored-openssl
pip install target/wheels/*.whl
maturin build --profile ci --features vendored-openssl --out dist
pip install dist/*.whl
- name: Run Python E2E response API tests
run: |

4
.github/workflows/release-docker-router.yml vendored
View File

@@ -4,7 +4,7 @@ on:
branches:
- main
paths:
- "sgl-router/py_src/sglang_router/version.py"
- "sgl-router/bindings/python/sglang_router/version.py"
workflow_dispatch:
jobs:
@@ -23,7 +23,7 @@ jobs:
- name: Build and Push
run: |
version=$(cat sgl-router/py_src/sglang_router/version.py | cut -d'"' -f2)
version=$(cat sgl-router/bindings/python/sglang_router/version.py | cut -d'"' -f2)
tag=v${version}
docker build . -f docker/router.Dockerfile -t lmsysorg/sglang-router:${tag} --no-cache

4
.github/workflows/release-pypi-router.yml vendored
View File

@@ -5,7 +5,7 @@ on:
branches:
- main
paths:
- sgl-router/pyproject.toml
- sgl-router/bindings/python/pyproject.toml
workflow_dispatch:
jobs:
@@ -77,6 +77,7 @@ jobs:
- name: Build wheels
uses: PyO3/maturin-action@v1
with:
working-directory: bindings/python
target: ${{ matrix.target }}
manylinux: ${{ matrix.manylinux || 'auto' }}
args: --release --out dist --features vendored-openssl --interpreter ${{ matrix.interpreter || '3.9 3.10 3.11 3.12 3.13 3.14' }}
@@ -136,6 +137,7 @@ jobs:
- name: Build SDist
uses: PyO3/maturin-action@v1
with:
working-directory: bindings/python
command: sdist
args: --out dist
rust-toolchain: stable

9
sgl-router/.coveragerc
View File

@@ -1,9 +0,0 @@
[run]
source = py_src/sglang_router
omit =
*/sglang_router/mini_lb.py
[report]
fail_under = 80
omit =
*/sglang_router/mini_lb.py

10
sgl-router/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "sglang_router_rs"
version = "0.0.0"
name = "sglang-router"
version = "0.2.2"
edition = "2021"
[features]
@@ -14,9 +14,8 @@ unused_qualifications = "warn"
[lib]
name = "sglang_router_rs"
# Pure Rust library: Just omit crate-type (defaults to rlib)
# Python/C binding + Rust library: Use ["cdylib", "rlib"]
crate-type = ["cdylib", "rlib"]
# Pure Rust library (rlib for the binary and other Rust dependents)
crate-type = ["rlib"]
[[bin]]
name = "sglang-router"
@@ -37,7 +36,6 @@ rand = "0.9.2"
reqwest = { version = "0.12.8", features = ["stream", "blocking", "json", "rustls-tls"], default-features = false }
futures-util = "0.3"
futures = "0.3"
pyo3 = { version = "0.27.1", features = ["extension-module", "abi3-py38"] }
dashmap = "6.1.0"
lru = "0.16.2"
blake3 = "1.5"

1
sgl-router/LICENSE Symbolic link
View File

@@ -0,0 +1 @@
../LICENSE

7
sgl-router/MANIFEST.in
View File

@@ -1,7 +0,0 @@
# Must include:
include Cargo.toml # Rust project configuration
include build.rs # Build script for protobuf generation
include ../LICENSE
recursive-include src *.rs # Rust source files
recursive-include src/proto *.proto # Protobuf definitions
recursive-include py_src *.py # Python source files

103
sgl-router/Makefile
View File

@@ -1,6 +1,14 @@
# Model Gateway Makefile
# Provides convenient shortcuts for common development tasks
# Python bindings directory
PYTHON_DIR := bindings/python
# Auto-detect CPU cores and cap at reasonable limit to avoid thread exhaustion
# Can be overridden: make python-dev JOBS=4
NPROC := $(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 8)
JOBS ?= $(shell echo $$(($(NPROC) > 16 ? 16 : $(NPROC))))
# Check if sccache is available and set RUSTC_WRAPPER accordingly
SCCACHE := $(shell which sccache 2>/dev/null)
ifdef SCCACHE
@@ -10,7 +18,8 @@ else
$(info sccache not found. Install it for faster builds: cargo install sccache)
endif
.PHONY: help bench bench-quick bench-baseline bench-compare test build clean
.PHONY: help build test clean docs check fmt dev-setup pre-commit setup-sccache sccache-stats sccache-clean sccache-stop \
python-dev python-build python-build-release python-install python-clean python-test python-check
help: ## Show this help message
@echo "Model Gateway Development Commands"
@@ -27,26 +36,6 @@ test: ## Run all tests
@echo "Running tests..."
@cargo test
bench: ## Run full benchmark suite
@echo "Running full benchmarks..."
@python3 scripts/run_benchmarks.py
bench-quick: ## Run quick benchmarks only
@echo "Running quick benchmarks..."
@python3 scripts/run_benchmarks.py --quick
bench-baseline: ## Save current performance as baseline
@echo "Saving performance baseline..."
@python3 scripts/run_benchmarks.py --save-baseline main
bench-compare: ## Compare with saved baseline
@echo "Comparing with baseline..."
@python3 scripts/run_benchmarks.py --compare-baseline main
bench-ci: ## Run benchmarks suitable for CI (quick mode)
@echo "Running CI benchmarks..."
@python3 scripts/run_benchmarks.py --quick
clean: ## Clean build artifacts
@echo "Cleaning build artifacts..."
@cargo clean
@@ -69,37 +58,9 @@ fmt: ## Format code with rustfmt
dev-setup: build test ## Set up development environment
@echo "Development environment ready!"
pre-commit: fmt check test bench-quick ## Run pre-commit checks
pre-commit: fmt check test ## Run pre-commit checks
@echo "Pre-commit checks passed!"
# Benchmark analysis shortcuts
bench-report: ## Open benchmark HTML report
@if [ -f "target/criterion/request_processing/report/index.html" ]; then \
echo "Opening benchmark report..."; \
if command -v xdg-open >/dev/null 2>&1; then \
xdg-open target/criterion/request_processing/report/index.html; \
elif command -v open >/dev/null 2>&1; then \
open target/criterion/request_processing/report/index.html; \
else \
echo "Please open target/criterion/request_processing/report/index.html in your browser"; \
fi \
else \
echo "No benchmark report found. Run 'make bench' first."; \
fi
bench-clean: ## Clean benchmark results
@echo "Cleaning benchmark results..."
@rm -rf target/criterion
# Performance monitoring
perf-monitor: ## Run continuous performance monitoring
@echo "Starting performance monitoring..."
@if command -v watch >/dev/null 2>&1; then \
watch -n 300 'make bench-quick'; \
else \
echo "Warning: 'watch' command not found. Install it or run 'make bench-quick' manually."; \
fi
# sccache management targets
setup-sccache: ## Install and configure sccache
@echo "Setting up sccache..."
@@ -129,3 +90,45 @@ sccache-stop: ## Stop the sccache server
else \
echo "sccache not installed"; \
fi
# Python bindings (maturin) targets
python-dev: ## Build Python bindings in development mode (fast, debug build)
@echo "Building Python bindings in development mode (using $(JOBS) parallel jobs with sccache)..."
@cd $(PYTHON_DIR) && CARGO_BUILD_JOBS=$(JOBS) maturin develop
python-build: ## Build Python wheel (release mode with vendored OpenSSL)
@echo "Building Python wheel (release, vendored OpenSSL, using $(JOBS) parallel jobs with sccache)..."
@cd $(PYTHON_DIR) && CARGO_BUILD_JOBS=$(JOBS) maturin build --release --out dist --features vendored-openssl
python-build-release: python-build ## Alias for python-build
python-install: python-build ## Build and install Python wheel
@echo "Installing Python wheel..."
@pip install --force-reinstall $(PYTHON_DIR)/dist/*.whl
@echo "Python package installed!"
python-clean: ## Clean Python build artifacts
@echo "Cleaning Python build artifacts..."
@rm -rf $(PYTHON_DIR)/dist/
@rm -rf $(PYTHON_DIR)/target/
@rm -rf $(PYTHON_DIR)/sglang_router.egg-info/
@rm -rf $(PYTHON_DIR)/sglang_router/__pycache__/
@find $(PYTHON_DIR) -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true
@find $(PYTHON_DIR) -name "*.pyc" -delete 2>/dev/null || true
@echo "Python build artifacts cleaned!"
python-test: ## Run Python tests
@echo "Running Python tests..."
@pytest py_test/ -v
python-check: ## Check Python package with twine
@echo "Checking Python package..."
@cd $(PYTHON_DIR) && CARGO_BUILD_JOBS=$(JOBS) maturin build --release --out dist --features vendored-openssl
@pip install twine 2>/dev/null || true
@twine check $(PYTHON_DIR)/dist/*
@echo "Python package check passed!"
# Combined shortcuts
dev: python-dev ## Quick development setup (build Python bindings in dev mode)
install: python-install ## Build and install everything

16
sgl-router/README.md
View File

@@ -70,20 +70,23 @@ pip install maturin
# Fast development mode (debug build, no wheel, instant)
# Uses system OpenSSL (requires libssl-dev/openssl-devel)
cd bindings/python
maturin develop
# Production build (optimized, creates wheel)
# Uses vendored OpenSSL (cross-platform compatibility)
maturin build --release --features vendored-openssl
cd bindings/python
maturin build --release --out dist --features vendored-openssl
pip install --force-reinstall dist/*.whl
# Development build with system OpenSSL (faster)
# Requires: apt install libssl-dev pkg-config (Ubuntu/Debian)
# or: yum install openssl-devel (RHEL/CentOS)
maturin build --release
cd bindings/python
maturin build --release --out dist
pip install --force-reinstall dist/*.whl
```
> **Note:** Use `maturin develop` for fast iteration during development (builds in debug mode and installs directly). Use `maturin build --release --features vendored-openssl` for production wheels with full optimizations (opt-level="z", lto="fat") and cross-platform compatibility. The package uses abi3 support for Python 3.8+ compatibility.
> **Note:** Python bindings are located in `bindings/python/` with their own Cargo.toml. Use `maturin develop` for fast iteration during development (builds in debug mode and installs directly). Use `maturin build --release --features vendored-openssl` for production wheels with full optimizations (opt-level="z", lto="fat") and cross-platform compatibility. The package uses abi3 support for Python 3.8+ compatibility.
## Quick Start
### Regular HTTP Routing
@@ -574,12 +577,13 @@ cargo build
cargo test
# Fast Python development (rebuilds and installs in debug mode)
maturin develop
cd bindings/python && maturin develop
# Run Python tests
pytest
cd ../.. # Back to sgl-router root
pytest py_test/
```
For production builds, use `maturin build --release` to create optimized wheels. During development, `maturin develop` rebuilds and installs instantly without creating wheel files. Use `python -m sglang_router.launch_server` to co-launch router and SGLang workers in small clusters for local validation.
For production builds, use `maturin build --release --out dist` from the `bindings/python/` directory to create optimized wheels. During development, `maturin develop` rebuilds and installs instantly without creating wheel files. Use `python -m sglang_router.launch_server` to co-launch router and SGLang workers in small clusters for local validation.
---

6
sgl-router/bindings/python/.coveragerc Normal file
View File

@@ -0,0 +1,6 @@
[run]
source = sglang_router
omit = */mini_lb.py
[report]
fail_under = 80

28
sgl-router/bindings/python/Cargo.toml Normal file
View File

@@ -0,0 +1,28 @@
[package]
name = "sglang-router-python"
version = "0.2.2"
edition = "2021"
[lib]
name = "sglang_router_rs"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.27.1", features = ["extension-module", "abi3-py38"] }
tokio = { version = "1.42.0", features = ["full"] }
[dependencies.sglang-router]
path = "../.."
default-features = true
package = "sglang-router"
[features]
default = ["pyo3/extension-module"]
vendored-openssl = ["sglang-router/vendored-openssl"]
[profile.ci]
inherits = "release"
opt-level = 2 # Lighter optimization (still fast runtime, much faster compile)
lto = "thin" # Thin LTO - good balance
codegen-units = 16 # More parallelization for faster builds
strip = true

9
sgl-router/bindings/python/MANIFEST.in Normal file
View File

@@ -0,0 +1,9 @@
# Must include:
include Cargo.toml # Python bindings Cargo configuration
include ../../Cargo.toml # Main Rust project configuration
include ../../build.rs # Build script for protobuf generation
include ../../LICENSE
recursive-include src *.rs # Python bindings wrapper
recursive-include ../../src *.rs # Main Rust source files
recursive-include ../../src/proto *.proto # Protobuf definitions
recursive-include sglang_router *.py # Python source files

71
sgl-router/bindings/python/README.md Normal file
View File

@@ -0,0 +1,71 @@
# SGLang Model Gateway Python Bindings
This directory contains the Python bindings for the SGLang Router, built using [maturin](https://github.com/PyO3/maturin) and [PyO3](https://github.com/PyO3/pyo3).
## Directory Structure
```
bindings/python/
├── src/ # Rust source code for Python bindings
│ └── lib.rs # PyO3 bindings implementation
├── sglang_router/ # Python source code
│ ├── __init__.py
│ ├── version.py
│ ├── launch_server.py
│ ├── launch_router.py
│ ├── router.py
│ ├── router_args.py
│ └── mini_lb.py
├── Cargo.toml # Rust package configuration for bindings
├── pyproject.toml # Python package configuration
├── setup.py # Setup configuration
├── MANIFEST.in # Package manifest
├── .coveragerc # Test coverage configuration
└── README.md # This file
```
## Building
### Development Build
```bash
# Install maturin
pip install maturin
# Build and install in development mode
cd sgl-router/bindings/python
maturin develop --features vendored-openssl
```
### Production Build
```bash
# Build wheel
cd sgl-router/bindings/python
maturin build --release --out dist --features vendored-openssl
# Install the built wheel
pip install dist/sglang_router-*.whl
```
## Testing
```bash
# Run Python tests
cd sgl-router
pytest py_test/
```
## Configuration
- **pyproject.toml**: Defines package metadata, dependencies, and build configuration
- **python-source**: Set to "." to indicate Python source is in the same directory as pyproject.toml
- **module-name**: `sglang_router.sglang_router_rs` - the Rust extension module name
## Notes
- The Rust bindings source code is located in `src/lib.rs`
- The bindings have their own `Cargo.toml` in this directory
- The main sglang-router library is located in `../../` and is used as a dependency
- The package includes both Python code and Rust extensions built with PyO3
- PyO3 types are prefixed with `Py` in Rust but exposed to Python without the prefix using the `name` attribute

2
sgl-router/pyproject.toml → sgl-router/bindings/python/pyproject.toml
View File

@@ -42,5 +42,5 @@ dev = [
[tool.maturin]
python-source = "py_src"
python-source = "."
module-name = "sglang_router.sglang_router_rs"

0
sgl-router/setup.py → sgl-router/bindings/python/setup.py
View File

0
sgl-router/py_src/sglang_router/__init__.py → sgl-router/bindings/python/sglang_router/__init__.py
View File

0
sgl-router/py_src/sglang_router/launch_router.py → sgl-router/bindings/python/sglang_router/launch_router.py
View File

0
sgl-router/py_src/sglang_router/launch_server.py → sgl-router/bindings/python/sglang_router/launch_server.py
View File

0
sgl-router/py_src/sglang_router/mini_lb.py → sgl-router/bindings/python/sglang_router/mini_lb.py
View File

0
sgl-router/py_src/sglang_router/router.py → sgl-router/bindings/python/sglang_router/router.py
View File

0
sgl-router/py_src/sglang_router/router_args.py → sgl-router/bindings/python/sglang_router/router_args.py
View File

0
sgl-router/py_src/sglang_router/version.py → sgl-router/bindings/python/sglang_router/version.py
View File

716
sgl-router/bindings/python/src/lib.rs Normal file
View File

@@ -0,0 +1,716 @@
use pyo3::prelude::*;
use sglang_router::*;
use std::collections::HashMap;
// Define the enums with PyO3 bindings
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum PolicyType {
Random,
RoundRobin,
CacheAware,
PowerOfTwo,
Bucket,
}
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum BackendType {
Sglang,
Openai,
}
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum HistoryBackendType {
Memory,
None,
Oracle,
Postgres,
}
#[pyclass]
#[derive(Clone, PartialEq)]
pub struct PyOracleConfig {
#[pyo3(get, set)]
pub wallet_path: Option<String>,
#[pyo3(get, set)]
pub connect_descriptor: Option<String>,
#[pyo3(get, set)]
pub username: Option<String>,
#[pyo3(get, set)]
pub password: Option<String>,
#[pyo3(get, set)]
pub pool_min: usize,
#[pyo3(get, set)]
pub pool_max: usize,
#[pyo3(get, set)]
pub pool_timeout_secs: u64,
}
impl std::fmt::Debug for PyOracleConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PyOracleConfig")
.field("wallet_path", &self.wallet_path)
.field("connect_descriptor", &"<redacted>")
.field("username", &self.username)
.field("password", &"<redacted>")
.field("pool_min", &self.pool_min)
.field("pool_max", &self.pool_max)
.field("pool_timeout_secs", &self.pool_timeout_secs)
.finish()
}
}
#[pymethods]
impl PyOracleConfig {
#[new]
#[pyo3(signature = (
password = None,
username = None,
connect_descriptor = None,
wallet_path = None,
pool_min = 1,
pool_max = 16,
pool_timeout_secs = 30,
))]
fn new(
password: Option<String>,
username: Option<String>,
connect_descriptor: Option<String>,
wallet_path: Option<String>,
pool_min: usize,
pool_max: usize,
pool_timeout_secs: u64,
) -> PyResult<Self> {
if pool_min == 0 {
return Err(pyo3::exceptions::PyValueError::new_err(
"pool_min must be at least 1",
));
}
if pool_max < pool_min {
return Err(pyo3::exceptions::PyValueError::new_err(
"pool_max must be >= pool_min",
));
}
Ok(PyOracleConfig {
wallet_path,
connect_descriptor,
username,
password,
pool_min,
pool_max,
pool_timeout_secs,
})
}
}
impl PyOracleConfig {
pub fn to_config_oracle(&self) -> config::OracleConfig {
config::OracleConfig {
wallet_path: self.wallet_path.clone(),
connect_descriptor: self.connect_descriptor.clone().unwrap_or_default(),
username: self.username.clone().unwrap_or_default(),
password: self.password.clone().unwrap_or_default(),
pool_min: self.pool_min,
pool_max: self.pool_max,
pool_timeout_secs: self.pool_timeout_secs,
}
}
}
#[pyclass]
#[derive(Debug, Clone, PartialEq)]
pub struct PyPostgresConfig {
#[pyo3(get, set)]
pub db_url: Option<String>,
#[pyo3(get, set)]
pub pool_max: usize,
}
#[pymethods]
impl PyPostgresConfig {
#[new]
#[pyo3(signature = (db_url = None,pool_max = 16,))]
fn new(db_url: Option<String>, pool_max: usize) -> PyResult<Self> {
Ok(PyPostgresConfig { db_url, pool_max })
}
}
impl PyPostgresConfig {
pub fn to_config_postgres(&self) -> config::PostgresConfig {
config::PostgresConfig {
db_url: self.db_url.clone().unwrap_or_default(),
pool_max: self.pool_max,
}
}
}
#[pyclass]
#[derive(Debug, Clone, PartialEq)]
struct Router {
host: String,
port: u16,
worker_urls: Vec<String>,
policy: PolicyType,
worker_startup_timeout_secs: u64,
worker_startup_check_interval: u64,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
eviction_interval_secs: u64,
max_tree_size: usize,
max_payload_size: usize,
dp_aware: bool,
api_key: Option<String>,
log_dir: Option<String>,
log_level: Option<String>,
service_discovery: bool,
selector: HashMap<String, String>,
service_discovery_port: u16,
service_discovery_namespace: Option<String>,
prefill_selector: HashMap<String, String>,
decode_selector: HashMap<String, String>,
bootstrap_port_annotation: String,
prometheus_port: Option<u16>,
prometheus_host: Option<String>,
request_timeout_secs: u64,
request_id_headers: Option<Vec<String>>,
pd_disaggregation: bool,
bucket_adjust_interval_secs: usize,
prefill_urls: Option<Vec<(String, Option<u16>)>>,
decode_urls: Option<Vec<String>>,
prefill_policy: Option<PolicyType>,
decode_policy: Option<PolicyType>,
max_concurrent_requests: i32,
cors_allowed_origins: Vec<String>,
retry_max_retries: u32,
retry_initial_backoff_ms: u64,
retry_max_backoff_ms: u64,
retry_backoff_multiplier: f32,
retry_jitter_factor: f32,
disable_retries: bool,
cb_failure_threshold: u32,
cb_success_threshold: u32,
cb_timeout_duration_secs: u64,
cb_window_duration_secs: u64,
disable_circuit_breaker: bool,
health_failure_threshold: u32,
health_success_threshold: u32,
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
rate_limit_tokens_per_second: Option<i32>,
connection_mode: core::ConnectionMode,
model_path: Option<String>,
tokenizer_path: Option<String>,
chat_template: Option<String>,
tokenizer_cache_enable_l0: bool,
tokenizer_cache_l0_max_entries: usize,
tokenizer_cache_enable_l1: bool,
tokenizer_cache_l1_max_memory: usize,
reasoning_parser: Option<String>,
tool_call_parser: Option<String>,
mcp_config_path: Option<String>,
backend: BackendType,
history_backend: HistoryBackendType,
oracle_config: Option<PyOracleConfig>,
postgres_config: Option<PyPostgresConfig>,
client_cert_path: Option<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
}
impl Router {
fn determine_connection_mode(worker_urls: &[String]) -> core::ConnectionMode {
for url in worker_urls {
if url.starts_with("grpc://") || url.starts_with("grpcs://") {
return core::ConnectionMode::Grpc { port: None };
}
}
core::ConnectionMode::Http
}
pub fn to_router_config(&self) -> config::ConfigResult<config::RouterConfig> {
use config::{
DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode,
};
let convert_policy = |policy: &PolicyType| -> ConfigPolicyConfig {
match policy {
PolicyType::Random => ConfigPolicyConfig::Random,
PolicyType::RoundRobin => ConfigPolicyConfig::RoundRobin,
PolicyType::CacheAware => ConfigPolicyConfig::CacheAware {
cache_threshold: self.cache_threshold,
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
eviction_interval_secs: self.eviction_interval_secs,
max_tree_size: self.max_tree_size,
},
PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
load_check_interval_secs: 5,
},
PolicyType::Bucket => ConfigPolicyConfig::Bucket {
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
bucket_adjust_interval_secs: self.bucket_adjust_interval_secs,
},
}
};
let mode = if self.enable_igw {
RoutingMode::Regular {
worker_urls: vec![],
}
} else if matches!(self.backend, BackendType::Openai) {
RoutingMode::OpenAI {
worker_urls: self.worker_urls.clone(),
}
} else if self.pd_disaggregation {
RoutingMode::PrefillDecode {
prefill_urls: self.prefill_urls.clone().unwrap_or_default(),
decode_urls: self.decode_urls.clone().unwrap_or_default(),
prefill_policy: self.prefill_policy.as_ref().map(convert_policy),
decode_policy: self.decode_policy.as_ref().map(convert_policy),
}
} else {
RoutingMode::Regular {
worker_urls: self.worker_urls.clone(),
}
};
let policy = convert_policy(&self.policy);
let discovery = if self.service_discovery {
Some(DiscoveryConfig {
enabled: true,
namespace: self.service_discovery_namespace.clone(),
port: self.service_discovery_port,
check_interval_secs: 60,
selector: self.selector.clone(),
prefill_selector: self.prefill_selector.clone(),
decode_selector: self.decode_selector.clone(),
bootstrap_port_annotation: self.bootstrap_port_annotation.clone(),
})
} else {
None
};
let metrics = match (self.prometheus_port, self.prometheus_host.as_ref()) {
(Some(port), Some(host)) => Some(MetricsConfig {
port,
host: host.clone(),
}),
_ => None,
};
let history_backend = match self.history_backend {
HistoryBackendType::Memory => config::HistoryBackend::Memory,
HistoryBackendType::None => config::HistoryBackend::None,
HistoryBackendType::Oracle => config::HistoryBackend::Oracle,
HistoryBackendType::Postgres => config::HistoryBackend::Postgres,
};
let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) {
self.oracle_config
.as_ref()
.map(|cfg| cfg.to_config_oracle())
} else {
None
};
let postgres_config = if matches!(self.history_backend, HistoryBackendType::Postgres) {
self.postgres_config
.as_ref()
.map(|cfg| cfg.to_config_postgres())
} else {
None
};
config::RouterConfig::builder()
.mode(mode)
.policy(policy)
.host(&self.host)
.port(self.port)
.connection_mode(self.connection_mode.clone())
.max_payload_size(self.max_payload_size)
.request_timeout_secs(self.request_timeout_secs)
.worker_startup_timeout_secs(self.worker_startup_timeout_secs)
.worker_startup_check_interval_secs(self.worker_startup_check_interval)
.max_concurrent_requests(self.max_concurrent_requests)
.queue_size(self.queue_size)
.queue_timeout_secs(self.queue_timeout_secs)
.cors_allowed_origins(self.cors_allowed_origins.clone())
.retry_config(config::RetryConfig {
max_retries: self.retry_max_retries,
initial_backoff_ms: self.retry_initial_backoff_ms,
max_backoff_ms: self.retry_max_backoff_ms,
backoff_multiplier: self.retry_backoff_multiplier,
jitter_factor: self.retry_jitter_factor,
})
.circuit_breaker_config(config::CircuitBreakerConfig {
failure_threshold: self.cb_failure_threshold,
success_threshold: self.cb_success_threshold,
timeout_duration_secs: self.cb_timeout_duration_secs,
window_duration_secs: self.cb_window_duration_secs,
})
.health_check_config(config::HealthCheckConfig {
failure_threshold: self.health_failure_threshold,
success_threshold: self.health_success_threshold,
timeout_secs: self.health_check_timeout_secs,
check_interval_secs: self.health_check_interval_secs,
endpoint: self.health_check_endpoint.clone(),
})
.tokenizer_cache(config::TokenizerCacheConfig {
enable_l0: self.tokenizer_cache_enable_l0,
l0_max_entries: self.tokenizer_cache_l0_max_entries,
enable_l1: self.tokenizer_cache_enable_l1,
l1_max_memory: self.tokenizer_cache_l1_max_memory,
})
.history_backend(history_backend)
.maybe_api_key(self.api_key.as_ref())
.maybe_discovery(discovery)
.maybe_metrics(metrics)
.maybe_log_dir(self.log_dir.as_ref())
.maybe_log_level(self.log_level.as_ref())
.maybe_request_id_headers(self.request_id_headers.clone())
.maybe_rate_limit_tokens_per_second(self.rate_limit_tokens_per_second)
.maybe_model_path(self.model_path.as_ref())
.maybe_tokenizer_path(self.tokenizer_path.as_ref())
.maybe_chat_template(self.chat_template.as_ref())
.maybe_oracle(oracle)
.maybe_postgres(postgres_config)
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
.dp_aware(self.dp_aware)
.retries(!self.disable_retries)
.circuit_breaker(!self.disable_circuit_breaker)
.igw(self.enable_igw)
.maybe_client_cert_and_key(
self.client_cert_path.as_ref(),
self.client_key_path.as_ref(),
)
.add_ca_certificates(self.ca_cert_paths.clone())
.build()
}
}
#[pymethods]
impl Router {
#[new]
#[pyo3(signature = (
worker_urls,
policy = PolicyType::RoundRobin,
host = String::from("0.0.0.0"),
port = 3001,
worker_startup_timeout_secs = 600,
worker_startup_check_interval = 30,
cache_threshold = 0.3,
balance_abs_threshold = 64,
balance_rel_threshold = 1.5,
eviction_interval_secs = 120,
max_tree_size = 2usize.pow(26),
max_payload_size = 512 * 1024 * 1024,
dp_aware = false,
api_key = None,
log_dir = None,
log_level = None,
service_discovery = false,
selector = HashMap::new(),
service_discovery_port = 80,
service_discovery_namespace = None,
prefill_selector = HashMap::new(),
decode_selector = HashMap::new(),
bootstrap_port_annotation = String::from("sglang.ai/bootstrap-port"),
prometheus_port = None,
prometheus_host = None,
request_timeout_secs = 1800,
request_id_headers = None,
pd_disaggregation = false,
bucket_adjust_interval_secs = 5,
prefill_urls = None,
decode_urls = None,
prefill_policy = None,
decode_policy = None,
max_concurrent_requests = -1,
cors_allowed_origins = vec![],
retry_max_retries = 5,
retry_initial_backoff_ms = 50,
retry_max_backoff_ms = 30_000,
retry_backoff_multiplier = 1.5,
retry_jitter_factor = 0.2,
disable_retries = false,
cb_failure_threshold = 10,
cb_success_threshold = 3,
cb_timeout_duration_secs = 60,
cb_window_duration_secs = 120,
disable_circuit_breaker = false,
health_failure_threshold = 3,
health_success_threshold = 2,
health_check_timeout_secs = 5,
health_check_interval_secs = 60,
health_check_endpoint = String::from("/health"),
enable_igw = false,
queue_size = 100,
queue_timeout_secs = 60,
rate_limit_tokens_per_second = None,
model_path = None,
tokenizer_path = None,
chat_template = None,
tokenizer_cache_enable_l0 = false,
tokenizer_cache_l0_max_entries = 10000,
tokenizer_cache_enable_l1 = false,
tokenizer_cache_l1_max_memory = 52428800,
reasoning_parser = None,
tool_call_parser = None,
mcp_config_path = None,
backend = BackendType::Sglang,
history_backend = HistoryBackendType::Memory,
oracle_config = None,
postgres_config = None,
client_cert_path = None,
client_key_path = None,
ca_cert_paths = vec![],
))]
#[allow(clippy::too_many_arguments)]
fn new(
worker_urls: Vec<String>,
policy: PolicyType,
host: String,
port: u16,
worker_startup_timeout_secs: u64,
worker_startup_check_interval: u64,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
eviction_interval_secs: u64,
max_tree_size: usize,
max_payload_size: usize,
dp_aware: bool,
api_key: Option<String>,
log_dir: Option<String>,
log_level: Option<String>,
service_discovery: bool,
selector: HashMap<String, String>,
service_discovery_port: u16,
service_discovery_namespace: Option<String>,
prefill_selector: HashMap<String, String>,
decode_selector: HashMap<String, String>,
bootstrap_port_annotation: String,
prometheus_port: Option<u16>,
prometheus_host: Option<String>,
request_timeout_secs: u64,
request_id_headers: Option<Vec<String>>,
pd_disaggregation: bool,
bucket_adjust_interval_secs: usize,
prefill_urls: Option<Vec<(String, Option<u16>)>>,
decode_urls: Option<Vec<String>>,
prefill_policy: Option<PolicyType>,
decode_policy: Option<PolicyType>,
max_concurrent_requests: i32,
cors_allowed_origins: Vec<String>,
retry_max_retries: u32,
retry_initial_backoff_ms: u64,
retry_max_backoff_ms: u64,
retry_backoff_multiplier: f32,
retry_jitter_factor: f32,
disable_retries: bool,
cb_failure_threshold: u32,
cb_success_threshold: u32,
cb_timeout_duration_secs: u64,
cb_window_duration_secs: u64,
disable_circuit_breaker: bool,
health_failure_threshold: u32,
health_success_threshold: u32,
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
rate_limit_tokens_per_second: Option<i32>,
model_path: Option<String>,
tokenizer_path: Option<String>,
chat_template: Option<String>,
tokenizer_cache_enable_l0: bool,
tokenizer_cache_l0_max_entries: usize,
tokenizer_cache_enable_l1: bool,
tokenizer_cache_l1_max_memory: usize,
reasoning_parser: Option<String>,
tool_call_parser: Option<String>,
mcp_config_path: Option<String>,
backend: BackendType,
history_backend: HistoryBackendType,
oracle_config: Option<PyOracleConfig>,
postgres_config: Option<PyPostgresConfig>,
client_cert_path: Option<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
) -> PyResult<Self> {
let mut all_urls = worker_urls.clone();
if let Some(ref prefill_urls) = prefill_urls {
for (url, _) in prefill_urls {
all_urls.push(url.clone());
}
}
if let Some(ref decode_urls) = decode_urls {
all_urls.extend(decode_urls.clone());
}
let connection_mode = Self::determine_connection_mode(&all_urls);
Ok(Router {
host,
port,
worker_urls,
policy,
worker_startup_timeout_secs,
worker_startup_check_interval,
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
eviction_interval_secs,
max_tree_size,
max_payload_size,
dp_aware,
api_key,
log_dir,
log_level,
service_discovery,
selector,
service_discovery_port,
service_discovery_namespace,
prefill_selector,
decode_selector,
bootstrap_port_annotation,
prometheus_port,
prometheus_host,
request_timeout_secs,
request_id_headers,
pd_disaggregation,
bucket_adjust_interval_secs,
prefill_urls,
decode_urls,
prefill_policy,
decode_policy,
max_concurrent_requests,
cors_allowed_origins,
retry_max_retries,
retry_initial_backoff_ms,
retry_max_backoff_ms,
retry_backoff_multiplier,
retry_jitter_factor,
disable_retries,
cb_failure_threshold,
cb_success_threshold,
cb_timeout_duration_secs,
cb_window_duration_secs,
disable_circuit_breaker,
health_failure_threshold,
health_success_threshold,
health_check_timeout_secs,
health_check_interval_secs,
health_check_endpoint,
enable_igw,
queue_size,
queue_timeout_secs,
rate_limit_tokens_per_second,
connection_mode,
model_path,
tokenizer_path,
chat_template,
tokenizer_cache_enable_l0,
tokenizer_cache_l0_max_entries,
tokenizer_cache_enable_l1,
tokenizer_cache_l1_max_memory,
reasoning_parser,
tool_call_parser,
mcp_config_path,
backend,
history_backend,
oracle_config,
postgres_config,
client_cert_path,
client_key_path,
ca_cert_paths,
})
}
fn start(&self) -> PyResult<()> {
use metrics::PrometheusConfig;
let router_config = self.to_router_config().map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!("Configuration error: {}", e))
})?;
router_config.validate().map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!(
"Configuration validation failed: {}",
e
))
})?;
let service_discovery_config = if self.service_discovery {
Some(service_discovery::ServiceDiscoveryConfig {
enabled: true,
selector: self.selector.clone(),
check_interval: std::time::Duration::from_secs(60),
port: self.service_discovery_port,
namespace: self.service_discovery_namespace.clone(),
pd_mode: self.pd_disaggregation,
prefill_selector: self.prefill_selector.clone(),
decode_selector: self.decode_selector.clone(),
bootstrap_port_annotation: self.bootstrap_port_annotation.clone(),
})
} else {
None
};
let prometheus_config = Some(PrometheusConfig {
port: self.prometheus_port.unwrap_or(29000),
host: self
.prometheus_host
.clone()
.unwrap_or_else(|| "127.0.0.1".to_string()),
});
let runtime = tokio::runtime::Runtime::new()
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
runtime.block_on(async move {
server::startup(server::ServerConfig {
host: self.host.clone(),
port: self.port,
router_config,
max_payload_size: self.max_payload_size,
log_dir: self.log_dir.clone(),
log_level: self.log_level.clone(),
service_discovery_config,
prometheus_config,
request_timeout_secs: self.request_timeout_secs,
request_id_headers: self.request_id_headers.clone(),
})
.await
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
})
}
}
#[pymodule]
fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PolicyType>()?;
m.add_class::<BackendType>()?;
m.add_class::<HistoryBackendType>()?;
m.add_class::<PyOracleConfig>()?;
m.add_class::<PyPostgresConfig>()?;
m.add_class::<Router>()?;
Ok(())
}

6
sgl-router/py_test/conftest.py
View File

@@ -2,14 +2,14 @@ import sys
from importlib.util import find_spec
from pathlib import Path
# Only add py_src to path if the wheel is not installed (for local development)
# Only add bindings/python to path if the wheel is not installed (for local development)
# This ensures CI tests use the installed wheel which contains the Rust extension
_ROOT = Path(__file__).resolve().parents[1]
_SRC = _ROOT / "py_src"
_SRC = _ROOT / "bindings" / "python"
# Check if sglang_router is already installed with the Rust extension
_wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
# Only add py_src if wheel is not installed (development mode)
# Only add bindings/python if wheel is not installed (development mode)
if not _wheel_installed and str(_SRC) not in sys.path:
sys.path.insert(0, str(_SRC))

4
sgl-router/py_test/e2e_grpc/conftest.py
View File

@@ -9,9 +9,9 @@ from pathlib import Path
import pytest # noqa: F401
# Ensure router py_src is importable
# Ensure router bindings/python is importable
_ROUTER_ROOT = Path(__file__).resolve().parents[2]
_ROUTER_SRC = _ROUTER_ROOT / "py_src"
_ROUTER_SRC = _ROUTER_ROOT / "bindings" / "python"
if str(_ROUTER_SRC) not in sys.path:
sys.path.insert(0, str(_ROUTER_SRC))

716
sgl-router/src/lib.rs
View File

@@ -1,8 +1,6 @@
use pyo3::prelude::*;
pub mod app_context;
pub mod config;
pub mod logging;
use std::collections::HashMap;
pub mod core;
pub mod data_connector;
@@ -19,717 +17,3 @@ pub mod server;
pub mod service_discovery;
pub mod tokenizer;
pub mod tool_parser;
use crate::metrics::PrometheusConfig;
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum PolicyType {
Random,
RoundRobin,
CacheAware,
PowerOfTwo,
Bucket,
}
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum BackendType {
Sglang,
Openai,
}
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum HistoryBackendType {
Memory,
None,
Oracle,
Postgres,
}
#[pyclass]
#[derive(Clone, PartialEq)]
pub struct PyOracleConfig {
#[pyo3(get, set)]
pub wallet_path: Option<String>,
#[pyo3(get, set)]
pub connect_descriptor: Option<String>,
#[pyo3(get, set)]
pub username: Option<String>,
#[pyo3(get, set)]
pub password: Option<String>,
#[pyo3(get, set)]
pub pool_min: usize,
#[pyo3(get, set)]
pub pool_max: usize,
#[pyo3(get, set)]
pub pool_timeout_secs: u64,
}
impl std::fmt::Debug for PyOracleConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PyOracleConfig")
.field("wallet_path", &self.wallet_path)
.field("connect_descriptor", &"<redacted>")
.field("username", &self.username)
.field("password", &"<redacted>")
.field("pool_min", &self.pool_min)
.field("pool_max", &self.pool_max)
.field("pool_timeout_secs", &self.pool_timeout_secs)
.finish()
}
}
#[pymethods]
impl PyOracleConfig {
#[new]
#[pyo3(signature = (
password = None,
username = None,
connect_descriptor = None,
wallet_path = None,
pool_min = 1,
pool_max = 16,
pool_timeout_secs = 30,
))]
fn new(
password: Option<String>,
username: Option<String>,
connect_descriptor: Option<String>,
wallet_path: Option<String>,
pool_min: usize,
pool_max: usize,
pool_timeout_secs: u64,
) -> PyResult<Self> {
if pool_min == 0 {
return Err(pyo3::exceptions::PyValueError::new_err(
"pool_min must be at least 1",
));
}
if pool_max < pool_min {
return Err(pyo3::exceptions::PyValueError::new_err(
"pool_max must be >= pool_min",
));
}
Ok(PyOracleConfig {
wallet_path,
connect_descriptor,
username,
password,
pool_min,
pool_max,
pool_timeout_secs,
})
}
}
impl PyOracleConfig {
fn to_config_oracle(&self) -> config::OracleConfig {
// Simple conversion - validation happens later in validate_oracle()
config::OracleConfig {
wallet_path: self.wallet_path.clone(),
connect_descriptor: self.connect_descriptor.clone().unwrap_or_default(),
username: self.username.clone().unwrap_or_default(),
password: self.password.clone().unwrap_or_default(),
pool_min: self.pool_min,
pool_max: self.pool_max,
pool_timeout_secs: self.pool_timeout_secs,
}
}
}
#[pyclass]
#[derive(Debug, Clone, PartialEq)]
pub struct PyPostgresConfig {
#[pyo3(get, set)]
pub db_url: Option<String>,
#[pyo3(get, set)]
pub pool_max: usize,
}
#[pymethods]
impl PyPostgresConfig {
#[new]
#[pyo3(signature = (db_url = None,pool_max = 16,))]
fn new(db_url: Option<String>, pool_max: usize) -> PyResult<Self> {
Ok(PyPostgresConfig { db_url, pool_max })
}
}
impl PyPostgresConfig {
fn to_config_postgres(&self) -> config::PostgresConfig {
config::PostgresConfig {
db_url: self.db_url.clone().unwrap_or_default(),
pool_max: self.pool_max,
}
}
}
#[pyclass]
#[derive(Debug, Clone, PartialEq)]
struct Router {
host: String,
port: u16,
worker_urls: Vec<String>,
policy: PolicyType,
worker_startup_timeout_secs: u64,
worker_startup_check_interval: u64,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
eviction_interval_secs: u64,
max_tree_size: usize,
max_payload_size: usize,
dp_aware: bool,
api_key: Option<String>,
log_dir: Option<String>,
log_level: Option<String>,
service_discovery: bool,
selector: HashMap<String, String>,
service_discovery_port: u16,
service_discovery_namespace: Option<String>,
prefill_selector: HashMap<String, String>,
decode_selector: HashMap<String, String>,
bootstrap_port_annotation: String,
prometheus_port: Option<u16>,
prometheus_host: Option<String>,
request_timeout_secs: u64,
request_id_headers: Option<Vec<String>>,
pd_disaggregation: bool,
// Takes effect in PD mode and when policy = bucket
bucket_adjust_interval_secs: usize,
prefill_urls: Option<Vec<(String, Option<u16>)>>,
decode_urls: Option<Vec<String>>,
prefill_policy: Option<PolicyType>,
decode_policy: Option<PolicyType>,
max_concurrent_requests: i32,
cors_allowed_origins: Vec<String>,
retry_max_retries: u32,
retry_initial_backoff_ms: u64,
retry_max_backoff_ms: u64,
retry_backoff_multiplier: f32,
retry_jitter_factor: f32,
disable_retries: bool,
cb_failure_threshold: u32,
cb_success_threshold: u32,
cb_timeout_duration_secs: u64,
cb_window_duration_secs: u64,
disable_circuit_breaker: bool,
health_failure_threshold: u32,
health_success_threshold: u32,
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
rate_limit_tokens_per_second: Option<i32>,
connection_mode: core::ConnectionMode,
model_path: Option<String>,
tokenizer_path: Option<String>,
chat_template: Option<String>,
tokenizer_cache_enable_l0: bool,
tokenizer_cache_l0_max_entries: usize,
tokenizer_cache_enable_l1: bool,
tokenizer_cache_l1_max_memory: usize,
reasoning_parser: Option<String>,
tool_call_parser: Option<String>,
mcp_config_path: Option<String>,
backend: BackendType,
history_backend: HistoryBackendType,
oracle_config: Option<PyOracleConfig>,
postgres_config: Option<PyPostgresConfig>,
client_cert_path: Option<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
}
impl Router {
/// Determine connection mode from worker URLs
fn determine_connection_mode(worker_urls: &[String]) -> core::ConnectionMode {
for url in worker_urls {
if url.starts_with("grpc://") || url.starts_with("grpcs://") {
return core::ConnectionMode::Grpc { port: None };
}
}
core::ConnectionMode::Http
}
pub fn to_router_config(&self) -> config::ConfigResult<config::RouterConfig> {
use config::{
DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode,
};
let convert_policy = |policy: &PolicyType| -> ConfigPolicyConfig {
match policy {
PolicyType::Random => ConfigPolicyConfig::Random,
PolicyType::RoundRobin => ConfigPolicyConfig::RoundRobin,
PolicyType::CacheAware => ConfigPolicyConfig::CacheAware {
cache_threshold: self.cache_threshold,
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
eviction_interval_secs: self.eviction_interval_secs,
max_tree_size: self.max_tree_size,
},
PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
load_check_interval_secs: 5,
},
PolicyType::Bucket => ConfigPolicyConfig::Bucket {
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
bucket_adjust_interval_secs: self.bucket_adjust_interval_secs,
},
}
};
let mode = if self.enable_igw {
RoutingMode::Regular {
worker_urls: vec![],
}
} else if matches!(self.backend, BackendType::Openai) {
RoutingMode::OpenAI {
worker_urls: self.worker_urls.clone(),
}
} else if self.pd_disaggregation {
RoutingMode::PrefillDecode {
prefill_urls: self.prefill_urls.clone().unwrap_or_default(),
decode_urls: self.decode_urls.clone().unwrap_or_default(),
prefill_policy: self.prefill_policy.as_ref().map(convert_policy),
decode_policy: self.decode_policy.as_ref().map(convert_policy),
}
} else {
RoutingMode::Regular {
worker_urls: self.worker_urls.clone(),
}
};
let policy = convert_policy(&self.policy);
let discovery = if self.service_discovery {
Some(DiscoveryConfig {
enabled: true,
namespace: self.service_discovery_namespace.clone(),
port: self.service_discovery_port,
check_interval_secs: 60,
selector: self.selector.clone(),
prefill_selector: self.prefill_selector.clone(),
decode_selector: self.decode_selector.clone(),
bootstrap_port_annotation: self.bootstrap_port_annotation.clone(),
})
} else {
None
};
let metrics = match (self.prometheus_port, self.prometheus_host.as_ref()) {
(Some(port), Some(host)) => Some(MetricsConfig {
port,
host: host.clone(),
}),
_ => None,
};
let history_backend = match self.history_backend {
HistoryBackendType::Memory => config::HistoryBackend::Memory,
HistoryBackendType::None => config::HistoryBackend::None,
HistoryBackendType::Oracle => config::HistoryBackend::Oracle,
HistoryBackendType::Postgres => config::HistoryBackend::Postgres,
};
let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) {
self.oracle_config
.as_ref()
.map(|cfg| cfg.to_config_oracle())
} else {
None
};
let postgres_config = if matches!(self.history_backend, HistoryBackendType::Postgres) {
self.postgres_config
.as_ref()
.map(|cfg| cfg.to_config_postgres())
} else {
None
};
config::RouterConfig::builder()
.mode(mode)
.policy(policy)
.host(&self.host)
.port(self.port)
.connection_mode(self.connection_mode.clone())
.max_payload_size(self.max_payload_size)
.request_timeout_secs(self.request_timeout_secs)
.worker_startup_timeout_secs(self.worker_startup_timeout_secs)
.worker_startup_check_interval_secs(self.worker_startup_check_interval)
.max_concurrent_requests(self.max_concurrent_requests)
.queue_size(self.queue_size)
.queue_timeout_secs(self.queue_timeout_secs)
.cors_allowed_origins(self.cors_allowed_origins.clone())
.retry_config(config::RetryConfig {
max_retries: self.retry_max_retries,
initial_backoff_ms: self.retry_initial_backoff_ms,
max_backoff_ms: self.retry_max_backoff_ms,
backoff_multiplier: self.retry_backoff_multiplier,
jitter_factor: self.retry_jitter_factor,
})
.circuit_breaker_config(config::CircuitBreakerConfig {
failure_threshold: self.cb_failure_threshold,
success_threshold: self.cb_success_threshold,
timeout_duration_secs: self.cb_timeout_duration_secs,
window_duration_secs: self.cb_window_duration_secs,
})
.health_check_config(config::HealthCheckConfig {
failure_threshold: self.health_failure_threshold,
success_threshold: self.health_success_threshold,
timeout_secs: self.health_check_timeout_secs,
check_interval_secs: self.health_check_interval_secs,
endpoint: self.health_check_endpoint.clone(),
})
.tokenizer_cache(config::TokenizerCacheConfig {
enable_l0: self.tokenizer_cache_enable_l0,
l0_max_entries: self.tokenizer_cache_l0_max_entries,
enable_l1: self.tokenizer_cache_enable_l1,
l1_max_memory: self.tokenizer_cache_l1_max_memory,
})
.history_backend(history_backend)
.maybe_api_key(self.api_key.as_ref())
.maybe_discovery(discovery)
.maybe_metrics(metrics)
.maybe_log_dir(self.log_dir.as_ref())
.maybe_log_level(self.log_level.as_ref())
.maybe_request_id_headers(self.request_id_headers.clone())
.maybe_rate_limit_tokens_per_second(self.rate_limit_tokens_per_second)
.maybe_model_path(self.model_path.as_ref())
.maybe_tokenizer_path(self.tokenizer_path.as_ref())
.maybe_chat_template(self.chat_template.as_ref())
.maybe_oracle(oracle)
.maybe_postgres(postgres_config)
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
.dp_aware(self.dp_aware)
.retries(!self.disable_retries)
.circuit_breaker(!self.disable_circuit_breaker)
.igw(self.enable_igw)
.maybe_client_cert_and_key(
self.client_cert_path.as_ref(),
self.client_key_path.as_ref(),
)
.add_ca_certificates(self.ca_cert_paths.clone())
.build()
}
}
#[pymethods]
impl Router {
#[new]
#[pyo3(signature = (
worker_urls,
policy = PolicyType::RoundRobin,
host = String::from("0.0.0.0"),
port = 3001,
worker_startup_timeout_secs = 600,
worker_startup_check_interval = 30,
cache_threshold = 0.3,
balance_abs_threshold = 64,
balance_rel_threshold = 1.5,
eviction_interval_secs = 120,
max_tree_size = 2usize.pow(26),
max_payload_size = 512 * 1024 * 1024,
dp_aware = false,
api_key = None,
log_dir = None,
log_level = None,
service_discovery = false,
selector = HashMap::new(),
service_discovery_port = 80,
service_discovery_namespace = None,
prefill_selector = HashMap::new(),
decode_selector = HashMap::new(),
bootstrap_port_annotation = String::from("sglang.ai/bootstrap-port"),
prometheus_port = None,
prometheus_host = None,
request_timeout_secs = 1800,
request_id_headers = None,
pd_disaggregation = false,
bucket_adjust_interval_secs = 5,
prefill_urls = None,
decode_urls = None,
prefill_policy = None,
decode_policy = None,
max_concurrent_requests = -1,
cors_allowed_origins = vec![],
retry_max_retries = 5,
retry_initial_backoff_ms = 50,
retry_max_backoff_ms = 30_000,
retry_backoff_multiplier = 1.5,
retry_jitter_factor = 0.2,
disable_retries = false,
cb_failure_threshold = 10,
cb_success_threshold = 3,
cb_timeout_duration_secs = 60,
cb_window_duration_secs = 120,
disable_circuit_breaker = false,
health_failure_threshold = 3,
health_success_threshold = 2,
health_check_timeout_secs = 5,
health_check_interval_secs = 60,
health_check_endpoint = String::from("/health"),
enable_igw = false,
queue_size = 100,
queue_timeout_secs = 60,
rate_limit_tokens_per_second = None,
model_path = None,
tokenizer_path = None,
chat_template = None,
tokenizer_cache_enable_l0 = false,
tokenizer_cache_l0_max_entries = 10000,
tokenizer_cache_enable_l1 = false,
tokenizer_cache_l1_max_memory = 52428800,
reasoning_parser = None,
tool_call_parser = None,
mcp_config_path = None,
backend = BackendType::Sglang,
history_backend = HistoryBackendType::Memory,
oracle_config = None,
postgres_config = None,
client_cert_path = None,
client_key_path = None,
ca_cert_paths = vec![],
))]
#[allow(clippy::too_many_arguments)]
fn new(
worker_urls: Vec<String>,
policy: PolicyType,
host: String,
port: u16,
worker_startup_timeout_secs: u64,
worker_startup_check_interval: u64,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
eviction_interval_secs: u64,
max_tree_size: usize,
max_payload_size: usize,
dp_aware: bool,
api_key: Option<String>,
log_dir: Option<String>,
log_level: Option<String>,
service_discovery: bool,
selector: HashMap<String, String>,
service_discovery_port: u16,
service_discovery_namespace: Option<String>,
prefill_selector: HashMap<String, String>,
decode_selector: HashMap<String, String>,
bootstrap_port_annotation: String,
prometheus_port: Option<u16>,
prometheus_host: Option<String>,
request_timeout_secs: u64,
request_id_headers: Option<Vec<String>>,
pd_disaggregation: bool,
bucket_adjust_interval_secs: usize,
prefill_urls: Option<Vec<(String, Option<u16>)>>,
decode_urls: Option<Vec<String>>,
prefill_policy: Option<PolicyType>,
decode_policy: Option<PolicyType>,
max_concurrent_requests: i32,
cors_allowed_origins: Vec<String>,
retry_max_retries: u32,
retry_initial_backoff_ms: u64,
retry_max_backoff_ms: u64,
retry_backoff_multiplier: f32,
retry_jitter_factor: f32,
disable_retries: bool,
cb_failure_threshold: u32,
cb_success_threshold: u32,
cb_timeout_duration_secs: u64,
cb_window_duration_secs: u64,
disable_circuit_breaker: bool,
health_failure_threshold: u32,
health_success_threshold: u32,
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
rate_limit_tokens_per_second: Option<i32>,
model_path: Option<String>,
tokenizer_path: Option<String>,
chat_template: Option<String>,
tokenizer_cache_enable_l0: bool,
tokenizer_cache_l0_max_entries: usize,
tokenizer_cache_enable_l1: bool,
tokenizer_cache_l1_max_memory: usize,
reasoning_parser: Option<String>,
tool_call_parser: Option<String>,
mcp_config_path: Option<String>,
backend: BackendType,
history_backend: HistoryBackendType,
oracle_config: Option<PyOracleConfig>,
postgres_config: Option<PyPostgresConfig>,
client_cert_path: Option<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
) -> PyResult<Self> {
let mut all_urls = worker_urls.clone();
if let Some(ref prefill_urls) = prefill_urls {
for (url, _) in prefill_urls {
all_urls.push(url.clone());
}
}
if let Some(ref decode_urls) = decode_urls {
all_urls.extend(decode_urls.clone());
}
let connection_mode = Self::determine_connection_mode(&all_urls);
Ok(Router {
host,
port,
worker_urls,
policy,
worker_startup_timeout_secs,
worker_startup_check_interval,
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
eviction_interval_secs,
max_tree_size,
max_payload_size,
dp_aware,
api_key,
log_dir,
log_level,
service_discovery,
selector,
service_discovery_port,
service_discovery_namespace,
prefill_selector,
decode_selector,
bootstrap_port_annotation,
prometheus_port,
prometheus_host,
request_timeout_secs,
request_id_headers,
pd_disaggregation,
bucket_adjust_interval_secs,
prefill_urls,
decode_urls,
prefill_policy,
decode_policy,
max_concurrent_requests,
cors_allowed_origins,
retry_max_retries,
retry_initial_backoff_ms,
retry_max_backoff_ms,
retry_backoff_multiplier,
retry_jitter_factor,
disable_retries,
cb_failure_threshold,
cb_success_threshold,
cb_timeout_duration_secs,
cb_window_duration_secs,
disable_circuit_breaker,
health_failure_threshold,
health_success_threshold,
health_check_timeout_secs,
health_check_interval_secs,
health_check_endpoint,
enable_igw,
queue_size,
queue_timeout_secs,
rate_limit_tokens_per_second,
connection_mode,
model_path,
tokenizer_path,
chat_template,
tokenizer_cache_enable_l0,
tokenizer_cache_l0_max_entries,
tokenizer_cache_enable_l1,
tokenizer_cache_l1_max_memory,
reasoning_parser,
tool_call_parser,
mcp_config_path,
backend,
history_backend,
oracle_config,
postgres_config,
client_cert_path,
client_key_path,
ca_cert_paths,
})
}
fn start(&self) -> PyResult<()> {
let router_config = self.to_router_config().map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!("Configuration error: {}", e))
})?;
router_config.validate().map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!(
"Configuration validation failed: {}",
e
))
})?;
let service_discovery_config = if self.service_discovery {
Some(service_discovery::ServiceDiscoveryConfig {
enabled: true,
selector: self.selector.clone(),
check_interval: std::time::Duration::from_secs(60),
port: self.service_discovery_port,
namespace: self.service_discovery_namespace.clone(),
pd_mode: self.pd_disaggregation,
prefill_selector: self.prefill_selector.clone(),
decode_selector: self.decode_selector.clone(),
bootstrap_port_annotation: self.bootstrap_port_annotation.clone(),
})
} else {
None
};
let prometheus_config = Some(PrometheusConfig {
port: self.prometheus_port.unwrap_or(29000),
host: self
.prometheus_host
.clone()
.unwrap_or_else(|| "127.0.0.1".to_string()),
});
let runtime = tokio::runtime::Runtime::new()
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
runtime.block_on(async move {
server::startup(server::ServerConfig {
host: self.host.clone(),
port: self.port,
router_config,
max_payload_size: self.max_payload_size,
log_dir: self.log_dir.clone(),
log_level: self.log_level.clone(),
service_discovery_config,
prometheus_config,
request_timeout_secs: self.request_timeout_secs,
request_id_headers: self.request_id_headers.clone(),
})
.await
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
})
}
}
#[pymodule]
fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PolicyType>()?;
m.add_class::<BackendType>()?;
m.add_class::<HistoryBackendType>()?;
m.add_class::<PyOracleConfig>()?;
m.add_class::<PyPostgresConfig>()?;
m.add_class::<Router>()?;
Ok(())
}