diff --git a/.github/workflows/nightly-release-router.yml b/.github/workflows/nightly-release-router.yml index 97b31976e..74b4e90ef 100644 --- a/.github/workflows/nightly-release-router.yml +++ b/.github/workflows/nightly-release-router.yml @@ -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 diff --git a/.github/workflows/pr-test-pd-router.yml b/.github/workflows/pr-test-pd-router.yml index 03e313985..063503a36 100644 --- a/.github/workflows/pr-test-pd-router.yml +++ b/.github/workflows/pr-test-pd-router.yml @@ -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 diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 4f6e21cf7..c834ceb1c 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -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: | diff --git a/.github/workflows/release-docker-router.yml b/.github/workflows/release-docker-router.yml index 78fde711a..e8722f073 100644 --- a/.github/workflows/release-docker-router.yml +++ b/.github/workflows/release-docker-router.yml @@ -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 diff --git a/.github/workflows/release-pypi-router.yml b/.github/workflows/release-pypi-router.yml index 636076a5f..df5bc6fe8 100644 --- a/.github/workflows/release-pypi-router.yml +++ b/.github/workflows/release-pypi-router.yml @@ -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 diff --git a/sgl-router/.coveragerc b/sgl-router/.coveragerc deleted file mode 100644 index f908584fe..000000000 --- a/sgl-router/.coveragerc +++ /dev/null @@ -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 diff --git a/sgl-router/Cargo.toml b/sgl-router/Cargo.toml index 8d5f44dd0..84c6d1663 100644 --- a/sgl-router/Cargo.toml +++ b/sgl-router/Cargo.toml @@ -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" diff --git a/sgl-router/LICENSE b/sgl-router/LICENSE new file mode 120000 index 000000000..ea5b60640 --- /dev/null +++ b/sgl-router/LICENSE @@ -0,0 +1 @@ +../LICENSE \ No newline at end of file diff --git a/sgl-router/MANIFEST.in b/sgl-router/MANIFEST.in deleted file mode 100644 index 460524b50..000000000 --- a/sgl-router/MANIFEST.in +++ /dev/null @@ -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 diff --git a/sgl-router/Makefile b/sgl-router/Makefile index 60cc72e9c..9336ad0fa 100644 --- a/sgl-router/Makefile +++ b/sgl-router/Makefile @@ -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 diff --git a/sgl-router/README.md b/sgl-router/README.md index ccd5a4210..3eede1d0c 100644 --- a/sgl-router/README.md +++ b/sgl-router/README.md @@ -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. --- diff --git a/sgl-router/bindings/python/.coveragerc b/sgl-router/bindings/python/.coveragerc new file mode 100644 index 000000000..4a066e7b2 --- /dev/null +++ b/sgl-router/bindings/python/.coveragerc @@ -0,0 +1,6 @@ +[run] +source = sglang_router +omit = */mini_lb.py + +[report] +fail_under = 80 diff --git a/sgl-router/bindings/python/Cargo.toml b/sgl-router/bindings/python/Cargo.toml new file mode 100644 index 000000000..d2054d76d --- /dev/null +++ b/sgl-router/bindings/python/Cargo.toml @@ -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 diff --git a/sgl-router/bindings/python/MANIFEST.in b/sgl-router/bindings/python/MANIFEST.in new file mode 100644 index 000000000..6c65950fa --- /dev/null +++ b/sgl-router/bindings/python/MANIFEST.in @@ -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 diff --git a/sgl-router/bindings/python/README.md b/sgl-router/bindings/python/README.md new file mode 100644 index 000000000..41848c45d --- /dev/null +++ b/sgl-router/bindings/python/README.md @@ -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 diff --git a/sgl-router/pyproject.toml b/sgl-router/bindings/python/pyproject.toml similarity index 98% rename from sgl-router/pyproject.toml rename to sgl-router/bindings/python/pyproject.toml index 595dd2faf..62575218a 100644 --- a/sgl-router/pyproject.toml +++ b/sgl-router/bindings/python/pyproject.toml @@ -42,5 +42,5 @@ dev = [ [tool.maturin] -python-source = "py_src" +python-source = "." module-name = "sglang_router.sglang_router_rs" diff --git a/sgl-router/setup.py b/sgl-router/bindings/python/setup.py similarity index 100% rename from sgl-router/setup.py rename to sgl-router/bindings/python/setup.py diff --git a/sgl-router/py_src/sglang_router/__init__.py b/sgl-router/bindings/python/sglang_router/__init__.py similarity index 100% rename from sgl-router/py_src/sglang_router/__init__.py rename to sgl-router/bindings/python/sglang_router/__init__.py diff --git a/sgl-router/py_src/sglang_router/launch_router.py b/sgl-router/bindings/python/sglang_router/launch_router.py similarity index 100% rename from sgl-router/py_src/sglang_router/launch_router.py rename to sgl-router/bindings/python/sglang_router/launch_router.py diff --git a/sgl-router/py_src/sglang_router/launch_server.py b/sgl-router/bindings/python/sglang_router/launch_server.py similarity index 100% rename from sgl-router/py_src/sglang_router/launch_server.py rename to sgl-router/bindings/python/sglang_router/launch_server.py diff --git a/sgl-router/py_src/sglang_router/mini_lb.py b/sgl-router/bindings/python/sglang_router/mini_lb.py similarity index 100% rename from sgl-router/py_src/sglang_router/mini_lb.py rename to sgl-router/bindings/python/sglang_router/mini_lb.py diff --git a/sgl-router/py_src/sglang_router/router.py b/sgl-router/bindings/python/sglang_router/router.py similarity index 100% rename from sgl-router/py_src/sglang_router/router.py rename to sgl-router/bindings/python/sglang_router/router.py diff --git a/sgl-router/py_src/sglang_router/router_args.py b/sgl-router/bindings/python/sglang_router/router_args.py similarity index 100% rename from sgl-router/py_src/sglang_router/router_args.py rename to sgl-router/bindings/python/sglang_router/router_args.py diff --git a/sgl-router/py_src/sglang_router/version.py b/sgl-router/bindings/python/sglang_router/version.py similarity index 100% rename from sgl-router/py_src/sglang_router/version.py rename to sgl-router/bindings/python/sglang_router/version.py diff --git a/sgl-router/bindings/python/src/lib.rs b/sgl-router/bindings/python/src/lib.rs new file mode 100644 index 000000000..e145ca1f9 --- /dev/null +++ b/sgl-router/bindings/python/src/lib.rs @@ -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, + #[pyo3(get, set)] + pub connect_descriptor: Option, + #[pyo3(get, set)] + pub username: Option, + #[pyo3(get, set)] + pub password: Option, + #[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", &"") + .field("username", &self.username) + .field("password", &"") + .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, + username: Option, + connect_descriptor: Option, + wallet_path: Option, + pool_min: usize, + pool_max: usize, + pool_timeout_secs: u64, + ) -> PyResult { + 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, + + #[pyo3(get, set)] + pub pool_max: usize, +} + +#[pymethods] +impl PyPostgresConfig { + #[new] + #[pyo3(signature = (db_url = None,pool_max = 16,))] + fn new(db_url: Option, pool_max: usize) -> PyResult { + 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, + 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, + log_dir: Option, + log_level: Option, + service_discovery: bool, + selector: HashMap, + service_discovery_port: u16, + service_discovery_namespace: Option, + prefill_selector: HashMap, + decode_selector: HashMap, + bootstrap_port_annotation: String, + prometheus_port: Option, + prometheus_host: Option, + request_timeout_secs: u64, + request_id_headers: Option>, + pd_disaggregation: bool, + bucket_adjust_interval_secs: usize, + prefill_urls: Option)>>, + decode_urls: Option>, + prefill_policy: Option, + decode_policy: Option, + max_concurrent_requests: i32, + cors_allowed_origins: Vec, + 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, + connection_mode: core::ConnectionMode, + model_path: Option, + tokenizer_path: Option, + chat_template: Option, + 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, + tool_call_parser: Option, + mcp_config_path: Option, + backend: BackendType, + history_backend: HistoryBackendType, + oracle_config: Option, + postgres_config: Option, + client_cert_path: Option, + client_key_path: Option, + ca_cert_paths: Vec, +} + +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 { + 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, + 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, + log_dir: Option, + log_level: Option, + service_discovery: bool, + selector: HashMap, + service_discovery_port: u16, + service_discovery_namespace: Option, + prefill_selector: HashMap, + decode_selector: HashMap, + bootstrap_port_annotation: String, + prometheus_port: Option, + prometheus_host: Option, + request_timeout_secs: u64, + request_id_headers: Option>, + pd_disaggregation: bool, + bucket_adjust_interval_secs: usize, + prefill_urls: Option)>>, + decode_urls: Option>, + prefill_policy: Option, + decode_policy: Option, + max_concurrent_requests: i32, + cors_allowed_origins: Vec, + 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, + model_path: Option, + tokenizer_path: Option, + chat_template: Option, + 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, + tool_call_parser: Option, + mcp_config_path: Option, + backend: BackendType, + history_backend: HistoryBackendType, + oracle_config: Option, + postgres_config: Option, + client_cert_path: Option, + client_key_path: Option, + ca_cert_paths: Vec, + ) -> PyResult { + 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::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/sgl-router/py_test/conftest.py b/sgl-router/py_test/conftest.py index b3c6dc66a..fe4fc6994 100644 --- a/sgl-router/py_test/conftest.py +++ b/sgl-router/py_test/conftest.py @@ -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)) diff --git a/sgl-router/py_test/e2e_grpc/conftest.py b/sgl-router/py_test/e2e_grpc/conftest.py index 7eda72ea8..b39d0d2c2 100644 --- a/sgl-router/py_test/e2e_grpc/conftest.py +++ b/sgl-router/py_test/e2e_grpc/conftest.py @@ -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)) diff --git a/sgl-router/src/lib.rs b/sgl-router/src/lib.rs index 8dfcc083e..6c66f02d9 100644 --- a/sgl-router/src/lib.rs +++ b/sgl-router/src/lib.rs @@ -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, - #[pyo3(get, set)] - pub connect_descriptor: Option, - #[pyo3(get, set)] - pub username: Option, - #[pyo3(get, set)] - pub password: Option, - #[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", &"") - .field("username", &self.username) - .field("password", &"") - .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, - username: Option, - connect_descriptor: Option, - wallet_path: Option, - pool_min: usize, - pool_max: usize, - pool_timeout_secs: u64, - ) -> PyResult { - 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, - - #[pyo3(get, set)] - pub pool_max: usize, -} - -#[pymethods] -impl PyPostgresConfig { - #[new] - #[pyo3(signature = (db_url = None,pool_max = 16,))] - fn new(db_url: Option, pool_max: usize) -> PyResult { - 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, - 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, - log_dir: Option, - log_level: Option, - service_discovery: bool, - selector: HashMap, - service_discovery_port: u16, - service_discovery_namespace: Option, - prefill_selector: HashMap, - decode_selector: HashMap, - bootstrap_port_annotation: String, - prometheus_port: Option, - prometheus_host: Option, - request_timeout_secs: u64, - request_id_headers: Option>, - pd_disaggregation: bool, - // Takes effect in PD mode and when policy = bucket - bucket_adjust_interval_secs: usize, - prefill_urls: Option)>>, - decode_urls: Option>, - prefill_policy: Option, - decode_policy: Option, - max_concurrent_requests: i32, - cors_allowed_origins: Vec, - 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, - connection_mode: core::ConnectionMode, - model_path: Option, - tokenizer_path: Option, - chat_template: Option, - 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, - tool_call_parser: Option, - mcp_config_path: Option, - backend: BackendType, - history_backend: HistoryBackendType, - oracle_config: Option, - postgres_config: Option, - client_cert_path: Option, - client_key_path: Option, - ca_cert_paths: Vec, -} - -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 { - 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, - 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, - log_dir: Option, - log_level: Option, - service_discovery: bool, - selector: HashMap, - service_discovery_port: u16, - service_discovery_namespace: Option, - prefill_selector: HashMap, - decode_selector: HashMap, - bootstrap_port_annotation: String, - prometheus_port: Option, - prometheus_host: Option, - request_timeout_secs: u64, - request_id_headers: Option>, - pd_disaggregation: bool, - bucket_adjust_interval_secs: usize, - prefill_urls: Option)>>, - decode_urls: Option>, - prefill_policy: Option, - decode_policy: Option, - max_concurrent_requests: i32, - cors_allowed_origins: Vec, - 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, - model_path: Option, - tokenizer_path: Option, - chat_template: Option, - 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, - tool_call_parser: Option, - mcp_config_path: Option, - backend: BackendType, - history_backend: HistoryBackendType, - oracle_config: Option, - postgres_config: Option, - client_cert_path: Option, - client_key_path: Option, - ca_cert_paths: Vec, - ) -> PyResult { - 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::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - Ok(()) -}