EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

sglangv0.5.2 & support Qwen3-Next-80B-A3B-Instruct

Browse Source
This commit is contained in:
maxiao1
2025-09-13 17:00:20 +08:00
commit 118f1fc726
2037 changed files with 515371 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
sgl-router/.cargo/config.toml Normal file
View File

@@ -0,0 +1,15 @@
[build]
rustflags = []
incremental = true
[target.aarch64-apple-darwin]
rustflags = [
"-C", "link-arg=-undefined",
"-C", "link-arg=dynamic_lookup",
]
[target.x86_64-apple-darwin]
rustflags = [
"-C", "link-arg=-undefined",
"-C", "link-arg=dynamic_lookup",
]

9
sgl-router/.coveragerc Normal file
View File

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

121
sgl-router/Cargo.toml Normal file
View File

@@ -0,0 +1,121 @@
[package]
name = "sglang_router_rs"
version = "0.0.0"
edition = "2021"
[features]
default = ["grpc-client"]
grpc-client = []
grpc-server = []
[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"]
[[bin]]
name = "sglang-router"
path = "src/main.rs"
[dependencies]
clap = { version = "4", features = ["derive"] }
axum = { version = "0.8.4", features = ["macros", "ws", "tracing"] }
tower = { version = "0.5", features = ["full"] }
tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "limit", "request-id", "util"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
bytes = "1.8.0"
rand = "0.9.2"
reqwest = { version = "0.12.8", features = ["stream", "blocking", "json"] }
futures-util = "0.3"
futures = "0.3"
pyo3 = { version = "0.25.1", features = ["extension-module"] }
dashmap = "6.1.0"
http = "1.1.0"
tokio = { version = "1.42.0", features = ["full"] }
async-trait = "0.1"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "chrono"] }
tracing-log = "0.2"
tracing-appender = "0.2.3"
chrono = "0.4"
kube = { version = "1.1.0", features = ["runtime", "derive"] }
k8s-openapi = { version = "0.25.0", features = ["v1_33"] }
metrics = "0.24.2"
metrics-exporter-prometheus = "0.17.0"
uuid = { version = "1.10", features = ["v4", "serde"] }
thiserror = "2.0.12"
regex = "1.10"
url = "2.5.4"
tokio-stream = { version = "0.1", features = ["sync"] }
anyhow = "1.0"
tokenizers = { version = "0.22.0" }
tiktoken-rs = { version = "0.7.0" }
minijinja = { version = "2.0" }
rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
hf-hub = { version = "0.4.3", features = ["tokio"] }
rmcp = { version = "0.6.3", features = ["client", "server",
"transport-child-process",
"transport-sse-client-reqwest",
"transport-streamable-http-client-reqwest",
"transport-streamable-http-server",
"transport-streamable-http-server-session",
"reqwest",
"auth"] }
serde_yaml = "0.9"
# gRPC and Protobuf dependencies
tonic = { version = "0.12", features = ["tls", "gzip", "transport"] }
prost = "0.13"
prost-types = "0.13"
deadpool = { version = "0.12", features = ["managed", "rt_tokio_1"] }
backoff = { version = "0.4", features = ["tokio"] }
strum = { version = "0.26", features = ["derive"] }
once_cell = "1.21.3"
[build-dependencies]
tonic-build = "0.12"
prost-build = "0.13"
[dev-dependencies]
criterion = { version = "0.5", features = ["html_reports"] }
tower = { version = "0.5", features = ["util"] }
http-body-util = "0.1"
portpicker = "0.1"
tempfile = "3.8"
lazy_static = "1.4"
[[bench]]
name = "request_processing"
harness = false
path = "benches/request_processing.rs"
[[bench]]
name = "tokenizer_benchmark"
harness = false
path = "benches/tokenizer_benchmark.rs"
[[bench]]
name = "tool_parser_benchmark"
harness = false
path = "benches/tool_parser_benchmark.rs"
[profile.release]
lto = "thin"
codegen-units = 1
[profile.dev]
opt-level = 0
debug = true
split-debuginfo = "unpacked"
incremental = true
[profile.dev.build-override]
opt-level = 3
codegen-units = 1
[profile.dev-opt]
inherits = "dev"
opt-level = 1

5
sgl-router/MANIFEST.in Normal file
View File

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

92
sgl-router/Makefile Normal file
View File

@@ -0,0 +1,92 @@
# SGLang Router Makefile
# Provides convenient shortcuts for common development tasks
.PHONY: help bench bench-quick bench-baseline bench-compare test build clean
help: ## Show this help message
@echo "SGLang Router Development Commands"
@echo "=================================="
@echo ""
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-20s\033[0m %s\n", $$1, $$2}'
@echo ""
build: ## Build the project in release mode
@echo "Building SGLang Router..."
@cargo build --release
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
docs: ## Generate and open documentation
@echo "Generating documentation..."
@cargo doc --open
check: ## Run cargo check and clippy
@echo "Running cargo check..."
@cargo check
@echo "Running clippy..."
@cargo clippy
fmt: ## Format code with rustfmt
@echo "Formatting code..."
@cargo fmt
# Development workflow shortcuts
dev-setup: build test ## Set up development environment
@echo "Development environment ready!"
pre-commit: fmt check test bench-quick ## 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

402
sgl-router/README.md Normal file
View File

@@ -0,0 +1,402 @@
# SGLang Router
SGLang router is a standalone Rust module that enables data parallelism across SGLang instances, providing high-performance request routing and advanced load balancing. The router supports multiple load balancing algorithms including cache-aware, power of two, random, and round robin, and acts as a specialized load balancer for prefill-decode disaggregated serving architectures.
## Documentation
- **User Guide**: [docs.sglang.ai/advanced_features/router.html](https://docs.sglang.ai/advanced_features/router.html)
## Quick Start
### Prerequisites
**Rust and Cargo:**
```bash
# Install rustup (Rust installer and version manager)
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
# Follow the installation prompts, then reload your shell
source $HOME/.cargo/env
# Verify installation
rustc --version
cargo --version
```
**Python with pip installed**
### Installation
#### Option A: Build and Install Wheel (Recommended)
```bash
# Install build dependencies
pip install setuptools-rust wheel build
# Build the wheel package
python -m build
# Install the generated wheel
pip install dist/*.whl
# One-liner for development (rebuild + install)
python -m build && pip install --force-reinstall dist/*.whl
```
#### Option B: Development Mode
```bash
# Currently broken
pip install -e .
```
⚠️ **Warning**: Editable installs may suffer performance degradation. Use wheel builds for performance testing.
### Basic Usage
```bash
# Build Rust components
cargo build
```
#### Using the Rust Binary Directly (Alternative to Python)
```bash
# Build the Rust binary
cargo build --release
# Launch router with worker URLs in regular mode
./target/release/sglang-router \
--worker-urls http://worker1:8000 http://worker2:8000
# Or use cargo run
cargo run --release -- \
--worker-urls http://worker1:8000 http://worker2:8000
```
#### Launch Router with Python (Original Method)
```bash
# Launch router with worker URLs
python -m sglang_router.launch_router \
--worker-urls http://worker1:8000 http://worker2:8000
```
#### Launch Router with Worker URLs in prefill-decode mode
```bash
# Note that the prefill and decode URLs must be provided in the following format:
# http://<ip>:<port> for decode nodes
# http://<ip>:<port> bootstrap-port for prefill nodes, where bootstrap-port is optional
# Using Rust binary directly
./target/release/sglang-router \
--pd-disaggregation \
--policy cache_aware \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 9002 \
--prefill http://127.0.0.3:30003 9003 \
--prefill http://127.0.0.4:30004 9004 \
--decode http://127.0.0.5:30005 \
--decode http://127.0.0.6:30006 \
--decode http://127.0.0.7:30007 \
--host 0.0.0.0 \
--port 8080
# Or using Python launcher
python -m sglang_router.launch_router \
--pd-disaggregation \
--policy cache_aware \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 9002 \
--prefill http://127.0.0.3:30003 9003 \
--prefill http://127.0.0.4:30004 9004 \
--decode http://127.0.0.5:30005 \
--decode http://127.0.0.6:30006 \
--decode http://127.0.0.7:30007 \
--host 0.0.0.0 \
--port 8080
````
## Configuration
### Logging
Enable structured logging with optional file output:
```python
from sglang_router import Router
# Console logging (default)
router = Router(worker_urls=["http://worker1:8000", "http://worker2:8000"])
# File logging enabled
router = Router(
worker_urls=["http://worker1:8000", "http://worker2:8000"],
log_dir="./logs" # Daily log files created here
)
```
Set log level with `--log-level` flag ([documentation](https://docs.sglang.ai/backend/server_arguments.html#logging)).
### Metrics
Prometheus metrics endpoint available at `127.0.0.1:29000` by default.
```bash
# Custom metrics configuration
python -m sglang_router.launch_router \
--worker-urls http://localhost:8080 http://localhost:8081 \
--prometheus-host 0.0.0.0 \
--prometheus-port 9000
```
### Retries and Circuit Breakers
- Retries (regular router) are enabled by default with exponential backoff and jitter. You can tune them via CLI:
```bash
python -m sglang_router.launch_router \
--worker-urls http://localhost:8080 http://localhost:8081 \
--retry-max-retries 3 \
--retry-initial-backoff-ms 100 \
--retry-max-backoff-ms 10000 \
--retry-backoff-multiplier 2.0 \
--retry-jitter-factor 0.1
```
- Circuit Breaker defaults protect workers and auto-recover. Tune thresholds/timeouts:
```bash
python -m sglang_router.launch_router \
--worker-urls http://localhost:8080 http://localhost:8081 \
--cb-failure-threshold 5 \
--cb-success-threshold 2 \
--cb-timeout-duration-secs 30 \
--cb-window-duration-secs 60
```
Behavior summary:
- Closed → Open after N consecutive failures (failure-threshold)
- Open → HalfOpen after timeout (timeout-duration-secs)
- HalfOpen → Closed after M consecutive successes (success-threshold)
- Any failure in HalfOpen reopens immediately
Retry predicate (regular router): retry on 408/429/500/502/503/504, otherwise return immediately. Backoff/jitter observed between attempts.
### Request ID Tracking
Track requests across distributed systems with configurable headers:
```bash
# Use custom request ID headers
python -m sglang_router.launch_router \
--worker-urls http://localhost:8080 \
--request-id-headers x-trace-id x-request-id
```
Default headers: `x-request-id`, `x-correlation-id`, `x-trace-id`, `request-id`
## Advanced Features
### Kubernetes Service Discovery
Automatic worker discovery and management in Kubernetes environments.
#### Basic Service Discovery
```bash
python -m sglang_router.launch_router \
--service-discovery \
--selector app=sglang-worker role=inference \
--service-discovery-namespace default
```
#### PD (Prefill-Decode) Mode
For disaggregated prefill/decode routing:
```bash
python -m sglang_router.launch_router \
--pd-disaggregation \
--policy cache_aware \
--service-discovery \
--prefill-selector app=sglang component=prefill \
--decode-selector app=sglang component=decode \
--service-discovery-namespace sglang-system
# With separate routing policies:
python -m sglang_router.launch_router \
--pd-disaggregation \
--prefill-policy cache_aware \
--decode-policy power_of_two \
--service-discovery \
--prefill-selector app=sglang component=prefill \
--decode-selector app=sglang component=decode \
--service-discovery-namespace sglang-system
# in lws case, such as tp16(1 leader pod, 1 worker pod)
python -m sglang_router.launch_router \
--pd-disaggregation \
--policy cache_aware \
--service-discovery \
--prefill-selector app=sglang component=prefill role=leader\
--decode-selector app=sglang component=decode role=leader\
--service-discovery-namespace sglang-system
```
#### Kubernetes Pod Configuration
**Prefill Server Pod:**
```yaml
apiVersion: v1
kind: Pod
metadata:
name: sglang-prefill-1
labels:
app: sglang
component: prefill
annotations:
sglang.ai/bootstrap-port: "9001" # Optional: Bootstrap port
spec:
containers:
- name: sglang
image: lmsys/sglang:latest
ports:
- containerPort: 8000 # Main API port
- containerPort: 9001 # Optional: Bootstrap port
```
**Decode Server Pod:**
```yaml
apiVersion: v1
kind: Pod
metadata:
name: sglang-decode-1
labels:
app: sglang
component: decode
spec:
containers:
- name: sglang
image: lmsys/sglang:latest
ports:
- containerPort: 8000
```
#### RBAC Configuration
**Namespace-scoped (recommended):**
```yaml
apiVersion: v1
kind: ServiceAccount
metadata:
name: sglang-router
namespace: sglang-system
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
namespace: sglang-system
name: sglang-router
rules:
- apiGroups: [""]
resources: ["pods"]
verbs: ["get", "list", "watch"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
name: sglang-router
namespace: sglang-system
subjects:
- kind: ServiceAccount
name: sglang-router
namespace: sglang-system
roleRef:
kind: Role
name: sglang-router
apiGroup: rbac.authorization.k8s.io
```
#### Complete PD Example
```bash
python -m sglang_router.launch_router \
--pd-disaggregation \
--policy cache_aware \
--service-discovery \
--prefill-selector app=sglang component=prefill environment=production \
--decode-selector app=sglang component=decode environment=production \
--service-discovery-namespace production \
--host 0.0.0.0 \
--port 8080 \
--prometheus-host 0.0.0.0 \
--prometheus-port 9090
```
### Command Line Arguments Reference
#### Service Discovery
- `--service-discovery`: Enable Kubernetes service discovery
- `--service-discovery-port`: Port for worker URLs (default: 8000)
- `--service-discovery-namespace`: Kubernetes namespace to watch
- `--selector`: Label selectors for regular mode (format: `key1=value1 key2=value2`)
#### PD Mode
- `--pd-disaggregation`: Enable Prefill-Decode disaggregated mode
- `--prefill`: Initial prefill server (format: `URL BOOTSTRAP_PORT`)
- `--decode`: Initial decode server URL
- `--prefill-selector`: Label selector for prefill pods
- `--decode-selector`: Label selector for decode pods
- `--policy`: Routing policy (`cache_aware`, `random`, `power_of_two`, `round_robin`)
- `--prefill-policy`: Separate routing policy for prefill nodes (optional, overrides `--policy` for prefill)
- `--decode-policy`: Separate routing policy for decode nodes (optional, overrides `--policy` for decode)
## Development
### Build Process
```bash
# Build Rust project
cargo build
# Build Python binding (see Installation section above)
```
**Note**: When modifying Rust code, you must rebuild the wheel for changes to take effect.
### Troubleshooting
**VSCode Rust Analyzer Issues:**
Set `rust-analyzer.linkedProjects` to the absolute path of `Cargo.toml`:
```json
{
"rust-analyzer.linkedProjects": ["/workspaces/sglang/sgl-router/Cargo.toml"]
}
```
### CI/CD Pipeline
The continuous integration pipeline includes comprehensive testing, benchmarking, and publishing:
#### Build & Test
1. **Build Wheels**: Uses `cibuildwheel` for manylinux x86_64 packages
2. **Build Source Distribution**: Creates source distribution for pip fallback
3. **Rust HTTP Server Benchmarking**: Performance testing of router overhead
4. **Basic Inference Testing**: End-to-end validation through the router
5. **PD Disaggregation Testing**: Benchmark and sanity checks for prefill-decode load balancing
#### Publishing
- **PyPI Publishing**: Wheels and source distributions are published only when the version changes in `pyproject.toml`
- **Container Images**: Docker images published using `/docker/Dockerfile.router`
## Features
- **High Performance**: Rust-based routing with connection pooling and optimized request handling
- **Advanced Load Balancing**: Multiple algorithms including:
- **Cache-Aware**: Intelligent routing based on cache locality for optimal performance
- **Power of Two**: Chooses the less loaded of two randomly selected workers
- **Random**: Distributes requests randomly across available workers
- **Round Robin**: Sequential distribution across workers in rotation
- **Prefill-Decode Disaggregation**: Specialized load balancing for separated prefill and decode servers
- **Service Discovery**: Automatic Kubernetes worker discovery and health management
- **Monitoring**: Comprehensive Prometheus metrics and structured logging
- **Scalability**: Handles thousands of concurrent connections with efficient resource utilization

692
sgl-router/benches/request_processing.rs Normal file
View File

@@ -0,0 +1,692 @@
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use serde_json::{from_str, to_string, to_value, to_vec};
use std::time::Instant;
use sglang_router_rs::core::{BasicWorker, Worker, WorkerType};
use sglang_router_rs::protocols::spec::{
ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateParameters, GenerateRequest,
SamplingParams, StringOrArray, UserMessageContent,
};
use sglang_router_rs::routers::http::pd_types::{
generate_room_id, get_hostname, RequestWithBootstrap,
};
fn create_test_worker() -> BasicWorker {
BasicWorker::new(
"http://test-server:8000".to_string(),
WorkerType::Prefill {
bootstrap_port: Some(5678),
},
)
}
// Helper function to get bootstrap info from worker
fn get_bootstrap_info(worker: &BasicWorker) -> (String, Option<u16>) {
let hostname = get_hostname(worker.url());
let bootstrap_port = match worker.worker_type() {
WorkerType::Prefill { bootstrap_port } => bootstrap_port,
_ => None,
};
(hostname, bootstrap_port)
}
/// Create a default GenerateRequest for benchmarks with minimal fields set
fn default_generate_request() -> GenerateRequest {
GenerateRequest {
text: None,
prompt: None,
input_ids: None,
stream: false,
parameters: None,
sampling_params: None,
return_logprob: false,
// SGLang Extensions
lora_path: None,
session_params: None,
return_hidden_states: false,
rid: None,
}
}
/// Create a default ChatCompletionRequest for benchmarks with minimal fields set
fn default_chat_completion_request() -> ChatCompletionRequest {
ChatCompletionRequest {
model: String::new(),
messages: vec![],
max_tokens: None,
max_completion_tokens: None,
temperature: None,
top_p: None,
n: None,
stream: false,
stream_options: None,
stop: None,
presence_penalty: None,
frequency_penalty: None,
logit_bias: None,
logprobs: false,
top_logprobs: None,
user: None,
response_format: None,
seed: None,
tools: None,
tool_choice: None,
parallel_tool_calls: None,
function_call: None,
functions: None,
// SGLang Extensions
top_k: None,
min_p: None,
min_tokens: None,
repetition_penalty: None,
regex: None,
ebnf: None,
stop_token_ids: None,
no_stop_trim: false,
ignore_eos: false,
continue_final_message: false,
skip_special_tokens: true,
// SGLang Extensions
lora_path: None,
session_params: None,
separate_reasoning: true,
stream_reasoning: true,
chat_template_kwargs: None,
return_hidden_states: false,
}
}
/// Create a default CompletionRequest for benchmarks with minimal fields set
fn default_completion_request() -> CompletionRequest {
CompletionRequest {
model: String::new(),
prompt: StringOrArray::String(String::new()),
suffix: None,
max_tokens: None,
temperature: None,
top_p: None,
n: None,
stream: false,
stream_options: None,
logprobs: None,
echo: false,
stop: None,
presence_penalty: None,
frequency_penalty: None,
best_of: None,
logit_bias: None,
user: None,
seed: None,
// SGLang Extensions
top_k: None,
min_p: None,
min_tokens: None,
repetition_penalty: None,
regex: None,
ebnf: None,
json_schema: None,
stop_token_ids: None,
no_stop_trim: false,
ignore_eos: false,
skip_special_tokens: true,
// SGLang Extensions
lora_path: None,
session_params: None,
return_hidden_states: false,
other: serde_json::Map::new(),
}
}
// Sample request data for benchmarks
fn create_sample_generate_request() -> GenerateRequest {
GenerateRequest {
text: Some("Write a story about artificial intelligence".to_string()),
parameters: Some(GenerateParameters {
max_new_tokens: Some(100),
temperature: Some(0.8),
top_p: Some(0.9),
top_k: Some(50),
repetition_penalty: Some(1.0),
..Default::default()
}),
sampling_params: Some(SamplingParams {
temperature: Some(0.8),
top_p: Some(0.9),
top_k: Some(50),
frequency_penalty: Some(0.0),
presence_penalty: Some(0.0),
repetition_penalty: Some(1.0),
..Default::default()
}),
..default_generate_request()
}
}
fn create_sample_chat_completion_request() -> ChatCompletionRequest {
ChatCompletionRequest {
model: "gpt-3.5-turbo".to_string(),
messages: vec![
ChatMessage::System {
role: "system".to_string(),
content: "You are a helpful assistant".to_string(),
name: None,
},
ChatMessage::User {
role: "user".to_string(),
content: UserMessageContent::Text(
"Explain quantum computing in simple terms".to_string(),
),
name: None,
},
],
max_tokens: Some(150),
max_completion_tokens: Some(150),
temperature: Some(0.7),
top_p: Some(1.0),
n: Some(1),
presence_penalty: Some(0.0),
frequency_penalty: Some(0.0),
parallel_tool_calls: Some(true),
..default_chat_completion_request()
}
}
fn create_sample_completion_request() -> CompletionRequest {
CompletionRequest {
model: "text-davinci-003".to_string(),
prompt: StringOrArray::String("Complete this sentence: The future of AI is".to_string()),
max_tokens: Some(50),
temperature: Some(0.8),
top_p: Some(1.0),
n: Some(1),
presence_penalty: Some(0.0),
frequency_penalty: Some(0.0),
best_of: Some(1),
..default_completion_request()
}
}
fn create_large_chat_completion_request() -> ChatCompletionRequest {
let mut messages = vec![ChatMessage::System {
role: "system".to_string(),
content: "You are a helpful assistant with extensive knowledge.".to_string(),
name: None,
}];
// Add many user/assistant pairs to simulate a long conversation
for i in 0..50 {
messages.push(ChatMessage::User {
role: "user".to_string(),
content: UserMessageContent::Text(format!("Question {}: What do you think about topic number {} which involves complex reasoning about multiple interconnected systems and their relationships?", i, i)),
name: None,
});
messages.push(ChatMessage::Assistant {
role: "assistant".to_string(),
content: Some(format!("Answer {}: This is a detailed response about topic {} that covers multiple aspects and provides comprehensive analysis of the interconnected systems you mentioned.", i, i)),
name: None,
tool_calls: None,
function_call: None,
reasoning_content: None,
});
}
ChatCompletionRequest {
model: "gpt-4".to_string(),
messages,
max_tokens: Some(1000),
max_completion_tokens: Some(1000),
temperature: Some(0.7),
top_p: Some(0.95),
n: Some(1),
presence_penalty: Some(0.1),
frequency_penalty: Some(0.1),
top_logprobs: Some(5),
user: Some("benchmark_user".to_string()),
seed: Some(42),
parallel_tool_calls: Some(true),
..default_chat_completion_request()
}
}
// Benchmark JSON serialization
fn bench_json_serialization(c: &mut Criterion) {
let mut group = c.benchmark_group("json_serialization");
let generate_req = create_sample_generate_request();
let chat_req = create_sample_chat_completion_request();
let completion_req = create_sample_completion_request();
let large_chat_req = create_large_chat_completion_request();
group.bench_function("generate_request", |b| {
b.iter(|| {
let json = to_string(black_box(&generate_req)).unwrap();
black_box(json);
});
});
group.bench_function("chat_completion_request", |b| {
b.iter(|| {
let json = to_string(black_box(&chat_req)).unwrap();
black_box(json);
});
});
group.bench_function("completion_request", |b| {
b.iter(|| {
let json = to_string(black_box(&completion_req)).unwrap();
black_box(json);
});
});
group.bench_function("large_chat_completion_request", |b| {
b.iter(|| {
let json = to_string(black_box(&large_chat_req)).unwrap();
black_box(json);
});
});
group.bench_function("generate_request_to_bytes", |b| {
b.iter(|| {
let bytes = to_vec(black_box(&generate_req)).unwrap();
black_box(bytes);
});
});
group.finish();
}
// Benchmark JSON deserialization
fn bench_json_deserialization(c: &mut Criterion) {
let mut group = c.benchmark_group("json_deserialization");
let generate_json = to_string(&create_sample_generate_request()).unwrap();
let chat_json = to_string(&create_sample_chat_completion_request()).unwrap();
let completion_json = to_string(&create_sample_completion_request()).unwrap();
let large_chat_json = to_string(&create_large_chat_completion_request()).unwrap();
group.bench_function("generate_request", |b| {
b.iter(|| {
let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap();
black_box(req);
});
});
group.bench_function("chat_completion_request", |b| {
b.iter(|| {
let req: ChatCompletionRequest = from_str(black_box(&chat_json)).unwrap();
black_box(req);
});
});
group.bench_function("completion_request", |b| {
b.iter(|| {
let req: CompletionRequest = from_str(black_box(&completion_json)).unwrap();
black_box(req);
});
});
group.bench_function("large_chat_completion_request", |b| {
b.iter(|| {
let req: ChatCompletionRequest = from_str(black_box(&large_chat_json)).unwrap();
black_box(req);
});
});
group.finish();
}
// Benchmark bootstrap injection (replaces request adaptation)
fn bench_bootstrap_injection(c: &mut Criterion) {
let mut group = c.benchmark_group("bootstrap_injection");
let generate_req = create_sample_generate_request();
let chat_req = create_sample_chat_completion_request();
let completion_req = create_sample_completion_request();
let large_chat_req = create_large_chat_completion_request();
let worker = create_test_worker();
let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
group.bench_function("generate_bootstrap_injection", |b| {
b.iter(|| {
let request_with_bootstrap = RequestWithBootstrap {
original: &generate_req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let json = to_value(black_box(&request_with_bootstrap)).unwrap();
black_box(json);
});
});
group.bench_function("chat_completion_bootstrap_injection", |b| {
b.iter(|| {
let request_with_bootstrap = RequestWithBootstrap {
original: &chat_req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let json = to_value(black_box(&request_with_bootstrap)).unwrap();
black_box(json);
});
});
group.bench_function("completion_bootstrap_injection", |b| {
b.iter(|| {
let request_with_bootstrap = RequestWithBootstrap {
original: &completion_req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let json = to_value(black_box(&request_with_bootstrap)).unwrap();
black_box(json);
});
});
group.bench_function("large_chat_completion_bootstrap_injection", |b| {
b.iter(|| {
let request_with_bootstrap = RequestWithBootstrap {
original: &large_chat_req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let json = to_value(black_box(&request_with_bootstrap)).unwrap();
black_box(json);
});
});
group.finish();
}
// Benchmark direct JSON routing (replaces regular routing)
fn bench_direct_json_routing(c: &mut Criterion) {
let mut group = c.benchmark_group("direct_json_routing");
let generate_req = create_sample_generate_request();
let chat_req = create_sample_chat_completion_request();
let completion_req = create_sample_completion_request();
group.bench_function("generate_to_json", |b| {
b.iter(|| {
let json = to_value(black_box(&generate_req)).unwrap();
black_box(json);
});
});
group.bench_function("generate_to_json_string", |b| {
b.iter(|| {
let json = to_string(black_box(&generate_req)).unwrap();
black_box(json);
});
});
group.bench_function("generate_to_bytes", |b| {
b.iter(|| {
let bytes = to_vec(black_box(&generate_req)).unwrap();
black_box(bytes);
});
});
group.bench_function("chat_completion_to_json", |b| {
b.iter(|| {
let json = to_value(black_box(&chat_req)).unwrap();
black_box(json);
});
});
group.bench_function("chat_completion_to_json_string", |b| {
b.iter(|| {
let json = to_string(black_box(&chat_req)).unwrap();
black_box(json);
});
});
group.bench_function("completion_to_json", |b| {
b.iter(|| {
let json = to_value(black_box(&completion_req)).unwrap();
black_box(json);
});
});
group.finish();
}
// Benchmark throughput with different request sizes
fn bench_throughput_by_size(c: &mut Criterion) {
let mut group = c.benchmark_group("throughput_by_size");
// Create requests of different sizes
let small_generate = GenerateRequest {
text: Some("Hi".to_string()),
..default_generate_request()
};
let medium_generate = GenerateRequest {
text: Some("Write a medium length story about AI".repeat(10)),
..default_generate_request()
};
let large_generate = GenerateRequest {
text: Some("Write a very long and detailed story about artificial intelligence and its impact on society".repeat(100)),
..default_generate_request()
};
let worker = create_test_worker();
let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
for (name, req) in [
("small", &small_generate),
("medium", &medium_generate),
("large", &large_generate),
] {
let json = to_string(req).unwrap();
let size_bytes = json.len();
let hostname_clone = hostname.clone();
group.throughput(Throughput::Bytes(size_bytes as u64));
group.bench_with_input(BenchmarkId::new("serialize", name), &req, |b, req| {
b.iter(|| {
let json = to_string(black_box(req)).unwrap();
black_box(json);
});
});
group.bench_with_input(
BenchmarkId::new("deserialize", name),
&json,
|b, json_str| {
b.iter(|| {
let req: GenerateRequest = black_box(from_str(json_str)).unwrap();
black_box(req);
});
},
);
group.bench_with_input(
BenchmarkId::new("bootstrap_inject", name),
&req,
move |b, req| {
let hostname = hostname_clone.clone();
b.iter(|| {
let request_with_bootstrap = RequestWithBootstrap {
original: req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let json = to_value(&request_with_bootstrap).unwrap();
black_box(json);
});
},
);
}
group.finish();
}
// Benchmark full round-trip: deserialize -> inject bootstrap -> serialize
fn bench_full_round_trip(c: &mut Criterion) {
let mut group = c.benchmark_group("full_round_trip");
let generate_json = to_string(&create_sample_generate_request()).unwrap();
let chat_json = to_string(&create_sample_chat_completion_request()).unwrap();
let completion_json = to_string(&create_sample_completion_request()).unwrap();
let worker = create_test_worker();
let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
group.bench_function("generate_openai_to_pd_pipeline", |b| {
b.iter(|| {
// Deserialize OpenAI request
let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap();
// Create wrapper with bootstrap fields
let request_with_bootstrap = RequestWithBootstrap {
original: &req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
// Serialize final request
let pd_json = to_string(&request_with_bootstrap).unwrap();
black_box(pd_json);
});
});
group.bench_function("chat_completion_openai_to_pd_pipeline", |b| {
b.iter(|| {
let req: ChatCompletionRequest = from_str(black_box(&chat_json)).unwrap();
let request_with_bootstrap = RequestWithBootstrap {
original: &req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let pd_json = to_string(&request_with_bootstrap).unwrap();
black_box(pd_json);
});
});
group.bench_function("completion_openai_to_pd_pipeline", |b| {
b.iter(|| {
let req: CompletionRequest = from_str(black_box(&completion_json)).unwrap();
let request_with_bootstrap = RequestWithBootstrap {
original: &req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let pd_json = to_string(&request_with_bootstrap).unwrap();
black_box(pd_json);
});
});
group.bench_function("generate_direct_json_pipeline", |b| {
b.iter(|| {
// Deserialize OpenAI request
let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap();
// Convert to JSON for direct routing (no bootstrap injection)
let routing_json = to_value(&req).unwrap();
let json_string = to_string(&routing_json).unwrap();
black_box(json_string);
});
});
group.finish();
}
fn benchmark_summary(c: &mut Criterion) {
let group = c.benchmark_group("benchmark_summary");
println!("\nSGLang Router Performance Benchmark Suite");
println!("=============================================");
// Quick performance overview
let generate_req = create_sample_generate_request();
let worker = create_test_worker();
println!("\nQuick Performance Overview:");
// Measure serialization
let start = Instant::now();
for _ in 0..1000 {
let _ = black_box(to_string(&generate_req).unwrap());
}
let serialize_time = start.elapsed().as_nanos() / 1000;
println!(" * Serialization (avg): {:>8} ns/req", serialize_time);
// Measure deserialization
let json = to_string(&generate_req).unwrap();
let start = Instant::now();
for _ in 0..1000 {
let _: GenerateRequest = black_box(from_str(&json).unwrap());
}
let deserialize_time = start.elapsed().as_nanos() / 1000;
println!(
" * Deserialization (avg): {:>8} ns/req",
deserialize_time
);
// Measure bootstrap injection (replaces adaptation)
let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
let start = Instant::now();
for _ in 0..1000 {
let request_with_bootstrap = RequestWithBootstrap {
original: &generate_req,
bootstrap_host: hostname.clone(),
bootstrap_port,
bootstrap_room: generate_room_id(),
};
let _ = black_box(to_value(&request_with_bootstrap).unwrap());
}
let inject_time = start.elapsed().as_nanos() / 1000;
println!(" * Bootstrap Injection (avg): {:>6} ns/req", inject_time);
// Calculate ratios
let total_pipeline = serialize_time + deserialize_time + inject_time;
println!(" * Total Pipeline (avg): {:>8} ns/req", total_pipeline);
println!("\nPerformance Insights:");
if deserialize_time > serialize_time * 2 {
println!(" • Deserialization is significantly faster than serialization");
}
if inject_time < serialize_time / 10 {
println!(
" • Bootstrap injection overhead is negligible ({:.1}% of serialization)",
(inject_time as f64 / serialize_time as f64) * 100.0
);
}
if total_pipeline < 100_000 {
println!(" • Total pipeline latency is excellent (< 100μs)");
}
println!("\nSimplification Benefits:");
println!(" • Eliminated complex type conversion layer");
println!(" • Reduced memory allocations");
println!(" • Automatic field preservation (no manual mapping)");
println!(" • Direct JSON manipulation improves performance");
println!("\nRecommendations:");
if serialize_time > deserialize_time {
println!(" • Focus optimization efforts on serialization rather than deserialization");
}
println!(" • PD mode overhead is minimal - safe to use for latency-sensitive workloads");
println!(" • Consider batching small requests to improve overall throughput");
println!("\n{}", "=".repeat(50));
group.finish();
}
criterion_group!(
benches,
benchmark_summary,
bench_json_serialization,
bench_json_deserialization,
bench_bootstrap_injection,
bench_direct_json_routing,
bench_throughput_by_size,
bench_full_round_trip
);
criterion_main!(benches);

1400
sgl-router/benches/tokenizer_benchmark.rs Normal file
View File

File diff suppressed because it is too large Load Diff

848
sgl-router/benches/tool_parser_benchmark.rs Normal file
View File

@@ -0,0 +1,848 @@
//! Comprehensive tool parser benchmark for measuring performance under various scenarios
//!
//! This benchmark tests:
//! - Single parser parsing performance
//! - Registry creation overhead
//! - Concurrent parsing with shared parsers
//! - Streaming vs complete parsing
//! - Different model formats (JSON, Mistral, Qwen, Pythonic, etc.)
use criterion::{black_box, criterion_group, BenchmarkId, Criterion, Throughput};
use sglang_router_rs::tool_parser::{
registry::ParserRegistry, state::ParseState, types::StreamResult,
};
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};
use tokio::runtime::Runtime;
// Test data for different parser formats - realistic complex examples
const JSON_SIMPLE: &str = r#"{"name": "code_interpreter", "arguments": "{\"language\": \"python\", \"code\": \"import numpy as np\\nimport matplotlib.pyplot as plt\\n\\n# Generate sample data\\nx = np.linspace(0, 10, 100)\\ny = np.sin(x) * np.exp(-x/10)\\n\\n# Create the plot\\nplt.figure(figsize=(10, 6))\\nplt.plot(x, y, 'b-', linewidth=2)\\nplt.grid(True)\\nplt.xlabel('Time (s)')\\nplt.ylabel('Amplitude')\\nplt.title('Damped Oscillation')\\nplt.show()\"}"}"#;
const JSON_ARRAY: &str = r#"[{"name": "web_search", "arguments": "{\"query\": \"latest developments in quantum computing 2024\", \"num_results\": 10, \"search_type\": \"news\", \"date_range\": \"2024-01-01:2024-12-31\", \"exclude_domains\": [\"reddit.com\", \"facebook.com\"], \"language\": \"en\"}"}, {"name": "analyze_sentiment", "arguments": "{\"text\": \"The breakthrough in quantum error correction represents a significant milestone. Researchers are optimistic about practical applications within the next decade.\", \"granularity\": \"sentence\", \"aspects\": [\"technology\", \"timeline\", \"impact\"], \"confidence_threshold\": 0.85}"}, {"name": "create_summary", "arguments": "{\"content_ids\": [\"doc_1234\", \"doc_5678\", \"doc_9012\"], \"max_length\": 500, \"style\": \"technical\", \"include_citations\": true}"}]"#;
const JSON_WITH_PARAMS: &str = r#"{"name": "database_query", "parameters": {"connection_string": "postgresql://user:pass@localhost:5432/analytics", "query": "SELECT customer_id, COUNT(*) as order_count, SUM(total_amount) as lifetime_value, AVG(order_amount) as avg_order_value FROM orders WHERE created_at >= '2024-01-01' GROUP BY customer_id HAVING COUNT(*) > 5 ORDER BY lifetime_value DESC LIMIT 100", "timeout_ms": 30000, "read_consistency": "strong", "partition_key": "customer_id"}}"#;
const MISTRAL_FORMAT: &str = r#"I'll help you analyze the sales data and create visualizations. Let me start by querying the database and then create some charts.
[TOOL_CALLS] [{"name": "sql_query", "arguments": {"database": "sales_analytics", "query": "WITH monthly_sales AS (SELECT DATE_TRUNC('month', order_date) as month, SUM(total_amount) as revenue, COUNT(DISTINCT customer_id) as unique_customers, COUNT(*) as total_orders FROM orders WHERE order_date >= CURRENT_DATE - INTERVAL '12 months' GROUP BY DATE_TRUNC('month', order_date)) SELECT month, revenue, unique_customers, total_orders, LAG(revenue) OVER (ORDER BY month) as prev_month_revenue, (revenue - LAG(revenue) OVER (ORDER BY month)) / LAG(revenue) OVER (ORDER BY month) * 100 as growth_rate FROM monthly_sales ORDER BY month DESC", "format": "json", "timeout": 60000}}]
Based on the query results, I can see interesting trends in your sales data."#;
const MISTRAL_MULTI: &str = r#"Let me help you with a comprehensive analysis of your application's performance.
[TOOL_CALLS] [{"name": "get_metrics", "arguments": {"service": "api-gateway", "metrics": ["latency_p50", "latency_p95", "latency_p99", "error_rate", "requests_per_second"], "start_time": "2024-01-01T00:00:00Z", "end_time": "2024-01-01T23:59:59Z", "aggregation": "5m", "filters": {"environment": "production", "region": "us-east-1"}}}, {"name": "analyze_logs", "arguments": {"log_group": "/aws/lambda/process-orders", "query": "fields @timestamp, @message, @requestId, duration | filter @message like /ERROR/ | stats count() by bin(@timestamp, 5m) as time_window", "start_time": "2024-01-01T00:00:00Z", "end_time": "2024-01-01T23:59:59Z", "limit": 1000}}, {"name": "get_traces", "arguments": {"service": "order-processing", "operation": "ProcessOrder", "min_duration_ms": 1000, "max_results": 100, "include_downstream": true}}]
Now let me create a comprehensive report based on this data."#;
const QWEN_FORMAT: &str = r#"Let me search for information about machine learning frameworks and their performance benchmarks.
<tool_call>
{"name": "academic_search", "arguments": {"query": "transformer architecture optimization techniques GPU inference latency reduction", "databases": ["arxiv", "ieee", "acm"], "year_range": [2020, 2024], "citation_count_min": 10, "include_code": true, "page_size": 25, "sort_by": "relevance"}}
</tool_call>
I found several interesting papers on optimization techniques."#;
const QWEN_MULTI: &str = r#"I'll help you set up a complete data pipeline for your analytics system.
<tool_call>
{"name": "create_data_pipeline", "arguments": {"name": "customer_analytics_etl", "source": {"type": "kafka", "config": {"bootstrap_servers": "kafka1:9092,kafka2:9092", "topic": "customer_events", "consumer_group": "analytics_consumer", "auto_offset_reset": "earliest"}}, "transformations": [{"type": "filter", "condition": "event_type IN ('purchase', 'signup', 'churn')"}, {"type": "aggregate", "window": "1h", "group_by": ["customer_id", "event_type"], "metrics": ["count", "sum(amount)"]}], "destination": {"type": "bigquery", "dataset": "analytics", "table": "customer_metrics", "write_mode": "append"}}}
</tool_call>
<tool_call>
{"name": "schedule_job", "arguments": {"job_id": "customer_analytics_etl", "schedule": "0 */4 * * *", "timezone": "UTC", "retry_policy": {"max_attempts": 3, "backoff_multiplier": 2, "max_backoff": 3600}, "notifications": {"on_failure": ["ops-team@company.com"], "on_success": null}, "monitoring": {"sla_minutes": 30, "alert_threshold": 0.95}}}
</tool_call>
<tool_call>
{"name": "create_dashboard", "arguments": {"title": "Customer Analytics Dashboard", "widgets": [{"type": "time_series", "title": "Customer Acquisition", "query": "SELECT DATE(timestamp) as date, COUNT(DISTINCT customer_id) as new_customers FROM analytics.customer_metrics WHERE event_type = 'signup' GROUP BY date ORDER BY date", "visualization": "line"}, {"type": "metric", "title": "Total Revenue", "query": "SELECT SUM(amount) as total FROM analytics.customer_metrics WHERE event_type = 'purchase' AND DATE(timestamp) = CURRENT_DATE()", "format": "currency"}, {"type": "table", "title": "Top Customers", "query": "SELECT customer_id, COUNT(*) as purchases, SUM(amount) as total_spent FROM analytics.customer_metrics WHERE event_type = 'purchase' GROUP BY customer_id ORDER BY total_spent DESC LIMIT 10"}], "refresh_interval": 300}}
</tool_call>
The data pipeline has been configured and the dashboard is ready."#;
const LLAMA_FORMAT: &str = r#"<|python_tag|>{"name": "execute_code", "arguments": "{\"code\": \"import pandas as pd\\nimport numpy as np\\nfrom sklearn.model_selection import train_test_split\\nfrom sklearn.ensemble import RandomForestClassifier\\nfrom sklearn.metrics import classification_report, confusion_matrix\\nimport joblib\\n\\n# Load and preprocess data\\ndf = pd.read_csv('/data/customer_churn.csv')\\nprint(f'Dataset shape: {df.shape}')\\nprint(f'Missing values: {df.isnull().sum().sum()}')\\n\\n# Feature engineering\\ndf['tenure_months'] = pd.to_datetime('today') - pd.to_datetime(df['signup_date'])\\ndf['tenure_months'] = df['tenure_months'].dt.days // 30\\ndf['avg_monthly_spend'] = df['total_spend'] / df['tenure_months'].clip(lower=1)\\n\\n# Prepare features and target\\nfeature_cols = ['tenure_months', 'avg_monthly_spend', 'support_tickets', 'product_usage_hours', 'feature_adoption_score']\\nX = df[feature_cols]\\ny = df['churned']\\n\\n# Split and train\\nX_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)\\nrf_model = RandomForestClassifier(n_estimators=100, max_depth=10, min_samples_split=5, random_state=42)\\nrf_model.fit(X_train, y_train)\\n\\n# Evaluate\\ny_pred = rf_model.predict(X_test)\\nprint('Classification Report:')\\nprint(classification_report(y_test, y_pred))\\n\\n# Save model\\njoblib.dump(rf_model, '/models/churn_predictor_v1.pkl')\\nprint('Model saved successfully!')\"}"}"#;
const PYTHONIC_FORMAT: &str = r#"[retrieve_context(query="How do transformer models handle long-range dependencies in natural language processing tasks?", index="ml_knowledge_base", top_k=5, similarity_threshold=0.75, rerank=True, include_metadata=True, filters={"category": "deep_learning", "year": {"$gte": 2020}})]"#;
const PYTHONIC_MULTI: &str = r#"[fetch_api_data(endpoint="https://api.weather.com/v1/forecast", params={"lat": 37.7749, "lon": -122.4194, "units": "metric", "days": 7, "hourly": True}, headers={"API-Key": "${WEATHER_API_KEY}"}, timeout=30, retry_count=3), process_weather_data(data="${response}", extract_fields=["temperature", "humidity", "precipitation", "wind_speed", "uv_index"], aggregation="daily", calculate_trends=True), generate_report(data="${processed_data}", template="weather_forecast", format="html", include_charts=True, language="en")]"#;
const DEEPSEEK_FORMAT: &str = r#"I'll analyze your codebase and identify potential security vulnerabilities.
🤔[{"name": "scan_repository", "arguments": {"repo_path": "/src/application", "scan_types": ["security", "dependencies", "secrets", "code_quality"], "file_patterns": ["*.py", "*.js", "*.java", "*.go"], "exclude_dirs": ["node_modules", ".git", "vendor", "build"], "vulnerability_databases": ["cve", "nvd", "ghsa"], "min_severity": "medium", "check_dependencies": true, "deep_scan": true, "parallel_workers": 8}}]
Let me examine the scan results and provide recommendations."#;
const KIMIK2_FORMAT: &str = r#"⍼validate_and_deploy⍁{"deployment_config": {"application": "payment-service", "version": "2.3.1", "environment": "staging", "region": "us-west-2", "deployment_strategy": "blue_green", "health_check": {"endpoint": "/health", "interval": 30, "timeout": 5, "healthy_threshold": 2, "unhealthy_threshold": 3}, "rollback_on_failure": true, "canary_config": {"percentage": 10, "duration_minutes": 30, "metrics": ["error_rate", "latency_p99", "success_rate"], "thresholds": {"error_rate": 0.01, "latency_p99": 500, "success_rate": 0.99}}, "pre_deployment_hooks": ["run_tests", "security_scan", "backup_database"], "post_deployment_hooks": ["smoke_tests", "notify_team", "update_documentation"]}}"#;
const GLM4_FORMAT: &str = r#"<tool>
analyze_customer_behavior
<parameter>dataset_id=customer_interactions_2024</parameter>
<parameter>analysis_type=cohort_retention</parameter>
<parameter>cohort_definition=signup_month</parameter>
<parameter>retention_periods=[1, 7, 14, 30, 60, 90, 180, 365]</parameter>
<parameter>segment_by=["acquisition_channel", "pricing_tier", "industry", "company_size"]</parameter>
<parameter>metrics=["active_users", "revenue", "feature_usage", "engagement_score"]</parameter>
<parameter>statistical_tests=["chi_square", "anova", "trend_analysis"]</parameter>
<parameter>visualization_types=["heatmap", "line_chart", "funnel", "sankey"]</parameter>
<parameter>export_format=dashboard</parameter>
<parameter>confidence_level=0.95</parameter>
</tool>"#;
const STEP3_FORMAT: &str = r#"<step.tML version="0.1">
<call>
<name>orchestrate_ml_pipeline</name>
<parameters>
<parameter name="pipeline_name">fraud_detection_model_v3</parameter>
<parameter name="data_source">s3://ml-datasets/transactions/2024/</parameter>
<parameter name="preprocessing_steps">
<step order="1" type="clean">{"remove_duplicates": true, "handle_missing": "interpolate", "outlier_method": "isolation_forest"}</step>
<step order="2" type="feature_engineering">{"create_ratios": true, "time_features": ["hour", "day_of_week", "month"], "aggregations": ["mean", "std", "max"]}</step>
<step order="3" type="normalize">{"method": "robust_scaler", "clip_outliers": true}</step>
</parameter>
<parameter name="model_config">{"algorithm": "xgboost", "hyperparameters": {"n_estimators": 500, "max_depth": 8, "learning_rate": 0.01, "subsample": 0.8}, "cross_validation": {"method": "stratified_kfold", "n_splits": 5}}</parameter>
<parameter name="evaluation_metrics">["auc_roc", "precision_recall", "f1", "confusion_matrix"]</parameter>
<parameter name="deployment_target">sagemaker_endpoint</parameter>
<parameter name="monitoring_config">{"drift_detection": true, "performance_threshold": 0.92, "alert_emails": ["ml-team@company.com"]}</parameter>
</parameters>
</call>
</step.tML>"#;
const GPT_OSS_FORMAT: &str = r#"<Channel.vector_search>{"collection": "technical_documentation", "query_embedding": [0.0234, -0.1456, 0.0891, 0.2341, -0.0567, 0.1234, 0.0456, -0.0789, 0.1567, 0.0234, -0.1123, 0.0678, 0.2345, -0.0456, 0.0891, 0.1234, -0.0567, 0.0789, 0.1456, -0.0234, 0.0891, 0.1567, -0.0678, 0.0345, 0.1234, -0.0456, 0.0789, 0.1891, -0.0234, 0.0567, 0.1345, -0.0891], "top_k": 10, "similarity_metric": "cosine", "filters": {"language": "en", "last_updated": {"$gte": "2023-01-01"}, "categories": {"$in": ["api", "sdk", "integration"]}}, "include_metadata": true, "rerank_with_cross_encoder": true}</Channel.vector_search>"#;
// Large test data for stress testing
fn generate_large_json(num_tools: usize) -> String {
let mut tools = Vec::new();
for i in 0..num_tools {
tools.push(format!(
r#"{{"name": "tool_{}", "arguments": {{"param1": "value{}", "param2": {}, "param3": true}}}}"#,
i, i, i
));
}
format!("[{}]", tools.join(", "))
}
// Global results storage
lazy_static::lazy_static! {
static ref BENCHMARK_RESULTS: Mutex<BTreeMap<String, String>> = Mutex::new(BTreeMap::new());
}
fn add_result(category: &str, result: String) {
let mut results = BENCHMARK_RESULTS.lock().unwrap();
let index = results.len();
results.insert(format!("{:03}_{}", index, category), result);
}
fn bench_registry_creation(c: &mut Criterion) {
let mut group = c.benchmark_group("registry_creation");
let printed = Arc::new(AtomicBool::new(false));
group.bench_function("new_registry", |b| {
let printed_clone = printed.clone();
b.iter_custom(|iters| {
let start = Instant::now();
for _ in 0..iters {
let registry = black_box(ParserRegistry::new());
// Force evaluation to prevent optimization
black_box(registry.list_parsers());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let time_per_op = duration.as_micros() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>12.0} | {:>12.1}µs | {:>15}",
"Registry Creation", ops_per_sec, time_per_op, "N/A"
);
add_result("registry", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
group.finish();
}
fn bench_parser_lookup(c: &mut Criterion) {
let registry = Arc::new(ParserRegistry::new());
let models = vec![
"gpt-4",
"mistral-large",
"qwen-72b",
"llama-3.2",
"deepseek-v3",
"unknown-model",
];
let mut group = c.benchmark_group("parser_lookup");
for model in models {
let printed = Arc::new(AtomicBool::new(false));
let registry_clone = registry.clone();
group.bench_function(model, |b| {
let printed_clone = printed.clone();
let registry = registry_clone.clone();
b.iter_custom(|iters| {
let start = Instant::now();
for _ in 0..iters {
let parser = black_box(registry.get_parser(model));
// Force evaluation
black_box(parser.is_some());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let time_per_op = duration.as_nanos() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>12.0} | {:>12.1}ns | {:>15}",
format!("Lookup {}", model),
ops_per_sec,
time_per_op,
if registry.get_parser(model).is_some() {
"Found"
} else {
"Fallback"
}
);
add_result("lookup", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
}
group.finish();
}
fn bench_complete_parsing(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let registry = Arc::new(ParserRegistry::new());
let test_cases = vec![
("json_simple", "json", JSON_SIMPLE),
("json_array", "json", JSON_ARRAY),
("json_params", "json", JSON_WITH_PARAMS),
("mistral_single", "mistral", MISTRAL_FORMAT),
("mistral_multi", "mistral", MISTRAL_MULTI),
("qwen_single", "qwen", QWEN_FORMAT),
("qwen_multi", "qwen", QWEN_MULTI),
("llama", "llama", LLAMA_FORMAT),
("pythonic_single", "pythonic", PYTHONIC_FORMAT),
("pythonic_multi", "pythonic", PYTHONIC_MULTI),
("deepseek", "deepseek", DEEPSEEK_FORMAT),
("kimik2", "kimik2", KIMIK2_FORMAT),
("glm4", "glm4_moe", GLM4_FORMAT),
("step3", "step3", STEP3_FORMAT),
("gpt_oss", "gpt_oss", GPT_OSS_FORMAT),
];
let mut group = c.benchmark_group("complete_parsing");
for (name, parser_name, input) in test_cases {
let printed = Arc::new(AtomicBool::new(false));
let registry_clone = registry.clone();
let input_len = input.len();
group.throughput(Throughput::Bytes(input_len as u64));
group.bench_function(name, |b| {
let printed_clone = printed.clone();
let registry = registry_clone.clone();
let rt = rt.handle().clone();
b.iter_custom(|iters| {
let parser = registry.get_parser(parser_name).expect("Parser not found");
let start = Instant::now();
for _ in 0..iters {
let parser = parser.clone();
let result = rt.block_on(async { parser.parse_complete(input).await });
black_box(result.unwrap());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let bytes_per_sec = (iters as f64 * input_len as f64) / duration.as_secs_f64();
let time_per_op = duration.as_micros() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>10} | {:>12.0} | {:>12.0} | {:>10.1}µs",
name, input_len, ops_per_sec, bytes_per_sec, time_per_op
);
add_result("complete", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
}
group.finish();
}
fn bench_streaming_parsing(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let registry = Arc::new(ParserRegistry::new());
// Streaming test with chunked input
let chunks = vec![
r#"{"na"#,
r#"me": "sear"#,
r#"ch", "argu"#,
r#"ments": {"qu"#,
r#"ery": "rust prog"#,
r#"ramming", "li"#,
r#"mit": 10, "off"#,
r#"set": 0}"#,
r#"}"#,
];
let mut group = c.benchmark_group("streaming_parsing");
let printed = Arc::new(AtomicBool::new(false));
group.bench_function("json_streaming", |b| {
let printed_clone = printed.clone();
let registry = registry.clone();
let rt = rt.handle().clone();
b.iter_custom(|iters| {
let parser = registry.get_parser("json").expect("Parser not found");
let start = Instant::now();
for _ in 0..iters {
let parser = parser.clone();
let mut state = ParseState::new();
let mut complete_tools = Vec::new();
rt.block_on(async {
for chunk in &chunks {
if let StreamResult::ToolComplete(tool) =
parser.parse_incremental(chunk, &mut state).await.unwrap()
{
complete_tools.push(tool);
}
}
});
black_box(complete_tools);
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let time_per_op = duration.as_micros() as f64 / iters as f64;
let chunks_per_sec = (iters as f64 * chunks.len() as f64) / duration.as_secs_f64();
let result = format!(
"{:<25} | {:>10} | {:>12.0} | {:>12.0} | {:>10.1}µs",
"JSON Streaming",
chunks.len(),
ops_per_sec,
chunks_per_sec,
time_per_op
);
add_result("streaming", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
group.finish();
}
fn bench_concurrent_parsing(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let registry = Arc::new(ParserRegistry::new());
let parser = registry.get_parser("json").expect("Parser not found");
let thread_counts = vec![1, 2, 4, 8, 16, 32];
let operations_per_thread = 100;
let mut group = c.benchmark_group("concurrent_parsing");
group.measurement_time(Duration::from_secs(3));
for num_threads in thread_counts {
let printed = Arc::new(AtomicBool::new(false));
let parser_clone = parser.clone();
group.bench_with_input(
BenchmarkId::from_parameter(num_threads),
&num_threads,
|b, &threads| {
let printed_clone = printed.clone();
let parser = parser_clone.clone();
let rt = rt.handle().clone();
b.iter_custom(|_iters| {
let total_operations = Arc::new(AtomicU64::new(0));
let total_parsed = Arc::new(AtomicU64::new(0));
let start = Instant::now();
let handles: Vec<_> = (0..threads)
.map(|_thread_id| {
let parser = parser.clone();
let total_ops = total_operations.clone();
let total_p = total_parsed.clone();
let rt = rt.clone();
thread::spawn(move || {
let test_inputs = [JSON_SIMPLE, JSON_ARRAY, JSON_WITH_PARAMS];
for i in 0..operations_per_thread {
let input = test_inputs[i % test_inputs.len()];
let result =
rt.block_on(async { parser.parse_complete(input).await });
if let Ok(tools) = result {
total_p.fetch_add(tools.len() as u64, Ordering::Relaxed);
}
}
total_ops
.fetch_add(operations_per_thread as u64, Ordering::Relaxed);
})
})
.collect();
for handle in handles {
handle.join().unwrap();
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let total_ops = total_operations.load(Ordering::Relaxed);
let total_p = total_parsed.load(Ordering::Relaxed);
let ops_per_sec = total_ops as f64 / duration.as_secs_f64();
let tools_per_sec = total_p as f64 / duration.as_secs_f64();
let result = format!(
"{:<25} | {:>10} | {:>12.0} | {:>12.0} | {:>10}",
format!("{}_threads", threads),
total_ops,
ops_per_sec,
tools_per_sec,
threads
);
add_result("concurrent", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
},
);
}
group.finish();
}
fn bench_large_payloads(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let registry = Arc::new(ParserRegistry::new());
let parser = registry.get_parser("json").expect("Parser not found");
let sizes = vec![1, 10, 50, 100, 500];
let mut group = c.benchmark_group("large_payloads");
for size in sizes {
let large_json = generate_large_json(size);
let input_len = large_json.len();
let printed = Arc::new(AtomicBool::new(false));
let parser_clone = parser.clone();
group.throughput(Throughput::Bytes(input_len as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &num_tools| {
let printed_clone = printed.clone();
let parser = parser_clone.clone();
let rt = rt.handle().clone();
let input = &large_json;
b.iter_custom(|iters| {
let start = Instant::now();
for _ in 0..iters {
let parser = parser.clone();
let result = rt.block_on(async { parser.parse_complete(input).await });
black_box(result.unwrap());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let bytes_per_sec = (iters as f64 * input_len as f64) / duration.as_secs_f64();
let time_per_op = duration.as_millis() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>10} | {:>10} | {:>12.0} | {:>12.0} | {:>10.1}ms",
format!("{}_tools", num_tools),
num_tools,
input_len,
ops_per_sec,
bytes_per_sec,
time_per_op
);
add_result("large", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
}
group.finish();
}
fn bench_parser_reuse(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("parser_reuse");
// Benchmark creating new registry each time
let printed_new = Arc::new(AtomicBool::new(false));
group.bench_function("new_registry_each_time", |b| {
let printed_clone = printed_new.clone();
let rt = rt.handle().clone();
b.iter_custom(|iters| {
let start = Instant::now();
for _ in 0..iters {
let registry = ParserRegistry::new();
let parser = registry.get_parser("json").unwrap();
let result = rt.block_on(async { parser.parse_complete(JSON_SIMPLE).await });
black_box(result.unwrap());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let time_per_op = duration.as_micros() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>12.0} | {:>12.1}µs | {:>15}",
"New Registry Each Time", ops_per_sec, time_per_op, "Baseline"
);
add_result("reuse", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
// Benchmark reusing registry
let printed_reuse = Arc::new(AtomicBool::new(false));
let shared_registry = Arc::new(ParserRegistry::new());
group.bench_function("reuse_registry", |b| {
let printed_clone = printed_reuse.clone();
let registry = shared_registry.clone();
let rt = rt.handle().clone();
b.iter_custom(|iters| {
let parser = registry.get_parser("json").unwrap();
let start = Instant::now();
for _ in 0..iters {
let parser = parser.clone();
let result = rt.block_on(async { parser.parse_complete(JSON_SIMPLE).await });
black_box(result.unwrap());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let time_per_op = duration.as_micros() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>12.0} | {:>12.1}µs | {:>15}",
"Reuse Registry", ops_per_sec, time_per_op, "Optimized"
);
add_result("reuse", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
// Benchmark reusing parser
let printed_parser = Arc::new(AtomicBool::new(false));
let shared_parser = shared_registry.get_parser("json").unwrap();
group.bench_function("reuse_parser", |b| {
let printed_clone = printed_parser.clone();
let parser = shared_parser.clone();
let rt = rt.handle().clone();
b.iter_custom(|iters| {
let start = Instant::now();
for _ in 0..iters {
let parser = parser.clone();
let result = rt.block_on(async { parser.parse_complete(JSON_SIMPLE).await });
black_box(result.unwrap());
}
let duration = start.elapsed();
if !printed_clone.load(Ordering::Relaxed) {
let ops_per_sec = iters as f64 / duration.as_secs_f64();
let time_per_op = duration.as_micros() as f64 / iters as f64;
let result = format!(
"{:<25} | {:>12.0} | {:>12.1}µs | {:>15}",
"Reuse Parser", ops_per_sec, time_per_op, "Best"
);
add_result("reuse", result);
printed_clone.store(true, Ordering::Relaxed);
}
duration
});
});
group.finish();
}
fn bench_latency_distribution(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let registry = Arc::new(ParserRegistry::new());
let test_cases = vec![
("json", JSON_SIMPLE),
("mistral", MISTRAL_FORMAT),
("qwen", QWEN_FORMAT),
("pythonic", PYTHONIC_FORMAT),
];
let mut group = c.benchmark_group("latency");
for (parser_name, input) in test_cases {
let printed = Arc::new(AtomicBool::new(false));
let registry_clone = registry.clone();
group.bench_function(parser_name, |b| {
let printed_clone = printed.clone();
let registry = registry_clone.clone();
let rt = rt.handle().clone();
b.iter_custom(|iters| {
let parser = registry.get_parser(parser_name).expect("Parser not found");
let total_duration = if !printed_clone.load(Ordering::Relaxed) {
let mut latencies = Vec::new();
// Warm up
for _ in 0..100 {
let parser = parser.clone();
rt.block_on(async { parser.parse_complete(input).await })
.unwrap();
}
// Measure for statistics
for _ in 0..1000 {
let parser = parser.clone();
let start = Instant::now();
rt.block_on(async { parser.parse_complete(input).await })
.unwrap();
let latency = start.elapsed();
latencies.push(latency);
}
latencies.sort();
let p50 = latencies[latencies.len() / 2];
let p95 = latencies[latencies.len() * 95 / 100];
let p99 = latencies[latencies.len() * 99 / 100];
let max = latencies.last().unwrap();
let result = format!(
"{:<25} | {:>10.1} | {:>10.1} | {:>10.1} | {:>10.1} | {:>10}",
parser_name,
p50.as_micros() as f64,
p95.as_micros() as f64,
p99.as_micros() as f64,
max.as_micros() as f64,
1000
);
add_result("latency", result);
printed_clone.store(true, Ordering::Relaxed);
// Return median for consistency
p50 * iters as u32
} else {
// Regular benchmark iterations
let start = Instant::now();
for _ in 0..iters {
let parser = parser.clone();
rt.block_on(async { parser.parse_complete(input).await })
.unwrap();
}
start.elapsed()
};
total_duration
});
});
}
group.finish();
}
// Print final summary table
fn print_summary() {
println!("\n{}", "=".repeat(120));
println!("TOOL PARSER BENCHMARK SUMMARY");
println!("{}", "=".repeat(120));
let results = BENCHMARK_RESULTS.lock().unwrap();
let mut current_category = String::new();
for (key, value) in results.iter() {
let category = key.split('_').skip(1).collect::<Vec<_>>().join("_");
if category != current_category {
current_category = category.clone();
// Print section header based on category
println!("\n{}", "-".repeat(120));
match category.as_str() {
"registry" => {
println!("REGISTRY OPERATIONS");
println!(
"{:<25} | {:>12} | {:>12} | {:>15}",
"Operation", "Ops/sec", "Time/op", "Notes"
);
}
"lookup" => {
println!("PARSER LOOKUP PERFORMANCE");
println!(
"{:<25} | {:>12} | {:>12} | {:>15}",
"Model", "Lookups/sec", "Time/lookup", "Result"
);
}
"complete" => {
println!("COMPLETE PARSING PERFORMANCE");
println!(
"{:<25} | {:>10} | {:>12} | {:>12} | {:>12}",
"Parser Format", "Size(B)", "Ops/sec", "Bytes/sec", "Time/op"
);
}
"streaming" => {
println!("STREAMING PARSING PERFORMANCE");
println!(
"{:<25} | {:>10} | {:>12} | {:>12} | {:>12}",
"Parser", "Chunks", "Ops/sec", "Chunks/sec", "Time/op"
);
}
"concurrent" => {
println!("CONCURRENT PARSING");
println!(
"{:<25} | {:>10} | {:>12} | {:>12} | {:>10}",
"Configuration", "Total Ops", "Ops/sec", "Tools/sec", "Threads"
);
}
"large" => {
println!("LARGE PAYLOAD PARSING");
println!(
"{:<25} | {:>10} | {:>10} | {:>12} | {:>12} | {:>12}",
"Payload", "Tools", "Size(B)", "Ops/sec", "Bytes/sec", "Time/op"
);
}
"reuse" => {
println!("PARSER REUSE COMPARISON");
println!(
"{:<25} | {:>12} | {:>12} | {:>15}",
"Strategy", "Ops/sec", "Time/op", "Performance"
);
}
"latency" => {
println!("LATENCY DISTRIBUTION");
println!(
"{:<25} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10}",
"Parser", "P50(µs)", "P95(µs)", "P99(µs)", "Max(µs)", "Samples"
);
}
_ => {}
}
println!("{}", "-".repeat(120));
}
println!("{}", value);
}
println!("\n{}", "=".repeat(120));
// Print performance analysis
println!("\nPERFORMANCE ANALYSIS:");
println!("{}", "-".repeat(120));
// Calculate and display key metrics
if let Some(new_registry) = results.get("007_reuse") {
if let Some(reuse_parser) = results.get("009_reuse") {
// Extract ops/sec values
let new_ops: f64 = new_registry
.split('|')
.nth(1)
.and_then(|s| s.trim().parse().ok())
.unwrap_or(0.0);
let reuse_ops: f64 = reuse_parser
.split('|')
.nth(1)
.and_then(|s| s.trim().parse().ok())
.unwrap_or(0.0);
if new_ops > 0.0 && reuse_ops > 0.0 {
let improvement = (reuse_ops / new_ops - 1.0) * 100.0;
println!("Parser Reuse Improvement: {:.1}% faster", improvement);
if improvement < 100.0 {
println!("⚠️ WARNING: Parser reuse improvement is lower than expected!");
println!(" Expected: >100% improvement with singleton pattern");
println!(" Actual: {:.1}% improvement", improvement);
println!(" Recommendation: Implement global singleton registry");
}
}
}
}
println!("{}", "=".repeat(120));
}
fn run_benchmarks(c: &mut Criterion) {
bench_registry_creation(c);
bench_parser_lookup(c);
bench_complete_parsing(c);
bench_streaming_parsing(c);
bench_concurrent_parsing(c);
bench_large_payloads(c);
bench_parser_reuse(c);
bench_latency_distribution(c);
// Print summary at the end
print_summary();
}
criterion_group!(benches, run_benchmarks);
criterion::criterion_main!(benches);

35
sgl-router/build.rs Normal file
View File

@@ -0,0 +1,35 @@
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Only regenerate if the proto file changes
println!("cargo:rerun-if-changed=src/proto/sglang_scheduler.proto");
// Configure protobuf compilation with custom settings
let config = prost_build::Config::new();
// Skip serde for types that use prost_types::Struct
// These cause conflicts and we don't need serde for all generated types
// Configure tonic-build for gRPC code generation
tonic_build::configure()
// Generate both client and server code
.build_server(true)
.build_client(true)
// Add a module-level attribute for documentation and clippy warnings
.server_mod_attribute(
"sglang.grpc.scheduler",
"#[allow(unused, clippy::mixed_attributes_style)]",
)
.client_mod_attribute(
"sglang.grpc.scheduler",
"#[allow(unused, clippy::mixed_attributes_style)]",
)
// Compile the proto file with the custom config
.compile_protos_with_config(
config,
&["src/proto/sglang_scheduler.proto"],
&["src/proto"],
)?;
println!("cargo:warning=Protobuf compilation completed successfully");
Ok(())
}

3
sgl-router/py_src/sglang_router/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
from sglang_router.version import __version__
__all__ = ["__version__"]

109
sgl-router/py_src/sglang_router/launch_router.py Normal file
View File

@@ -0,0 +1,109 @@
import argparse
import logging
import sys
from typing import List, Optional
import setproctitle
from sglang_router.mini_lb import MiniLoadBalancer
from sglang_router.router_args import RouterArgs
logger = logging.getLogger("router")
try:
from sglang_router.router import Router
except ImportError:
Router = None
logger.warning(
"Rust Router is not installed, only python MiniLB (debugging only) is available"
)
def launch_router(args: argparse.Namespace) -> Optional[Router]:
"""
Launch the SGLang router with the configuration from parsed arguments.
Args:
args: Namespace object containing router configuration
Can be either raw argparse.Namespace or converted RouterArgs
Returns:
Router instance if successful, None if failed
"""
setproctitle.setproctitle("sglang::router")
try:
# Convert to RouterArgs if needed
if not isinstance(args, RouterArgs):
router_args = RouterArgs.from_cli_args(args)
else:
router_args = args
if router_args.mini_lb:
mini_lb = MiniLoadBalancer(router_args)
mini_lb.start()
else:
if Router is None:
raise RuntimeError("Rust Router is not installed")
router_args._validate_router_args()
router = Router.from_args(router_args)
router.start()
except Exception as e:
logger.error(f"Error starting router: {e}")
raise e
class CustomHelpFormatter(
argparse.RawDescriptionHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
):
"""Custom formatter that preserves both description formatting and shows defaults"""
pass
def parse_router_args(args: List[str]) -> RouterArgs:
"""Parse command line arguments and return RouterArgs instance."""
parser = argparse.ArgumentParser(
description="""SGLang Router - High-performance request distribution across worker nodes
Usage:
This launcher enables starting a router with individual worker instances. It is useful for
multi-node setups or when you want to start workers and router separately.
Examples:
# Regular mode
python -m sglang_router.launch_router --worker-urls http://worker1:8000 http://worker2:8000
# PD disaggregated mode with same policy for both
python -m sglang_router.launch_router --pd-disaggregation \\
--prefill http://prefill1:8000 9000 --prefill http://prefill2:8000 \\
--decode http://decode1:8001 --decode http://decode2:8001 \\
--policy cache_aware
# PD mode with optional bootstrap ports
python -m sglang_router.launch_router --pd-disaggregation \\
--prefill http://prefill1:8000 9000 \\ # With bootstrap port
--prefill http://prefill2:8000 none \\ # Explicitly no bootstrap port
--prefill http://prefill3:8000 \\ # Defaults to no bootstrap port
--decode http://decode1:8001 --decode http://decode2:8001
# PD mode with different policies for prefill and decode
python -m sglang_router.launch_router --pd-disaggregation \\
--prefill http://prefill1:8000 --prefill http://prefill2:8000 \\
--decode http://decode1:8001 --decode http://decode2:8001 \\
--prefill-policy cache_aware --decode-policy power_of_two
""",
formatter_class=CustomHelpFormatter,
)
RouterArgs.add_cli_args(parser, use_router_prefix=False)
return RouterArgs.from_cli_args(parser.parse_args(args), use_router_prefix=False)
def main() -> None:
router_args = parse_router_args(sys.argv[1:])
launch_router(router_args)
if __name__ == "__main__":
main()

203
sgl-router/py_src/sglang_router/launch_server.py Normal file
View File

@@ -0,0 +1,203 @@
import argparse
import copy
import logging
import multiprocessing as mp
import os
import random
import signal
import sys
import time
from typing import List
import requests
from setproctitle import setproctitle
from sglang_router.launch_router import RouterArgs, launch_router
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_port_available
def setup_logger():
logger = logging.getLogger("router")
logger.setLevel(logging.INFO)
formatter = logging.Formatter(
"[Router (Python)] %(asctime)s - %(levelname)s - %(message)s - %(filename)s:%(lineno)d",
datefmt="%Y-%m-%d %H:%M:%S",
)
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)
return logger
logger = setup_logger()
# Create new process group
def run_server(server_args, dp_rank):
"""
Note:
1. Without os.setpgrp(), all processes share the same PGID. When you press Ctrl+C, the terminal sends SIGINT to all processes in the group simultaneously.
This can cause leaf processes to terminate first, which messes up the cleaning order and produces orphaned processes.
Terminal (PGID=100)
└── Main Python Process (PGID=100)
└── Server Process 1 (PGID=100)
└── Scheduler 1
└── Detokenizer 1
└── Server Process 2 (PGID=100)
└── Scheduler 2
└── Detokenizer 2
2. With os.setpgrp(), the main Python process and its children are in a separate group. Now:
Terminal (PGID=100)
└── Main Python Process (PGID=200)
└── Server Process 1 (PGID=300)
└── Scheduler 1
└── Detokenizer 1
└── Server Process 2 (PGID=400)
└── Scheduler 2
└── Detokenizer 2
"""
# create new process group
os.setpgrp()
setproctitle("sglang::server")
# Set SGLANG_DP_RANK environment variable
os.environ["SGLANG_DP_RANK"] = str(dp_rank)
launch_server(server_args)
def launch_server_process(
server_args: ServerArgs, worker_port: int, dp_id: int
) -> mp.Process:
"""Launch a single server process with the given args and port."""
server_args = copy.deepcopy(server_args)
server_args.port = worker_port
server_args.base_gpu_id = dp_id * server_args.tp_size
server_args.dp_size = 1
proc = mp.Process(target=run_server, args=(server_args, dp_id))
proc.start()
return proc
def wait_for_server_health(host: str, port: int, timeout: int = 300) -> bool:
"""Wait for server to be healthy by checking /health endpoint."""
start_time = time.perf_counter()
url = f"http://{host}:{port}/health"
while time.perf_counter() - start_time < timeout:
try:
response = requests.get(url, timeout=5)
if response.status_code == 200:
return True
except requests.exceptions.RequestException:
pass
time.sleep(1)
return False
def find_available_ports(base_port: int, count: int) -> List[int]:
"""Find consecutive available ports starting from base_port."""
available_ports = []
current_port = base_port
while len(available_ports) < count:
if is_port_available(current_port):
available_ports.append(current_port)
current_port += random.randint(100, 1000)
return available_ports
def cleanup_processes(processes: List[mp.Process]):
for process in processes:
logger.info(f"Terminating process group {process.pid}")
try:
os.killpg(process.pid, signal.SIGTERM)
except ProcessLookupError:
# Process group may already be terminated
pass
# Wait for processes to terminate
for process in processes:
process.join(timeout=5)
if process.is_alive():
logger.warning(
f"Process {process.pid} did not terminate gracefully, forcing kill"
)
try:
os.killpg(process.pid, signal.SIGKILL)
except ProcessLookupError:
pass
logger.info("All process groups terminated")
def main():
# CUDA runtime isn't fork-safe, which can lead to subtle bugs or crashes
mp.set_start_method("spawn")
parser = argparse.ArgumentParser(
description="Launch SGLang router and server processes"
)
ServerArgs.add_cli_args(parser)
RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
parser.add_argument(
"--router-dp-worker-base-port",
type=int,
default=31000,
help="Base port number for data parallel workers",
)
# No extra retry/CB flags here; RouterArgs.add_cli_args already defines them with router- prefix
args = parser.parse_args()
server_args = ServerArgs.from_cli_args(args)
router_args = RouterArgs.from_cli_args(args, use_router_prefix=True)
# Find available ports for workers
worker_ports = find_available_ports(
args.router_dp_worker_base_port, server_args.dp_size
)
# Start server processes
server_processes = []
for i, worker_port in enumerate(worker_ports):
logger.info(f"Launching DP server process {i} on port {worker_port}")
proc = launch_server_process(server_args, worker_port, i)
server_processes.append(proc)
signal.signal(signal.SIGINT, lambda sig, frame: cleanup_processes(server_processes))
signal.signal(
signal.SIGTERM, lambda sig, frame: cleanup_processes(server_processes)
)
signal.signal(
signal.SIGQUIT, lambda sig, frame: cleanup_processes(server_processes)
)
# Update router args with worker URLs
router_args.worker_urls = [
f"http://{server_args.host}:{port}" for port in worker_ports
]
# Start the router
try:
launch_router(router_args)
except Exception as e:
logger.error(f"Failed to start router: {e}")
cleanup_processes(server_processes)
sys.exit(1)
if __name__ == "__main__":
main()

395
sgl-router/py_src/sglang_router/mini_lb.py Normal file
View File

@@ -0,0 +1,395 @@
"""
Minimal HTTP load balancer for prefill and decode servers for testing.
"""
import asyncio
import ipaddress
import logging
import random
import urllib
from http import HTTPStatus
from itertools import chain
from typing import Optional
import aiohttp
import orjson
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
from sglang_router.router_args import RouterArgs
logger = logging.getLogger(__name__)
AIOHTTP_STREAM_READ_CHUNK_SIZE = (
1024 * 64
) # 64KB, to prevent aiohttp's "Chunk too big" error
def maybe_wrap_ipv6_address(address: str) -> str:
try:
ipaddress.IPv6Address(address)
return f"[{address}]"
except ValueError:
return address
class MiniLoadBalancer:
def __init__(
self,
router_args: RouterArgs,
):
self._validate_router_args(router_args)
self.host = router_args.host
self.port = router_args.port
self.timeout = router_args.request_timeout_secs
self.prefill_urls = [url[0] for url in router_args.prefill_urls]
self.prefill_bootstrap_ports = [url[1] for url in router_args.prefill_urls]
self.decode_urls = router_args.decode_urls
def _validate_router_args(self, router_args: RouterArgs):
logger.warning(
"\x1b[33mMiniLB is only for debugging purposes, it only supports random policy!\033[0m"
)
# NOTE: too many arguments unsupported, just validate some important ones
if router_args.policy != "random":
logger.warning("[MiniLB] Overriding policy to random")
router_args.policy = "random"
if not router_args.pd_disaggregation:
raise ValueError("MiniLB only supports PD disaggregation mode")
if len(router_args.prefill_urls) == 0 or len(router_args.decode_urls) == 0:
raise ValueError(
"MiniLB requires at least one prefill and one decode server"
)
def start(self):
global lb
lb = self
uvicorn.run(app, host=self.host, port=self.port)
def select_pair(self):
assert len(self.prefill_urls) > 0, "No prefill servers available"
assert len(self.decode_urls) > 0, "No decode servers available"
pidx = random.randint(0, len(self.prefill_urls) - 1)
didx = random.randint(0, len(self.decode_urls) - 1)
return (
self.prefill_urls[pidx],
self.prefill_bootstrap_ports[pidx],
self.decode_urls[didx],
)
async def generate(
self, modified_request, prefill_server, decode_server, endpoint
) -> ORJSONResponse:
assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
async with aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(
total=self.timeout
) # Add timeout for request reliability
) as session:
tasks = [
session.post(f"{prefill_server}/{endpoint}", json=modified_request),
session.post(f"{decode_server}/{endpoint}", json=modified_request),
]
# Wait for both responses to complete. Prefill should end first.
prefill_response, decode_response = await asyncio.gather(*tasks)
if "return_logprob" in modified_request:
prefill_json = await prefill_response.json()
ret_json = await decode_response.json()
# merge `meta_info.input_token_logprobs` from prefill to decode
if "meta_info" in ret_json:
if "input_token_logprobs" in ret_json["meta_info"]:
ret_json["meta_info"]["input_token_logprobs"] = (
prefill_json["meta_info"]["input_token_logprobs"]
+ ret_json["meta_info"]["input_token_logprobs"]
)
else:
ret_json = await decode_response.json()
return ORJSONResponse(
content=ret_json,
status_code=decode_response.status,
)
async def generate_stream(
self, modified_request, prefill_server, decode_server, endpoint="generate"
):
assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
async def stream_results():
async with aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(
total=self.timeout
) # Add timeout for request reliability
) as session:
# Create the tasks for both prefill and decode requests
tasks = [
session.post(f"{prefill_server}/{endpoint}", json=modified_request),
session.post(f"{decode_server}/{endpoint}", json=modified_request),
]
# Wait for both responses to complete. Since this is streaming, they return immediately.
prefill_response, decode_response = await asyncio.gather(*tasks)
if modified_request.get("return_logprob", False):
prefill_chunks = []
async for chunk in prefill_response.content:
prefill_chunks.append(chunk)
first_prefill_chunk = (
prefill_chunks[0].decode("utf-8")[5:].strip("\n")
)
first_prefill_chunk_json = orjson.loads(first_prefill_chunk)
async for chunk in decode_response.content:
# Note: This is inefficient
# merge prefill input_token_logprobs, output_token_logprobs to decode
decoded_chunk = chunk.decode("utf-8")
if (
decoded_chunk
and decoded_chunk.startswith("data:")
and "[DONE]" not in decoded_chunk
):
ret_json = orjson.loads(decoded_chunk[5:].strip("\n"))
ret_json["meta_info"]["input_token_logprobs"] = (
first_prefill_chunk_json["meta_info"][
"input_token_logprobs"
]
+ ret_json["meta_info"]["input_token_logprobs"]
)
yield b"data: " + orjson.dumps(ret_json) + b"\n\n"
else:
yield chunk
else:
async for chunk in decode_response.content.iter_chunked(
AIOHTTP_STREAM_READ_CHUNK_SIZE
):
yield chunk
return StreamingResponse(
stream_results(),
media_type="text/event-stream",
)
app = FastAPI()
lb: Optional[MiniLoadBalancer] = None
@app.get("/health")
async def health_check():
return Response(status_code=200)
@app.get("/health_generate")
async def health_generate():
async with aiohttp.ClientSession() as session:
# Create the tasks
tasks = []
for server in chain(lb.prefill_urls, lb.decode_urls):
tasks.append(session.get(f"{server}/health_generate"))
for i, response in enumerate(asyncio.as_completed(tasks)):
await response
return Response(status_code=200)
@app.post("/flush_cache")
async def flush_cache():
async with aiohttp.ClientSession() as session:
# Create the tasks
tasks = []
for server in chain(lb.prefill_urls, lb.decode_urls):
tasks.append(session.post(f"{server}/flush_cache"))
for i, response in enumerate(asyncio.as_completed(tasks)):
await response
return Response(status_code=200)
@app.get("/get_server_info")
async def get_server_info():
prefill_infos = []
decode_infos = []
all_internal_states = []
async with aiohttp.ClientSession() as session:
for server in lb.prefill_urls:
server_info = await session.get(f"{server}/get_server_info")
prefill_infos.append(await server_info.json())
for server in lb.decode_urls:
server_info = await session.get(f"{server}/get_server_info")
info_json = await server_info.json()
decode_infos.append(info_json)
# Extract internal_states from decode servers
if "internal_states" in info_json:
all_internal_states.extend(info_json["internal_states"])
# Return format expected by bench_one_batch_server.py
if all_internal_states:
return {
"internal_states": all_internal_states,
"prefill": prefill_infos,
"decode": decode_infos,
}
else:
# Fallback with dummy data if no internal states found
return {
"internal_states": [
{
"last_gen_throughput": 0.0,
"avg_spec_accept_length": None,
}
],
"prefill": prefill_infos,
"decode": decode_infos,
}
@app.get("/get_model_info")
async def get_model_info():
if not lb or not lb.prefill_urls:
raise HTTPException(
status_code=HTTPStatus.SERVICE_UNAVAILABLE,
detail="There is no server registered",
)
target_server_url = lb.prefill_urls[0]
endpoint_url = f"{target_server_url}/get_model_info"
async with aiohttp.ClientSession() as session:
try:
async with session.get(endpoint_url) as response:
if response.status != 200:
error_text = await response.text()
raise HTTPException(
status_code=HTTPStatus.BAD_GATEWAY,
detail=(
f"Failed to get model info from {target_server_url}"
f"Status: {response.status}, Response: {error_text}"
),
)
model_info_json = await response.json()
return ORJSONResponse(content=model_info_json)
except aiohttp.ClientError as e:
raise HTTPException(
status_code=HTTPStatus.SERVICE_UNAVAILABLE,
detail=f"Failed to get model info from backend",
)
@app.post("/generate")
async def handle_generate_request(request_data: dict):
prefill_server, bootstrap_port, decode_server = lb.select_pair()
# Parse and transform prefill_server for bootstrap data
parsed_url = urllib.parse.urlparse(prefill_server)
hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
modified_request = request_data.copy()
batch_size = _get_request_batch_size(modified_request)
if batch_size is not None:
modified_request.update(
{
"bootstrap_host": [hostname] * batch_size,
"bootstrap_port": [bootstrap_port] * batch_size,
"bootstrap_room": [
_generate_bootstrap_room() for _ in range(batch_size)
],
}
)
else:
modified_request.update(
{
"bootstrap_host": hostname,
"bootstrap_port": bootstrap_port,
"bootstrap_room": _generate_bootstrap_room(),
}
)
if request_data.get("stream", False):
return await lb.generate_stream(
modified_request, prefill_server, decode_server, "generate"
)
else:
return await lb.generate(
modified_request, prefill_server, decode_server, "generate"
)
async def _forward_to_backend(request_data: dict, endpoint_name: str):
prefill_server, bootstrap_port, decode_server = lb.select_pair()
# Parse and transform prefill_server for bootstrap data
parsed_url = urllib.parse.urlparse(prefill_server)
hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
modified_request = request_data.copy()
modified_request.update(
{
"bootstrap_host": hostname,
"bootstrap_port": bootstrap_port,
"bootstrap_room": _generate_bootstrap_room(),
}
)
if request_data.get("stream", False):
return await lb.generate_stream(
modified_request,
prefill_server,
decode_server,
endpoint=endpoint_name,
)
else:
return await lb.generate(
modified_request,
prefill_server,
decode_server,
endpoint=endpoint_name,
)
@app.post("/v1/chat/completions")
async def handle_chat_completion_request(request_data: dict):
return await _forward_to_backend(request_data, "v1/chat/completions")
@app.post("/v1/completions")
async def handle_completion_request(request_data: dict):
return await _forward_to_backend(request_data, "v1/completions")
def _generate_bootstrap_room():
return random.randint(0, 2**63 - 1)
# We may utilize `GenerateReqInput`'s logic later
def _get_request_batch_size(request):
if (text := request.get("text")) is not None:
return None if isinstance(text, str) else len(text)
if (input_ids := request.get("input_ids")) is not None:
return None if isinstance(input_ids[0], int) else len(input_ids)
return None
@app.get("/v1/models")
async def get_models():
prefill_server = lb.prefill_urls[0] # Get the first prefill server
async with aiohttp.ClientSession() as session:
try:
response = await session.get(f"{prefill_server}/v1/models")
if response.status != 200:
raise HTTPException(
status_code=response.status,
detail=f"Prefill server error: Status {response.status}",
)
return ORJSONResponse(content=await response.json())
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))

129
sgl-router/py_src/sglang_router/router.py Normal file
View File

@@ -0,0 +1,129 @@
from typing import Optional
from sglang_router.router_args import RouterArgs
from sglang_router_rs import PolicyType
from sglang_router_rs import Router as _Router
def policy_from_str(policy_str: Optional[str]) -> PolicyType:
"""Convert policy string to PolicyType enum."""
if policy_str is None:
return None
policy_map = {
"random": PolicyType.Random,
"round_robin": PolicyType.RoundRobin,
"cache_aware": PolicyType.CacheAware,
"power_of_two": PolicyType.PowerOfTwo,
}
return policy_map[policy_str]
class Router:
"""
A high-performance router for distributing requests across worker nodes.
Args:
worker_urls: List of URLs for worker nodes that will handle requests. Each URL should include
the protocol, host, and port (e.g., ['http://worker1:8000', 'http://worker2:8000'])
policy: Load balancing policy to use. Options:
- PolicyType.Random: Randomly select workers
- PolicyType.RoundRobin: Distribute requests in round-robin fashion
- PolicyType.CacheAware: Distribute requests based on cache state and load balance
- PolicyType.PowerOfTwo: Select best of two random workers based on load (PD mode only)
host: Host address to bind the router server. Default: '127.0.0.1'
port: Port number to bind the router server. Default: 3001
worker_startup_timeout_secs: Timeout in seconds for worker startup. Default: 300
worker_startup_check_interval: Interval in seconds between checks for worker initialization. Default: 10
cache_threshold: Cache threshold (0.0-1.0) for cache-aware routing. Routes to cached worker
if the match rate exceeds threshold, otherwise routes to the worker with the smallest
tree. Default: 0.5
balance_abs_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold
AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 32
balance_rel_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold
AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 1.0001
eviction_interval_secs: Interval in seconds between cache eviction operations in cache-aware
routing. Default: 60
max_payload_size: Maximum payload size in bytes. Default: 256MB
max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24
dp_aware: Enable data parallelism aware schedule. Default: False
api_key: The api key used for the authorization with the worker.
Useful when the dp aware scheduling strategy is enabled.
Default: None
log_dir: Directory to store log files. If None, logs are only output to console. Default: None
log_level: Logging level. Options: 'debug', 'info', 'warning', 'error', 'critical'.
service_discovery: Enable Kubernetes service discovery. When enabled, the router will
automatically discover worker pods based on the selector. Default: False
selector: Dictionary mapping of label keys to values for Kubernetes pod selection.
Example: {"app": "sglang-worker"}. Default: {}
service_discovery_port: Port to use for service discovery. The router will generate
worker URLs using this port. Default: 80
service_discovery_namespace: Kubernetes namespace to watch for pods. If not provided,
watches pods across all namespaces (requires cluster-wide permissions). Default: None
prefill_selector: Dictionary mapping of label keys to values for Kubernetes pod selection
for prefill servers (PD mode only). Default: {}
decode_selector: Dictionary mapping of label keys to values for Kubernetes pod selection
for decode servers (PD mode only). Default: {}
prometheus_port: Port to expose Prometheus metrics. Default: None
prometheus_host: Host address to bind the Prometheus metrics server. Default: None
pd_disaggregation: Enable PD (Prefill-Decode) disaggregated mode. Default: False
prefill_urls: List of (url, bootstrap_port) tuples for prefill servers (PD mode only)
decode_urls: List of URLs for decode servers (PD mode only)
prefill_policy: Specific load balancing policy for prefill nodes (PD mode only).
If not specified, uses the main policy. Default: None
decode_policy: Specific load balancing policy for decode nodes (PD mode only).
If not specified, uses the main policy. Default: None
request_id_headers: List of HTTP headers to check for request IDs. If not specified,
uses common defaults: ['x-request-id', 'x-correlation-id', 'x-trace-id', 'request-id'].
Example: ['x-my-request-id', 'x-custom-trace-id']. Default: None
bootstrap_port_annotation: Kubernetes annotation name for bootstrap port (PD mode).
Default: 'sglang.ai/bootstrap-port'
request_timeout_secs: Request timeout in seconds. Default: 600
max_concurrent_requests: Maximum number of concurrent requests allowed for rate limiting. Default: 256
queue_size: Queue size for pending requests when max concurrent limit reached (0 = no queue, return 429 immediately). Default: 100
queue_timeout_secs: Maximum time (in seconds) a request can wait in queue before timing out. Default: 60
rate_limit_tokens_per_second: Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests. Default: None
cors_allowed_origins: List of allowed origins for CORS. Empty list allows all origins. Default: []
health_failure_threshold: Number of consecutive health check failures before marking worker unhealthy. Default: 3
health_success_threshold: Number of consecutive health check successes before marking worker healthy. Default: 2
health_check_timeout_secs: Timeout in seconds for health check requests. Default: 5
health_check_interval_secs: Interval in seconds between runtime health checks. Default: 60
health_check_endpoint: Health check endpoint path. Default: '/health'
model_path: Model path for loading tokenizer (HuggingFace model ID or local path). Default: None
tokenizer_path: Explicit tokenizer path (overrides model_path tokenizer if provided). Default: None
"""
def __init__(self, router: _Router):
self._router = router
@staticmethod
def from_args(args: RouterArgs) -> "Router":
"""Create a router from a RouterArgs instance."""
args_dict = vars(args)
# Convert RouterArgs to _Router parameters
args_dict["worker_urls"] = (
[]
if args_dict["service_discovery"] or args_dict["pd_disaggregation"]
else args_dict["worker_urls"]
)
args_dict["policy"] = policy_from_str(args_dict["policy"])
args_dict["prefill_urls"] = (
args_dict["prefill_urls"] if args_dict["pd_disaggregation"] else None
)
args_dict["decode_urls"] = (
args_dict["decode_urls"] if args_dict["pd_disaggregation"] else None
)
args_dict["prefill_policy"] = policy_from_str(args_dict["prefill_policy"])
args_dict["decode_policy"] = policy_from_str(args_dict["decode_policy"])
# remoge mini_lb parameter
args_dict.pop("mini_lb")
return Router(_Router(**args_dict))
def start(self) -> None:
"""Start the router server.
This method blocks until the server is shut down.
"""
self._router.start()

577
sgl-router/py_src/sglang_router/router_args.py Normal file
View File

@@ -0,0 +1,577 @@
import argparse
import dataclasses
import logging
from typing import Dict, List, Optional
logger = logging.getLogger(__name__)
@dataclasses.dataclass
class RouterArgs:
# Worker configuration
worker_urls: List[str] = dataclasses.field(default_factory=list)
host: str = "127.0.0.1"
port: int = 30000
# PD-specific configuration
mini_lb: bool = False
pd_disaggregation: bool = False # Enable PD disaggregated mode
prefill_urls: List[tuple] = dataclasses.field(
default_factory=list
) # List of (url, bootstrap_port)
decode_urls: List[str] = dataclasses.field(default_factory=list)
# Routing policy
policy: str = "cache_aware"
prefill_policy: Optional[str] = None # Specific policy for prefill nodes in PD mode
decode_policy: Optional[str] = None # Specific policy for decode nodes in PD mode
worker_startup_timeout_secs: int = 600
worker_startup_check_interval: int = 30
cache_threshold: float = 0.3
balance_abs_threshold: int = 64
balance_rel_threshold: float = 1.5
eviction_interval_secs: int = 120
max_tree_size: int = 2**26
max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches
dp_aware: bool = False
api_key: Optional[str] = None
log_dir: Optional[str] = None
log_level: Optional[str] = None
# Service discovery configuration
service_discovery: bool = False
selector: Dict[str, str] = dataclasses.field(default_factory=dict)
service_discovery_port: int = 80
service_discovery_namespace: Optional[str] = None
# PD service discovery configuration
prefill_selector: Dict[str, str] = dataclasses.field(default_factory=dict)
decode_selector: Dict[str, str] = dataclasses.field(default_factory=dict)
bootstrap_port_annotation: str = "sglang.ai/bootstrap-port"
# Prometheus configuration
prometheus_port: Optional[int] = None
prometheus_host: Optional[str] = None
# Request ID headers configuration
request_id_headers: Optional[List[str]] = None
# Request timeout in seconds
request_timeout_secs: int = 1800
# Max concurrent requests for rate limiting
max_concurrent_requests: int = 256
# Queue size for pending requests when max concurrent limit reached
queue_size: int = 100
# Maximum time (in seconds) a request can wait in queue before timing out
queue_timeout_secs: int = 60
# Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests
rate_limit_tokens_per_second: Optional[int] = None
# CORS allowed origins
cors_allowed_origins: List[str] = dataclasses.field(default_factory=list)
# Retry configuration
retry_max_retries: int = 5
retry_initial_backoff_ms: int = 50
retry_max_backoff_ms: int = 30_000
retry_backoff_multiplier: float = 1.5
retry_jitter_factor: float = 0.2
disable_retries: bool = False
# Health check configuration
health_failure_threshold: int = 3
health_success_threshold: int = 2
health_check_timeout_secs: int = 5
health_check_interval_secs: int = 60
health_check_endpoint: str = "/health"
# Circuit breaker configuration
cb_failure_threshold: int = 10
cb_success_threshold: int = 3
cb_timeout_duration_secs: int = 60
cb_window_duration_secs: int = 120
disable_circuit_breaker: bool = False
# Tokenizer configuration
model_path: Optional[str] = None
tokenizer_path: Optional[str] = None
@staticmethod
def add_cli_args(
parser: argparse.ArgumentParser,
use_router_prefix: bool = False,
exclude_host_port: bool = False,
):
"""
Add router-specific arguments to an argument parser.
Args:
parser: The argument parser to add arguments to
use_router_prefix: If True, prefix all arguments with 'router-' to avoid conflicts
exclude_host_port: If True, don't add host and port arguments (used when inheriting from server)
"""
prefix = "router-" if use_router_prefix else ""
# Worker configuration
if not exclude_host_port:
parser.add_argument(
"--host",
type=str,
default=RouterArgs.host,
help="Host address to bind the router server",
)
parser.add_argument(
"--port",
type=int,
default=RouterArgs.port,
help="Port number to bind the router server",
)
parser.add_argument(
"--worker-urls",
type=str,
nargs="*",
default=[],
help="List of worker URLs (e.g., http://worker1:8000 http://worker2:8000)",
)
# Routing policy configuration
parser.add_argument(
f"--{prefix}policy",
type=str,
default=RouterArgs.policy,
choices=["random", "round_robin", "cache_aware", "power_of_two"],
help="Load balancing policy to use. In PD mode, this is used for both prefill and decode unless overridden",
)
parser.add_argument(
f"--{prefix}prefill-policy",
type=str,
default=None,
choices=["random", "round_robin", "cache_aware", "power_of_two"],
help="Specific policy for prefill nodes in PD mode. If not specified, uses the main policy",
)
parser.add_argument(
f"--{prefix}decode-policy",
type=str,
default=None,
choices=["random", "round_robin", "cache_aware", "power_of_two"],
help="Specific policy for decode nodes in PD mode. If not specified, uses the main policy",
)
# PD-specific arguments
parser.add_argument(
f"--{prefix}mini-lb",
action="store_true",
help="Enable MiniLB",
)
parser.add_argument(
f"--{prefix}pd-disaggregation",
action="store_true",
help="Enable PD (Prefill-Decode) disaggregated mode",
)
parser.add_argument(
f"--{prefix}prefill",
nargs="+",
action="append",
help="Prefill server URL and optional bootstrap port. Can be specified multiple times. "
"Format: --prefill URL [BOOTSTRAP_PORT]. "
"BOOTSTRAP_PORT can be a port number, 'none', or omitted (defaults to none).",
)
parser.add_argument(
f"--{prefix}decode",
nargs=1,
action="append",
metavar=("URL",),
help="Decode server URL. Can be specified multiple times.",
)
parser.add_argument(
f"--{prefix}worker-startup-timeout-secs",
type=int,
default=RouterArgs.worker_startup_timeout_secs,
help="Timeout in seconds for worker startup",
)
parser.add_argument(
f"--{prefix}worker-startup-check-interval",
type=int,
default=RouterArgs.worker_startup_check_interval,
help="Interval in seconds between checks for worker startup",
)
parser.add_argument(
f"--{prefix}cache-threshold",
type=float,
default=RouterArgs.cache_threshold,
help="Cache threshold (0.0-1.0) for cache-aware routing",
)
parser.add_argument(
f"--{prefix}balance-abs-threshold",
type=int,
default=RouterArgs.balance_abs_threshold,
help="Load balancing is triggered when (max_load - min_load) > abs_threshold AND max_load > min_load * rel_threshold. Otherwise, use cache aware",
)
parser.add_argument(
f"--{prefix}balance-rel-threshold",
type=float,
default=RouterArgs.balance_rel_threshold,
help="Load balancing is triggered when (max_load - min_load) > abs_threshold AND max_load > min_load * rel_threshold. Otherwise, use cache aware",
)
parser.add_argument(
f"--{prefix}eviction-interval-secs",
type=int,
default=RouterArgs.eviction_interval_secs,
help="Interval in seconds between cache eviction operations",
)
parser.add_argument(
f"--{prefix}max-tree-size",
type=int,
default=RouterArgs.max_tree_size,
help="Maximum size of the approximation tree for cache-aware routing",
)
parser.add_argument(
f"--{prefix}max-payload-size",
type=int,
default=RouterArgs.max_payload_size,
help="Maximum payload size in bytes",
)
parser.add_argument(
f"--{prefix}dp-aware",
action="store_true",
help="Enable data parallelism aware schedule",
)
parser.add_argument(
f"--{prefix}api-key",
type=str,
default=None,
help="The api key used for the authorization with the worker. Useful when the dp aware scheduling strategy is enaled.",
)
parser.add_argument(
f"--{prefix}log-dir",
type=str,
default=None,
help="Directory to store log files. If not specified, logs are only output to console.",
)
parser.add_argument(
f"--{prefix}log-level",
type=str,
default="info",
choices=["debug", "info", "warning", "error", "critical"],
help="Set the logging level. If not specified, defaults to INFO.",
)
parser.add_argument(
f"--{prefix}service-discovery",
action="store_true",
help="Enable Kubernetes service discovery",
)
parser.add_argument(
f"--{prefix}selector",
type=str,
nargs="+",
default={},
help="Label selector for Kubernetes service discovery (format: key1=value1 key2=value2)",
)
parser.add_argument(
f"--{prefix}service-discovery-port",
type=int,
default=RouterArgs.service_discovery_port,
help="Port to use for discovered worker pods",
)
parser.add_argument(
f"--{prefix}service-discovery-namespace",
type=str,
help="Kubernetes namespace to watch for pods. If not provided, watches all namespaces (requires cluster-wide permissions)",
)
parser.add_argument(
f"--{prefix}prefill-selector",
type=str,
nargs="+",
default={},
help="Label selector for prefill server pods in PD mode (format: key1=value1 key2=value2)",
)
parser.add_argument(
f"--{prefix}decode-selector",
type=str,
nargs="+",
default={},
help="Label selector for decode server pods in PD mode (format: key1=value1 key2=value2)",
)
# Prometheus configuration
parser.add_argument(
f"--{prefix}prometheus-port",
type=int,
default=29000,
help="Port to expose Prometheus metrics. If not specified, Prometheus metrics are disabled",
)
parser.add_argument(
f"--{prefix}prometheus-host",
type=str,
default="127.0.0.1",
help="Host address to bind the Prometheus metrics server",
)
parser.add_argument(
f"--{prefix}request-id-headers",
type=str,
nargs="*",
help="Custom HTTP headers to check for request IDs (e.g., x-request-id x-trace-id). If not specified, uses common defaults.",
)
parser.add_argument(
f"--{prefix}request-timeout-secs",
type=int,
default=RouterArgs.request_timeout_secs,
help="Request timeout in seconds",
)
# Retry configuration
parser.add_argument(
f"--{prefix}retry-max-retries",
type=int,
default=RouterArgs.retry_max_retries,
)
parser.add_argument(
f"--{prefix}retry-initial-backoff-ms",
type=int,
default=RouterArgs.retry_initial_backoff_ms,
)
parser.add_argument(
f"--{prefix}retry-max-backoff-ms",
type=int,
default=RouterArgs.retry_max_backoff_ms,
)
parser.add_argument(
f"--{prefix}retry-backoff-multiplier",
type=float,
default=RouterArgs.retry_backoff_multiplier,
)
parser.add_argument(
f"--{prefix}retry-jitter-factor",
type=float,
default=RouterArgs.retry_jitter_factor,
)
parser.add_argument(
f"--{prefix}disable-retries",
action="store_true",
help="Disable retries (equivalent to setting retry_max_retries=1)",
)
# Circuit breaker configuration
parser.add_argument(
f"--{prefix}cb-failure-threshold",
type=int,
default=RouterArgs.cb_failure_threshold,
)
parser.add_argument(
f"--{prefix}cb-success-threshold",
type=int,
default=RouterArgs.cb_success_threshold,
)
parser.add_argument(
f"--{prefix}cb-timeout-duration-secs",
type=int,
default=RouterArgs.cb_timeout_duration_secs,
)
parser.add_argument(
f"--{prefix}cb-window-duration-secs",
type=int,
default=RouterArgs.cb_window_duration_secs,
)
parser.add_argument(
f"--{prefix}disable-circuit-breaker",
action="store_true",
help="Disable circuit breaker (equivalent to setting cb_failure_threshold to u32::MAX)",
)
# Health check configuration
parser.add_argument(
f"--{prefix}health-failure-threshold",
type=int,
default=RouterArgs.health_failure_threshold,
help="Number of consecutive health check failures before marking worker unhealthy",
)
parser.add_argument(
f"--{prefix}health-success-threshold",
type=int,
default=RouterArgs.health_success_threshold,
help="Number of consecutive health check successes before marking worker healthy",
)
parser.add_argument(
f"--{prefix}health-check-timeout-secs",
type=int,
default=RouterArgs.health_check_timeout_secs,
help="Timeout in seconds for health check requests",
)
parser.add_argument(
f"--{prefix}health-check-interval-secs",
type=int,
default=RouterArgs.health_check_interval_secs,
help="Interval in seconds between runtime health checks",
)
parser.add_argument(
f"--{prefix}health-check-endpoint",
type=str,
default=RouterArgs.health_check_endpoint,
help="Health check endpoint path",
)
parser.add_argument(
f"--{prefix}max-concurrent-requests",
type=int,
default=RouterArgs.max_concurrent_requests,
help="Maximum number of concurrent requests allowed (for rate limiting)",
)
parser.add_argument(
f"--{prefix}queue-size",
type=int,
default=RouterArgs.queue_size,
help="Queue size for pending requests when max concurrent limit reached (0 = no queue, return 429 immediately)",
)
parser.add_argument(
f"--{prefix}queue-timeout-secs",
type=int,
default=RouterArgs.queue_timeout_secs,
help="Maximum time (in seconds) a request can wait in queue before timing out",
)
parser.add_argument(
f"--{prefix}rate-limit-tokens-per-second",
type=int,
default=RouterArgs.rate_limit_tokens_per_second,
help="Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests",
)
parser.add_argument(
f"--{prefix}cors-allowed-origins",
type=str,
nargs="*",
default=[],
help="CORS allowed origins (e.g., http://localhost:3000 https://example.com)",
)
# Tokenizer configuration
parser.add_argument(
f"--{prefix}model-path",
type=str,
default=None,
help="Model path for loading tokenizer (HuggingFace model ID or local path)",
)
parser.add_argument(
f"--{prefix}tokenizer-path",
type=str,
default=None,
help="Explicit tokenizer path (overrides model_path tokenizer if provided)",
)
@classmethod
def from_cli_args(
cls, args: argparse.Namespace, use_router_prefix: bool = False
) -> "RouterArgs":
"""
Create RouterArgs instance from parsed command line arguments.
Args:
args: Parsed command line arguments
use_router_prefix: If True, look for arguments with 'router-' prefix
"""
prefix = "router_" if use_router_prefix else ""
cli_args_dict = vars(args)
args_dict = {}
for attr in dataclasses.fields(cls):
# Auto strip prefix from args
if f"{prefix}{attr.name}" in cli_args_dict:
args_dict[attr.name] = cli_args_dict[f"{prefix}{attr.name}"]
elif attr.name in cli_args_dict:
args_dict[attr.name] = cli_args_dict[attr.name]
# parse special arguments and remove "--prefill" and "--decode" from cli_args_dict
args_dict["prefill_urls"] = cls._parse_prefill_urls(
cli_args_dict.get(f"{prefix}prefill", None)
)
args_dict["decode_urls"] = cls._parse_decode_urls(
cli_args_dict.get(f"{prefix}decode", None)
)
args_dict["selector"] = cls._parse_selector(
cli_args_dict.get(f"{prefix}selector", None)
)
args_dict["prefill_selector"] = cls._parse_selector(
cli_args_dict.get(f"{prefix}prefill_selector", None)
)
args_dict["decode_selector"] = cls._parse_selector(
cli_args_dict.get(f"{prefix}decode_selector", None)
)
# Mooncake-specific annotation
args_dict["bootstrap_port_annotation"] = "sglang.ai/bootstrap-port"
return cls(**args_dict)
def _validate_router_args(self):
# Validate configuration based on mode
if self.pd_disaggregation:
# Validate PD configuration - skip URL requirements if using service discovery
if not self.service_discovery:
if not self.prefill_urls:
raise ValueError("PD disaggregation mode requires --prefill")
if not self.decode_urls:
raise ValueError("PD disaggregation mode requires --decode")
# Warn about policy usage in PD mode
if self.prefill_policy and self.decode_policy and self.policy:
logger.warning(
"Both --prefill-policy and --decode-policy are specified. "
"The main --policy flag will be ignored for PD mode."
)
elif self.prefill_policy and not self.decode_policy and self.policy:
logger.info(
f"Using --prefill-policy '{self.prefill_policy}' for prefill nodes "
f"and --policy '{self.policy}' for decode nodes."
)
elif self.decode_policy and not self.prefill_policy and self.policy:
logger.info(
f"Using --policy '{self.policy}' for prefill nodes "
f"and --decode-policy '{self.decode_policy}' for decode nodes."
)
@staticmethod
def _parse_selector(selector_list):
if not selector_list:
return {}
selector = {}
for item in selector_list:
if "=" in item:
key, value = item.split("=", 1)
selector[key] = value
return selector
@staticmethod
def _parse_prefill_urls(prefill_list):
"""Parse prefill URLs from --prefill arguments.
Format: --prefill URL [BOOTSTRAP_PORT]
Example:
--prefill http://prefill1:8080 9000 # With bootstrap port
--prefill http://prefill2:8080 none # Explicitly no bootstrap port
--prefill http://prefill3:8080 # Defaults to no bootstrap port
"""
if not prefill_list:
return []
prefill_urls = []
for prefill_args in prefill_list:
url = prefill_args[0]
# Handle optional bootstrap port
if len(prefill_args) >= 2:
bootstrap_port_str = prefill_args[1]
# Handle 'none' as None
if bootstrap_port_str.lower() == "none":
bootstrap_port = None
else:
try:
bootstrap_port = int(bootstrap_port_str)
except ValueError:
raise ValueError(
f"Invalid bootstrap port: {bootstrap_port_str}. Must be a number or 'none'"
)
else:
# No bootstrap port specified, default to None
bootstrap_port = None
prefill_urls.append((url, bootstrap_port))
return prefill_urls
@staticmethod
def _parse_decode_urls(decode_list):
"""Parse decode URLs from --decode arguments.
Format: --decode URL
Example: --decode http://decode1:8081 --decode http://decode2:8081
"""
if not decode_list:
return []
# decode_list is a list of single-element lists due to nargs=1
return [url[0] for url in decode_list]

1
sgl-router/py_src/sglang_router/version.py Normal file
View File

@@ -0,0 +1 @@
__version__ = "0.1.9"

1
sgl-router/py_test/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Test package root for router Python tests."""

8
sgl-router/py_test/conftest.py Normal file
View File

@@ -0,0 +1,8 @@
import sys
from pathlib import Path
# Ensure local sources in py_src are importable ahead of any installed package
_ROOT = Path(__file__).resolve().parents[1]
_SRC = _ROOT / "py_src"
if str(_SRC) not in sys.path:
sys.path.insert(0, str(_SRC))

512
sgl-router/py_test/e2e/conftest.py Normal file
View File

@@ -0,0 +1,512 @@
import json
import logging
import os
import shutil
import signal
import socket
import subprocess
import time
from pathlib import Path
from types import SimpleNamespace
from typing import Callable, Optional
from urllib.parse import urlparse
import pytest
import requests
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
)
logger = logging.getLogger(__name__)
def _find_available_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _parse_url(base_url: str) -> tuple[str, str]:
"""Parse a base URL and return (host, port) as strings.
This is more robust than simple string splitting and supports different schemes
and URL shapes like trailing paths.
"""
parsed = urlparse(base_url)
return parsed.hostname or "127.0.0.1", (
str(parsed.port) if parsed.port is not None else ""
)
def _wait_router_health(base_url: str, timeout: float) -> None:
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{base_url}/health", timeout=5)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(2)
raise TimeoutError("Router failed to become healthy in time")
def _popen_launch_router(
model: str,
base_url: str,
dp_size: int,
timeout: float,
policy: str = "cache_aware",
) -> subprocess.Popen:
host, port = _parse_url(base_url)
prom_port = _find_available_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
"--dp",
str(dp_size),
"--router-policy",
policy,
"--allow-auto-truncate",
"--router-prometheus-port",
str(prom_port),
"--router-prometheus-host",
"127.0.0.1",
]
proc = subprocess.Popen(cmd)
_wait_router_health(base_url, timeout)
return proc
def _popen_launch_worker(
model: str,
base_url: str,
*,
dp_size: int | None = None,
api_key: str | None = None,
base_gpu_id: int | None = 0,
) -> subprocess.Popen:
host, port = _parse_url(base_url)
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
"--base-gpu-id",
str(base_gpu_id or 0),
]
if dp_size is not None:
cmd += ["--dp-size", str(dp_size)]
if api_key is not None:
cmd += ["--api-key", api_key]
return subprocess.Popen(cmd)
def _popen_launch_router_only(
base_url: str,
policy: str = "round_robin",
timeout: float = 120.0,
*,
dp_aware: bool = False,
api_key: str | None = None,
) -> subprocess.Popen:
host, port = _parse_url(base_url)
prom_port = _find_available_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
host,
"--port",
port,
"--policy",
policy,
]
if dp_aware:
cmd += ["--dp-aware"]
if api_key is not None:
cmd += ["--api-key", api_key]
cmd += [
"--prometheus-port",
str(prom_port),
"--prometheus-host",
"127.0.0.1",
]
proc = subprocess.Popen(cmd)
_wait_router_health(base_url, timeout)
return proc
def _terminate(proc: subprocess.Popen, timeout: float = 120) -> None:
if proc is None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
def _which(cmd: str) -> Optional[str]:
try:
return shutil.which(cmd)
except Exception as e:
logger.warning("shutil.which(%r) failed: %s", cmd, e)
return None
def _graceful_stop_popen(p: subprocess.Popen) -> None:
if p is None:
return
try:
if p.poll() is None:
p.terminate()
for _ in range(5):
if p.poll() is not None:
break
time.sleep(1)
if p.poll() is None:
p.kill()
except Exception as e:
logger.warning("Exception during graceful stop of popen: %s", e)
def _pid_alive(pid: int) -> bool:
try:
os.kill(pid, 0)
return True
except Exception:
return False
def _graceful_stop_pid(pid: int) -> None:
try:
if _pid_alive(pid):
try:
os.kill(pid, signal.SIGTERM)
except Exception:
pass
for _ in range(5):
if not _pid_alive(pid):
break
time.sleep(1)
if _pid_alive(pid):
try:
os.kill(pid, signal.SIGKILL)
except Exception:
pass
except Exception:
pass
def _graceful_stop_any(obj) -> None:
try:
if isinstance(obj, subprocess.Popen):
_graceful_stop_popen(obj)
return
if isinstance(obj, int):
_graceful_stop_pid(obj)
return
proc_obj = getattr(obj, "proc", None)
if isinstance(proc_obj, subprocess.Popen):
_graceful_stop_popen(proc_obj)
except Exception:
pass
@pytest.fixture(scope="session")
def genai_bench_runner() -> Callable[..., None]:
"""Provide a callable to run genai-bench and validate metrics.
Usage in tests:
def test(..., genai_bench_runner):
genai_bench_runner(router_url=..., model_path=..., experiment_folder=...)
"""
def _run(
*,
router_url: str,
model_path: str,
experiment_folder: str,
timeout_sec: int | None = None,
thresholds: dict | None = None,
extra_env: dict | None = None,
num_concurrency: int = 32,
traffic_scenario: str = "D(4000,100)",
max_requests_per_run: int | None = None,
clean_experiment: bool = True,
kill_procs: list | None = None,
drain_delay_sec: int = 6,
) -> None:
cli = _which("genai-bench")
if not cli:
pytest.fail(
"genai-bench CLI not found; please install it to run benchmarks"
)
# Clean previous experiment folder under current working directory
if clean_experiment:
exp_dir = Path.cwd() / experiment_folder
if exp_dir.exists():
shutil.rmtree(exp_dir, ignore_errors=True)
# Default requests per run if not provided
mrr = (
max_requests_per_run
if max_requests_per_run is not None
else num_concurrency * 3
)
cmd = [
cli,
"benchmark",
"--api-backend",
"openai",
"--api-base",
router_url,
"--api-key",
"dummy-token",
"--api-model-name",
model_path,
"--model-tokenizer",
model_path,
"--task",
"text-to-text",
"--num-concurrency",
str(num_concurrency),
"--traffic-scenario",
traffic_scenario,
"--max-requests-per-run",
str(mrr),
"--max-time-per-run",
"2",
"--experiment-folder-name",
experiment_folder,
"--experiment-base-dir",
str(Path.cwd()),
]
env = os.environ.copy()
if extra_env:
env.update(extra_env)
to = timeout_sec or int(os.environ.get("GENAI_BENCH_TEST_TIMEOUT", "120"))
proc = subprocess.Popen(
cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
)
stdout = stderr = ""
rc = None
try:
try:
stdout, stderr = proc.communicate(timeout=to)
except subprocess.TimeoutExpired:
# Simple: kill the CLI process if it doesn't exit in time
try:
proc.kill()
except Exception:
pass
stdout, stderr = proc.communicate()
rc = proc.returncode
# Prefer exact path under cwd; fallback to rglob search
base = Path.cwd()
direct = base / experiment_folder
candidates = [direct] if direct.is_dir() else []
if not candidates:
for p in base.rglob(experiment_folder):
if p.is_dir() and p.name == experiment_folder:
candidates = [p]
break
if not candidates:
raise AssertionError(
"Benchmark failed: experiment folder not found: "
f"{experiment_folder}\nExit code: {rc}\nSTDOUT (tail):\n{stdout[-1000:]}\nSTDERR (tail):\n{stderr[-1000:]}"
)
actual_folder = candidates[0]
json_files = [
p
for p in actual_folder.rglob("*.json")
if "experiment_metadata" not in p.name
]
if not json_files:
raise AssertionError(
"Benchmark failed: no JSON results found\n"
f"Exit code: {rc}\nSTDOUT (tail):\n{stdout[-1000:]}\nSTDERR (tail):\n{stderr[-1000:]}"
)
th = thresholds # None means "log only", no validation
for jf in json_files:
with jf.open("r") as f:
data = json.load(f)
stats = data.get("aggregated_metrics", {}).get("stats", {})
ttft_mean = float(stats.get("ttft", {}).get("mean", float("inf")))
e2e_latency_mean = float(
stats.get("e2e_latency", {}).get("mean", float("inf"))
)
input_tp_mean = float(stats.get("input_throughput", {}).get("mean", 0.0))
output_tp_mean = float(stats.get("output_throughput", {}).get("mean", 0.0))
logger.info(
"genai-bench[%s] %s ttft_mean=%.3fs e2e_latency_mean=%.3fs input_tp_mean=%.1f tok/s output_tp_mean=%.1f tok/s",
experiment_folder,
jf.name,
ttft_mean,
e2e_latency_mean,
input_tp_mean,
output_tp_mean,
)
if th is not None:
assert (
ttft_mean <= th["ttft_mean_max"]
), f"TTFT validation failed: {ttft_mean} > {th['ttft_mean_max']} (file={jf.name})"
assert (
e2e_latency_mean <= th["e2e_latency_mean_max"]
), f"E2E latency validation failed: {e2e_latency_mean} > {th['e2e_latency_mean_max']} (file={jf.name})"
assert (
input_tp_mean >= th["input_throughput_mean_min"]
), f"Input throughput validation failed: {input_tp_mean} < {th['input_throughput_mean_min']} (file={jf.name})"
assert (
output_tp_mean >= th["output_throughput_mean_min"]
), f"Output throughput validation failed: {output_tp_mean} < {th['output_throughput_mean_min']} (file={jf.name})"
finally:
# Always attempt to stop workers to avoid resource leakage
if kill_procs:
# Give router/workers a small grace period to finish any last drains
if drain_delay_sec > 0:
try:
time.sleep(drain_delay_sec)
except Exception:
pass
for p in kill_procs:
_graceful_stop_any(p)
try:
time.sleep(2)
except Exception:
pass
return _run
def pytest_configure(config):
config.addinivalue_line("markers", "e2e: mark as end-to-end test")
@pytest.fixture(scope="session")
def e2e_model() -> str:
# Always use the default test model
return DEFAULT_MODEL_NAME_FOR_TEST
@pytest.fixture
def e2e_router(e2e_model: str):
# Keep this available but tests below use router-only to avoid GPU contention
base_url = DEFAULT_URL_FOR_TEST
proc = _popen_launch_router(
e2e_model, base_url, dp_size=2, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture
def e2e_router_only_rr():
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_router_only(base_url, policy="round_robin")
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_primary_worker(e2e_model: str):
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_worker(e2e_model, base_url)
# Router health gate will handle worker readiness
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture
def e2e_router_only_rr_dp_aware_api():
"""Router-only with dp-aware enabled and an API key."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
api_key = "secret"
proc = _popen_launch_router_only(
base_url, policy="round_robin", timeout=180.0, dp_aware=True, api_key=api_key
)
try:
yield SimpleNamespace(proc=proc, url=base_url, api_key=api_key)
finally:
_terminate(proc)
@pytest.fixture
def e2e_worker_dp2_api(e2e_model: str, e2e_router_only_rr_dp_aware_api):
"""Worker with dp-size=2 and the same API key as the dp-aware router."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
api_key = e2e_router_only_rr_dp_aware_api.api_key
proc = _popen_launch_worker(e2e_model, base_url, dp_size=2, api_key=api_key)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_two_workers_dp2(e2e_model: str):
"""Launch two workers, each with dp_size=2, mapped to GPUs [0,1] and [2,3]."""
workers = []
try:
# Worker A on GPUs 0-1
port_a = _find_available_port()
url_a = f"http://127.0.0.1:{port_a}"
proc_a = _popen_launch_worker(e2e_model, url_a, dp_size=2, base_gpu_id=0)
workers.append(SimpleNamespace(proc=proc_a, url=url_a))
# Worker B on GPUs 2-3
port_b = _find_available_port()
url_b = f"http://127.0.0.1:{port_b}"
proc_b = _popen_launch_worker(e2e_model, url_b, dp_size=2, base_gpu_id=2)
workers.append(SimpleNamespace(proc=proc_b, url=url_b))
yield workers
finally:
for w in workers:
_terminate(w.proc)

262
sgl-router/py_test/e2e/test_pd_router.py Normal file
View File

@@ -0,0 +1,262 @@
import logging
import os
import socket
import subprocess
import time
from types import SimpleNamespace
from typing import Optional
import pytest
import requests
from sglang.test.run_eval import run_eval
logger = logging.getLogger(__name__)
def _find_available_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _wait_health(url: str, timeout: float = 180.0) -> None:
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{url}/health", timeout=5)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(1)
raise TimeoutError(f"Service at {url} failed to become healthy in time")
def _detect_ib_device() -> Optional[str]:
"""Return first active IB device name (e.g., mlx5_0) or None if unavailable."""
# Fast check that ibv_devinfo exists
try:
subprocess.run(
["ibv_devinfo", "-l"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=1,
)
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
for i in range(12):
dev = f"mlx5_{i}"
try:
res = subprocess.run(
["ibv_devinfo", dev],
capture_output=True,
text=True,
timeout=2,
)
if res.returncode == 0 and ("state:" in res.stdout):
for line in res.stdout.splitlines():
if "state:" in line and "PORT_ACTIVE" in line:
return dev
except Exception:
pass
return None
def _popen_launch_prefill_worker(
model: str,
bootstrap_port: int,
ib_device: Optional[str] = None,
base_gpu_id: int = 0,
) -> SimpleNamespace:
port = _find_available_port()
url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--disaggregation-mode",
"prefill",
"--host",
"127.0.0.1",
"--port",
str(port),
"--disaggregation-bootstrap-port",
str(bootstrap_port),
"--base-gpu-id",
str(base_gpu_id),
]
if ib_device:
cmd += ["--disaggregation-ib-device", ib_device]
proc = subprocess.Popen(cmd)
_wait_health(url, timeout=300.0)
return SimpleNamespace(proc=proc, url=url, bootstrap_port=bootstrap_port)
def _popen_launch_decode_worker(
model: str, ib_device: Optional[str] = None, base_gpu_id: int = 0
) -> SimpleNamespace:
port = _find_available_port()
url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--disaggregation-mode",
"decode",
"--host",
"127.0.0.1",
"--port",
str(port),
"--base-gpu-id",
str(base_gpu_id),
]
if ib_device:
cmd += ["--disaggregation-ib-device", ib_device]
proc = subprocess.Popen(cmd)
_wait_health(url, timeout=300.0)
return SimpleNamespace(proc=proc, url=url)
def _terminate(proc: subprocess.Popen, timeout: float = 120) -> None:
if proc is None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
@pytest.fixture(scope="module")
def pd_cluster(e2e_model: str):
"""Start 2 prefill + 2 decode workers and one PD router, once per module."""
# Environment capability checks: require sgl_kernel and GPU backend
try:
import sgl_kernel # noqa: F401
except Exception as e: # pragma: no cover - environment dependent
pytest.fail(f"PD e2e requires sgl_kernel but it is not available: {e}")
try:
import torch # noqa: F401
except Exception as e: # pragma: no cover - environment dependent
pytest.fail(
f"PD e2e requires torch but it is not available or misconfigured: {e}"
)
if not torch.cuda.is_available(): # pragma: no cover - environment dependent
pytest.fail("PD e2e requires CUDA backend, but CUDA is not available")
workers: list[SimpleNamespace] = []
router_proc = None
try:
ib_device = _detect_ib_device()
# Launch 4 workers across 4 GPUs: prefill on 0,1 and decode on 2,3
pf1 = _popen_launch_prefill_worker(
e2e_model,
bootstrap_port=_find_available_port(),
ib_device=ib_device,
base_gpu_id=0,
)
pf2 = _popen_launch_prefill_worker(
e2e_model,
bootstrap_port=_find_available_port(),
ib_device=ib_device,
base_gpu_id=1,
)
dc1 = _popen_launch_decode_worker(e2e_model, ib_device=ib_device, base_gpu_id=2)
dc2 = _popen_launch_decode_worker(e2e_model, ib_device=ib_device, base_gpu_id=3)
prefills = [pf1, pf2]
decodes = [dc1, dc2]
workers.extend(prefills + decodes)
# PD router with two prefill and two decode endpoints
rport = _find_available_port()
router_url = f"http://127.0.0.1:{rport}"
pport = _find_available_port()
prefill = [(pf.url, pf.bootstrap_port) for pf in prefills]
decode = [dc.url for dc in decodes]
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(rport),
"--policy",
"round_robin",
"--pd-disaggregation",
]
for url, bport in prefill:
cmd += ["--prefill", url, str(bport)]
for url in decode:
cmd += ["--decode", url]
cmd += [
"--prometheus-port",
str(pport),
"--prometheus-host",
"127.0.0.1",
]
router_proc = subprocess.Popen(cmd)
_wait_health(router_url, timeout=180.0)
yield SimpleNamespace(
router_url=router_url, workers=workers, router_proc=router_proc
)
finally:
if router_proc is not None:
_terminate(router_proc)
for w in workers:
_terminate(w.proc)
@pytest.mark.e2e
def test_pd_mmlu(e2e_model: str, pd_cluster):
"""
Launch 4 workers, start a PD router (2 prefill + 2 decode), then run MMLU.
"""
args = SimpleNamespace(
base_url=pd_cluster.router_url,
model=e2e_model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert metrics["score"] >= 0.65
@pytest.mark.e2e
def test_pd_genai_bench(e2e_model: str, pd_cluster, genai_bench_runner):
"""
Launch 4 workers, start a PD router (2 prefill + 2 decode), then run a
short genai-bench benchmark and validate aggregate metrics.
"""
# Run genai-bench against the shared router
policy_label = "benchmark_round_robin_pd"
genai_bench_runner(
router_url=pd_cluster.router_url,
model_path=e2e_model,
experiment_folder=policy_label,
thresholds={
"ttft_mean_max": 12,
"e2e_latency_mean_max": 15,
"input_throughput_mean_min": 400,
"output_throughput_mean_min": 20,
},
kill_procs=pd_cluster.workers,
)

169
sgl-router/py_test/e2e/test_regular_router.py Normal file
View File

@@ -0,0 +1,169 @@
import threading
import time
from types import SimpleNamespace
import pytest
import requests
from sglang.test.run_eval import run_eval
@pytest.mark.e2e
def test_mmlu(e2e_router_only_rr, e2e_two_workers_dp2, e2e_model):
# Attach two dp=2 workers (total 4 GPUs) to a fresh router-only instance
base = e2e_router_only_rr.url
for w in e2e_two_workers_dp2:
r = requests.post(f"{base}/add_worker", params={"url": w.url}, timeout=180)
r.raise_for_status()
args = SimpleNamespace(
base_url=base,
model=e2e_model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert metrics["score"] >= 0.65
@pytest.mark.e2e
def test_genai_bench(
e2e_router_only_rr, e2e_two_workers_dp2, e2e_model, genai_bench_runner
):
"""Attach a worker to the regular router and run a short genai-bench."""
base = e2e_router_only_rr.url
for w in e2e_two_workers_dp2:
r = requests.post(f"{base}/add_worker", params={"url": w.url}, timeout=180)
r.raise_for_status()
genai_bench_runner(
router_url=base,
model_path=e2e_model,
experiment_folder="benchmark_round_robin_regular",
thresholds={
"ttft_mean_max": 6,
"e2e_latency_mean_max": 14,
"input_throughput_mean_min": 1000,
"output_throughput_mean_min": 12,
},
kill_procs=e2e_two_workers_dp2,
)
@pytest.mark.e2e
def test_add_and_remove_worker_live(e2e_router_only_rr, e2e_primary_worker, e2e_model):
base = e2e_router_only_rr.url
worker_url = e2e_primary_worker.url
r = requests.post(f"{base}/add_worker", params={"url": worker_url}, timeout=180)
r.raise_for_status()
with requests.Session() as s:
for i in range(8):
r = s.post(
f"{base}/v1/completions",
json={
"model": e2e_model,
"prompt": f"x{i}",
"max_tokens": 1,
"stream": False,
},
timeout=120,
)
r.raise_for_status()
# Remove the worker
r = requests.post(f"{base}/remove_worker", params={"url": worker_url}, timeout=60)
r.raise_for_status()
@pytest.mark.e2e
def test_lazy_fault_tolerance_live(e2e_router_only_rr, e2e_primary_worker, e2e_model):
base = e2e_router_only_rr.url
worker = e2e_primary_worker
r = requests.post(f"{base}/add_worker", params={"url": worker.url}, timeout=180)
r.raise_for_status()
def killer():
time.sleep(10)
try:
worker.proc.terminate()
except Exception:
pass
t = threading.Thread(target=killer, daemon=True)
t.start()
args = SimpleNamespace(
base_url=base,
model=e2e_model,
eval_name="mmlu",
num_examples=32,
num_threads=16,
temperature=0.0,
)
metrics = run_eval(args)
assert 0.0 <= metrics["score"] <= 1.0
@pytest.mark.e2e
def test_dp_aware_worker_expansion_and_api_key(
e2e_model,
e2e_router_only_rr_dp_aware_api,
e2e_worker_dp2_api,
):
"""
Launch a router-only instance in dp_aware mode and a single worker with dp_size=2
and API key protection. Verify expansion, auth enforcement, and basic eval.
"""
import os
router_url = e2e_router_only_rr_dp_aware_api.url
worker_url = e2e_worker_dp2_api.url
api_key = e2e_router_only_rr_dp_aware_api.api_key
# Attach worker; router should expand to dp_size logical workers
r = requests.post(
f"{router_url}/add_worker", params={"url": worker_url}, timeout=180
)
r.raise_for_status()
r = requests.get(f"{router_url}/list_workers", timeout=30)
r.raise_for_status()
urls = r.json().get("urls", [])
assert len(urls) == 2
assert set(urls) == {f"{worker_url}@0", f"{worker_url}@1"}
# Verify API key enforcement path-through
# 1) Without Authorization -> 401 from backend
r = requests.post(
f"{router_url}/v1/completions",
json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
timeout=60,
)
assert r.status_code == 401
# 2) With correct Authorization -> 200
r = requests.post(
f"{router_url}/v1/completions",
json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
headers={"Authorization": f"Bearer {api_key}"},
timeout=60,
)
assert r.status_code == 200
# Finally, run MMLU eval through the router with auth
os.environ["OPENAI_API_KEY"] = api_key
args = SimpleNamespace(
base_url=router_url,
model=e2e_model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert metrics["score"] >= 0.65

1
sgl-router/py_test/fixtures/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Shared fixtures for router integration tests."""

252
sgl-router/py_test/fixtures/mock_worker.py Normal file
View File

@@ -0,0 +1,252 @@
"""
Lightweight mock worker HTTP server for router integration tests.
Implements minimal endpoints used by the router:
- GET /health, /health_generate
- POST /generate, /v1/completions, /v1/chat/completions
- POST /flush_cache
- GET /get_server_info, /get_model_info, /v1/models
Behavior knobs are controlled via CLI flags to simulate failures, latency, and load.
"""
import argparse
import asyncio
import json
import os
import random
import signal
import sys
import time
from contextlib import asynccontextmanager
from typing import Optional
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
# Global state (per-process)
_inflight = 0
_failures_seen = 0
def _parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--host", default="127.0.0.1")
p.add_argument("--port", type=int, required=True)
p.add_argument("--worker-id", default=None)
p.add_argument("--latency-ms", type=int, default=0)
p.add_argument("--timeout", action="store_true")
p.add_argument("--status-code", type=int, default=200)
p.add_argument("--fail-first-n", type=int, default=0)
p.add_argument("--random-fail-rate", type=float, default=0.0)
p.add_argument("--require-api-key", action="store_true")
p.add_argument("--api-key", default=None)
p.add_argument("--max-payload-bytes", type=int, default=10 * 1024 * 1024)
p.add_argument("--stream", action="store_true")
p.add_argument("--dp-size", type=int, default=1)
p.add_argument("--crash-on-request", action="store_true")
p.add_argument("--health-fail-after-ms", type=int, default=0)
return p.parse_args()
def _extract_worker_id(args: argparse.Namespace) -> str:
if args.worker_id:
return str(args.worker_id)
# default to port (unique enough for tests)
return f"worker-{args.port}"
def create_app(args: argparse.Namespace) -> FastAPI:
app = FastAPI()
worker_id = _extract_worker_id(args)
start_ts = time.time()
crashed = {"done": False}
async def maybe_delay():
if args.latency_ms > 0:
await asyncio.sleep(args.latency_ms / 1000.0)
def should_fail() -> Optional[int]:
global _failures_seen
# Fail first N requests (500)
if args.fail_first_n > 0 and _failures_seen < args.fail_first_n:
_failures_seen += 1
return 500
# Random failure probability (500)
if args.random_fail_rate > 0.0 and random.random() < args.random_fail_rate:
return 500
# Forced status code override (non-200) for all responses
if args.status_code != 200:
return int(args.status_code)
return None
def check_api_key(request: Request):
if not args.require_api_key:
return
auth = request.headers.get("Authorization")
if not auth or not auth.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Unauthorized")
key = auth.split(" ", 1)[1]
if args.api_key and key != args.api_key:
raise HTTPException(status_code=401, detail="Unauthorized")
@asynccontextmanager
async def track_inflight():
global _inflight
_inflight += 1
try:
yield
finally:
_inflight -= 1
@app.get("/health")
async def health():
if (
args.health_fail_after_ms
and (time.time() - start_ts) * 1000.0 >= args.health_fail_after_ms
):
return PlainTextResponse("bad", status_code=500)
return PlainTextResponse("ok", status_code=200)
@app.get("/health_generate")
async def health_generate():
return PlainTextResponse("ok", status_code=200)
@app.post("/flush_cache")
async def flush_cache():
return PlainTextResponse("ok", status_code=200)
@app.get("/get_model_info")
async def get_model_info():
return JSONResponse({"model": "mock", "vocab_size": 32000})
@app.get("/v1/models")
async def list_models():
return JSONResponse({"data": [{"id": "mock", "object": "model"}]})
@app.get("/get_server_info")
async def get_server_info(request: Request):
# Enforce API key on server info when required (used by dp_aware probing)
check_api_key(request)
return JSONResponse(
{
"worker_id": worker_id,
"load_in_flight": _inflight,
"cache": {"size": 0, "hit_rate": 0.0},
"dp_size": int(args.dp_size),
}
)
@app.get("/get_load")
async def get_load():
return JSONResponse({"load": _inflight})
def make_json_response(obj: dict, status_code: int = 200) -> JSONResponse:
resp = JSONResponse(obj, status_code=status_code)
resp.headers["X-Worker-Id"] = worker_id
return resp
async def handle_text_request(request: Request):
# Authorization
check_api_key(request)
# Payload limit
body = await request.body()
if len(body) > args.max_payload_bytes:
return make_json_response({"error": "payload too large"}, status_code=413)
# Simulate crash on first request
if args.crash_on_request and not crashed["done"]:
crashed["done"] = True
os._exit(1)
# Optional timeout (simulate hang)
if args.timeout:
await asyncio.sleep(3600)
# Optional latency
await maybe_delay()
# Optional failures
fail_code = should_fail()
if fail_code is not None and fail_code != 200:
return make_json_response(
{"error": f"mock failure {fail_code}"}, status_code=fail_code
)
# Build response echoing minimal shape
try:
data = await request.json()
except (json.JSONDecodeError, ValueError):
data = {}
now = time.time()
ret = {
"id": f"cmpl-{int(now*1000)}",
"object": "text_completion",
"created": int(now),
"model": "mock",
"choices": [
{
"text": "ok",
"index": 0,
"finish_reason": "stop",
}
],
"worker_id": worker_id,
"echo": data,
}
return make_json_response(ret, status_code=200)
async def handle_stream_request(request: Request):
check_api_key(request)
async def gen():
# minimal 2-chunk stream then [DONE]
for i in range(2):
await asyncio.sleep(0.01)
chunk = {
"choices": [{"delta": {"content": "x"}}],
"worker_id": worker_id,
}
yield f"data: {json.dumps(chunk)}\n\n"
yield "data: [DONE]\n\n"
headers = {"X-Worker-Id": worker_id}
return StreamingResponse(gen(), media_type="text/event-stream", headers=headers)
@app.post("/generate")
async def generate(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
@app.post("/v1/completions")
async def completions(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
return app
def main() -> None:
args = _parse_args()
app = create_app(args)
# Handle SIGTERM gracefully for fast test teardown
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
if __name__ == "__main__":
main()

8
sgl-router/py_test/fixtures/ports.py Normal file
View File

@@ -0,0 +1,8 @@
import socket
def find_free_port() -> int:
"""Return an available TCP port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]

158
sgl-router/py_test/fixtures/router_manager.py Normal file
View File

@@ -0,0 +1,158 @@
import subprocess
import time
from dataclasses import dataclass
from typing import Dict, List, Optional
import requests
from .ports import find_free_port
@dataclass
class ProcHandle:
process: subprocess.Popen
url: str
class RouterManager:
"""Helper to spawn a router process and interact with admin endpoints."""
def __init__(self):
self._children: List[subprocess.Popen] = []
def start_router(
self,
worker_urls: Optional[List[str]] = None,
policy: str = "round_robin",
port: Optional[int] = None,
extra: Optional[Dict] = None,
# PD options
pd_disaggregation: bool = False,
prefill_urls: Optional[List[tuple]] = None,
decode_urls: Optional[List[str]] = None,
prefill_policy: Optional[str] = None,
decode_policy: Optional[str] = None,
) -> ProcHandle:
worker_urls = worker_urls or []
port = port or find_free_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(port),
"--policy",
policy,
]
# Avoid Prometheus port collisions by assigning a free port per router
prom_port = find_free_port()
cmd.extend(
["--prometheus-port", str(prom_port), "--prometheus-host", "127.0.0.1"]
)
if worker_urls:
cmd.extend(["--worker-urls", *worker_urls])
# PD routing configuration
if pd_disaggregation:
cmd.append("--pd-disaggregation")
if prefill_urls:
for url, bport in prefill_urls:
if bport is None:
cmd.extend(["--prefill", url, "none"])
else:
cmd.extend(["--prefill", url, str(bport)])
if decode_urls:
for url in decode_urls:
cmd.extend(["--decode", url])
if prefill_policy:
cmd.extend(["--prefill-policy", prefill_policy])
if decode_policy:
cmd.extend(["--decode-policy", decode_policy])
# Map supported extras to CLI flags (subset for integration)
if extra:
flag_map = {
"max_payload_size": "--max-payload-size",
"dp_aware": "--dp-aware",
"api_key": "--api-key",
# Health/monitoring
"worker_startup_check_interval": "--worker-startup-check-interval",
# Cache-aware tuning
"cache_threshold": "--cache-threshold",
"balance_abs_threshold": "--balance-abs-threshold",
"balance_rel_threshold": "--balance-rel-threshold",
# Retry
"retry_max_retries": "--retry-max-retries",
"retry_initial_backoff_ms": "--retry-initial-backoff-ms",
"retry_max_backoff_ms": "--retry-max-backoff-ms",
"retry_backoff_multiplier": "--retry-backoff-multiplier",
"retry_jitter_factor": "--retry-jitter-factor",
"disable_retries": "--disable-retries",
# Circuit breaker
"cb_failure_threshold": "--cb-failure-threshold",
"cb_success_threshold": "--cb-success-threshold",
"cb_timeout_duration_secs": "--cb-timeout-duration-secs",
"cb_window_duration_secs": "--cb-window-duration-secs",
"disable_circuit_breaker": "--disable-circuit-breaker",
# Rate limiting
"max_concurrent_requests": "--max-concurrent-requests",
"queue_size": "--queue-size",
"queue_timeout_secs": "--queue-timeout-secs",
"rate_limit_tokens_per_second": "--rate-limit-tokens-per-second",
}
for k, v in extra.items():
if v is None:
continue
flag = flag_map.get(k)
if not flag:
continue
if isinstance(v, bool):
if v:
cmd.append(flag)
else:
cmd.extend([flag, str(v)])
proc = subprocess.Popen(cmd)
self._children.append(proc)
url = f"http://127.0.0.1:{port}"
self._wait_health(url)
return ProcHandle(process=proc, url=url)
def _wait_health(self, base_url: str, timeout: float = 30.0):
start = time.time()
with requests.Session() as s:
while time.time() - start < timeout:
try:
r = s.get(f"{base_url}/health", timeout=2)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(0.2)
raise TimeoutError(f"Router at {base_url} did not become healthy")
def add_worker(self, base_url: str, worker_url: str) -> None:
r = requests.post(f"{base_url}/add_worker", params={"url": worker_url})
assert r.status_code == 200, f"add_worker failed: {r.status_code} {r.text}"
def remove_worker(self, base_url: str, worker_url: str) -> None:
r = requests.post(f"{base_url}/remove_worker", params={"url": worker_url})
assert r.status_code == 200, f"remove_worker failed: {r.status_code} {r.text}"
def list_workers(self, base_url: str) -> list[str]:
r = requests.get(f"{base_url}/list_workers")
assert r.status_code == 200, f"list_workers failed: {r.status_code} {r.text}"
data = r.json()
return data.get("urls", [])
def stop_all(self):
for p in self._children:
if p.poll() is None:
p.terminate()
try:
p.wait(timeout=5)
except subprocess.TimeoutExpired:
p.kill()
self._children.clear()

1
sgl-router/py_test/integration/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Integration test package for the router."""

109
sgl-router/py_test/integration/conftest.py Normal file
View File

@@ -0,0 +1,109 @@
import os
import subprocess
import time
from pathlib import Path
from typing import Dict, Iterable, List, Tuple
import pytest
import requests
from ..fixtures.ports import find_free_port
from ..fixtures.router_manager import RouterManager
def pytest_configure(config):
config.addinivalue_line("markers", "integration: mark as router integration test")
@pytest.fixture
def router_manager() -> Iterable[RouterManager]:
mgr = RouterManager()
try:
yield mgr
finally:
mgr.stop_all()
def _spawn_mock_worker(args: List[str]) -> Tuple[subprocess.Popen, str, str]:
repo_root = Path(__file__).resolve().parents[2]
script = repo_root / "py_test" / "fixtures" / "mock_worker.py"
port = find_free_port()
worker_id = f"worker-{port}"
base_cmd = [
"python3",
str(script),
"--port",
str(port),
"--worker-id",
worker_id,
]
cmd = base_cmd + args
proc = subprocess.Popen(cmd)
url = f"http://127.0.0.1:{port}"
_wait_health(url)
return proc, url, worker_id
def _wait_health(url: str, timeout: float = 10.0):
start = time.time()
with requests.Session() as s:
while time.time() - start < timeout:
try:
r = s.get(f"{url}/health", timeout=1)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(0.1)
raise TimeoutError(f"Mock worker at {url} did not become healthy")
@pytest.fixture
def mock_worker():
"""Start a single healthy mock worker; yields (process, url, worker_id)."""
proc, url, worker_id = _spawn_mock_worker([])
try:
yield proc, url, worker_id
finally:
if proc.poll() is None:
proc.terminate()
try:
proc.wait(timeout=3)
except subprocess.TimeoutExpired:
proc.kill()
@pytest.fixture
def mock_workers():
"""Factory to start N workers with custom args.
Usage:
procs, urls, ids = mock_workers(n=3, args=["--latency-ms", "5"]) # same args for all
...
"""
procs: List[subprocess.Popen] = []
def _start(n: int, args: List[str] | None = None):
args = args or []
new_procs: List[subprocess.Popen] = []
urls: List[str] = []
ids: List[str] = []
for _ in range(n):
p, url, wid = _spawn_mock_worker(args)
procs.append(p)
new_procs.append(p)
urls.append(url)
ids.append(wid)
return new_procs, urls, ids
try:
yield _start
finally:
for p in procs:
if p.poll() is None:
p.terminate()
try:
p.wait(timeout=3)
except subprocess.TimeoutExpired:
p.kill()

1
sgl-router/py_test/integration/load_balancing/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Load balancing integration tests."""

73
sgl-router/py_test/integration/load_balancing/test_cache_aware.py Normal file
View File

@@ -0,0 +1,73 @@
import collections
import concurrent.futures
import uuid
import pytest
import requests
@pytest.mark.integration
def test_cache_aware_affinity(mock_workers, router_manager):
# Two workers; same prompt should stick to one due to cache tree
_, urls, ids = mock_workers(n=2)
rh = router_manager.start_router(worker_urls=urls, policy="cache_aware")
counts = collections.Counter()
with requests.Session() as s:
for i in range(12):
r = s.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "repeated prompt for cache",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
counts[wid] += 1
# Expect strong skew toward one worker (tree match); majority > 80%
top = max(counts.values())
assert top >= 10, counts
@pytest.mark.integration
def test_cache_aware_diverse_prompts_balances(mock_workers, router_manager):
# Add latency so concurrent requests overlap and influence load-based selection
_, urls, ids = mock_workers(n=3, args=["--latency-ms", "30"])
rh = router_manager.start_router(
worker_urls=urls,
policy="cache_aware",
extra={
"cache_threshold": 0.99,
"balance_abs_threshold": 0,
"balance_rel_threshold": 1.0,
},
)
counts = collections.Counter()
def call(i):
# Use diverse, unrelated prompts to avoid prefix matches entirely
prompt = str(uuid.uuid4())
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": prompt,
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
assert r.status_code == 200
return r.headers.get("X-Worker-Id") or r.json().get("worker_id")
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex:
for wid in ex.map(call, range(40)):
counts[wid] += 1
# Expect participation of at least two workers
assert sum(1 for v in counts.values() if v > 0) >= 2, counts

89
sgl-router/py_test/integration/load_balancing/test_power_of_two.py Normal file
View File

@@ -0,0 +1,89 @@
import collections
import concurrent.futures
import time
import pytest
import requests
@pytest.mark.integration
def test_power_of_two_prefers_less_loaded(mock_workers, router_manager):
# Start two workers: one slow (higher inflight), one fast
# Router monitors /get_load and Power-of-Two uses cached loads to choose
# Start one slow and one fast worker using the fixture factory
procs_slow, urls_slow, ids_slow = mock_workers(n=1, args=["--latency-ms", "200"])
procs_fast, urls_fast, ids_fast = mock_workers(n=1, args=["--latency-ms", "0"])
procs = procs_slow + procs_fast
urls = urls_slow + urls_fast
ids = ids_slow + ids_fast
slow_id = ids_slow[0]
rh = router_manager.start_router(
worker_urls=urls,
policy="power_of_two",
extra={"worker_startup_check_interval": 1},
)
# Prime: fire a burst to create measurable load on slow worker, then wait for monitor tick
def _prime_call(i):
try:
requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"warm-{i}",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
except Exception:
pass
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex:
list(ex.map(_prime_call, range(128)))
time.sleep(2)
# Apply direct background load on the slow worker to amplify load diff
def _direct_load(i):
try:
requests.post(
f"{slow_url}/v1/completions",
json={
"model": "test-model",
"prompt": f"bg-{i}",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
except Exception:
pass
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex:
list(ex.map(_direct_load, range(128)))
time.sleep(1)
def call(i):
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"p{i}",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
assert r.status_code == 200
return r.headers.get("X-Worker-Id") or r.json().get("worker_id")
counts = collections.Counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex:
for wid in ex.map(call, range(200)):
counts[wid] += 1
# Expect the slow worker (higher latency/inflight) to receive fewer requests
fast_worker_id = [i for i in ids if i != slow_id][0]
assert counts[slow_id] < counts[fast_worker_id], counts

33
sgl-router/py_test/integration/load_balancing/test_random.py Normal file
View File

@@ -0,0 +1,33 @@
import collections
import math
import pytest
import requests
@pytest.mark.integration
def test_random_distribution(mock_workers, router_manager):
procs, urls, ids = mock_workers(n=4)
rh = router_manager.start_router(worker_urls=urls, policy="random")
counts = collections.Counter()
N = 200
with requests.Session() as s:
for i in range(N):
r = s.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"p{i}",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
counts[wid] += 1
# simple statistical tolerance: each worker should be within ±50% of mean
mean = N / len(ids)
for wid in ids:
assert 0.5 * mean <= counts[wid] <= 1.5 * mean, counts

34
sgl-router/py_test/integration/load_balancing/test_round_robin.py Normal file
View File

@@ -0,0 +1,34 @@
import collections
import time
import pytest
import requests
@pytest.mark.integration
def test_round_robin_distribution(mock_workers, router_manager):
procs, urls, ids = mock_workers(n=3)
rh = router_manager.start_router(worker_urls=urls, policy="round_robin")
counts = collections.Counter()
with requests.Session() as s:
for i in range(30):
r = s.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"hello {i}",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
assert wid in ids
counts[wid] += 1
# Expect near-even distribution across 3 workers
# 30 requests -> ideally 10 each; allow small tolerance ±3
for wid in ids:
assert 7 <= counts[wid] <= 13, counts

38
sgl-router/py_test/integration/test_api_auth.py Normal file
View File

@@ -0,0 +1,38 @@
import pytest
import requests
@pytest.mark.integration
def test_router_api_key_enforcement(router_manager, mock_workers):
# Start backend requiring API key; router should forward Authorization header transparently
_, urls, _ = mock_workers(
n=1, args=["--require-api-key", "--api-key", "correct_api_key"]
)
rh = router_manager.start_router(
worker_urls=urls,
policy="round_robin",
extra={},
)
# No auth -> 401
r = requests.post(
f"{rh.url}/v1/completions",
json={"model": "test-model", "prompt": "x", "max_tokens": 1, "stream": False},
)
assert r.status_code == 401
# Invalid auth -> 401
r = requests.post(
f"{rh.url}/v1/completions",
json={"model": "test-model", "prompt": "x", "max_tokens": 1, "stream": False},
headers={"Authorization": "Bearer wrong"},
)
assert r.status_code == 401
# Correct auth -> 200
r = requests.post(
f"{rh.url}/v1/completions",
json={"model": "test-model", "prompt": "x", "max_tokens": 1, "stream": False},
headers={"Authorization": "Bearer correct_api_key"},
)
assert r.status_code == 200

191
sgl-router/py_test/integration/test_circuit_breaker.py Normal file
View File

@@ -0,0 +1,191 @@
import time
import pytest
import requests
@pytest.mark.integration
def test_circuit_breaker_opens_and_recovers(router_manager, mock_workers):
# A single worker that fails first 3 requests, then succeeds
_, [wurl], _ = mock_workers(n=1, args=["--fail-first-n", "3"]) # fails first 3
rh = router_manager.start_router(
worker_urls=[wurl],
policy="round_robin",
extra={
"cb_failure_threshold": 3,
"cb_success_threshold": 2,
"cb_timeout_duration_secs": 3,
"cb_window_duration_secs": 10,
"disable_retries": True,
},
)
def post_once():
return requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "trigger",
"max_tokens": 1,
"stream": False,
},
timeout=3,
)
saw_503 = False
for _ in range(8):
r = post_once()
if r.status_code == 503:
saw_503 = True
break
assert saw_503, "circuit breaker did not open to return 503"
time.sleep(4)
r1 = post_once()
r2 = post_once()
assert r1.status_code == 200 and r2.status_code == 200
@pytest.mark.integration
def test_circuit_breaker_half_open_failure_reopens(router_manager, mock_workers):
_, [wurl], _ = mock_workers(n=1, args=["--status-code", "500"]) # always fail
rh = router_manager.start_router(
worker_urls=[wurl],
policy="round_robin",
extra={
"cb_failure_threshold": 2,
"cb_success_threshold": 2,
"cb_timeout_duration_secs": 2,
"cb_window_duration_secs": 5,
"disable_retries": True,
},
)
def post_once():
return requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "x",
"max_tokens": 1,
"stream": False,
},
timeout=3,
)
opened = False
for _ in range(8):
r = post_once()
if r.status_code == 503:
opened = True
break
assert opened, "circuit breaker did not open"
time.sleep(3)
r = post_once()
assert r.status_code == 500
r2 = post_once()
assert r2.status_code == 503
@pytest.mark.integration
def test_circuit_breaker_disable_flag(router_manager, mock_workers):
_, [wurl], _ = mock_workers(n=1, args=["--status-code", "500"]) # always fail
rh = router_manager.start_router(
worker_urls=[wurl],
policy="round_robin",
extra={
"disable_circuit_breaker": True,
"disable_retries": True,
},
)
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "x",
"max_tokens": 1,
"stream": False,
},
timeout=3,
)
assert r.status_code == 500
@pytest.mark.integration
def test_circuit_breaker_per_worker_isolation(router_manager, mock_workers):
_, [fail_url], _ = mock_workers(n=1, args=["--status-code", "500"]) # always fail
_, [ok_url], _ = mock_workers(n=1)
rh = router_manager.start_router(
worker_urls=[fail_url, ok_url],
policy="round_robin",
extra={
"cb_failure_threshold": 2,
"cb_success_threshold": 1,
"cb_timeout_duration_secs": 2,
"cb_window_duration_secs": 10,
"disable_retries": True,
},
)
def post_once():
return requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "y",
"max_tokens": 1,
"stream": False,
},
timeout=3,
)
failures = 0
successes_after_open = 0
opened = False
for _ in range(30):
r = post_once()
if not opened:
if r.status_code == 500:
failures += 1
if failures >= 2:
_ = post_once()
_ = post_once()
opened = True
else:
if r.status_code == 200:
successes_after_open += 1
else:
assert False, f"Unexpected non-200 after CB open: {r.status_code}"
assert opened and successes_after_open >= 5
@pytest.mark.integration
def test_circuit_breaker_with_retries(router_manager, mock_workers):
_, [fail_url], _ = mock_workers(n=1, args=["--status-code", "500"]) # always fail
_, [ok_url], _ = mock_workers(n=1)
rh = router_manager.start_router(
worker_urls=[fail_url, ok_url],
policy="round_robin",
extra={
"retry_max_retries": 3,
"retry_initial_backoff_ms": 10,
"retry_max_backoff_ms": 50,
"cb_failure_threshold": 2,
"cb_success_threshold": 1,
"cb_timeout_duration_secs": 2,
"cb_window_duration_secs": 10,
},
)
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "z",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
assert r.status_code == 200

36
sgl-router/py_test/integration/test_fault_tolerance.py Normal file
View File

@@ -0,0 +1,36 @@
import concurrent.futures
import subprocess
import time
import pytest
import requests
@pytest.mark.integration
def test_worker_crash_reroute_with_retries(router_manager, mock_workers):
# Start one healthy and one that will crash on first request
_, [ok_url], _ = mock_workers(n=1)
_, [crash_url], _ = mock_workers(n=1, args=["--crash-on-request"])
rh = router_manager.start_router(
worker_urls=[crash_url, ok_url],
policy="round_robin",
extra={
"retry_max_retries": 3,
"retry_initial_backoff_ms": 10,
"retry_max_backoff_ms": 50,
},
)
# A single request should succeed via retry to the healthy worker
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "crash",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
assert r.status_code == 200
# mock_workers fixture handles cleanup

33
sgl-router/py_test/integration/test_payload_size.py Normal file
View File

@@ -0,0 +1,33 @@
import pytest
import requests
@pytest.mark.integration
def test_payload_size_limit(router_manager, mock_workers):
# Start one backend and a router with a 1MB payload limit
_, urls, _ = mock_workers(n=1)
rh = router_manager.start_router(
worker_urls=urls,
policy="round_robin",
extra={"max_payload_size": 1 * 1024 * 1024}, # 1MB
)
# Payload just under 1MB should succeed
payload_small = {
"model": "test-model",
"prompt": "x" * int(0.5 * 1024 * 1024), # ~0.5MB
"max_tokens": 1,
"stream": False,
}
r = requests.post(f"{rh.url}/v1/completions", json=payload_small)
assert r.status_code == 200
# Payload over 1MB should fail with 413
payload_large = {
"model": "test-model",
"prompt": "x" * int(1.2 * 1024 * 1024), # ~1.2MB
"max_tokens": 1,
"stream": False,
}
r = requests.post(f"{rh.url}/v1/completions", json=payload_large)
assert r.status_code == 413

127
sgl-router/py_test/integration/test_pd_routing.py Normal file
View File

@@ -0,0 +1,127 @@
import collections
import concurrent.futures
import subprocess
import time
import pytest
import requests
@pytest.mark.integration
def test_pd_power_of_two_decode_attribution(router_manager, mock_workers):
# Start two prefill and three decode mock workers via fixture
_, prefill_urls_raw, prefill_ids = mock_workers(n=2)
_, decode_urls_raw, decode_ids_list = mock_workers(n=3)
prefill_urls = [(u, None) for u in prefill_urls_raw]
decode_urls = list(decode_urls_raw)
decode_ids = set(decode_ids_list)
rh = router_manager.start_router(
policy="power_of_two",
pd_disaggregation=True,
prefill_urls=prefill_urls,
decode_urls=decode_urls,
extra={"worker_startup_check_interval": 1},
)
counts = collections.Counter()
with requests.Session() as s:
for i in range(30):
r = s.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"p{i}",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
assert wid in decode_ids
counts[wid] += 1
assert sum(1 for v in counts.values() if v > 0) >= 2
@pytest.mark.integration
def test_pd_power_of_two_skews_to_faster_decode(router_manager, mock_workers):
# Start two prefill workers (fast)
_, prefill_urls_raw, _ = mock_workers(n=2)
# Start two decode workers: one slow, one fast
_, [decode_slow_url], [slow_id] = mock_workers(
n=1, args=["--latency-ms", "300"]
) # slower decode
_, [decode_fast_url], [fast_id] = mock_workers(n=1)
decode_urls_raw = [decode_slow_url, decode_fast_url]
prefill_urls = [(u, None) for u in prefill_urls_raw]
decode_urls = list(decode_urls_raw)
rh = router_manager.start_router(
policy="power_of_two",
pd_disaggregation=True,
prefill_urls=prefill_urls,
decode_urls=decode_urls,
extra={"worker_startup_check_interval": 1},
)
def _prime_call(i):
try:
requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"warm-{i}",
"max_tokens": 1,
"stream": False,
},
timeout=8,
)
except Exception:
pass
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex:
list(ex.map(_prime_call, range(128)))
time.sleep(2)
def _direct_decode_load(i):
try:
requests.post(
f"{decode_slow_url}/v1/completions",
json={
"model": "test-model",
"prompt": f"bg-{i}",
"max_tokens": 1,
"stream": False,
},
timeout=8,
)
except Exception:
pass
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex:
list(ex.map(_direct_decode_load, range(128)))
time.sleep(1)
def call(i):
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"p{i}",
"max_tokens": 1,
"stream": False,
},
timeout=8,
)
assert r.status_code == 200
return r.headers.get("X-Worker-Id") or r.json().get("worker_id")
counts = collections.Counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex:
for wid in ex.map(call, range(200)):
counts[wid] += 1
assert counts[slow_id] < counts[fast_id], counts

91
sgl-router/py_test/integration/test_rate_limiting.py Normal file
View File

@@ -0,0 +1,91 @@
import concurrent.futures
import time
import pytest
import requests
@pytest.mark.integration
def test_rate_limit_and_queue(router_manager, mock_workers):
# One fast backend
_, urls, _ = mock_workers(n=1)
rh = router_manager.start_router(
worker_urls=urls,
policy="round_robin",
extra={
"max_concurrent_requests": 2,
"queue_size": 0, # no queue -> immediate 429 when limit exceeded
},
)
def call_once(i):
try:
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"p{i}",
"max_tokens": 1,
"stream": False,
},
timeout=3,
)
return r.status_code
except Exception:
return 599
# Fire a burst of concurrent requests
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex:
results = list(ex.map(call_once, range(16)))
# Expect some to succeed and some to be rate limited (429)
assert any(code == 200 for code in results)
assert any(code == 429 for code in results)
@pytest.mark.integration
def test_rate_limit_queue_and_timeout(router_manager, mock_workers):
# Slow backend: ~2s per request ensures queue wait > timeout
_, urls, _ = mock_workers(n=1, args=["--latency-ms", "2000"]) # 2.0s per request
# Allow 1 concurrent, queue up to 1, with 1s queue timeout
rh = router_manager.start_router(
worker_urls=urls,
policy="round_robin",
extra={
"max_concurrent_requests": 1,
"queue_size": 1,
"queue_timeout_secs": 1,
},
)
def call_once(i):
try:
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"q{i}",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
return r.status_code
except Exception:
return 599
# Fire 4 concurrent requests: 1 runs (~2s), 1 queued (times out at 1s -> 408), 2 overflow -> 429
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex:
results = list(ex.map(call_once, range(4)))
# We expect:
# - Some 200s (processed)
# - At least one 408 (queued too long and timed out)
# - Remaining non-200s are either 429 (queue overflow) or additional 408s depending on scheduling
assert any(code == 200 for code in results)
assert any(code == 408 for code in results), results
non200 = [c for c in results if c != 200]
assert len(non200) >= 2 and all(c in (408, 429) for c in non200), results

65
sgl-router/py_test/integration/test_retries.py Normal file
View File

@@ -0,0 +1,65 @@
import concurrent.futures
import subprocess
import time
import pytest
import requests
@pytest.mark.integration
def test_retry_reroutes_to_healthy_worker(router_manager, mock_workers):
# Worker A always 500; Worker B healthy
# Worker A always 500; Worker B/C healthy
_, [url_a], [id_a] = mock_workers(n=1, args=["--status-code", "500"]) # fail
_, [url_b], [id_b] = mock_workers(n=1)
_, [url_c], [id_c] = mock_workers(n=1)
rh = router_manager.start_router(
worker_urls=[url_a, url_b, url_c],
policy="round_robin",
extra={
"retry_max_retries": 3,
"retry_initial_backoff_ms": 10,
"retry_max_backoff_ms": 50,
},
)
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "x",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
assert wid == id_b # should have retried onto healthy worker
# mock_workers fixture handles cleanup
@pytest.mark.integration
def test_disable_retries_surfaces_failure(router_manager, mock_workers):
# Single failing worker, retries disabled -> should return 500
_, [url], [wid] = mock_workers(n=1, args=["--status-code", "500"]) # always fail
rh = router_manager.start_router(
worker_urls=[url],
policy="round_robin",
extra={
"disable_retries": True,
},
)
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "x",
"max_tokens": 1,
"stream": False,
},
timeout=5,
)
assert r.status_code == 500
# mock_workers fixture handles cleanup

36
sgl-router/py_test/integration/test_service_discovery_shim.py Normal file
View File

@@ -0,0 +1,36 @@
import pytest
import requests
@pytest.mark.integration
def test_discovery_shim_add_remove(router_manager, mock_workers):
# Start router without workers
rh = router_manager.start_router(worker_urls=[], policy="round_robin")
# Initially empty
urls = router_manager.list_workers(rh.url)
assert urls == []
# Add a worker (simulate discovery event)
_, [wurl], [wid] = mock_workers(n=1)
router_manager.add_worker(rh.url, wurl)
urls = router_manager.list_workers(rh.url)
assert wurl in urls
# Can serve a request
r = requests.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": "hi",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
# Remove worker (simulate pod deletion)
router_manager.remove_worker(rh.url, wurl)
urls = router_manager.list_workers(rh.url)
assert wurl not in urls
# mock_workers fixture handles cleanup

61
sgl-router/py_test/integration/test_worker_management.py Normal file
View File

@@ -0,0 +1,61 @@
import collections
import subprocess
import time
import pytest
import requests
@pytest.mark.integration
def test_add_and_remove_worker(mock_worker, router_manager, mock_workers):
# Start with a single worker
proc1, url1, id1 = mock_worker
rh = router_manager.start_router(worker_urls=[url1], policy="round_robin")
# Add a second worker
procs2, urls2, ids2 = mock_workers(n=1)
url2 = urls2[0]
id2 = ids2[0]
router_manager.add_worker(rh.url, url2)
# Send some requests and ensure both workers are seen
seen = set()
with requests.Session() as s:
for i in range(20):
r = s.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"x{i}",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
seen.add(wid)
if len(seen) == 2:
break
assert id1 in seen and id2 in seen
# Now remove the second worker
router_manager.remove_worker(rh.url, url2)
# After removal, subsequent requests should only come from first worker
with requests.Session() as s:
for i in range(10):
r = s.post(
f"{rh.url}/v1/completions",
json={
"model": "test-model",
"prompt": f"y{i}",
"max_tokens": 1,
"stream": False,
},
)
assert r.status_code == 200
wid = r.headers.get("X-Worker-Id") or r.json().get("worker_id")
assert wid == id1
# mock_workers fixture handles cleanup

7
sgl-router/py_test/unit/__init__.py Normal file
View File

@@ -0,0 +1,7 @@
"""
Unit tests for sglang_router.
This package contains fast, isolated unit tests for Python components
of the SGLang router. These tests focus on testing individual functions
and classes in isolation without starting actual router instances.
"""

628
sgl-router/py_test/unit/test_arg_parser.py Normal file
View File

@@ -0,0 +1,628 @@
"""
Unit tests for argument parsing functionality in sglang_router.
These tests focus on testing the argument parsing logic in isolation,
without starting actual router instances.
"""
import argparse
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from sglang_router.launch_router import RouterArgs, parse_router_args
from sglang_router.router import policy_from_str
class TestRouterArgs:
"""Test RouterArgs dataclass and its methods."""
def test_default_values(self):
"""Test that RouterArgs has correct default values."""
args = RouterArgs()
# Test basic defaults
assert args.host == "127.0.0.1"
assert args.port == 30000
assert args.policy == "cache_aware"
assert args.worker_urls == []
assert args.pd_disaggregation is False
assert args.prefill_urls == []
assert args.decode_urls == []
# Test PD-specific defaults
assert args.prefill_policy is None
assert args.decode_policy is None
# Test service discovery defaults
assert args.service_discovery is False
assert args.selector == {}
assert args.service_discovery_port == 80
assert args.service_discovery_namespace is None
# Test retry and circuit breaker defaults
assert args.retry_max_retries == 5
assert args.cb_failure_threshold == 10
assert args.disable_retries is False
assert args.disable_circuit_breaker is False
def test_parse_selector_valid(self):
"""Test parsing valid selector arguments."""
# Test single key-value pair
result = RouterArgs._parse_selector(["app=worker"])
assert result == {"app": "worker"}
# Test multiple key-value pairs
result = RouterArgs._parse_selector(["app=worker", "env=prod", "version=v1"])
assert result == {"app": "worker", "env": "prod", "version": "v1"}
# Test empty list
result = RouterArgs._parse_selector([])
assert result == {}
# Test None
result = RouterArgs._parse_selector(None)
assert result == {}
def test_parse_selector_invalid(self):
"""Test parsing invalid selector arguments."""
# Test malformed selector (no equals sign)
result = RouterArgs._parse_selector(["app"])
assert result == {}
# Test multiple equals signs (should use first one)
result = RouterArgs._parse_selector(["app=worker=extra"])
assert result == {"app": "worker=extra"}
def test_parse_prefill_urls_valid(self):
"""Test parsing valid prefill URL arguments."""
# Test with bootstrap port
result = RouterArgs._parse_prefill_urls([["http://prefill1:8000", "9000"]])
assert result == [("http://prefill1:8000", 9000)]
# Test with 'none' bootstrap port
result = RouterArgs._parse_prefill_urls([["http://prefill1:8000", "none"]])
assert result == [("http://prefill1:8000", None)]
# Test without bootstrap port
result = RouterArgs._parse_prefill_urls([["http://prefill1:8000"]])
assert result == [("http://prefill1:8000", None)]
# Test multiple prefill URLs
result = RouterArgs._parse_prefill_urls(
[
["http://prefill1:8000", "9000"],
["http://prefill2:8000", "none"],
["http://prefill3:8000"],
]
)
expected = [
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
("http://prefill3:8000", None),
]
assert result == expected
# Test empty list
result = RouterArgs._parse_prefill_urls([])
assert result == []
# Test None
result = RouterArgs._parse_prefill_urls(None)
assert result == []
def test_parse_prefill_urls_invalid(self):
"""Test parsing invalid prefill URL arguments."""
# Test invalid bootstrap port
with pytest.raises(ValueError, match="Invalid bootstrap port"):
RouterArgs._parse_prefill_urls([["http://prefill1:8000", "invalid"]])
def test_parse_decode_urls_valid(self):
"""Test parsing valid decode URL arguments."""
# Test single decode URL
result = RouterArgs._parse_decode_urls([["http://decode1:8001"]])
assert result == ["http://decode1:8001"]
# Test multiple decode URLs
result = RouterArgs._parse_decode_urls(
[["http://decode1:8001"], ["http://decode2:8001"]]
)
assert result == ["http://decode1:8001", "http://decode2:8001"]
# Test empty list
result = RouterArgs._parse_decode_urls([])
assert result == []
# Test None
result = RouterArgs._parse_decode_urls(None)
assert result == []
def test_from_cli_args_basic(self):
"""Test creating RouterArgs from basic CLI arguments."""
args = SimpleNamespace(
host="0.0.0.0",
port=30001,
worker_urls=["http://worker1:8000", "http://worker2:8000"],
policy="round_robin",
prefill=None,
decode=None,
router_policy="round_robin",
router_pd_disaggregation=False,
router_prefill_policy=None,
router_decode_policy=None,
router_worker_startup_timeout_secs=300,
router_worker_startup_check_interval=15,
router_cache_threshold=0.7,
router_balance_abs_threshold=128,
router_balance_rel_threshold=2.0,
router_eviction_interval=180,
router_max_tree_size=2**28,
router_max_payload_size=1024 * 1024 * 1024, # 1GB
router_dp_aware=True,
router_api_key="test-key",
router_log_dir="/tmp/logs",
router_log_level="debug",
router_service_discovery=True,
router_selector=["app=worker", "env=test"],
router_service_discovery_port=8080,
router_service_discovery_namespace="default",
router_prefill_selector=["app=prefill"],
router_decode_selector=["app=decode"],
router_prometheus_port=29000,
router_prometheus_host="0.0.0.0",
router_request_id_headers=["x-request-id", "x-trace-id"],
router_request_timeout_secs=1200,
router_max_concurrent_requests=512,
router_queue_size=200,
router_queue_timeout_secs=120,
router_rate_limit_tokens_per_second=100,
router_cors_allowed_origins=["http://localhost:3000"],
router_retry_max_retries=3,
router_retry_initial_backoff_ms=100,
router_retry_max_backoff_ms=10000,
router_retry_backoff_multiplier=2.0,
router_retry_jitter_factor=0.1,
router_cb_failure_threshold=5,
router_cb_success_threshold=2,
router_cb_timeout_duration_secs=30,
router_cb_window_duration_secs=60,
router_disable_retries=False,
router_disable_circuit_breaker=False,
router_health_failure_threshold=2,
router_health_success_threshold=1,
router_health_check_timeout_secs=3,
router_health_check_interval_secs=30,
router_health_check_endpoint="/healthz",
)
router_args = RouterArgs.from_cli_args(args, use_router_prefix=True)
# Test basic configuration
assert router_args.host == "0.0.0.0"
assert router_args.port == 30001
assert router_args.worker_urls == ["http://worker1:8000", "http://worker2:8000"]
assert router_args.policy == "round_robin"
# Test PD configuration
assert router_args.pd_disaggregation is False
assert router_args.prefill_urls == []
assert router_args.decode_urls == []
# Test service discovery
assert router_args.service_discovery is True
assert router_args.selector == {"app": "worker", "env": "test"}
assert router_args.service_discovery_port == 8080
assert router_args.service_discovery_namespace == "default"
assert router_args.prefill_selector == {"app": "prefill"}
assert router_args.decode_selector == {"app": "decode"}
# Test other configurations
assert router_args.dp_aware is True
assert router_args.api_key == "test-key"
assert router_args.log_dir == "/tmp/logs"
assert router_args.log_level == "debug"
assert router_args.prometheus_port == 29000
assert router_args.prometheus_host == "0.0.0.0"
assert router_args.request_id_headers == ["x-request-id", "x-trace-id"]
assert router_args.request_timeout_secs == 1200
assert router_args.max_concurrent_requests == 512
assert router_args.queue_size == 200
assert router_args.queue_timeout_secs == 120
assert router_args.rate_limit_tokens_per_second == 100
assert router_args.cors_allowed_origins == ["http://localhost:3000"]
# Test retry configuration
assert router_args.retry_max_retries == 3
assert router_args.retry_initial_backoff_ms == 100
assert router_args.retry_max_backoff_ms == 10000
assert router_args.retry_backoff_multiplier == 2.0
assert router_args.retry_jitter_factor == 0.1
# Test circuit breaker configuration
assert router_args.cb_failure_threshold == 5
assert router_args.cb_success_threshold == 2
assert router_args.cb_timeout_duration_secs == 30
assert router_args.cb_window_duration_secs == 60
assert router_args.disable_retries is False
assert router_args.disable_circuit_breaker is False
# Test health check configuration
assert router_args.health_failure_threshold == 2
assert router_args.health_success_threshold == 1
assert router_args.health_check_timeout_secs == 3
assert router_args.health_check_interval_secs == 30
assert router_args.health_check_endpoint == "/healthz"
# Note: model_path and tokenizer_path are not available in current RouterArgs
def test_from_cli_args_pd_mode(self):
"""Test creating RouterArgs from CLI arguments in PD mode."""
args = SimpleNamespace(
host="127.0.0.1",
port=30000,
worker_urls=[],
policy="cache_aware",
prefill=[
["http://prefill1:8000", "9000"],
["http://prefill2:8000", "none"],
],
decode=[["http://decode1:8001"], ["http://decode2:8001"]],
router_prefill=[
["http://prefill1:8000", "9000"],
["http://prefill2:8000", "none"],
],
router_decode=[["http://decode1:8001"], ["http://decode2:8001"]],
router_policy="cache_aware",
router_pd_disaggregation=True,
router_prefill_policy="power_of_two",
router_decode_policy="round_robin",
# Include all required fields with defaults
router_worker_startup_timeout_secs=600,
router_worker_startup_check_interval=30,
router_cache_threshold=0.3,
router_balance_abs_threshold=64,
router_balance_rel_threshold=1.5,
router_eviction_interval=120,
router_max_tree_size=2**26,
router_max_payload_size=512 * 1024 * 1024,
router_dp_aware=False,
router_api_key=None,
router_log_dir=None,
router_log_level=None,
router_service_discovery=False,
router_selector=None,
router_service_discovery_port=80,
router_service_discovery_namespace=None,
router_prefill_selector=None,
router_decode_selector=None,
router_prometheus_port=None,
router_prometheus_host=None,
router_request_id_headers=None,
router_request_timeout_secs=1800,
router_max_concurrent_requests=256,
router_queue_size=100,
router_queue_timeout_secs=60,
router_rate_limit_tokens_per_second=None,
router_cors_allowed_origins=[],
router_retry_max_retries=5,
router_retry_initial_backoff_ms=50,
router_retry_max_backoff_ms=30000,
router_retry_backoff_multiplier=1.5,
router_retry_jitter_factor=0.2,
router_cb_failure_threshold=10,
router_cb_success_threshold=3,
router_cb_timeout_duration_secs=60,
router_cb_window_duration_secs=120,
router_disable_retries=False,
router_disable_circuit_breaker=False,
router_health_failure_threshold=3,
router_health_success_threshold=2,
router_health_check_timeout_secs=5,
router_health_check_interval_secs=60,
router_health_check_endpoint="/health",
)
router_args = RouterArgs.from_cli_args(args, use_router_prefix=True)
# Test PD configuration
assert router_args.pd_disaggregation is True
assert router_args.prefill_urls == [
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
]
assert router_args.decode_urls == ["http://decode1:8001", "http://decode2:8001"]
assert router_args.prefill_policy == "power_of_two"
assert router_args.decode_policy == "round_robin"
assert router_args.policy == "cache_aware" # Main policy still set
def test_from_cli_args_without_prefix(self):
"""Test creating RouterArgs from CLI arguments without router prefix."""
args = SimpleNamespace(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000"],
policy="random",
prefill=None,
decode=None,
pd_disaggregation=False,
prefill_policy=None,
decode_policy=None,
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=120,
max_tree_size=2**26,
max_payload_size=512 * 1024 * 1024,
dp_aware=False,
api_key=None,
log_dir=None,
log_level=None,
service_discovery=False,
selector=None,
service_discovery_port=80,
service_discovery_namespace=None,
prefill_selector=None,
decode_selector=None,
prometheus_port=None,
prometheus_host=None,
request_id_headers=None,
request_timeout_secs=1800,
max_concurrent_requests=256,
queue_size=100,
queue_timeout_secs=60,
rate_limit_tokens_per_second=None,
cors_allowed_origins=[],
retry_max_retries=5,
retry_initial_backoff_ms=50,
retry_max_backoff_ms=30000,
retry_backoff_multiplier=1.5,
retry_jitter_factor=0.2,
cb_failure_threshold=10,
cb_success_threshold=3,
cb_timeout_duration_secs=60,
cb_window_duration_secs=120,
disable_retries=False,
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="/health",
model_path=None,
tokenizer_path=None,
)
router_args = RouterArgs.from_cli_args(args, use_router_prefix=False)
assert router_args.host == "127.0.0.1"
assert router_args.port == 30000
assert router_args.worker_urls == ["http://worker1:8000"]
assert router_args.policy == "random"
assert router_args.pd_disaggregation is False
class TestPolicyFromStr:
"""Test policy string to enum conversion."""
def test_valid_policies(self):
"""Test conversion of valid policy strings."""
from sglang_router_rs import PolicyType
assert policy_from_str("random") == PolicyType.Random
assert policy_from_str("round_robin") == PolicyType.RoundRobin
assert policy_from_str("cache_aware") == PolicyType.CacheAware
assert policy_from_str("power_of_two") == PolicyType.PowerOfTwo
def test_invalid_policy(self):
"""Test conversion of invalid policy string."""
with pytest.raises(KeyError):
policy_from_str("invalid_policy")
class TestParseRouterArgs:
"""Test the parse_router_args function."""
def test_parse_basic_args(self):
"""Test parsing basic router arguments."""
args = [
"--host",
"0.0.0.0",
"--port",
"30001",
"--worker-urls",
"http://worker1:8000",
"http://worker2:8000",
"--policy",
"round_robin",
]
router_args = parse_router_args(args)
assert router_args.host == "0.0.0.0"
assert router_args.port == 30001
assert router_args.worker_urls == ["http://worker1:8000", "http://worker2:8000"]
assert router_args.policy == "round_robin"
def test_parse_pd_args(self):
"""Test parsing PD disaggregated mode arguments."""
args = [
"--pd-disaggregation",
"--prefill",
"http://prefill1:8000",
"9000",
"--prefill",
"http://prefill2:8000",
"none",
"--decode",
"http://decode1:8001",
"--decode",
"http://decode2:8001",
"--prefill-policy",
"power_of_two",
"--decode-policy",
"round_robin",
]
router_args = parse_router_args(args)
assert router_args.pd_disaggregation is True
assert router_args.prefill_urls == [
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
]
assert router_args.decode_urls == ["http://decode1:8001", "http://decode2:8001"]
assert router_args.prefill_policy == "power_of_two"
assert router_args.decode_policy == "round_robin"
def test_parse_service_discovery_args(self):
"""Test parsing service discovery arguments."""
args = [
"--service-discovery",
"--selector",
"app=worker",
"env=prod",
"--service-discovery-port",
"8080",
"--service-discovery-namespace",
"default",
]
router_args = parse_router_args(args)
assert router_args.service_discovery is True
assert router_args.selector == {"app": "worker", "env": "prod"}
assert router_args.service_discovery_port == 8080
assert router_args.service_discovery_namespace == "default"
def test_parse_retry_and_circuit_breaker_args(self):
"""Test parsing retry and circuit breaker arguments."""
args = [
"--retry-max-retries",
"3",
"--retry-initial-backoff-ms",
"100",
"--retry-max-backoff-ms",
"10000",
"--retry-backoff-multiplier",
"2.0",
"--retry-jitter-factor",
"0.1",
"--disable-retries",
"--cb-failure-threshold",
"5",
"--cb-success-threshold",
"2",
"--cb-timeout-duration-secs",
"30",
"--cb-window-duration-secs",
"60",
"--disable-circuit-breaker",
]
router_args = parse_router_args(args)
# Test retry configuration
assert router_args.retry_max_retries == 3
assert router_args.retry_initial_backoff_ms == 100
assert router_args.retry_max_backoff_ms == 10000
assert router_args.retry_backoff_multiplier == 2.0
assert router_args.retry_jitter_factor == 0.1
assert router_args.disable_retries is True
# Test circuit breaker configuration
assert router_args.cb_failure_threshold == 5
assert router_args.cb_success_threshold == 2
assert router_args.cb_timeout_duration_secs == 30
assert router_args.cb_window_duration_secs == 60
assert router_args.disable_circuit_breaker is True
def test_parse_rate_limiting_args(self):
"""Test parsing rate limiting arguments."""
args = [
"--max-concurrent-requests",
"512",
"--queue-size",
"200",
"--queue-timeout-secs",
"120",
"--rate-limit-tokens-per-second",
"100",
]
router_args = parse_router_args(args)
assert router_args.max_concurrent_requests == 512
assert router_args.queue_size == 200
assert router_args.queue_timeout_secs == 120
assert router_args.rate_limit_tokens_per_second == 100
def test_parse_health_check_args(self):
"""Test parsing health check arguments."""
args = [
"--health-failure-threshold",
"2",
"--health-success-threshold",
"1",
"--health-check-timeout-secs",
"3",
"--health-check-interval-secs",
"30",
"--health-check-endpoint",
"/healthz",
]
router_args = parse_router_args(args)
assert router_args.health_failure_threshold == 2
assert router_args.health_success_threshold == 1
assert router_args.health_check_timeout_secs == 3
assert router_args.health_check_interval_secs == 30
assert router_args.health_check_endpoint == "/healthz"
def test_parse_cors_args(self):
"""Test parsing CORS arguments."""
args = [
"--cors-allowed-origins",
"http://localhost:3000",
"https://example.com",
]
router_args = parse_router_args(args)
assert router_args.cors_allowed_origins == [
"http://localhost:3000",
"https://example.com",
]
def test_parse_tokenizer_args(self):
"""Test parsing tokenizer arguments."""
# Note: model-path and tokenizer-path arguments are not available in current implementation
# This test is skipped until those arguments are added
pytest.skip("Tokenizer arguments not available in current implementation")
def test_parse_invalid_args(self):
"""Test parsing invalid arguments."""
# Test invalid policy
with pytest.raises(SystemExit):
parse_router_args(["--policy", "invalid_policy"])
# Test invalid bootstrap port
with pytest.raises(ValueError, match="Invalid bootstrap port"):
parse_router_args(
[
"--pd-disaggregation",
"--prefill",
"http://prefill1:8000",
"invalid_port",
]
)
def test_help_output(self):
"""Test that help output is generated correctly."""
with pytest.raises(SystemExit) as exc_info:
parse_router_args(["--help"])
# SystemExit with code 0 indicates help was displayed
assert exc_info.value.code == 0

421
sgl-router/py_test/unit/test_router_config.py Normal file
View File

@@ -0,0 +1,421 @@
"""
Unit tests for router configuration validation and setup.
These tests focus on testing the router configuration logic in isolation,
including validation of configuration parameters and their interactions.
"""
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from sglang_router.launch_router import RouterArgs, launch_router
from sglang_router.router import policy_from_str
from sglang_router_rs import PolicyType
class TestRouterConfigValidation:
"""Test router configuration validation logic."""
def test_valid_basic_config(self):
"""Test that a valid basic configuration passes validation."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000", "http://worker2:8000"],
policy="cache_aware",
)
# Should not raise any exceptions
assert args.host == "127.0.0.1"
assert args.port == 30000
assert args.worker_urls == ["http://worker1:8000", "http://worker2:8000"]
assert args.policy == "cache_aware"
def test_valid_pd_config(self):
"""Test that a valid PD configuration passes validation."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
pd_disaggregation=True,
prefill_urls=[
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
],
decode_urls=["http://decode1:8001", "http://decode2:8001"],
policy="cache_aware",
)
assert args.pd_disaggregation is True
assert args.prefill_urls == [
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
]
assert args.decode_urls == ["http://decode1:8001", "http://decode2:8001"]
assert args.policy == "cache_aware"
def test_pd_config_without_urls_raises_error(self):
"""Test that PD mode without URLs raises validation error."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=False,
)
# This should raise an error when trying to launch
with pytest.raises(
ValueError, match="PD disaggregation mode requires --prefill"
):
launch_router(args)
def test_pd_config_with_service_discovery_allows_empty_urls(self):
"""Test that PD mode with service discovery allows empty URLs."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=True,
)
# Should not raise validation error when service discovery is enabled
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_regular_mode_without_workers_allows_empty_urls(self):
"""Test that regular mode allows empty worker URLs."""
args = RouterArgs(worker_urls=[], service_discovery=False)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_cache_threshold_validation(self):
"""Test cache threshold validation."""
# Valid cache threshold
args = RouterArgs(cache_threshold=0.5)
assert args.cache_threshold == 0.5
# Edge cases
args = RouterArgs(cache_threshold=0.0)
assert args.cache_threshold == 0.0
args = RouterArgs(cache_threshold=1.0)
assert args.cache_threshold == 1.0
def test_balance_threshold_validation(self):
"""Test load balancing threshold validation."""
# Valid thresholds
args = RouterArgs(balance_abs_threshold=64, balance_rel_threshold=1.5)
assert args.balance_abs_threshold == 64
assert args.balance_rel_threshold == 1.5
# Edge cases
args = RouterArgs(balance_abs_threshold=0, balance_rel_threshold=1.0)
assert args.balance_abs_threshold == 0
assert args.balance_rel_threshold == 1.0
def test_timeout_validation(self):
"""Test timeout parameter validation."""
# Valid timeouts
args = RouterArgs(
worker_startup_timeout_secs=600,
worker_startup_check_interval=30,
request_timeout_secs=1800,
queue_timeout_secs=60,
)
assert args.worker_startup_timeout_secs == 600
assert args.worker_startup_check_interval == 30
assert args.request_timeout_secs == 1800
assert args.queue_timeout_secs == 60
def test_retry_config_validation(self):
"""Test retry configuration validation."""
# Valid retry config
args = RouterArgs(
retry_max_retries=5,
retry_initial_backoff_ms=50,
retry_max_backoff_ms=30000,
retry_backoff_multiplier=1.5,
retry_jitter_factor=0.2,
disable_retries=False,
)
assert args.retry_max_retries == 5
assert args.retry_initial_backoff_ms == 50
assert args.retry_max_backoff_ms == 30000
assert args.retry_backoff_multiplier == 1.5
assert args.retry_jitter_factor == 0.2
assert args.disable_retries is False
def test_circuit_breaker_config_validation(self):
"""Test circuit breaker configuration validation."""
# Valid circuit breaker config
args = RouterArgs(
cb_failure_threshold=10,
cb_success_threshold=3,
cb_timeout_duration_secs=60,
cb_window_duration_secs=120,
disable_circuit_breaker=False,
)
assert args.cb_failure_threshold == 10
assert args.cb_success_threshold == 3
assert args.cb_timeout_duration_secs == 60
assert args.cb_window_duration_secs == 120
assert args.disable_circuit_breaker is False
def test_health_check_config_validation(self):
"""Test health check configuration validation."""
# Valid health check config
args = RouterArgs(
health_failure_threshold=3,
health_success_threshold=2,
health_check_timeout_secs=5,
health_check_interval_secs=60,
health_check_endpoint="/health",
)
assert args.health_failure_threshold == 3
assert args.health_success_threshold == 2
assert args.health_check_timeout_secs == 5
assert args.health_check_interval_secs == 60
assert args.health_check_endpoint == "/health"
def test_rate_limiting_config_validation(self):
"""Test rate limiting configuration validation."""
# Valid rate limiting config
args = RouterArgs(
max_concurrent_requests=256,
queue_size=100,
queue_timeout_secs=60,
rate_limit_tokens_per_second=100,
)
assert args.max_concurrent_requests == 256
assert args.queue_size == 100
assert args.queue_timeout_secs == 60
assert args.rate_limit_tokens_per_second == 100
def test_service_discovery_config_validation(self):
"""Test service discovery configuration validation."""
# Valid service discovery config
args = RouterArgs(
service_discovery=True,
selector={"app": "worker", "env": "prod"},
service_discovery_port=8080,
service_discovery_namespace="default",
)
assert args.service_discovery is True
assert args.selector == {"app": "worker", "env": "prod"}
assert args.service_discovery_port == 8080
assert args.service_discovery_namespace == "default"
def test_pd_service_discovery_config_validation(self):
"""Test PD service discovery configuration validation."""
# Valid PD service discovery config
args = RouterArgs(
pd_disaggregation=True,
service_discovery=True,
prefill_selector={"app": "prefill"},
decode_selector={"app": "decode"},
bootstrap_port_annotation="sglang.ai/bootstrap-port",
)
assert args.pd_disaggregation is True
assert args.service_discovery is True
assert args.prefill_selector == {"app": "prefill"}
assert args.decode_selector == {"app": "decode"}
assert args.bootstrap_port_annotation == "sglang.ai/bootstrap-port"
def test_prometheus_config_validation(self):
"""Test Prometheus configuration validation."""
# Valid Prometheus config
args = RouterArgs(prometheus_port=29000, prometheus_host="127.0.0.1")
assert args.prometheus_port == 29000
assert args.prometheus_host == "127.0.0.1"
def test_cors_config_validation(self):
"""Test CORS configuration validation."""
# Valid CORS config
args = RouterArgs(
cors_allowed_origins=["http://localhost:3000", "https://example.com"]
)
assert args.cors_allowed_origins == [
"http://localhost:3000",
"https://example.com",
]
def test_tokenizer_config_validation(self):
"""Test tokenizer configuration validation."""
# Note: model_path and tokenizer_path are not available in current RouterArgs
pytest.skip("Tokenizer configuration not available in current implementation")
def test_dp_aware_config_validation(self):
"""Test data parallelism aware configuration validation."""
# Valid DP aware config
args = RouterArgs(dp_aware=True, api_key="test-api-key")
assert args.dp_aware is True
assert args.api_key == "test-api-key"
def test_request_id_headers_validation(self):
"""Test request ID headers configuration validation."""
# Valid request ID headers config
args = RouterArgs(
request_id_headers=["x-request-id", "x-trace-id", "x-correlation-id"]
)
assert args.request_id_headers == [
"x-request-id",
"x-trace-id",
"x-correlation-id",
]
def test_policy_consistency_validation(self):
"""Test policy consistency validation in PD mode."""
# Test with both prefill and decode policies specified
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy="power_of_two",
decode_policy="round_robin",
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_policy_fallback_validation(self):
"""Test policy fallback validation in PD mode."""
# Test with only prefill policy specified
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy="power_of_two",
decode_policy=None,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_policy_enum_conversion(self):
"""Test policy string to enum conversion."""
# Test all valid policy conversions
assert policy_from_str("random") == PolicyType.Random
assert policy_from_str("round_robin") == PolicyType.RoundRobin
assert policy_from_str("cache_aware") == PolicyType.CacheAware
assert policy_from_str("power_of_two") == PolicyType.PowerOfTwo
def test_invalid_policy_enum_conversion(self):
"""Test invalid policy string to enum conversion."""
with pytest.raises(KeyError):
policy_from_str("invalid_policy")
def test_config_immutability(self):
"""Test that configuration objects are properly immutable."""
args = RouterArgs(
host="127.0.0.1", port=30000, worker_urls=["http://worker1:8000"]
)
# Test that we can't modify the configuration after creation
# (This is more of a design test - dataclasses are mutable by default)
original_host = args.host
args.host = "0.0.0.0"
assert args.host == "0.0.0.0" # Dataclasses are mutable
assert args.host != original_host
def test_config_defaults_consistency(self):
"""Test that configuration defaults are consistent."""
args1 = RouterArgs()
args2 = RouterArgs()
# Both instances should have the same defaults
assert args1.host == args2.host
assert args1.port == args2.port
assert args1.policy == args2.policy
assert args1.worker_urls == args2.worker_urls
assert args1.pd_disaggregation == args2.pd_disaggregation
def test_config_serialization(self):
"""Test that configuration can be serialized/deserialized."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000"],
policy="cache_aware",
cache_threshold=0.5,
)
# Test that we can access all attributes
assert hasattr(args, "host")
assert hasattr(args, "port")
assert hasattr(args, "worker_urls")
assert hasattr(args, "policy")
assert hasattr(args, "cache_threshold")
def test_config_with_none_values(self):
"""Test configuration with None values."""
args = RouterArgs(
api_key=None,
log_dir=None,
log_level=None,
prometheus_port=None,
prometheus_host=None,
request_id_headers=None,
rate_limit_tokens_per_second=None,
service_discovery_namespace=None,
)
# All None values should be preserved
assert args.api_key is None
assert args.log_dir is None
assert args.log_level is None
assert args.prometheus_port is None
assert args.prometheus_host is None
assert args.request_id_headers is None
assert args.rate_limit_tokens_per_second is None
assert args.service_discovery_namespace is None
def test_config_with_empty_lists(self):
"""Test configuration with empty lists."""
args = RouterArgs(
worker_urls=[], prefill_urls=[], decode_urls=[], cors_allowed_origins=[]
)
# All empty lists should be preserved
assert args.worker_urls == []
assert args.prefill_urls == []
assert args.decode_urls == []
assert args.cors_allowed_origins == []
def test_config_with_empty_dicts(self):
"""Test configuration with empty dictionaries."""
args = RouterArgs(selector={}, prefill_selector={}, decode_selector={})
# All empty dictionaries should be preserved
assert args.selector == {}
assert args.prefill_selector == {}
assert args.decode_selector == {}

1053
sgl-router/py_test/unit/test_startup_sequence.py Normal file
View File

File diff suppressed because it is too large Load Diff

506
sgl-router/py_test/unit/test_validation.py Normal file
View File

@@ -0,0 +1,506 @@
"""
Unit tests for validation logic in sglang_router.
These tests focus on testing the validation logic in isolation,
including parameter validation, URL validation, and configuration validation.
"""
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from sglang_router.launch_router import RouterArgs, launch_router
class TestURLValidation:
"""Test URL validation logic."""
def test_valid_worker_urls(self):
"""Test validation of valid worker URLs."""
valid_urls = [
"http://worker1:8000",
"https://worker2:8000",
"http://localhost:8000",
"http://127.0.0.1:8000",
"http://192.168.1.100:8000",
"http://worker.example.com:8000",
]
for url in valid_urls:
args = RouterArgs(worker_urls=[url])
# Should not raise any validation errors
assert url in args.worker_urls
def test_valid_prefill_urls(self):
"""Test validation of valid prefill URLs."""
valid_prefill_urls = [
("http://prefill1:8000", 9000),
("https://prefill2:8000", None),
("http://localhost:8000", 9000),
("http://127.0.0.1:8000", None),
]
for url, bootstrap_port in valid_prefill_urls:
args = RouterArgs(prefill_urls=[(url, bootstrap_port)])
# Should not raise any validation errors
assert (url, bootstrap_port) in args.prefill_urls
def test_valid_decode_urls(self):
"""Test validation of valid decode URLs."""
valid_decode_urls = [
"http://decode1:8001",
"https://decode2:8001",
"http://localhost:8001",
"http://127.0.0.1:8001",
]
for url in valid_decode_urls:
args = RouterArgs(decode_urls=[url])
# Should not raise any validation errors
assert url in args.decode_urls
def test_malformed_urls(self):
"""Test handling of malformed URLs."""
# Note: The current implementation doesn't validate URL format
# This test documents the current behavior
malformed_urls = [
"not-a-url",
"ftp://worker1:8000", # Wrong protocol
"http://", # Missing host
":8000", # Missing protocol and host
"http://worker1", # Missing port
]
for url in malformed_urls:
args = RouterArgs(worker_urls=[url])
# Currently, malformed URLs are accepted
# This might be something to improve in the future
assert url in args.worker_urls
class TestPortValidation:
"""Test port validation logic."""
def test_valid_ports(self):
"""Test validation of valid port numbers."""
valid_ports = [1, 80, 8000, 30000, 65535]
for port in valid_ports:
args = RouterArgs(port=port)
assert args.port == port
def test_invalid_ports(self):
"""Test handling of invalid port numbers."""
# Note: The current implementation doesn't validate port ranges
# This test documents the current behavior
invalid_ports = [0, -1, 65536, 70000]
for port in invalid_ports:
args = RouterArgs(port=port)
# Currently, invalid ports are accepted
# This might be something to improve in the future
assert args.port == port
def test_bootstrap_port_validation(self):
"""Test validation of bootstrap ports in PD mode."""
valid_bootstrap_ports = [1, 80, 9000, 30000, 65535, None]
for bootstrap_port in valid_bootstrap_ports:
args = RouterArgs(prefill_urls=[("http://prefill1:8000", bootstrap_port)])
assert args.prefill_urls[0][1] == bootstrap_port
class TestParameterValidation:
"""Test parameter validation logic."""
def test_cache_threshold_validation(self):
"""Test cache threshold parameter validation."""
# Valid cache thresholds
valid_thresholds = [0.0, 0.1, 0.5, 0.9, 1.0]
for threshold in valid_thresholds:
args = RouterArgs(cache_threshold=threshold)
assert args.cache_threshold == threshold
def test_balance_threshold_validation(self):
"""Test load balancing threshold parameter validation."""
# Valid absolute thresholds
valid_abs_thresholds = [0, 1, 32, 64, 128, 1000]
for threshold in valid_abs_thresholds:
args = RouterArgs(balance_abs_threshold=threshold)
assert args.balance_abs_threshold == threshold
# Valid relative thresholds
valid_rel_thresholds = [1.0, 1.1, 1.5, 2.0, 10.0]
for threshold in valid_rel_thresholds:
args = RouterArgs(balance_rel_threshold=threshold)
assert args.balance_rel_threshold == threshold
def test_timeout_validation(self):
"""Test timeout parameter validation."""
# Valid timeouts
valid_timeouts = [1, 30, 60, 300, 600, 1800, 3600]
for timeout in valid_timeouts:
args = RouterArgs(
worker_startup_timeout_secs=timeout,
worker_startup_check_interval=timeout,
request_timeout_secs=timeout,
queue_timeout_secs=timeout,
)
assert args.worker_startup_timeout_secs == timeout
assert args.worker_startup_check_interval == timeout
assert args.request_timeout_secs == timeout
assert args.queue_timeout_secs == timeout
def test_retry_parameter_validation(self):
"""Test retry parameter validation."""
# Valid retry parameters
valid_retry_counts = [0, 1, 3, 5, 10]
for count in valid_retry_counts:
args = RouterArgs(retry_max_retries=count)
assert args.retry_max_retries == count
# Valid backoff parameters
valid_backoff_ms = [1, 50, 100, 1000, 30000]
for backoff in valid_backoff_ms:
args = RouterArgs(
retry_initial_backoff_ms=backoff, retry_max_backoff_ms=backoff
)
assert args.retry_initial_backoff_ms == backoff
assert args.retry_max_backoff_ms == backoff
# Valid multiplier parameters
valid_multipliers = [1.0, 1.5, 2.0, 3.0]
for multiplier in valid_multipliers:
args = RouterArgs(retry_backoff_multiplier=multiplier)
assert args.retry_backoff_multiplier == multiplier
# Valid jitter parameters
valid_jitter = [0.0, 0.1, 0.2, 0.5]
for jitter in valid_jitter:
args = RouterArgs(retry_jitter_factor=jitter)
assert args.retry_jitter_factor == jitter
def test_circuit_breaker_parameter_validation(self):
"""Test circuit breaker parameter validation."""
# Valid failure thresholds
valid_failure_thresholds = [1, 3, 5, 10, 20]
for threshold in valid_failure_thresholds:
args = RouterArgs(cb_failure_threshold=threshold)
assert args.cb_failure_threshold == threshold
# Valid success thresholds
valid_success_thresholds = [1, 2, 3, 5]
for threshold in valid_success_thresholds:
args = RouterArgs(cb_success_threshold=threshold)
assert args.cb_success_threshold == threshold
# Valid timeout durations
valid_timeouts = [10, 30, 60, 120, 300]
for timeout in valid_timeouts:
args = RouterArgs(
cb_timeout_duration_secs=timeout, cb_window_duration_secs=timeout
)
assert args.cb_timeout_duration_secs == timeout
assert args.cb_window_duration_secs == timeout
def test_health_check_parameter_validation(self):
"""Test health check parameter validation."""
# Valid failure thresholds
valid_failure_thresholds = [1, 2, 3, 5, 10]
for threshold in valid_failure_thresholds:
args = RouterArgs(health_failure_threshold=threshold)
assert args.health_failure_threshold == threshold
# Valid success thresholds
valid_success_thresholds = [1, 2, 3, 5]
for threshold in valid_success_thresholds:
args = RouterArgs(health_success_threshold=threshold)
assert args.health_success_threshold == threshold
# Valid timeouts and intervals
valid_times = [1, 5, 10, 30, 60, 120]
for time_val in valid_times:
args = RouterArgs(
health_check_timeout_secs=time_val, health_check_interval_secs=time_val
)
assert args.health_check_timeout_secs == time_val
assert args.health_check_interval_secs == time_val
def test_rate_limiting_parameter_validation(self):
"""Test rate limiting parameter validation."""
# Valid concurrent request limits
valid_limits = [1, 10, 64, 256, 512, 1000]
for limit in valid_limits:
args = RouterArgs(max_concurrent_requests=limit)
assert args.max_concurrent_requests == limit
# Valid queue sizes
valid_queue_sizes = [0, 10, 50, 100, 500, 1000]
for size in valid_queue_sizes:
args = RouterArgs(queue_size=size)
assert args.queue_size == size
# Valid token rates
valid_rates = [1, 10, 50, 100, 500, 1000]
for rate in valid_rates:
args = RouterArgs(rate_limit_tokens_per_second=rate)
assert args.rate_limit_tokens_per_second == rate
def test_tree_size_validation(self):
"""Test tree size parameter validation."""
# Valid tree sizes (powers of 2)
valid_sizes = [2**10, 2**20, 2**24, 2**26, 2**28, 2**30]
for size in valid_sizes:
args = RouterArgs(max_tree_size=size)
assert args.max_tree_size == size
def test_payload_size_validation(self):
"""Test payload size parameter validation."""
# Valid payload sizes
valid_sizes = [
1024, # 1KB
1024 * 1024, # 1MB
10 * 1024 * 1024, # 10MB
100 * 1024 * 1024, # 100MB
512 * 1024 * 1024, # 512MB
1024 * 1024 * 1024, # 1GB
]
for size in valid_sizes:
args = RouterArgs(max_payload_size=size)
assert args.max_payload_size == size
class TestConfigurationValidation:
"""Test configuration validation logic."""
def test_pd_mode_validation(self):
"""Test PD mode configuration validation."""
# Valid PD configuration
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", 9000)],
decode_urls=["http://decode1:8001"],
)
assert args.pd_disaggregation is True
assert len(args.prefill_urls) > 0
assert len(args.decode_urls) > 0
def test_service_discovery_validation(self):
"""Test service discovery configuration validation."""
# Valid service discovery configuration
args = RouterArgs(
service_discovery=True,
selector={"app": "worker", "env": "prod"},
service_discovery_port=8080,
service_discovery_namespace="default",
)
assert args.service_discovery is True
assert args.selector == {"app": "worker", "env": "prod"}
assert args.service_discovery_port == 8080
assert args.service_discovery_namespace == "default"
def test_pd_service_discovery_validation(self):
"""Test PD service discovery configuration validation."""
# Valid PD service discovery configuration
args = RouterArgs(
pd_disaggregation=True,
service_discovery=True,
prefill_selector={"app": "prefill"},
decode_selector={"app": "decode"},
)
assert args.pd_disaggregation is True
assert args.service_discovery is True
assert args.prefill_selector == {"app": "prefill"}
assert args.decode_selector == {"app": "decode"}
def test_policy_validation(self):
"""Test policy configuration validation."""
# Valid policies
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two"]
for policy in valid_policies:
args = RouterArgs(policy=policy)
assert args.policy == policy
def test_pd_policy_validation(self):
"""Test PD policy configuration validation."""
# Valid PD policies
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two"]
for prefill_policy in valid_policies:
for decode_policy in valid_policies:
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
prefill_policy=prefill_policy,
decode_policy=decode_policy,
)
assert args.prefill_policy == prefill_policy
assert args.decode_policy == decode_policy
def test_cors_validation(self):
"""Test CORS configuration validation."""
# Valid CORS origins
valid_origins = [
[],
["http://localhost:3000"],
["https://example.com"],
["http://localhost:3000", "https://example.com"],
["*"], # Wildcard (if supported)
]
for origins in valid_origins:
args = RouterArgs(cors_allowed_origins=origins)
assert args.cors_allowed_origins == origins
def test_logging_validation(self):
"""Test logging configuration validation."""
# Valid log levels
valid_log_levels = ["debug", "info", "warning", "error", "critical"]
for level in valid_log_levels:
args = RouterArgs(log_level=level)
assert args.log_level == level
def test_prometheus_validation(self):
"""Test Prometheus configuration validation."""
# Valid Prometheus configuration
args = RouterArgs(prometheus_port=29000, prometheus_host="127.0.0.1")
assert args.prometheus_port == 29000
assert args.prometheus_host == "127.0.0.1"
def test_tokenizer_validation(self):
"""Test tokenizer configuration validation."""
# Note: model_path and tokenizer_path are not available in current RouterArgs
pytest.skip("Tokenizer configuration not available in current implementation")
def test_request_id_headers_validation(self):
"""Test request ID headers configuration validation."""
# Valid request ID headers
valid_headers = [
["x-request-id"],
["x-request-id", "x-trace-id"],
["x-request-id", "x-trace-id", "x-correlation-id"],
["custom-header"],
]
for headers in valid_headers:
args = RouterArgs(request_id_headers=headers)
assert args.request_id_headers == headers
class TestLaunchValidation:
"""Test launch-time validation logic."""
def test_pd_mode_requires_urls(self):
"""Test that PD mode requires prefill and decode URLs."""
# PD mode without URLs should fail
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=False,
)
with pytest.raises(
ValueError, match="PD disaggregation mode requires --prefill"
):
launch_router(args)
def test_pd_mode_with_service_discovery_allows_empty_urls(self):
"""Test that PD mode with service discovery allows empty URLs."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=True,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_regular_mode_allows_empty_worker_urls(self):
"""Test that regular mode allows empty worker URLs."""
args = RouterArgs(worker_urls=[], service_discovery=False)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_launch_with_valid_config(self):
"""Test launching with valid configuration."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000"],
policy="cache_aware",
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_launch_with_pd_config(self):
"""Test launching with valid PD configuration."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", 9000)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_launch_with_service_discovery_config(self):
"""Test launching with valid service discovery configuration."""
args = RouterArgs(
service_discovery=True,
selector={"app": "worker"},
service_discovery_port=8080,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()

31
sgl-router/pyproject.toml Normal file
View File

@@ -0,0 +1,31 @@
[build-system]
requires = ["setuptools>=45", "wheel", "setuptools-rust>=1.5.2"]
build-backend = "setuptools.build_meta"
[project]
name = "sglang-router"
version = "0.1.9"
description = "High-performance Rust-based load balancer for SGLang with multiple routing algorithms and prefill-decode disaggregation support"
authors = [{name = "Byron Hsu", email = "byronhsu1230@gmail.com"}]
requires-python = ">=3.8"
readme = "README.md"
license = { text = "Apache-2.0" }
classifiers = [
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Rust",
"Programming Language :: Python :: 3",
]
[project.optional-dependencies]
dev = [
"requests>=2.25.0",
]
# https://github.com/PyO3/setuptools-rust?tab=readme-ov-file
[tool.setuptools.packages]
find = { where = ["py_src"] }
# workaround for https://github.com/pypa/twine/issues/1216
[tool.setuptools]
license-files = []

5
sgl-router/pytest.ini Normal file
View File

@@ -0,0 +1,5 @@
[pytest]
testpaths = py_test
python_files = test_*.py
python_classes = Test*
python_functions = test_*

250
sgl-router/scripts/run_benchmarks.py Executable file
View File

@@ -0,0 +1,250 @@
#!/usr/bin/env python3
"""
SGLang Router Benchmark Runner
A Python script to run Rust benchmarks with various options and modes.
Replaces the shell script for better maintainability and integration.
"""
import argparse
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional
class BenchmarkRunner:
"""Handles running Rust benchmarks for the SGLang router."""
def __init__(self, project_root: str):
self.project_root = Path(project_root)
self.timestamp = time.strftime("%a %b %d %H:%M:%S UTC %Y", time.gmtime())
def run_command(
self, cmd: List[str], capture_output: bool = False
) -> subprocess.CompletedProcess:
"""Run a command and handle errors."""
try:
if capture_output:
result = subprocess.run(
cmd, capture_output=True, text=True, cwd=self.project_root
)
else:
result = subprocess.run(cmd, cwd=self.project_root)
return result
except subprocess.CalledProcessError as e:
print(f"Error running command: {' '.join(cmd)}")
print(f"Exit code: {e.returncode}")
sys.exit(1)
def print_header(self):
"""Print the benchmark runner header."""
print("SGLang Router Benchmark Runner")
print("=" * 30)
print(f"Project: {self.project_root.absolute()}")
print(f"Timestamp: {self.timestamp}")
print()
def build_release(self):
"""Build the project in release mode."""
print("Building in release mode...")
result = self.run_command(["cargo", "build", "--release", "--quiet"])
if result.returncode != 0:
print("Failed to build in release mode")
sys.exit(1)
def run_benchmarks(
self,
quick_mode: bool = False,
save_baseline: Optional[str] = None,
compare_baseline: Optional[str] = None,
) -> str:
"""Run benchmarks with specified options."""
bench_args = ["cargo", "bench", "--bench", "request_processing"]
if quick_mode:
bench_args.append("benchmark_summary")
print("Running quick benchmarks...")
else:
print("Running full benchmark suite...")
# Note: Criterion baselines are handled via target directory structure
# For now, we'll implement baseline functionality via file copying
if save_baseline:
print(f"Will save results as baseline: {save_baseline}")
if compare_baseline:
print(f"Will compare with baseline: {compare_baseline}")
print(f"Executing: {' '.join(bench_args)}")
result = self.run_command(bench_args, capture_output=True)
if result.returncode != 0:
print("Benchmark execution failed!")
print("STDOUT:", result.stdout)
print("STDERR:", result.stderr)
sys.exit(1)
# Handle baseline saving after successful run
if save_baseline:
self._save_baseline(save_baseline, result.stdout)
return result.stdout
def _save_baseline(self, filename: str, output: str):
"""Save benchmark results to a file as baseline."""
filepath = self.project_root / filename
with open(filepath, "w") as f:
f.write(output)
print(f"Baseline saved to: {filepath}")
def parse_benchmark_results(self, output: str) -> Dict[str, str]:
"""Parse benchmark output to extract performance metrics."""
results = {}
# Look for performance overview section
lines = output.split("\n")
parsing_overview = False
for line in lines:
line = line.strip()
if "Quick Performance Overview:" in line:
parsing_overview = True
continue
if parsing_overview and line.startswith("* "):
# Parse lines like "* Serialization (avg): 481 ns/req"
if "Serialization (avg):" in line:
results["serialization_time"] = self._extract_time(line)
elif "Deserialization (avg):" in line:
results["deserialization_time"] = self._extract_time(line)
elif "Bootstrap Injection (avg):" in line:
results["bootstrap_injection_time"] = self._extract_time(line)
elif "Total Pipeline (avg):" in line:
results["total_time"] = self._extract_time(line)
# Stop parsing after the overview section
if parsing_overview and line.startswith("Performance Insights:"):
break
return results
def _extract_time(self, line: str) -> str:
"""Extract time value from a benchmark line."""
# Extract number followed by ns/req
import re
match = re.search(r"(\d+)\s*ns/req", line)
return match.group(1) if match else "N/A"
def validate_thresholds(self, results: Dict[str, str]) -> bool:
"""Validate benchmark results against performance thresholds."""
thresholds = {
"serialization_time": 2000, # 2μs max
"deserialization_time": 2000, # 2μs max
"bootstrap_injection_time": 5000, # 5μs max
"total_time": 10000, # 10μs max
}
all_passed = True
print("\nPerformance Threshold Validation:")
print("=" * 35)
for metric, threshold in thresholds.items():
if metric in results and results[metric] != "N/A":
try:
value = int(results[metric])
passed = value <= threshold
status = "✓ PASS" if passed else "✗ FAIL"
print(f"{metric:20}: {value:>6}ns <= {threshold:>6}ns {status}")
if not passed:
all_passed = False
except ValueError:
print(f"{metric:20}: Invalid value: {results[metric]}")
all_passed = False
else:
print(f"{metric:20}: No data available")
all_passed = False
print()
if all_passed:
print("All performance thresholds passed!")
else:
print("Some performance thresholds failed!")
return all_passed
def save_results_to_file(
self, results: Dict[str, str], filename: str = "benchmark_results.env"
):
"""Save benchmark results to a file for CI consumption."""
filepath = self.project_root / filename
with open(filepath, "w") as f:
for key, value in results.items():
f.write(f"{key}={value}\n")
print(f"Results saved to: {filepath}")
def main():
parser = argparse.ArgumentParser(description="Run SGLang router benchmarks")
parser.add_argument(
"--quick", action="store_true", help="Run quick benchmarks (summary only)"
)
parser.add_argument(
"--save-baseline", type=str, help="Save benchmark results as baseline"
)
parser.add_argument(
"--compare-baseline", type=str, help="Compare with saved baseline"
)
parser.add_argument(
"--validate-thresholds",
action="store_true",
help="Validate results against performance thresholds",
)
parser.add_argument(
"--save-results", action="store_true", help="Save results to file for CI"
)
args = parser.parse_args()
# Determine project root (script is in scripts/ subdirectory)
script_dir = Path(__file__).parent
project_root = script_dir.parent
runner = BenchmarkRunner(str(project_root))
runner.print_header()
# Build in release mode
runner.build_release()
# Run benchmarks
output = runner.run_benchmarks(
quick_mode=args.quick,
save_baseline=args.save_baseline,
compare_baseline=args.compare_baseline,
)
# Print the raw output
print(output)
# Parse and validate results if requested
if args.validate_thresholds or args.save_results:
results = runner.parse_benchmark_results(output)
if args.save_results:
runner.save_results_to_file(results)
if args.validate_thresholds:
passed = runner.validate_thresholds(results)
if not passed:
print("Validation failed - performance regression detected!")
sys.exit(1)
print("\nBenchmark run completed successfully!")
if __name__ == "__main__":
main()

21
sgl-router/setup.py Normal file
View File

@@ -0,0 +1,21 @@
import os
from setuptools import setup
from setuptools_rust import Binding, RustExtension
no_rust = os.environ.get("SGLANG_ROUTER_BUILD_NO_RUST") == "1"
rust_extensions = []
if not no_rust:
rust_extensions.append(
RustExtension(
target="sglang_router_rs",
path="Cargo.toml",
binding=Binding.PyO3,
)
)
setup(
rust_extensions=rust_extensions,
zip_safe=False,
)

28
sgl-router/src/config/mod.rs Normal file
View File

@@ -0,0 +1,28 @@
pub mod types;
pub mod validation;
pub use types::*;
pub use validation::*;
/// Configuration errors
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
#[error("Validation failed: {reason}")]
ValidationFailed { reason: String },
#[error("Invalid value for field '{field}': {value} - {reason}")]
InvalidValue {
field: String,
value: String,
reason: String,
},
#[error("Incompatible configuration: {reason}")]
IncompatibleConfig { reason: String },
#[error("Missing required field: {field}")]
MissingRequired { field: String },
}
/// Result type for configuration operations
pub type ConfigResult<T> = Result<T, ConfigError>;

1241
sgl-router/src/config/types.rs Normal file
View File

File diff suppressed because it is too large Load Diff

776
sgl-router/src/config/validation.rs Normal file
View File

@@ -0,0 +1,776 @@
use super::*;
/// Configuration validator
pub struct ConfigValidator;
impl ConfigValidator {
/// Validate a complete router configuration
pub fn validate(config: &RouterConfig) -> ConfigResult<()> {
// Check if service discovery is enabled
let has_service_discovery = config.discovery.as_ref().is_some_and(|d| d.enabled);
Self::validate_mode(&config.mode, has_service_discovery)?;
Self::validate_policy(&config.policy)?;
Self::validate_server_settings(config)?;
if let Some(discovery) = &config.discovery {
Self::validate_discovery(discovery, &config.mode)?;
}
if let Some(metrics) = &config.metrics {
Self::validate_metrics(metrics)?;
}
Self::validate_compatibility(config)?;
// Validate effective retry/CB configs (respect disable flags)
let retry_cfg = config.effective_retry_config();
let cb_cfg = config.effective_circuit_breaker_config();
Self::validate_retry(&retry_cfg)?;
Self::validate_circuit_breaker(&cb_cfg)?;
Ok(())
}
/// Validate routing mode configuration
fn validate_mode(mode: &RoutingMode, has_service_discovery: bool) -> ConfigResult<()> {
match mode {
RoutingMode::Regular { worker_urls } => {
// Validate URLs if any are provided
if !worker_urls.is_empty() {
Self::validate_urls(worker_urls)?;
}
// Note: We allow empty worker URLs even without service discovery
// to let the router start and fail at runtime when routing requests.
// This matches legacy behavior and test expectations.
}
RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
prefill_policy,
decode_policy,
} => {
// Only require URLs if service discovery is disabled
if !has_service_discovery {
if prefill_urls.is_empty() {
return Err(ConfigError::ValidationFailed {
reason: "PD mode requires at least one prefill worker URL".to_string(),
});
}
if decode_urls.is_empty() {
return Err(ConfigError::ValidationFailed {
reason: "PD mode requires at least one decode worker URL".to_string(),
});
}
}
// Validate URLs if any are provided
if !prefill_urls.is_empty() {
let prefill_url_strings: Vec<String> =
prefill_urls.iter().map(|(url, _)| url.clone()).collect();
Self::validate_urls(&prefill_url_strings)?;
}
if !decode_urls.is_empty() {
Self::validate_urls(decode_urls)?;
}
// Validate bootstrap ports
for (_url, port) in prefill_urls {
if let Some(port) = port {
if *port == 0 {
return Err(ConfigError::InvalidValue {
field: "bootstrap_port".to_string(),
value: port.to_string(),
reason: "Port must be between 1 and 65535".to_string(),
});
}
}
}
// Validate optional prefill and decode policies
if let Some(p_policy) = prefill_policy {
Self::validate_policy(p_policy)?;
}
if let Some(d_policy) = decode_policy {
Self::validate_policy(d_policy)?;
}
}
RoutingMode::OpenAI { worker_urls } => {
// Require exactly one worker URL for OpenAI router
if worker_urls.len() != 1 {
return Err(ConfigError::ValidationFailed {
reason: "OpenAI mode requires exactly one --worker-urls entry".to_string(),
});
}
// Validate URL format
if let Err(e) = url::Url::parse(&worker_urls[0]) {
return Err(ConfigError::ValidationFailed {
reason: format!("Invalid OpenAI worker URL '{}': {}", &worker_urls[0], e),
});
}
}
}
Ok(())
}
/// Validate policy configuration
fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> {
match policy {
PolicyConfig::Random | PolicyConfig::RoundRobin => {
// No specific validation needed
}
PolicyConfig::CacheAware {
cache_threshold,
balance_abs_threshold: _,
balance_rel_threshold,
eviction_interval_secs,
max_tree_size,
} => {
if !(0.0..=1.0).contains(cache_threshold) {
return Err(ConfigError::InvalidValue {
field: "cache_threshold".to_string(),
value: cache_threshold.to_string(),
reason: "Must be between 0.0 and 1.0".to_string(),
});
}
if *balance_rel_threshold < 1.0 {
return Err(ConfigError::InvalidValue {
field: "balance_rel_threshold".to_string(),
value: balance_rel_threshold.to_string(),
reason: "Must be >= 1.0".to_string(),
});
}
if *eviction_interval_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "eviction_interval_secs".to_string(),
value: eviction_interval_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
if *max_tree_size == 0 {
return Err(ConfigError::InvalidValue {
field: "max_tree_size".to_string(),
value: max_tree_size.to_string(),
reason: "Must be > 0".to_string(),
});
}
}
PolicyConfig::PowerOfTwo {
load_check_interval_secs,
} => {
if *load_check_interval_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "load_check_interval_secs".to_string(),
value: load_check_interval_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
}
}
Ok(())
}
/// Validate server configuration
fn validate_server_settings(config: &RouterConfig) -> ConfigResult<()> {
if config.port == 0 {
return Err(ConfigError::InvalidValue {
field: "port".to_string(),
value: config.port.to_string(),
reason: "Port must be > 0".to_string(),
});
}
if config.max_payload_size == 0 {
return Err(ConfigError::InvalidValue {
field: "max_payload_size".to_string(),
value: config.max_payload_size.to_string(),
reason: "Must be > 0".to_string(),
});
}
if config.request_timeout_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "request_timeout_secs".to_string(),
value: config.request_timeout_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
if config.worker_startup_timeout_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "worker_startup_timeout_secs".to_string(),
value: config.worker_startup_timeout_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
if config.worker_startup_check_interval_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "worker_startup_check_interval_secs".to_string(),
value: config.worker_startup_check_interval_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
Ok(())
}
/// Validate service discovery configuration
fn validate_discovery(discovery: &DiscoveryConfig, mode: &RoutingMode) -> ConfigResult<()> {
if !discovery.enabled {
return Ok(()); // No validation needed if disabled
}
if discovery.port == 0 {
return Err(ConfigError::InvalidValue {
field: "discovery.port".to_string(),
value: discovery.port.to_string(),
reason: "Port must be > 0".to_string(),
});
}
if discovery.check_interval_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "discovery.check_interval_secs".to_string(),
value: discovery.check_interval_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
// Validate selectors based on mode
match mode {
RoutingMode::Regular { .. } => {
if discovery.selector.is_empty() {
return Err(ConfigError::ValidationFailed {
reason: "Regular mode with service discovery requires a non-empty selector"
.to_string(),
});
}
}
RoutingMode::PrefillDecode { .. } => {
if discovery.prefill_selector.is_empty() && discovery.decode_selector.is_empty() {
return Err(ConfigError::ValidationFailed {
reason: "PD mode with service discovery requires at least one non-empty selector (prefill or decode)".to_string(),
});
}
}
RoutingMode::OpenAI { .. } => {
// OpenAI mode doesn't use service discovery
return Err(ConfigError::ValidationFailed {
reason: "OpenAI mode does not support service discovery".to_string(),
});
}
}
Ok(())
}
/// Validate metrics configuration
fn validate_metrics(metrics: &MetricsConfig) -> ConfigResult<()> {
if metrics.port == 0 {
return Err(ConfigError::InvalidValue {
field: "metrics.port".to_string(),
value: metrics.port.to_string(),
reason: "Port must be > 0".to_string(),
});
}
if metrics.host.is_empty() {
return Err(ConfigError::InvalidValue {
field: "metrics.host".to_string(),
value: metrics.host.clone(),
reason: "Host cannot be empty".to_string(),
});
}
Ok(())
}
/// Validate retry configuration
fn validate_retry(retry: &RetryConfig) -> ConfigResult<()> {
if retry.max_retries < 1 {
return Err(ConfigError::InvalidValue {
field: "retry.max_retries".to_string(),
value: retry.max_retries.to_string(),
reason: "Must be >= 1 (set to 1 to effectively disable retries)".to_string(),
});
}
if retry.initial_backoff_ms == 0 {
return Err(ConfigError::InvalidValue {
field: "retry.initial_backoff_ms".to_string(),
value: retry.initial_backoff_ms.to_string(),
reason: "Must be > 0".to_string(),
});
}
if retry.max_backoff_ms < retry.initial_backoff_ms {
return Err(ConfigError::InvalidValue {
field: "retry.max_backoff_ms".to_string(),
value: retry.max_backoff_ms.to_string(),
reason: "Must be >= initial_backoff_ms".to_string(),
});
}
if retry.backoff_multiplier < 1.0 {
return Err(ConfigError::InvalidValue {
field: "retry.backoff_multiplier".to_string(),
value: retry.backoff_multiplier.to_string(),
reason: "Must be >= 1.0".to_string(),
});
}
if !(0.0..=1.0).contains(&retry.jitter_factor) {
return Err(ConfigError::InvalidValue {
field: "retry.jitter_factor".to_string(),
value: retry.jitter_factor.to_string(),
reason: "Must be between 0.0 and 1.0".to_string(),
});
}
Ok(())
}
/// Validate circuit breaker configuration
fn validate_circuit_breaker(cb: &CircuitBreakerConfig) -> ConfigResult<()> {
if cb.failure_threshold < 1 {
return Err(ConfigError::InvalidValue {
field: "circuit_breaker.failure_threshold".to_string(),
value: cb.failure_threshold.to_string(),
reason: "Must be >= 1 (set to u32::MAX to effectively disable CB)".to_string(),
});
}
if cb.success_threshold < 1 {
return Err(ConfigError::InvalidValue {
field: "circuit_breaker.success_threshold".to_string(),
value: cb.success_threshold.to_string(),
reason: "Must be >= 1".to_string(),
});
}
if cb.timeout_duration_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "circuit_breaker.timeout_duration_secs".to_string(),
value: cb.timeout_duration_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
if cb.window_duration_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "circuit_breaker.window_duration_secs".to_string(),
value: cb.window_duration_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}
Ok(())
}
/// Validate compatibility between different configuration sections
fn validate_compatibility(config: &RouterConfig) -> ConfigResult<()> {
// IGW mode is independent - skip other compatibility checks when enabled
if config.enable_igw {
return Ok(());
}
// Validate gRPC connection mode requires tokenizer configuration
if config.connection_mode == ConnectionMode::Grpc
&& config.tokenizer_path.is_none()
&& config.model_path.is_none()
{
return Err(ConfigError::ValidationFailed {
reason: "gRPC connection mode requires either --tokenizer-path or --model-path to be specified".to_string(),
});
}
// All policies are now supported for both router types thanks to the unified trait design
// No mode/policy restrictions needed anymore
// Check if service discovery is enabled for worker count validation
let has_service_discovery = config.discovery.as_ref().is_some_and(|d| d.enabled);
// Only validate worker counts if service discovery is disabled
if !has_service_discovery {
// Check if power-of-two policy makes sense with insufficient workers
if let PolicyConfig::PowerOfTwo { .. } = &config.policy {
let worker_count = config.mode.worker_count();
if worker_count < 2 {
return Err(ConfigError::IncompatibleConfig {
reason: "Power-of-two policy requires at least 2 workers".to_string(),
});
}
}
// For PD mode, validate that policies have sufficient workers
if let RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
prefill_policy,
decode_policy,
} = &config.mode
{
// Check power-of-two for prefill
if let Some(PolicyConfig::PowerOfTwo { .. }) = prefill_policy {
if prefill_urls.len() < 2 {
return Err(ConfigError::IncompatibleConfig {
reason: "Power-of-two policy for prefill requires at least 2 prefill workers".to_string(),
});
}
}
// Check power-of-two for decode
if let Some(PolicyConfig::PowerOfTwo { .. }) = decode_policy {
if decode_urls.len() < 2 {
return Err(ConfigError::IncompatibleConfig {
reason:
"Power-of-two policy for decode requires at least 2 decode workers"
.to_string(),
});
}
}
}
}
// Service discovery is conflict with dp_aware routing for now
// since it's not fully supported yet
if has_service_discovery && config.dp_aware {
return Err(ConfigError::IncompatibleConfig {
reason: "DP-aware routing is not compatible with service discovery".to_string(),
});
}
Ok(())
}
/// Validate URL format
fn validate_urls(urls: &[String]) -> ConfigResult<()> {
for url in urls {
if url.is_empty() {
return Err(ConfigError::InvalidValue {
field: "worker_url".to_string(),
value: url.clone(),
reason: "URL cannot be empty".to_string(),
});
}
if !url.starts_with("http://")
&& !url.starts_with("https://")
&& !url.starts_with("grpc://")
{
return Err(ConfigError::InvalidValue {
field: "worker_url".to_string(),
value: url.clone(),
reason: "URL must start with http://, https://, or grpc://".to_string(),
});
}
// Basic URL validation
match ::url::Url::parse(url) {
Ok(parsed) => {
// Additional validation
if parsed.host_str().is_none() {
return Err(ConfigError::InvalidValue {
field: "worker_url".to_string(),
value: url.clone(),
reason: "URL must have a valid host".to_string(),
});
}
}
Err(e) => {
return Err(ConfigError::InvalidValue {
field: "worker_url".to_string(),
value: url.clone(),
reason: format!("Invalid URL format: {}", e),
});
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_validate_regular_mode() {
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["http://worker:8000".to_string()],
},
PolicyConfig::Random,
);
assert!(ConfigValidator::validate(&config).is_ok());
}
#[test]
fn test_validate_empty_worker_urls() {
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec![],
},
PolicyConfig::Random,
);
// Empty worker URLs are now allowed to match legacy behavior
assert!(ConfigValidator::validate(&config).is_ok());
}
#[test]
fn test_validate_empty_worker_urls_with_service_discovery() {
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec![],
},
PolicyConfig::Random,
);
// Enable service discovery
config.discovery = Some(DiscoveryConfig {
enabled: true,
selector: vec![("app".to_string(), "test".to_string())]
.into_iter()
.collect(),
..Default::default()
});
// Should pass validation since service discovery is enabled
assert!(ConfigValidator::validate(&config).is_ok());
}
#[test]
fn test_validate_invalid_urls() {
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["invalid-url".to_string()],
},
PolicyConfig::Random,
);
assert!(ConfigValidator::validate(&config).is_err());
}
#[test]
fn test_validate_cache_aware_thresholds() {
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec![
"http://worker1:8000".to_string(),
"http://worker2:8000".to_string(),
],
},
PolicyConfig::CacheAware {
cache_threshold: 1.5, // Invalid: > 1.0
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
eviction_interval_secs: 60,
max_tree_size: 1000,
},
);
assert!(ConfigValidator::validate(&config).is_err());
}
#[test]
fn test_validate_cache_aware_single_worker() {
// Cache-aware with single worker should be allowed (even if not optimal)
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["http://worker1:8000".to_string()],
},
PolicyConfig::CacheAware {
cache_threshold: 0.5,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
eviction_interval_secs: 60,
max_tree_size: 1000,
},
);
assert!(ConfigValidator::validate(&config).is_ok());
}
#[test]
fn test_validate_pd_mode() {
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![("http://prefill:8000".to_string(), Some(8081))],
decode_urls: vec!["http://decode:8000".to_string()],
prefill_policy: None,
decode_policy: None,
},
PolicyConfig::Random,
);
assert!(ConfigValidator::validate(&config).is_ok());
}
#[test]
fn test_validate_roundrobin_with_pd_mode() {
// RoundRobin with PD mode is now supported
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![("http://prefill:8000".to_string(), None)],
decode_urls: vec!["http://decode:8000".to_string()],
prefill_policy: None,
decode_policy: None,
},
PolicyConfig::RoundRobin,
);
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_cache_aware_with_pd_mode() {
// CacheAware with PD mode is now supported
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![("http://prefill:8000".to_string(), None)],
decode_urls: vec!["http://decode:8000".to_string()],
prefill_policy: None,
decode_policy: None,
},
PolicyConfig::CacheAware {
cache_threshold: 0.5,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
eviction_interval_secs: 60,
max_tree_size: 1000,
},
);
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_power_of_two_with_regular_mode() {
// PowerOfTwo with Regular mode is now supported
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec![
"http://worker1:8000".to_string(),
"http://worker2:8000".to_string(),
],
},
PolicyConfig::PowerOfTwo {
load_check_interval_secs: 60,
},
);
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_pd_mode_with_separate_policies() {
// Test PD mode with different policies for prefill and decode
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![
("http://prefill1:8000".to_string(), None),
("http://prefill2:8000".to_string(), None),
],
decode_urls: vec![
"http://decode1:8000".to_string(),
"http://decode2:8000".to_string(),
],
prefill_policy: Some(PolicyConfig::CacheAware {
cache_threshold: 0.5,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
eviction_interval_secs: 60,
max_tree_size: 1000,
}),
decode_policy: Some(PolicyConfig::PowerOfTwo {
load_check_interval_secs: 60,
}),
},
PolicyConfig::Random, // Main policy as fallback
);
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_pd_mode_power_of_two_insufficient_workers() {
// Test that power-of-two policy requires at least 2 workers
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![("http://prefill1:8000".to_string(), None)], // Only 1 prefill
decode_urls: vec![
"http://decode1:8000".to_string(),
"http://decode2:8000".to_string(),
],
prefill_policy: Some(PolicyConfig::PowerOfTwo {
load_check_interval_secs: 60,
}), // Requires 2+ workers
decode_policy: None,
},
PolicyConfig::Random,
);
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
if let Err(e) = result {
assert!(e.to_string().contains("prefill requires at least 2"));
}
}
#[test]
fn test_validate_grpc_requires_tokenizer() {
// Test that gRPC connection mode requires tokenizer configuration
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
// Set connection mode to gRPC without tokenizer config
config.connection_mode = ConnectionMode::Grpc;
config.tokenizer_path = None;
config.model_path = None;
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
if let Err(e) = result {
assert!(e.to_string().contains("gRPC connection mode requires"));
}
}
#[test]
fn test_validate_grpc_with_model_path() {
// Test that gRPC works with model_path
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
config.connection_mode = ConnectionMode::Grpc;
config.model_path = Some("meta-llama/Llama-3-8B".to_string());
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_grpc_with_tokenizer_path() {
// Test that gRPC works with tokenizer_path
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
config.connection_mode = ConnectionMode::Grpc;
config.tokenizer_path = Some("/path/to/tokenizer.json".to_string());
let result = ConfigValidator::validate(&config);
assert!(result.is_ok());
}
}

555
sgl-router/src/core/circuit_breaker.rs Normal file
View File

@@ -0,0 +1,555 @@
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use tracing::info;
/// Circuit breaker configuration
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
/// Number of consecutive failures to open the circuit
pub failure_threshold: u32,
/// Success threshold to close circuit from half-open
pub success_threshold: u32,
/// Duration to wait before attempting half-open
pub timeout_duration: Duration,
/// Time window for failure counting
pub window_duration: Duration,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
success_threshold: 2,
timeout_duration: Duration::from_secs(30),
window_duration: Duration::from_secs(60),
}
}
}
/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
/// Normal operation - requests are allowed
Closed,
/// Circuit is open - requests are rejected
Open,
/// Testing if service has recovered - limited requests allowed
HalfOpen,
}
impl std::fmt::Display for CircuitState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CircuitState::Closed => write!(f, "Closed"),
CircuitState::Open => write!(f, "Open"),
CircuitState::HalfOpen => write!(f, "HalfOpen"),
}
}
}
/// Circuit breaker implementation
#[derive(Debug)]
pub struct CircuitBreaker {
state: Arc<RwLock<CircuitState>>,
consecutive_failures: Arc<AtomicU32>,
consecutive_successes: Arc<AtomicU32>,
total_failures: Arc<AtomicU64>,
total_successes: Arc<AtomicU64>,
last_failure_time: Arc<RwLock<Option<Instant>>>,
last_state_change: Arc<RwLock<Instant>>,
config: CircuitBreakerConfig,
}
impl CircuitBreaker {
/// Create a new circuit breaker with default configuration
pub fn new() -> Self {
Self::with_config(CircuitBreakerConfig::default())
}
/// Create a new circuit breaker with custom configuration
pub fn with_config(config: CircuitBreakerConfig) -> Self {
Self {
state: Arc::new(RwLock::new(CircuitState::Closed)),
consecutive_failures: Arc::new(AtomicU32::new(0)),
consecutive_successes: Arc::new(AtomicU32::new(0)),
total_failures: Arc::new(AtomicU64::new(0)),
total_successes: Arc::new(AtomicU64::new(0)),
last_failure_time: Arc::new(RwLock::new(None)),
last_state_change: Arc::new(RwLock::new(Instant::now())),
config,
}
}
/// Check if a request can be executed
pub fn can_execute(&self) -> bool {
// First check if we need to transition from Open to HalfOpen
self.check_and_update_state();
let state = *self.state.read().unwrap();
match state {
CircuitState::Closed => true,
CircuitState::Open => false,
CircuitState::HalfOpen => true, // Allow limited requests in half-open state
}
}
/// Get the current state
pub fn state(&self) -> CircuitState {
self.check_and_update_state();
*self.state.read().unwrap()
}
/// Record the outcome of a request
pub fn record_outcome(&self, success: bool) {
if success {
self.record_success();
} else {
self.record_failure();
}
}
/// Record a successful request
pub fn record_success(&self) {
self.total_successes.fetch_add(1, Ordering::Relaxed);
self.consecutive_failures.store(0, Ordering::Release);
let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1;
// Outcome-level metrics are recorded at the worker level where the worker label is known
let current_state = *self.state.read().unwrap();
match current_state {
CircuitState::HalfOpen => {
// Check if we've reached the success threshold to close the circuit
if successes >= self.config.success_threshold {
self.transition_to(CircuitState::Closed);
}
}
CircuitState::Closed => {
// Already closed, nothing to do
}
CircuitState::Open => {
// Shouldn't happen, but if it does, stay open
tracing::warn!("Success recorded while circuit is open");
}
}
}
/// Record a failed request
pub fn record_failure(&self) {
self.total_failures.fetch_add(1, Ordering::Relaxed);
self.consecutive_successes.store(0, Ordering::Release);
let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1;
// Outcome-level metrics are recorded at the worker level where the worker label is known
// Update last failure time
{
let mut last_failure = self.last_failure_time.write().unwrap();
*last_failure = Some(Instant::now());
}
let current_state = *self.state.read().unwrap();
match current_state {
CircuitState::Closed => {
// Check if we've reached the failure threshold to open the circuit
if failures >= self.config.failure_threshold {
self.transition_to(CircuitState::Open);
}
}
CircuitState::HalfOpen => {
// Single failure in half-open state reopens the circuit
self.transition_to(CircuitState::Open);
}
CircuitState::Open => {
// Already open, nothing to do
}
}
}
/// Check and update state based on timeout
fn check_and_update_state(&self) {
let current_state = *self.state.read().unwrap();
if current_state == CircuitState::Open {
// Check if timeout has expired
let last_change = *self.last_state_change.read().unwrap();
if last_change.elapsed() >= self.config.timeout_duration {
self.transition_to(CircuitState::HalfOpen);
}
}
}
/// Transition to a new state
fn transition_to(&self, new_state: CircuitState) {
let mut state = self.state.write().unwrap();
let old_state = *state;
if old_state != new_state {
*state = new_state;
// Update last state change time
let mut last_change = self.last_state_change.write().unwrap();
*last_change = Instant::now();
// Reset counters based on transition
match new_state {
CircuitState::Closed => {
self.consecutive_failures.store(0, Ordering::Release);
self.consecutive_successes.store(0, Ordering::Release);
}
CircuitState::Open => {
self.consecutive_successes.store(0, Ordering::Release);
}
CircuitState::HalfOpen => {
self.consecutive_failures.store(0, Ordering::Release);
self.consecutive_successes.store(0, Ordering::Release);
}
}
let from = match old_state {
CircuitState::Closed => "closed",
CircuitState::Open => "open",
CircuitState::HalfOpen => "half_open",
};
let to = match new_state {
CircuitState::Closed => "closed",
CircuitState::Open => "open",
CircuitState::HalfOpen => "half_open",
};
info!("Circuit breaker state transition: {} -> {}", from, to);
// Transition metrics are recorded at the worker level where the worker label is known
}
}
/// Get the number of consecutive failures
pub fn failure_count(&self) -> u32 {
self.consecutive_failures.load(Ordering::Acquire)
}
/// Get the number of consecutive successes
pub fn success_count(&self) -> u32 {
self.consecutive_successes.load(Ordering::Acquire)
}
/// Get total failures
pub fn total_failures(&self) -> u64 {
self.total_failures.load(Ordering::Relaxed)
}
/// Get total successes
pub fn total_successes(&self) -> u64 {
self.total_successes.load(Ordering::Relaxed)
}
/// Get time since last failure
pub fn time_since_last_failure(&self) -> Option<Duration> {
self.last_failure_time.read().unwrap().map(|t| t.elapsed())
}
/// Get time since last state change
pub fn time_since_last_state_change(&self) -> Duration {
self.last_state_change.read().unwrap().elapsed()
}
/// Check if the circuit is in a half-open state
pub fn is_half_open(&self) -> bool {
self.state() == CircuitState::HalfOpen
}
/// Record a test success (for health check probing)
pub fn record_test_success(&self) {
if self.is_half_open() {
self.record_success();
}
}
/// Record a test failure (for health check probing)
pub fn record_test_failure(&self) {
if self.is_half_open() {
self.record_failure();
}
}
/// Reset the circuit breaker to closed state
pub fn reset(&self) {
self.transition_to(CircuitState::Closed);
self.consecutive_failures.store(0, Ordering::Release);
self.consecutive_successes.store(0, Ordering::Release);
}
/// Force the circuit to open (for manual intervention)
pub fn force_open(&self) {
self.transition_to(CircuitState::Open);
}
/// Get circuit breaker statistics
pub fn stats(&self) -> CircuitBreakerStats {
CircuitBreakerStats {
state: self.state(),
consecutive_failures: self.failure_count(),
consecutive_successes: self.success_count(),
total_failures: self.total_failures(),
total_successes: self.total_successes(),
time_since_last_failure: self.time_since_last_failure(),
time_since_last_state_change: self.time_since_last_state_change(),
}
}
}
impl Clone for CircuitBreaker {
fn clone(&self) -> Self {
Self {
state: Arc::clone(&self.state),
consecutive_failures: Arc::clone(&self.consecutive_failures),
consecutive_successes: Arc::clone(&self.consecutive_successes),
total_failures: Arc::clone(&self.total_failures),
total_successes: Arc::clone(&self.total_successes),
last_failure_time: Arc::clone(&self.last_failure_time),
last_state_change: Arc::clone(&self.last_state_change),
config: self.config.clone(),
}
}
}
impl Default for CircuitBreaker {
fn default() -> Self {
Self::new()
}
}
/// Circuit breaker statistics
#[derive(Debug, Clone)]
pub struct CircuitBreakerStats {
pub state: CircuitState,
pub consecutive_failures: u32,
pub consecutive_successes: u32,
pub total_failures: u64,
pub total_successes: u64,
pub time_since_last_failure: Option<Duration>,
pub time_since_last_state_change: Duration,
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
#[test]
fn test_circuit_breaker_initial_state() {
let cb = CircuitBreaker::new();
assert_eq!(cb.state(), CircuitState::Closed);
assert!(cb.can_execute());
assert_eq!(cb.failure_count(), 0);
assert_eq!(cb.success_count(), 0);
}
#[test]
fn test_circuit_opens_on_threshold() {
let config = CircuitBreakerConfig {
failure_threshold: 3,
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
// Record failures up to threshold
assert_eq!(cb.state(), CircuitState::Closed);
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Closed);
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Closed);
cb.record_failure();
// Circuit should now be open
assert_eq!(cb.state(), CircuitState::Open);
assert!(!cb.can_execute());
assert_eq!(cb.failure_count(), 3);
}
#[test]
fn test_circuit_half_open_after_timeout() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
timeout_duration: Duration::from_millis(100),
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
// Open the circuit
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
// Wait for timeout
thread::sleep(Duration::from_millis(150));
// Circuit should be half-open
assert_eq!(cb.state(), CircuitState::HalfOpen);
assert!(cb.can_execute());
}
#[test]
fn test_circuit_closes_on_success_threshold() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
success_threshold: 2,
timeout_duration: Duration::from_millis(50),
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
// Open the circuit
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
// Wait for timeout
thread::sleep(Duration::from_millis(100));
assert_eq!(cb.state(), CircuitState::HalfOpen);
// Record successes
cb.record_success();
assert_eq!(cb.state(), CircuitState::HalfOpen);
cb.record_success();
// Circuit should now be closed
assert_eq!(cb.state(), CircuitState::Closed);
assert!(cb.can_execute());
}
#[test]
fn test_circuit_reopens_on_half_open_failure() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
timeout_duration: Duration::from_millis(50),
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
// Open the circuit
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
// Wait for timeout
thread::sleep(Duration::from_millis(100));
assert_eq!(cb.state(), CircuitState::HalfOpen);
// Record a failure in half-open state
cb.record_failure();
// Circuit should reopen immediately
assert_eq!(cb.state(), CircuitState::Open);
assert!(!cb.can_execute());
}
#[test]
fn test_success_resets_failure_count() {
let config = CircuitBreakerConfig {
failure_threshold: 3,
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
// Record some failures
cb.record_failure();
cb.record_failure();
assert_eq!(cb.failure_count(), 2);
// Success should reset failure count
cb.record_success();
assert_eq!(cb.failure_count(), 0);
assert_eq!(cb.success_count(), 1);
// Can now record more failures without opening
cb.record_failure();
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn test_manual_reset() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
// Open the circuit
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
// Manual reset
cb.reset();
assert_eq!(cb.state(), CircuitState::Closed);
assert_eq!(cb.failure_count(), 0);
assert_eq!(cb.success_count(), 0);
}
#[test]
fn test_force_open() {
let cb = CircuitBreaker::new();
assert_eq!(cb.state(), CircuitState::Closed);
cb.force_open();
assert_eq!(cb.state(), CircuitState::Open);
assert!(!cb.can_execute());
}
#[test]
fn test_stats() {
let config = CircuitBreakerConfig {
failure_threshold: 2,
..Default::default()
};
let cb = CircuitBreaker::with_config(config);
cb.record_success();
cb.record_failure();
cb.record_failure();
let stats = cb.stats();
assert_eq!(stats.state, CircuitState::Open);
assert_eq!(stats.consecutive_failures, 2);
assert_eq!(stats.consecutive_successes, 0);
assert_eq!(stats.total_failures, 2);
assert_eq!(stats.total_successes, 1);
}
#[test]
fn test_clone() {
let cb1 = CircuitBreaker::new();
cb1.record_failure();
let cb2 = cb1.clone();
assert_eq!(cb2.failure_count(), 1);
// Changes to cb1 affect cb2 (shared state)
cb1.record_failure();
assert_eq!(cb2.failure_count(), 2);
}
#[test]
fn test_thread_safety() {
use std::sync::Arc;
let cb = Arc::new(CircuitBreaker::new());
let mut handles = vec![];
// Spawn threads that record failures
for _ in 0..10 {
let cb_clone = Arc::clone(&cb);
let handle = thread::spawn(move || {
for _ in 0..100 {
cb_clone.record_failure();
}
});
handles.push(handle);
}
// Wait for all threads
for handle in handles {
handle.join().unwrap();
}
// Should have recorded 1000 failures
assert_eq!(cb.total_failures(), 1000);
}
}

240
sgl-router/src/core/error.rs Normal file
View File

@@ -0,0 +1,240 @@
//! Error types for the SGLang router core
//!
//! This module defines error types used throughout the router for worker operations.
use std::fmt;
/// Worker-related errors
#[derive(Debug)]
pub enum WorkerError {
/// Health check failed
HealthCheckFailed { url: String, reason: String },
/// Worker not found
WorkerNotFound { url: String },
/// Invalid worker configuration
InvalidConfiguration { message: String },
/// Network error
NetworkError { url: String, error: String },
/// Worker is at capacity
WorkerAtCapacity { url: String },
/// Invalid URL format
InvalidUrl { url: String },
}
impl fmt::Display for WorkerError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
WorkerError::HealthCheckFailed { url, reason } => {
write!(f, "Health check failed for worker {}: {}", url, reason)
}
WorkerError::WorkerNotFound { url } => {
write!(f, "Worker not found: {}", url)
}
WorkerError::InvalidConfiguration { message } => {
write!(f, "Invalid worker configuration: {}", message)
}
WorkerError::NetworkError { url, error } => {
write!(f, "Network error for worker {}: {}", url, error)
}
WorkerError::WorkerAtCapacity { url } => {
write!(f, "Worker at capacity: {}", url)
}
WorkerError::InvalidUrl { url } => {
write!(f, "Invalid URL format: {}", url)
}
}
}
}
impl std::error::Error for WorkerError {}
/// Result type for worker operations
pub type WorkerResult<T> = Result<T, WorkerError>;
/// Convert from reqwest errors to worker errors
impl From<reqwest::Error> for WorkerError {
fn from(err: reqwest::Error) -> Self {
WorkerError::NetworkError {
url: err.url().map(|u| u.to_string()).unwrap_or_default(),
error: err.to_string(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::error::Error;
#[test]
fn test_health_check_failed_display() {
let error = WorkerError::HealthCheckFailed {
url: "http://worker1:8080".to_string(),
reason: "Connection refused".to_string(),
};
assert_eq!(
error.to_string(),
"Health check failed for worker http://worker1:8080: Connection refused"
);
}
#[test]
fn test_worker_not_found_display() {
let error = WorkerError::WorkerNotFound {
url: "http://worker2:8080".to_string(),
};
assert_eq!(error.to_string(), "Worker not found: http://worker2:8080");
}
#[test]
fn test_invalid_configuration_display() {
let error = WorkerError::InvalidConfiguration {
message: "Missing port number".to_string(),
};
assert_eq!(
error.to_string(),
"Invalid worker configuration: Missing port number"
);
}
#[test]
fn test_network_error_display() {
let error = WorkerError::NetworkError {
url: "http://worker3:8080".to_string(),
error: "Timeout after 30s".to_string(),
};
assert_eq!(
error.to_string(),
"Network error for worker http://worker3:8080: Timeout after 30s"
);
}
#[test]
fn test_worker_at_capacity_display() {
let error = WorkerError::WorkerAtCapacity {
url: "http://worker4:8080".to_string(),
};
assert_eq!(error.to_string(), "Worker at capacity: http://worker4:8080");
}
#[test]
fn test_worker_error_implements_std_error() {
let error = WorkerError::WorkerNotFound {
url: "http://test".to_string(),
};
// Verify it implements Error trait
let _: &dyn Error = &error;
assert!(error.source().is_none());
}
#[test]
fn test_error_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<WorkerError>();
}
#[test]
fn test_worker_result_type_alias() {
// Test Ok variant
let result: WorkerResult<i32> = Ok(42);
assert!(matches!(result, Ok(42)));
// Test Err variant
let error = WorkerError::WorkerNotFound {
url: "test".to_string(),
};
let result: WorkerResult<i32> = Err(error);
assert!(result.is_err());
}
#[test]
fn test_empty_url_handling() {
// Test empty URLs in error variants
let error1 = WorkerError::HealthCheckFailed {
url: "".to_string(),
reason: "No connection".to_string(),
};
assert_eq!(
error1.to_string(),
"Health check failed for worker : No connection"
);
let error2 = WorkerError::NetworkError {
url: "".to_string(),
error: "DNS failure".to_string(),
};
assert_eq!(error2.to_string(), "Network error for worker : DNS failure");
let error3 = WorkerError::WorkerNotFound {
url: "".to_string(),
};
assert_eq!(error3.to_string(), "Worker not found: ");
}
#[test]
fn test_special_characters_in_messages() {
// Test with special characters
let error = WorkerError::InvalidConfiguration {
message: "Invalid JSON: {\"error\": \"test\"}".to_string(),
};
assert_eq!(
error.to_string(),
"Invalid worker configuration: Invalid JSON: {\"error\": \"test\"}"
);
// Test with unicode
let error2 = WorkerError::HealthCheckFailed {
url: "http://测试:8080".to_string(),
reason: "连接被拒绝".to_string(),
};
assert_eq!(
error2.to_string(),
"Health check failed for worker http://测试:8080: 连接被拒绝"
);
}
#[test]
fn test_very_long_error_messages() {
let long_message = "A".repeat(10000);
let error = WorkerError::InvalidConfiguration {
message: long_message.clone(),
};
let display = error.to_string();
assert!(display.contains(&long_message));
assert_eq!(
display.len(),
"Invalid worker configuration: ".len() + long_message.len()
);
}
// Mock reqwest error for testing conversion
#[test]
fn test_reqwest_error_conversion() {
// Test that NetworkError is the correct variant
let network_error = WorkerError::NetworkError {
url: "http://example.com".to_string(),
error: "connection timeout".to_string(),
};
match network_error {
WorkerError::NetworkError { url, error } => {
assert_eq!(url, "http://example.com");
assert_eq!(error, "connection timeout");
}
_ => panic!("Expected NetworkError variant"),
}
}
#[test]
fn test_error_equality() {
// WorkerError doesn't implement PartialEq, but we can test that
// the same error construction produces the same display output
let error1 = WorkerError::WorkerNotFound {
url: "http://test".to_string(),
};
let error2 = WorkerError::WorkerNotFound {
url: "http://test".to_string(),
};
assert_eq!(error1.to_string(), error2.to_string());
}
}

24
sgl-router/src/core/mod.rs Normal file
View File

@@ -0,0 +1,24 @@
//! Core abstractions for the SGLang router
//!
//! This module contains the fundamental types and traits used throughout the router:
//! - Worker trait and implementations
//! - Error types
//! - Circuit breaker for reliability
//! - Common utilities
pub mod circuit_breaker;
pub mod error;
pub mod retry;
pub mod token_bucket;
pub mod worker;
// Re-export commonly used types at the module level
pub use circuit_breaker::{
CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState,
};
pub use error::{WorkerError, WorkerResult};
pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor};
pub use worker::{
start_health_checker, BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig,
Worker, WorkerCollection, WorkerFactory, WorkerLoadGuard, WorkerType,
};

409
sgl-router/src/core/retry.rs Normal file
View File

@@ -0,0 +1,409 @@
use crate::config::types::RetryConfig;
use axum::http::StatusCode;
use axum::response::Response;
use rand::Rng;
use std::time::Duration;
use tracing::debug;
/// Check if an HTTP status code indicates a retryable error
pub fn is_retryable_status(status: StatusCode) -> bool {
matches!(
status,
StatusCode::REQUEST_TIMEOUT
| StatusCode::TOO_MANY_REQUESTS
| StatusCode::INTERNAL_SERVER_ERROR
| StatusCode::BAD_GATEWAY
| StatusCode::SERVICE_UNAVAILABLE
| StatusCode::GATEWAY_TIMEOUT
)
}
/// Computes exponential backoff with optional jitter.
#[derive(Debug, Clone)]
pub struct BackoffCalculator;
impl BackoffCalculator {
/// Calculate backoff delay for a given attempt index (0-based).
pub fn calculate_delay(config: &RetryConfig, attempt: u32) -> Duration {
// Base exponential backoff
let pow = config.backoff_multiplier.powi(attempt as i32);
let mut delay_ms = (config.initial_backoff_ms as f32 * pow) as u64;
if delay_ms > config.max_backoff_ms {
delay_ms = config.max_backoff_ms;
}
// Apply jitter in range [-j, +j]
let jitter = config.jitter_factor.clamp(0.0, 1.0);
if jitter > 0.0 {
let mut rng = rand::rng();
let jitter_scale: f32 = rng.random_range(-jitter..=jitter);
let jitter_ms = (delay_ms as f32 * jitter_scale)
.round()
.max(-(delay_ms as f32));
let adjusted = (delay_ms as i64 + jitter_ms as i64).max(0) as u64;
return Duration::from_millis(adjusted);
}
Duration::from_millis(delay_ms)
}
}
#[derive(Debug, thiserror::Error)]
pub enum RetryError {
#[error("no available workers")]
NoAvailableWorkers,
#[error("maximum retry attempts exceeded")]
MaxRetriesExceeded,
}
/// A thin async retry executor for generic operations.
#[derive(Debug, Clone, Default)]
pub struct RetryExecutor;
impl RetryExecutor {
/// Execute an async operation with retries and backoff.
/// The `operation` closure is invoked each attempt with the attempt index.
pub async fn execute_with_retry<F, Fut, T>(
config: &RetryConfig,
mut operation: F,
) -> Result<T, RetryError>
where
F: FnMut(u32) -> Fut,
Fut: std::future::Future<Output = Result<T, ()>>,
{
let max = config.max_retries.max(1);
let mut attempt: u32 = 0;
loop {
match operation(attempt).await {
Ok(val) => return Ok(val),
Err(_) => {
// Use the number of failures so far (0-indexed) to compute delay,
// so the first retry uses `initial_backoff_ms`.
let is_last = attempt + 1 >= max;
if is_last {
return Err(RetryError::MaxRetriesExceeded);
}
let delay = BackoffCalculator::calculate_delay(config, attempt);
attempt += 1; // advance to the next attempt after computing delay
tokio::time::sleep(delay).await;
}
}
}
}
/// Execute an operation that returns an HTTP Response with retries and backoff.
///
/// Usage pattern:
/// - `operation(attempt)`: perform one attempt (0-based). Construct and send the request,
/// then return the `Response`. Do any per-attempt bookkeeping (e.g., load tracking,
/// circuit-breaker outcome recording) inside this closure.
/// - `should_retry(&response, attempt)`: decide if the given response should be retried
/// (e.g., based on HTTP status). Returning false short-circuits and returns the response.
/// - `on_backoff(delay, next_attempt)`: called before sleeping between attempts.
/// Use this to record metrics.
/// - `on_exhausted()`: called when the executor has exhausted all retry attempts.
///
/// Example:
/// ```ignore
/// let resp = RetryExecutor::execute_response_with_retry(
/// &retry_cfg,
/// |attempt| async move {
/// let worker = select_cb_aware_worker()?;
/// let resp = send_request(worker).await;
/// worker.record_outcome(resp.status().is_success());
/// resp
/// },
/// |res, _| matches!(res.status(), StatusCode::REQUEST_TIMEOUT | StatusCode::TOO_MANY_REQUESTS | StatusCode::INTERNAL_SERVER_ERROR | StatusCode::BAD_GATEWAY | StatusCode::SERVICE_UNAVAILABLE | StatusCode::GATEWAY_TIMEOUT),
/// |delay, attempt| RouterMetrics::record_retry_backoff_duration(delay, attempt),
/// || RouterMetrics::record_retries_exhausted("/route"),
/// ).await;
/// ```
pub async fn execute_response_with_retry<Op, Fut, ShouldRetry, OnBackoff, OnExhausted>(
config: &RetryConfig,
mut operation: Op,
should_retry: ShouldRetry,
on_backoff: OnBackoff,
mut on_exhausted: OnExhausted,
) -> Response
where
Op: FnMut(u32) -> Fut,
Fut: std::future::Future<Output = Response>,
ShouldRetry: Fn(&Response, u32) -> bool,
OnBackoff: Fn(Duration, u32),
OnExhausted: FnMut(),
{
let max = config.max_retries.max(1);
let mut attempt: u32 = 0;
loop {
let response = operation(attempt).await;
let is_last = attempt + 1 >= max;
if !should_retry(&response, attempt) {
return response;
}
if is_last {
// Exhausted retries
on_exhausted();
return response;
}
// Backoff before next attempt
let next_attempt = attempt + 1;
// Compute delay based on the number of failures so far (0-indexed)
let delay = BackoffCalculator::calculate_delay(config, attempt);
debug!(
attempt = attempt,
next_attempt = next_attempt,
delay_ms = delay.as_millis() as u64,
"Retry backoff"
);
on_backoff(delay, next_attempt);
tokio::time::sleep(delay).await;
attempt = next_attempt;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
fn base_retry_config() -> RetryConfig {
RetryConfig {
max_retries: 3,
initial_backoff_ms: 1,
max_backoff_ms: 4,
backoff_multiplier: 2.0,
jitter_factor: 0.0,
}
}
#[test]
fn test_backoff_no_jitter_progression_and_cap() {
let cfg = RetryConfig {
max_retries: 10,
initial_backoff_ms: 100,
max_backoff_ms: 250,
backoff_multiplier: 2.0,
jitter_factor: 0.0,
};
// attempt=0 => 100ms
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 0),
Duration::from_millis(100)
);
// attempt=1 => 200ms
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 1),
Duration::from_millis(200)
);
// attempt=2 => 400ms -> capped to 250ms
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 2),
Duration::from_millis(250)
);
// large attempt still capped
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 10),
Duration::from_millis(250)
);
}
#[test]
fn test_backoff_with_jitter_within_bounds() {
let cfg = RetryConfig {
max_retries: 5,
initial_backoff_ms: 100,
max_backoff_ms: 10_000,
backoff_multiplier: 2.0,
jitter_factor: 0.5,
};
// attempt=2 => base 400ms, jitter in [0.5x, 1.5x]
let base = 400.0;
for _ in 0..50 {
let d = BackoffCalculator::calculate_delay(&cfg, 2).as_millis() as f32;
assert!(d >= base * 0.5 - 1.0 && d <= base * 1.5 + 1.0);
}
}
#[tokio::test]
async fn test_execute_with_retry_success_after_failures() {
let cfg = base_retry_config();
let remaining = Arc::new(AtomicU32::new(2));
let calls = Arc::new(AtomicU32::new(0));
let res: Result<u32, RetryError> = RetryExecutor::execute_with_retry(&cfg, {
let remaining = remaining.clone();
let calls = calls.clone();
move |_attempt| {
calls.fetch_add(1, Ordering::Relaxed);
let remaining = remaining.clone();
async move {
if remaining
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |v| v.checked_sub(1))
.is_ok()
{
Err(())
} else {
Ok(42u32)
}
}
}
})
.await;
assert!(res.is_ok());
assert_eq!(res.unwrap(), 42);
assert_eq!(calls.load(Ordering::Relaxed), 3); // 2 fails + 1 success
}
#[tokio::test]
async fn test_execute_with_retry_exhausted() {
let cfg = base_retry_config();
let calls = Arc::new(AtomicU32::new(0));
let res: Result<u32, RetryError> = RetryExecutor::execute_with_retry(&cfg, {
let calls = calls.clone();
move |_attempt| {
calls.fetch_add(1, Ordering::Relaxed);
async move { Err(()) }
}
})
.await;
assert!(matches!(res, Err(RetryError::MaxRetriesExceeded)));
assert_eq!(calls.load(Ordering::Relaxed), cfg.max_retries);
}
#[tokio::test]
async fn test_execute_response_with_retry_success_path_and_hooks() {
let cfg = base_retry_config();
let remaining = Arc::new(AtomicU32::new(2));
let calls = Arc::new(AtomicU32::new(0));
let backoffs = Arc::new(AtomicU32::new(0));
let exhausted = Arc::new(AtomicU32::new(0));
let response = RetryExecutor::execute_response_with_retry(
&cfg,
{
let remaining = remaining.clone();
let calls = calls.clone();
move |_attempt| {
calls.fetch_add(1, Ordering::Relaxed);
let remaining = remaining.clone();
async move {
if remaining
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |v| v.checked_sub(1))
.is_ok()
{
(StatusCode::SERVICE_UNAVAILABLE, "fail").into_response()
} else {
(StatusCode::OK, "ok").into_response()
}
}
}
},
|res, _attempt| !res.status().is_success(), // retry until success
{
let backoffs = backoffs.clone();
move |_delay, _next_attempt| {
backoffs.fetch_add(1, Ordering::Relaxed);
}
},
{
let exhausted = exhausted.clone();
move || {
exhausted.fetch_add(1, Ordering::Relaxed);
}
},
)
.await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(calls.load(Ordering::Relaxed), 3); // 2 fails + 1 success
assert_eq!(backoffs.load(Ordering::Relaxed), 2);
assert_eq!(exhausted.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn test_execute_response_with_retry_non_retryable_short_circuit() {
let cfg = base_retry_config();
let calls = Arc::new(AtomicU32::new(0));
let backoffs = Arc::new(AtomicU32::new(0));
let exhausted = Arc::new(AtomicU32::new(0));
let response = RetryExecutor::execute_response_with_retry(
&cfg,
{
let calls = calls.clone();
move |_attempt| {
calls.fetch_add(1, Ordering::Relaxed);
async move { (StatusCode::BAD_REQUEST, "bad").into_response() }
}
},
|_res, _attempt| false, // never retry
{
let backoffs = backoffs.clone();
move |_delay, _next_attempt| {
backoffs.fetch_add(1, Ordering::Relaxed);
}
},
{
let exhausted = exhausted.clone();
move || {
exhausted.fetch_add(1, Ordering::Relaxed);
}
},
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(calls.load(Ordering::Relaxed), 1);
assert_eq!(backoffs.load(Ordering::Relaxed), 0);
assert_eq!(exhausted.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn test_execute_response_with_retry_exhausted_hooks() {
let cfg = base_retry_config();
let calls = Arc::new(AtomicU32::new(0));
let backoffs = Arc::new(AtomicU32::new(0));
let exhausted = Arc::new(AtomicU32::new(0));
let response = RetryExecutor::execute_response_with_retry(
&cfg,
{
let calls = calls.clone();
move |_attempt| {
calls.fetch_add(1, Ordering::Relaxed);
async move { (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response() }
}
},
|_res, _attempt| true, // keep retrying
{
let backoffs = backoffs.clone();
move |_delay, _next_attempt| {
backoffs.fetch_add(1, Ordering::Relaxed);
}
},
{
let exhausted = exhausted.clone();
move || {
exhausted.fetch_add(1, Ordering::Relaxed);
}
},
)
.await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(calls.load(Ordering::Relaxed), cfg.max_retries);
assert_eq!(backoffs.load(Ordering::Relaxed), cfg.max_retries - 1);
assert_eq!(exhausted.load(Ordering::Relaxed), 1);
}
}

195
sgl-router/src/core/token_bucket.rs Normal file
View File

@@ -0,0 +1,195 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, Notify};
use tracing::{debug, trace};
/// Token bucket for rate limiting
///
/// This implementation provides:
/// - Smooth rate limiting with configurable refill rate
/// - Burst capacity handling
/// - Fair queuing for waiting requests
#[derive(Clone)]
pub struct TokenBucket {
inner: Arc<Mutex<TokenBucketInner>>,
notify: Arc<Notify>,
capacity: f64,
refill_rate: f64, // tokens per second
}
struct TokenBucketInner {
tokens: f64,
last_refill: Instant,
}
impl TokenBucket {
/// Create a new token bucket
///
/// # Arguments
/// * `capacity` - Maximum number of tokens (burst capacity)
/// * `refill_rate` - Tokens added per second
pub fn new(capacity: usize, refill_rate: usize) -> Self {
let capacity = capacity as f64;
let refill_rate = refill_rate as f64;
// Ensure refill_rate is not zero to prevent division by zero
let refill_rate = if refill_rate > 0.0 {
refill_rate
} else {
1.0 // Default to 1 token per second if zero
};
Self {
inner: Arc::new(Mutex::new(TokenBucketInner {
tokens: capacity, // Start full
last_refill: Instant::now(),
})),
notify: Arc::new(Notify::new()),
capacity,
refill_rate,
}
}
/// Try to acquire tokens immediately
pub async fn try_acquire(&self, tokens: f64) -> Result<(), ()> {
let mut inner = self.inner.lock().await;
// Refill tokens based on elapsed time
let now = Instant::now();
let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
let refill_amount = elapsed * self.refill_rate;
inner.tokens = (inner.tokens + refill_amount).min(self.capacity);
inner.last_refill = now;
trace!(
"Token bucket: {} tokens available, requesting {}",
inner.tokens,
tokens
);
if inner.tokens >= tokens {
inner.tokens -= tokens;
debug!(
"Token bucket: acquired {} tokens, {} remaining",
tokens, inner.tokens
);
Ok(())
} else {
Err(())
}
}
/// Acquire tokens, waiting if necessary
pub async fn acquire(&self, tokens: f64) -> Result<(), tokio::time::error::Elapsed> {
// First try to acquire immediately
if self.try_acquire(tokens).await.is_ok() {
return Ok(());
}
// Calculate wait time
let wait_time = {
let inner = self.inner.lock().await;
let tokens_needed = tokens - inner.tokens;
let wait_secs = tokens_needed / self.refill_rate;
Duration::from_secs_f64(wait_secs)
};
debug!(
"Token bucket: waiting {:?} for {} tokens",
wait_time, tokens
);
// Wait for tokens to be available
tokio::time::timeout(wait_time, async {
loop {
// Check if we can acquire now
if self.try_acquire(tokens).await.is_ok() {
return;
}
// Wait for notification or small interval
tokio::select! {
_ = self.notify.notified() => {},
_ = tokio::time::sleep(Duration::from_millis(10)) => {},
}
}
})
.await?;
Ok(())
}
/// Acquire tokens with custom timeout
pub async fn acquire_timeout(
&self,
tokens: f64,
timeout: Duration,
) -> Result<(), tokio::time::error::Elapsed> {
tokio::time::timeout(timeout, self.acquire(tokens)).await?
}
/// Return tokens to the bucket (for cancelled requests)
pub async fn return_tokens(&self, tokens: f64) {
let mut inner = self.inner.lock().await;
inner.tokens = (inner.tokens + tokens).min(self.capacity);
self.notify.notify_waiters();
debug!(
"Token bucket: returned {} tokens, {} available",
tokens, inner.tokens
);
}
/// Get current available tokens (for monitoring)
pub async fn available_tokens(&self) -> f64 {
let mut inner = self.inner.lock().await;
// Refill before checking
let now = Instant::now();
let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
let refill_amount = elapsed * self.refill_rate;
inner.tokens = (inner.tokens + refill_amount).min(self.capacity);
inner.last_refill = now;
inner.tokens
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_token_bucket_basic() {
let bucket = TokenBucket::new(10, 5); // 10 capacity, 5 per second
// Should succeed - bucket starts full
assert!(bucket.try_acquire(5.0).await.is_ok());
assert!(bucket.try_acquire(5.0).await.is_ok());
// Should fail - no tokens left
assert!(bucket.try_acquire(1.0).await.is_err());
// Wait for refill
tokio::time::sleep(Duration::from_millis(300)).await;
// Should have ~1.5 tokens now
assert!(bucket.try_acquire(1.0).await.is_ok());
}
#[tokio::test]
async fn test_token_bucket_refill() {
let bucket = TokenBucket::new(10, 10); // 10 capacity, 10 per second
// Use all tokens
assert!(bucket.try_acquire(10.0).await.is_ok());
// Wait for partial refill
tokio::time::sleep(Duration::from_millis(500)).await;
// Should have ~5 tokens
let available = bucket.available_tokens().await;
assert!((4.0..=6.0).contains(&available));
}
}

1980
sgl-router/src/core/worker.rs Normal file
View File

File diff suppressed because it is too large Load Diff

327
sgl-router/src/grpc/client.rs Normal file
View File

@@ -0,0 +1,327 @@
use std::time::Duration;
use tonic::{transport::Channel, Request};
use tracing::debug;
// Include the generated protobuf code
pub mod proto {
tonic::include_proto!("sglang.grpc.scheduler");
}
// The generated module structure depends on the package name in the .proto file
// package sglang.grpc.scheduler; generates a nested module structure
/// gRPC client for SGLang scheduler
pub struct SglangSchedulerClient {
client: proto::sglang_scheduler_client::SglangSchedulerClient<Channel>,
}
impl SglangSchedulerClient {
/// Create a new client and connect to the scheduler
pub async fn connect(endpoint: &str) -> Result<Self, Box<dyn std::error::Error>> {
debug!("Connecting to SGLang scheduler at {}", endpoint);
// Convert grpc:// to http:// for tonic
let http_endpoint = if endpoint.starts_with("grpc://") {
endpoint.replace("grpc://", "http://")
} else {
endpoint.to_string()
};
let channel = Channel::from_shared(http_endpoint)?
.timeout(Duration::from_secs(30))
.connect()
.await?;
let client = proto::sglang_scheduler_client::SglangSchedulerClient::new(channel);
Ok(Self { client })
}
/// Initialize the connection
pub async fn initialize(
&mut self,
client_id: String,
) -> Result<proto::InitializeResponse, Box<dyn std::error::Error>> {
let request = Request::new(proto::InitializeRequest {
client_id,
client_version: "0.1.0".to_string(),
mode: proto::initialize_request::Mode::Regular as i32,
});
let response = self.client.initialize(request).await?;
Ok(response.into_inner())
}
/// Submit a generation request (returns streaming response)
pub async fn generate_stream(
&mut self,
req: proto::GenerateRequest,
) -> Result<tonic::Streaming<proto::GenerateResponse>, Box<dyn std::error::Error>> {
let request = Request::new(req);
let response = self.client.generate(request).await?;
Ok(response.into_inner())
}
/// Perform health check
pub async fn health_check(
&mut self,
) -> Result<proto::HealthCheckResponse, Box<dyn std::error::Error>> {
debug!("Sending health check request");
let request = Request::new(proto::HealthCheckRequest {
include_detailed_metrics: false,
});
let response = self.client.health_check(request).await?;
debug!("Health check response received");
Ok(response.into_inner())
}
/// Abort a request
pub async fn abort_request(
&mut self,
request_id: String,
reason: String,
) -> Result<(), Box<dyn std::error::Error>> {
let request = Request::new(proto::AbortRequest { request_id, reason });
self.client.abort(request).await?;
Ok(())
}
/// Flush cache
pub async fn flush_cache(
&mut self,
flush_all: bool,
session_ids: &[String],
) -> Result<proto::FlushCacheResponse, Box<dyn std::error::Error>> {
let request = Request::new(proto::FlushCacheRequest {
flush_all,
session_ids: session_ids.to_vec(),
});
let response = self.client.flush_cache(request).await?;
Ok(response.into_inner())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_proto_types_compilation() {
// Test that protobuf types can be constructed
let init_req = proto::InitializeRequest {
client_id: "test-client".to_string(),
client_version: "0.1.0".to_string(),
mode: 0,
};
assert_eq!(init_req.client_id, "test-client");
assert_eq!(init_req.client_version, "0.1.0");
assert_eq!(init_req.mode, 0);
}
#[test]
fn test_generate_request_construction() {
let sampling_params = proto::SamplingParams {
temperature: 0.7,
max_new_tokens: 128,
top_p: 0.9,
top_k: 50,
stop: vec!["</s>".to_string()],
..Default::default()
};
let gen_req = proto::GenerateRequest {
request_id: "test-req-123".to_string(),
input: Some(proto::generate_request::Input::Text(
"Hello world".to_string(),
)),
sampling_params: Some(sampling_params),
return_logprob: true,
logprob_start_len: 0,
top_logprobs_num: 5,
..Default::default()
};
assert_eq!(gen_req.request_id, "test-req-123");
if let Some(proto::generate_request::Input::Text(text)) = &gen_req.input {
assert_eq!(text, "Hello world");
}
assert!(gen_req.return_logprob);
assert_eq!(gen_req.top_logprobs_num, 5);
let params = gen_req.sampling_params.unwrap();
assert_eq!(params.temperature, 0.7);
assert_eq!(params.max_new_tokens, 128);
assert_eq!(params.stop, vec!["</s>"]);
}
#[test]
fn test_health_check_request() {
let health_req = proto::HealthCheckRequest {
include_detailed_metrics: true,
};
assert!(health_req.include_detailed_metrics);
}
#[test]
fn test_abort_request_construction() {
let abort_req = proto::AbortRequest {
request_id: "req-456".to_string(),
reason: "User canceled".to_string(),
};
assert_eq!(abort_req.request_id, "req-456");
assert_eq!(abort_req.reason, "User canceled");
}
#[test]
fn test_flush_cache_request() {
let flush_req = proto::FlushCacheRequest {
flush_all: true,
session_ids: vec!["session1".to_string(), "session2".to_string()],
};
assert!(flush_req.flush_all);
assert_eq!(flush_req.session_ids.len(), 2);
assert_eq!(flush_req.session_ids[0], "session1");
}
#[test]
fn test_sampling_params_defaults() {
let params = proto::SamplingParams::default();
assert_eq!(params.temperature, 0.0);
assert_eq!(params.max_new_tokens, 0);
assert_eq!(params.top_p, 0.0);
assert_eq!(params.top_k, 0);
assert!(params.stop.is_empty());
}
#[test]
fn test_multimodal_inputs() {
let mm_inputs = proto::MultimodalInputs {
image_urls: vec!["http://example.com/image.jpg".to_string()],
video_urls: vec![],
audio_urls: vec![],
image_data: vec![],
video_data: vec![],
audio_data: vec![],
modalities: vec!["image".to_string()],
..Default::default()
};
assert_eq!(mm_inputs.image_urls.len(), 1);
assert_eq!(mm_inputs.image_urls[0], "http://example.com/image.jpg");
assert_eq!(mm_inputs.modalities[0], "image");
}
#[test]
fn test_session_params() {
let session_params = proto::SessionParams {
session_id: "sess-789".to_string(),
request_id: "req-101".to_string(),
offset: 100,
replace: true,
drop_previous_output: false,
};
assert_eq!(session_params.session_id, "sess-789");
assert_eq!(session_params.request_id, "req-101");
assert_eq!(session_params.offset, 100);
assert!(session_params.replace);
assert!(!session_params.drop_previous_output);
}
#[test]
fn test_embed_request() {
let embed_req = proto::EmbedRequest {
request_id: "embed-req-202".to_string(),
input: Some(proto::embed_request::Input::Text(
"This is a test sentence for embedding".to_string(),
)),
log_metrics: true,
data_parallel_rank: 0,
..Default::default()
};
assert_eq!(embed_req.request_id, "embed-req-202");
if let Some(proto::embed_request::Input::Text(text)) = &embed_req.input {
assert_eq!(text, "This is a test sentence for embedding");
}
assert!(embed_req.log_metrics);
assert_eq!(embed_req.data_parallel_rank, 0);
}
#[tokio::test]
async fn test_client_connect_invalid_endpoint() {
// Test connecting to an invalid endpoint should return error
let result = SglangSchedulerClient::connect("invalid://endpoint").await;
assert!(result.is_err());
}
#[test]
fn test_tokenized_input() {
let tokenized = proto::TokenizedInput {
original_text: "Hello world".to_string(),
input_ids: vec![1, 15043, 1917, 2],
};
assert_eq!(tokenized.original_text, "Hello world");
assert_eq!(tokenized.input_ids, vec![1, 15043, 1917, 2]);
}
// Test response type construction
#[test]
fn test_generate_stream_chunk() {
let chunk = proto::GenerateStreamChunk {
token_id: 1234,
text: " world".to_string(),
prompt_tokens: 5,
completion_tokens: 2,
cached_tokens: 3,
generation_time: 0.025,
queue_time: 10,
..Default::default()
};
assert_eq!(chunk.token_id, 1234);
assert_eq!(chunk.text, " world");
assert_eq!(chunk.prompt_tokens, 5);
assert_eq!(chunk.completion_tokens, 2);
assert_eq!(chunk.cached_tokens, 3);
assert_eq!(chunk.generation_time, 0.025);
assert_eq!(chunk.queue_time, 10);
}
#[test]
fn test_model_info() {
let model_info = proto::ModelInfo {
model_name: "Meta-Llama-3-8B-Instruct".to_string(),
max_context_length: 8192,
vocab_size: 128256,
supports_tool_calling: true,
supports_vision: false,
special_tokens: vec![
"<|begin_of_text|>".to_string(),
"<|end_of_text|>".to_string(),
],
model_type: "llama".to_string(),
num_layers: 32,
hidden_size: 4096,
num_attention_heads: 32,
num_key_value_heads: 8,
tokenizer_type: "llama".to_string(),
eos_token_ids: vec![128001, 128009],
pad_token_id: 128001,
bos_token_id: 128000,
};
assert_eq!(model_info.model_name, "Meta-Llama-3-8B-Instruct");
assert_eq!(model_info.max_context_length, 8192);
assert_eq!(model_info.vocab_size, 128256);
assert!(model_info.supports_tool_calling);
assert!(!model_info.supports_vision);
assert_eq!(model_info.special_tokens.len(), 2);
assert_eq!(model_info.num_layers, 32);
assert_eq!(model_info.eos_token_ids, vec![128001, 128009]);
}
}

8
sgl-router/src/grpc/mod.rs Normal file
View File

@@ -0,0 +1,8 @@
//! gRPC client module for communicating with SGLang scheduler
//!
//! This module provides a gRPC client implementation for the SGLang router.
pub mod client;
// Re-export the client
pub use client::{proto, SglangSchedulerClient};

508
sgl-router/src/lib.rs Normal file
View File

@@ -0,0 +1,508 @@
use pyo3::prelude::*;
pub mod config;
pub mod logging;
use std::collections::HashMap;
pub mod core;
#[cfg(feature = "grpc-client")]
pub mod grpc;
pub mod mcp;
pub mod metrics;
pub mod middleware;
pub mod policies;
pub mod protocols;
pub mod reasoning_parser;
pub mod routers;
pub mod server;
pub mod service_discovery;
pub mod tokenizer;
pub mod tool_parser;
pub mod tree;
use crate::metrics::PrometheusConfig;
#[pyclass(eq)]
#[derive(Clone, PartialEq, Debug)]
pub enum PolicyType {
Random,
RoundRobin,
CacheAware,
PowerOfTwo,
}
#[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,
prefill_urls: Option<Vec<(String, Option<u16>)>>,
decode_urls: Option<Vec<String>>,
prefill_policy: Option<PolicyType>,
decode_policy: Option<PolicyType>,
max_concurrent_requests: usize,
cors_allowed_origins: Vec<String>,
// Retry configuration
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,
// Circuit breaker configuration
cb_failure_threshold: u32,
cb_success_threshold: u32,
cb_timeout_duration_secs: u64,
cb_window_duration_secs: u64,
disable_circuit_breaker: bool,
// Health check configuration
health_failure_threshold: u32,
health_success_threshold: u32,
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
// IGW (Inference Gateway) configuration
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
rate_limit_tokens_per_second: Option<usize>,
// Connection mode (determined from worker URLs)
connection_mode: config::ConnectionMode,
// Model path for tokenizer
model_path: Option<String>,
// Explicit tokenizer path
tokenizer_path: Option<String>,
}
impl Router {
/// Determine connection mode from worker URLs
fn determine_connection_mode(worker_urls: &[String]) -> config::ConnectionMode {
// Only consider it gRPC if explicitly specified with grpc:// or grpcs:// scheme
for url in worker_urls {
if url.starts_with("grpc://") || url.starts_with("grpcs://") {
return config::ConnectionMode::Grpc;
}
}
// Default to HTTP for all other cases (including http://, https://, or no scheme)
config::ConnectionMode::Http
}
/// Convert PyO3 Router to RouterConfig
pub fn to_router_config(&self) -> config::ConfigResult<config::RouterConfig> {
use config::{
DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode,
};
// Convert policy helper function
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, // Default value
},
}
};
// Determine routing mode
let mode = if self.enable_igw {
// IGW mode - routing mode is not used in IGW, but we need to provide a placeholder
RoutingMode::Regular {
worker_urls: vec![],
}
} 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(),
}
};
// Convert main policy
let policy = convert_policy(&self.policy);
// Service discovery configuration
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
};
// Metrics configuration
let metrics = match (self.prometheus_port, self.prometheus_host.as_ref()) {
(Some(port), Some(host)) => Some(MetricsConfig {
port,
host: host.clone(),
}),
_ => None,
};
Ok(config::RouterConfig {
mode,
policy,
host: self.host.clone(),
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,
dp_aware: self.dp_aware,
api_key: self.api_key.clone(),
discovery,
metrics,
log_dir: self.log_dir.clone(),
log_level: self.log_level.clone(),
request_id_headers: self.request_id_headers.clone(),
max_concurrent_requests: self.max_concurrent_requests,
queue_size: self.queue_size,
queue_timeout_secs: self.queue_timeout_secs,
rate_limit_tokens_per_second: self.rate_limit_tokens_per_second,
cors_allowed_origins: self.cors_allowed_origins.clone(),
retry: 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::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,
},
disable_retries: self.disable_retries,
disable_circuit_breaker: self.disable_circuit_breaker,
health_check: 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(),
},
enable_igw: self.enable_igw,
model_path: self.model_path.clone(),
tokenizer_path: self.tokenizer_path.clone(),
})
}
}
#[pymethods]
impl Router {
#[new]
#[pyo3(signature = (
worker_urls,
policy = PolicyType::RoundRobin,
host = String::from("127.0.0.1"),
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, // 512MB default for large batches
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, // Add configurable request timeout
request_id_headers = None, // Custom request ID headers
pd_disaggregation = false, // New flag for PD mode
prefill_urls = None,
decode_urls = None,
prefill_policy = None,
decode_policy = None,
max_concurrent_requests = 256,
cors_allowed_origins = vec![],
// Retry defaults
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,
// Circuit breaker defaults
cb_failure_threshold = 10,
cb_success_threshold = 3,
cb_timeout_duration_secs = 60,
cb_window_duration_secs = 120,
disable_circuit_breaker = false,
// Health check defaults
health_failure_threshold = 3,
health_success_threshold = 2,
health_check_timeout_secs = 5,
health_check_interval_secs = 60,
health_check_endpoint = String::from("/health"),
// IGW defaults
enable_igw = false,
queue_size = 100,
queue_timeout_secs = 60,
rate_limit_tokens_per_second = None,
// Tokenizer defaults
model_path = None,
tokenizer_path = None,
))]
#[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,
prefill_urls: Option<Vec<(String, Option<u16>)>>,
decode_urls: Option<Vec<String>>,
prefill_policy: Option<PolicyType>,
decode_policy: Option<PolicyType>,
max_concurrent_requests: usize,
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<usize>,
model_path: Option<String>,
tokenizer_path: Option<String>,
) -> PyResult<Self> {
// Determine connection mode from worker URLs
let mut all_urls = worker_urls.clone();
// Add prefill URLs if in PD mode
if let Some(ref prefill_urls) = prefill_urls {
for (url, _) in prefill_urls {
all_urls.push(url.clone());
}
}
// Add decode URLs if in PD mode
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,
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,
})
}
fn start(&self) -> PyResult<()> {
// Convert to RouterConfig and validate
let router_config = self.to_router_config().map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!("Configuration error: {}", e))
})?;
// Validate the configuration
router_config.validate().map_err(|e| {
pyo3::exceptions::PyValueError::new_err(format!(
"Configuration validation failed: {}",
e
))
})?;
// Create service discovery config if enabled
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
};
// Create Prometheus config if enabled
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()),
});
// Use tokio runtime instead of actix-web System for better compatibility
let runtime = tokio::runtime::Runtime::new()
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
// Block on the async startup function
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::<Router>()?;
Ok(())
}

163
sgl-router/src/logging.rs Normal file
View File

@@ -0,0 +1,163 @@
use std::path::PathBuf;
use tracing::Level;
use tracing_appender::non_blocking::WorkerGuard;
use tracing_appender::rolling::{RollingFileAppender, Rotation};
use tracing_log::LogTracer;
use tracing_subscriber::fmt::time::ChronoUtc;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{EnvFilter, Layer};
/// Configuration for the logging system
#[derive(Debug, Clone)]
pub struct LoggingConfig {
/// Log level for the application (default: INFO)
pub level: Level,
/// Whether to use json format for logs (default: false)
pub json_format: bool,
/// Path to store log files. If None, logs will only go to stdout/stderr
pub log_dir: Option<String>,
/// Whether to colorize logs when output is a terminal (default: true)
pub colorize: bool,
/// Log file name to use if log_dir is specified (default: "sgl-router")
pub log_file_name: String,
/// Custom log targets to filter (default: "sglang_router_rs")
pub log_targets: Option<Vec<String>>,
}
impl Default for LoggingConfig {
fn default() -> Self {
Self {
level: Level::INFO,
json_format: false,
log_dir: None,
colorize: true,
log_file_name: "sgl-router".to_string(),
log_targets: Some(vec!["sglang_router_rs".to_string()]),
}
}
}
/// Guard that keeps the file appender worker thread alive
///
/// This must be kept in scope for the duration of the program
/// to ensure logs are properly written to files
#[allow(dead_code)]
pub struct LogGuard {
_file_guard: Option<WorkerGuard>,
}
/// Initialize the logging system with the given configuration
///
/// # Arguments
/// * `config` - Configuration for the logging system
///
/// # Returns
/// A LogGuard that must be kept alive for the duration of the program
///
/// # Panics
/// Will not panic, as initialization errors are handled gracefully
pub fn init_logging(config: LoggingConfig) -> LogGuard {
// Forward logs to tracing - ignore errors to allow for multiple initialization
let _ = LogTracer::init();
// Convert log level to filter string
let level_filter = match config.level {
Level::TRACE => "trace",
Level::DEBUG => "debug",
Level::INFO => "info",
Level::WARN => "warn",
Level::ERROR => "error",
};
// Create env filter
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
// Format: <target>=<level>,<target2>=<level2>,...
let filter_string = if let Some(targets) = &config.log_targets {
targets
.iter()
.enumerate()
.map(|(i, target)| {
if i > 0 {
format!(",{}={}", target, level_filter)
} else {
format!("{}={}", target, level_filter)
}
})
.collect::<String>()
} else {
format!("sglang_router_rs={}", level_filter)
};
EnvFilter::new(filter_string)
});
// Setup stdout/stderr layer
let mut layers = Vec::new();
// Standard timestamp format: YYYY-MM-DD HH:MM:SS
let time_format = "%Y-%m-%d %H:%M:%S".to_string();
// Configure the console stdout layer
let stdout_layer = tracing_subscriber::fmt::layer()
.with_ansi(config.colorize)
.with_file(true)
.with_line_number(true)
.with_timer(ChronoUtc::new(time_format.clone()));
let stdout_layer = if config.json_format {
stdout_layer.json().flatten_event(true).boxed()
} else {
stdout_layer.boxed()
};
layers.push(stdout_layer);
// Create a file appender if log_dir is specified
let mut file_guard = None;
if let Some(log_dir) = &config.log_dir {
let file_name = config.log_file_name.clone();
let log_dir = PathBuf::from(log_dir);
// Create log directory if it doesn't exist
if !log_dir.exists() {
if let Err(e) = std::fs::create_dir_all(&log_dir) {
eprintln!("Failed to create log directory: {}", e);
return LogGuard { _file_guard: None };
}
}
let file_appender = RollingFileAppender::new(Rotation::DAILY, log_dir, file_name);
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
file_guard = Some(guard);
let file_layer = tracing_subscriber::fmt::layer()
.with_ansi(false) // Never use ANSI colors in log files
.with_file(true)
.with_line_number(true)
.with_timer(ChronoUtc::new(time_format))
.with_writer(non_blocking);
let file_layer = if config.json_format {
file_layer.json().flatten_event(true).boxed()
} else {
file_layer.boxed()
};
layers.push(file_layer);
}
// Initialize the subscriber with all layers
// Use try_init to handle errors gracefully in case another subscriber is already set
let _ = tracing_subscriber::registry()
.with(env_filter)
.with(layers)
.try_init();
// Return the guard to keep the file appender worker thread alive
LogGuard {
_file_guard: file_guard,
}
}

636
sgl-router/src/main.rs Normal file
View File

@@ -0,0 +1,636 @@
use clap::{ArgAction, Parser, ValueEnum};
use sglang_router_rs::config::{
CircuitBreakerConfig, ConfigError, ConfigResult, ConnectionMode, DiscoveryConfig,
HealthCheckConfig, MetricsConfig, PolicyConfig, RetryConfig, RouterConfig, RoutingMode,
};
use sglang_router_rs::metrics::PrometheusConfig;
use sglang_router_rs::server::{self, ServerConfig};
use sglang_router_rs::service_discovery::ServiceDiscoveryConfig;
use std::collections::HashMap;
// Helper function to parse prefill arguments from command line
fn parse_prefill_args() -> Vec<(String, Option<u16>)> {
let args: Vec<String> = std::env::args().collect();
let mut prefill_entries = Vec::new();
let mut i = 0;
while i < args.len() {
if args[i] == "--prefill" && i + 1 < args.len() {
let url = args[i + 1].clone();
let bootstrap_port = if i + 2 < args.len() && !args[i + 2].starts_with("--") {
// Check if next arg is a port number
if let Ok(port) = args[i + 2].parse::<u16>() {
i += 1; // Skip the port argument
Some(port)
} else if args[i + 2].to_lowercase() == "none" {
i += 1; // Skip the "none" argument
None
} else {
None
}
} else {
None
};
prefill_entries.push((url, bootstrap_port));
i += 2; // Skip --prefill and URL
} else {
i += 1;
}
}
prefill_entries
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
pub enum Backend {
#[value(name = "sglang")]
Sglang,
#[value(name = "vllm")]
Vllm,
#[value(name = "trtllm")]
Trtllm,
#[value(name = "openai")]
Openai,
#[value(name = "anthropic")]
Anthropic,
}
impl std::fmt::Display for Backend {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let s = match self {
Backend::Sglang => "sglang",
Backend::Vllm => "vllm",
Backend::Trtllm => "trtllm",
Backend::Openai => "openai",
Backend::Anthropic => "anthropic",
};
write!(f, "{}", s)
}
}
#[derive(Parser, Debug)]
#[command(name = "sglang-router")]
#[command(about = "SGLang Router - High-performance request distribution across worker nodes")]
#[command(long_about = r#"
SGLang Router - High-performance request distribution across worker nodes
Usage:
This launcher enables starting a router with individual worker instances. It is useful for
multi-node setups or when you want to start workers and router separately.
Examples:
# Regular mode
sglang-router --worker-urls http://worker1:8000 http://worker2:8000
# PD disaggregated mode with same policy for both
sglang-router --pd-disaggregation \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 9002 \
--decode http://127.0.0.3:30003 \
--decode http://127.0.0.4:30004 \
--policy cache_aware
# PD mode with different policies for prefill and decode
sglang-router --pd-disaggregation \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 \
--decode http://127.0.0.3:30003 \
--decode http://127.0.0.4:30004 \
--prefill-policy cache_aware --decode-policy power_of_two
"#)]
struct CliArgs {
/// Host address to bind the router server
#[arg(long, default_value = "127.0.0.1")]
host: String,
/// Port number to bind the router server
#[arg(long, default_value_t = 30000)]
port: u16,
/// List of worker URLs (e.g., http://worker1:8000 http://worker2:8000)
#[arg(long, num_args = 0..)]
worker_urls: Vec<String>,
/// Load balancing policy to use
#[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two"])]
policy: String,
/// Enable PD (Prefill-Decode) disaggregated mode
#[arg(long, default_value_t = false)]
pd_disaggregation: bool,
/// Decode server URL (can be specified multiple times)
#[arg(long, action = ArgAction::Append)]
decode: Vec<String>,
/// Specific policy for prefill nodes in PD mode
#[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two"])]
prefill_policy: Option<String>,
/// Specific policy for decode nodes in PD mode
#[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two"])]
decode_policy: Option<String>,
/// Timeout in seconds for worker startup
#[arg(long, default_value_t = 600)]
worker_startup_timeout_secs: u64,
/// Interval in seconds between checks for worker startup
#[arg(long, default_value_t = 30)]
worker_startup_check_interval: u64,
/// Cache threshold (0.0-1.0) for cache-aware routing
#[arg(long, default_value_t = 0.3)]
cache_threshold: f32,
/// Absolute threshold for load balancing
#[arg(long, default_value_t = 64)]
balance_abs_threshold: usize,
/// Relative threshold for load balancing
#[arg(long, default_value_t = 1.5)]
balance_rel_threshold: f32,
/// Interval in seconds between cache eviction operations
#[arg(long, default_value_t = 120)]
eviction_interval: u64,
/// Maximum size of the approximation tree for cache-aware routing
#[arg(long, default_value_t = 67108864)] // 2^26
max_tree_size: usize,
/// Maximum payload size in bytes
#[arg(long, default_value_t = 536870912)] // 512MB
max_payload_size: usize,
/// Enable data parallelism aware schedule
#[arg(long, default_value_t = false)]
dp_aware: bool,
/// API key for worker authorization
#[arg(long)]
api_key: Option<String>,
/// Backend to route requests to (sglang, vllm, trtllm, openai, anthropic)
#[arg(long, value_enum, default_value_t = Backend::Sglang, alias = "runtime")]
backend: Backend,
/// Directory to store log files
#[arg(long)]
log_dir: Option<String>,
/// Set the logging level
#[arg(long, default_value = "info", value_parser = ["debug", "info", "warn", "error"])]
log_level: String,
/// Enable Kubernetes service discovery
#[arg(long, default_value_t = false)]
service_discovery: bool,
/// Label selector for Kubernetes service discovery (format: key1=value1 key2=value2)
#[arg(long, num_args = 0..)]
selector: Vec<String>,
/// Port to use for discovered worker pods
#[arg(long, default_value_t = 80)]
service_discovery_port: u16,
/// Kubernetes namespace to watch for pods
#[arg(long)]
service_discovery_namespace: Option<String>,
/// Label selector for prefill server pods in PD mode
#[arg(long, num_args = 0..)]
prefill_selector: Vec<String>,
/// Label selector for decode server pods in PD mode
#[arg(long, num_args = 0..)]
decode_selector: Vec<String>,
/// Port to expose Prometheus metrics
#[arg(long, default_value_t = 29000)]
prometheus_port: u16,
/// Host address to bind the Prometheus metrics server
#[arg(long, default_value = "127.0.0.1")]
prometheus_host: String,
/// Custom HTTP headers to check for request IDs
#[arg(long, num_args = 0..)]
request_id_headers: Vec<String>,
/// Request timeout in seconds
#[arg(long, default_value_t = 1800)]
request_timeout_secs: u64,
/// Maximum number of concurrent requests allowed
#[arg(long, default_value_t = 256)]
max_concurrent_requests: usize,
/// CORS allowed origins
#[arg(long, num_args = 0..)]
cors_allowed_origins: Vec<String>,
// Retry configuration
/// Maximum number of retries
#[arg(long, default_value_t = 5)]
retry_max_retries: u32,
/// Initial backoff in milliseconds for retries
#[arg(long, default_value_t = 50)]
retry_initial_backoff_ms: u64,
/// Maximum backoff in milliseconds for retries
#[arg(long, default_value_t = 30000)]
retry_max_backoff_ms: u64,
/// Backoff multiplier for exponential backoff
#[arg(long, default_value_t = 1.5)]
retry_backoff_multiplier: f32,
/// Jitter factor for retry backoff
#[arg(long, default_value_t = 0.2)]
retry_jitter_factor: f32,
/// Disable retries
#[arg(long, default_value_t = false)]
disable_retries: bool,
// Circuit breaker configuration
/// Number of failures before circuit breaker opens
#[arg(long, default_value_t = 10)]
cb_failure_threshold: u32,
/// Number of successes before circuit breaker closes
#[arg(long, default_value_t = 3)]
cb_success_threshold: u32,
/// Timeout duration in seconds for circuit breaker
#[arg(long, default_value_t = 60)]
cb_timeout_duration_secs: u64,
/// Window duration in seconds for circuit breaker
#[arg(long, default_value_t = 120)]
cb_window_duration_secs: u64,
/// Disable circuit breaker
#[arg(long, default_value_t = false)]
disable_circuit_breaker: bool,
// Health check configuration
/// Number of consecutive health check failures before marking worker unhealthy
#[arg(long, default_value_t = 3)]
health_failure_threshold: u32,
/// Number of consecutive health check successes before marking worker healthy
#[arg(long, default_value_t = 2)]
health_success_threshold: u32,
/// Timeout in seconds for health check requests
#[arg(long, default_value_t = 5)]
health_check_timeout_secs: u64,
/// Interval in seconds between runtime health checks
#[arg(long, default_value_t = 60)]
health_check_interval_secs: u64,
/// Health check endpoint path
#[arg(long, default_value = "/health")]
health_check_endpoint: String,
// IGW (Inference Gateway) configuration
/// Enable Inference Gateway mode
#[arg(long, default_value_t = false)]
enable_igw: bool,
// Tokenizer configuration
/// Model path for loading tokenizer (HuggingFace model ID or local path)
#[arg(long)]
model_path: Option<String>,
/// Explicit tokenizer path (overrides model_path tokenizer if provided)
#[arg(long)]
tokenizer_path: Option<String>,
}
impl CliArgs {
/// Determine connection mode from worker URLs
fn determine_connection_mode(worker_urls: &[String]) -> ConnectionMode {
// Only consider it gRPC if explicitly specified with grpc:// or grpcs:// scheme
for url in worker_urls {
if url.starts_with("grpc://") || url.starts_with("grpcs://") {
return ConnectionMode::Grpc;
}
}
// Default to HTTP for all other cases (including http://, https://, or no scheme)
ConnectionMode::Http
}
/// Parse selector strings into HashMap
fn parse_selector(selector_list: &[String]) -> HashMap<String, String> {
let mut map = HashMap::new();
for item in selector_list {
if let Some(eq_pos) = item.find('=') {
let key = item[..eq_pos].to_string();
let value = item[eq_pos + 1..].to_string();
map.insert(key, value);
}
}
map
}
/// Convert policy string to PolicyConfig
fn parse_policy(&self, policy_str: &str) -> PolicyConfig {
match policy_str {
"random" => PolicyConfig::Random,
"round_robin" => PolicyConfig::RoundRobin,
"cache_aware" => PolicyConfig::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,
max_tree_size: self.max_tree_size,
},
"power_of_two" => PolicyConfig::PowerOfTwo {
load_check_interval_secs: 5, // Default value
},
_ => PolicyConfig::RoundRobin, // Fallback
}
}
/// Convert CLI arguments to RouterConfig
fn to_router_config(
&self,
prefill_urls: Vec<(String, Option<u16>)>,
) -> ConfigResult<RouterConfig> {
// Determine routing mode
let mode = if self.enable_igw {
// IGW mode - routing mode is not used in IGW, but we need to provide a placeholder
RoutingMode::Regular {
worker_urls: vec![],
}
} else if matches!(self.backend, Backend::Openai) {
// OpenAI backend mode - use worker_urls as base(s)
RoutingMode::OpenAI {
worker_urls: self.worker_urls.clone(),
}
} else if self.pd_disaggregation {
let decode_urls = self.decode.clone();
// Validate PD configuration if not using service discovery
if !self.service_discovery && (prefill_urls.is_empty() || decode_urls.is_empty()) {
return Err(ConfigError::ValidationFailed {
reason: "PD disaggregation mode requires --prefill and --decode URLs when not using service discovery".to_string(),
});
}
RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
prefill_policy: self.prefill_policy.as_ref().map(|p| self.parse_policy(p)),
decode_policy: self.decode_policy.as_ref().map(|p| self.parse_policy(p)),
}
} else {
// Regular mode
if !self.service_discovery && self.worker_urls.is_empty() {
return Err(ConfigError::ValidationFailed {
reason: "Regular mode requires --worker-urls when not using service discovery"
.to_string(),
});
}
RoutingMode::Regular {
worker_urls: self.worker_urls.clone(),
}
};
// Main policy
let policy = self.parse_policy(&self.policy);
// Service discovery configuration
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::parse_selector(&self.selector),
prefill_selector: Self::parse_selector(&self.prefill_selector),
decode_selector: Self::parse_selector(&self.decode_selector),
bootstrap_port_annotation: "sglang.ai/bootstrap-port".to_string(),
})
} else {
None
};
// Metrics configuration
let metrics = Some(MetricsConfig {
port: self.prometheus_port,
host: self.prometheus_host.clone(),
});
// Determine connection mode from all worker URLs
let mut all_urls = Vec::new();
match &mode {
RoutingMode::Regular { worker_urls } => {
all_urls.extend(worker_urls.clone());
}
RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
..
} => {
for (url, _) in prefill_urls {
all_urls.push(url.clone());
}
all_urls.extend(decode_urls.clone());
}
RoutingMode::OpenAI { .. } => {
// For connection-mode detection, skip URLs; OpenAI forces HTTP below.
}
}
let connection_mode = match &mode {
RoutingMode::OpenAI { .. } => ConnectionMode::Http,
_ => Self::determine_connection_mode(&all_urls),
};
// Build RouterConfig
Ok(RouterConfig {
mode,
policy,
connection_mode,
host: self.host.clone(),
port: self.port,
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,
dp_aware: self.dp_aware,
api_key: self.api_key.clone(),
discovery,
metrics,
log_dir: self.log_dir.clone(),
log_level: Some(self.log_level.clone()),
request_id_headers: if self.request_id_headers.is_empty() {
None
} else {
Some(self.request_id_headers.clone())
},
max_concurrent_requests: self.max_concurrent_requests,
queue_size: 100, // Default queue size
queue_timeout_secs: 60, // Default timeout
cors_allowed_origins: self.cors_allowed_origins.clone(),
retry: 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: 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,
},
disable_retries: self.disable_retries,
disable_circuit_breaker: self.disable_circuit_breaker,
health_check: 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(),
},
enable_igw: self.enable_igw,
rate_limit_tokens_per_second: None,
model_path: self.model_path.clone(),
tokenizer_path: self.tokenizer_path.clone(),
})
}
/// Create ServerConfig from CLI args and RouterConfig
fn to_server_config(&self, router_config: RouterConfig) -> ServerConfig {
// Create service discovery config if enabled
let service_discovery_config = if self.service_discovery {
Some(ServiceDiscoveryConfig {
enabled: true,
selector: Self::parse_selector(&self.selector),
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::parse_selector(&self.prefill_selector),
decode_selector: Self::parse_selector(&self.decode_selector),
bootstrap_port_annotation: "sglang.ai/bootstrap-port".to_string(),
})
} else {
None
};
// Create Prometheus config
let prometheus_config = Some(PrometheusConfig {
port: self.prometheus_port,
host: self.prometheus_host.clone(),
});
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: Some(self.log_level.clone()),
service_discovery_config,
prometheus_config,
request_timeout_secs: self.request_timeout_secs,
request_id_headers: if self.request_id_headers.is_empty() {
None
} else {
Some(self.request_id_headers.clone())
},
}
}
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Parse prefill arguments manually before clap parsing
let prefill_urls = parse_prefill_args();
// Filter out prefill arguments and their values before passing to clap
let mut filtered_args: Vec<String> = Vec::new();
let raw_args: Vec<String> = std::env::args().collect();
let mut i = 0;
while i < raw_args.len() {
if raw_args[i] == "--prefill" && i + 1 < raw_args.len() {
// Skip --prefill and its URL
i += 2;
// Also skip bootstrap port if present
if i < raw_args.len()
&& !raw_args[i].starts_with("--")
&& (raw_args[i].parse::<u16>().is_ok() || raw_args[i].to_lowercase() == "none")
{
i += 1;
}
} else {
filtered_args.push(raw_args[i].clone());
i += 1;
}
}
// Parse CLI arguments with clap using filtered args
let cli_args = CliArgs::parse_from(filtered_args);
// Print startup info
println!("SGLang Router starting...");
println!("Host: {}:{}", cli_args.host, cli_args.port);
let mode_str = if cli_args.enable_igw {
"IGW (Inference Gateway)".to_string()
} else if matches!(cli_args.backend, Backend::Openai) {
"OpenAI Backend".to_string()
} else if cli_args.pd_disaggregation {
"PD Disaggregated".to_string()
} else {
format!("Regular ({})", cli_args.backend)
};
println!("Mode: {}", mode_str);
// Warn for runtimes that are parsed but not yet implemented
match cli_args.backend {
Backend::Vllm | Backend::Trtllm | Backend::Anthropic => {
println!(
"WARNING: runtime '{}' not implemented yet; falling back to regular routing. \
Provide --worker-urls or PD flags as usual.",
cli_args.backend
);
}
Backend::Sglang | Backend::Openai => {}
}
if !cli_args.enable_igw {
println!("Policy: {}", cli_args.policy);
if cli_args.pd_disaggregation && !prefill_urls.is_empty() {
println!("Prefill nodes: {:?}", prefill_urls);
println!("Decode nodes: {:?}", cli_args.decode);
}
}
// Convert to RouterConfig
let router_config = cli_args.to_router_config(prefill_urls)?;
// Validate configuration
router_config.validate()?;
// Create ServerConfig
let server_config = cli_args.to_server_config(router_config);
// Create a new runtime for the server (like Python binding does)
let runtime = tokio::runtime::Runtime::new()?;
// Block on the async startup function
runtime.block_on(async move { server::startup(server_config).await })?;
Ok(())
}

535
sgl-router/src/mcp/client_manager.rs Normal file
View File

@@ -0,0 +1,535 @@
use backoff::ExponentialBackoffBuilder;
use dashmap::DashMap;
use rmcp::{
model::{
CallToolRequestParam, GetPromptRequestParam, GetPromptResult, Prompt,
ReadResourceRequestParam, ReadResourceResult, Resource, Tool as McpTool,
},
service::RunningService,
transport::{
sse_client::SseClientConfig, streamable_http_client::StreamableHttpClientTransportConfig,
ConfigureCommandExt, SseClientTransport, StreamableHttpClientTransport, TokioChildProcess,
},
RoleClient, ServiceExt,
};
use serde::{Deserialize, Serialize};
use std::{borrow::Cow, collections::HashMap, time::Duration};
use crate::mcp::{
config::{McpConfig, McpServerConfig, McpTransport},
error::{McpError, McpResult},
};
/// Information about an available tool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInfo {
pub name: String,
pub description: String,
pub server: String,
pub parameters: Option<serde_json::Value>,
}
/// Information about an available prompt
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptInfo {
pub name: String,
pub description: Option<String>,
pub server: String,
pub arguments: Option<Vec<serde_json::Value>>,
}
/// Information about an available resource
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceInfo {
pub uri: String,
pub name: String,
pub description: Option<String>,
pub mime_type: Option<String>,
pub server: String,
}
/// Manages MCP client connections and tool execution
pub struct McpClientManager {
/// Map of server_name -> MCP client
clients: HashMap<String, RunningService<RoleClient, ()>>,
/// Map of tool_name -> (server_name, tool_definition)
tools: DashMap<String, (String, McpTool)>,
/// Map of prompt_name -> (server_name, prompt_definition)
prompts: DashMap<String, (String, Prompt)>,
/// Map of resource_uri -> (server_name, resource_definition)
resources: DashMap<String, (String, Resource)>,
}
impl McpClientManager {
/// Create a new manager and connect to all configured servers
pub async fn new(config: McpConfig) -> McpResult<Self> {
let mut mgr = Self {
clients: HashMap::new(),
tools: DashMap::new(),
prompts: DashMap::new(),
resources: DashMap::new(),
};
for server_config in config.servers {
match Self::connect_server(&server_config).await {
Ok(client) => {
mgr.load_server_inventory(&server_config.name, &client)
.await;
mgr.clients.insert(server_config.name.clone(), client);
}
Err(e) => {
tracing::error!(
"Failed to connect to server '{}': {}",
server_config.name,
e
);
}
}
}
if mgr.clients.is_empty() {
return Err(McpError::ConnectionFailed(
"Failed to connect to any MCP servers".to_string(),
));
}
Ok(mgr)
}
/// Discover and cache tools/prompts/resources for a connected server
async fn load_server_inventory(
&self,
server_name: &str,
client: &RunningService<RoleClient, ()>,
) {
// Tools
match client.peer().list_all_tools().await {
Ok(ts) => {
tracing::info!("Discovered {} tools from '{}'", ts.len(), server_name);
for t in ts {
if self.tools.contains_key(t.name.as_ref()) {
tracing::warn!(
"Tool '{}' from server '{}' is overwriting an existing tool.",
&t.name,
server_name
);
}
self.tools
.insert(t.name.to_string(), (server_name.to_string(), t));
}
}
Err(e) => tracing::warn!("Failed to list tools from '{}': {}", server_name, e),
}
// Prompts
match client.peer().list_all_prompts().await {
Ok(ps) => {
tracing::info!("Discovered {} prompts from '{}'", ps.len(), server_name);
for p in ps {
if self.prompts.contains_key(&p.name) {
tracing::warn!(
"Prompt '{}' from server '{}' is overwriting an existing prompt.",
&p.name,
server_name
);
}
self.prompts
.insert(p.name.clone(), (server_name.to_string(), p));
}
}
Err(e) => tracing::debug!("No prompts or failed to list on '{}': {}", server_name, e),
}
// Resources
match client.peer().list_all_resources().await {
Ok(rs) => {
tracing::info!("Discovered {} resources from '{}'", rs.len(), server_name);
for r in rs {
if self.resources.contains_key(&r.uri) {
tracing::warn!(
"Resource '{}' from server '{}' is overwriting an existing resource.",
&r.uri,
server_name
);
}
self.resources
.insert(r.uri.clone(), (server_name.to_string(), r));
}
}
Err(e) => tracing::debug!("No resources or failed to list on '{}': {}", server_name, e),
}
}
/// Connect to a single MCP server with retry logic for remote transports
async fn connect_server(config: &McpServerConfig) -> McpResult<RunningService<RoleClient, ()>> {
let needs_retry = matches!(
&config.transport,
McpTransport::Sse { .. } | McpTransport::Streamable { .. }
);
if needs_retry {
Self::connect_server_with_retry(config).await
} else {
Self::connect_server_impl(config).await
}
}
/// Connect with exponential backoff retry for remote servers
async fn connect_server_with_retry(
config: &McpServerConfig,
) -> McpResult<RunningService<RoleClient, ()>> {
let backoff = ExponentialBackoffBuilder::new()
.with_initial_interval(Duration::from_secs(1))
.with_max_interval(Duration::from_secs(30))
.with_max_elapsed_time(Some(Duration::from_secs(120)))
.build();
backoff::future::retry(backoff, || async {
match Self::connect_server_impl(config).await {
Ok(client) => Ok(client),
Err(e) => {
tracing::warn!("Failed to connect to '{}', retrying: {}", config.name, e);
Err(backoff::Error::transient(e))
}
}
})
.await
}
/// Internal implementation of server connection
async fn connect_server_impl(
config: &McpServerConfig,
) -> McpResult<RunningService<RoleClient, ()>> {
tracing::info!(
"Connecting to MCP server '{}' via {:?}",
config.name,
config.transport
);
match &config.transport {
McpTransport::Stdio {
command,
args,
envs,
} => {
let transport = TokioChildProcess::new(
tokio::process::Command::new(command).configure(|cmd| {
cmd.args(args)
.envs(envs.iter())
.stderr(std::process::Stdio::inherit());
}),
)
.map_err(|e| McpError::Transport(format!("create stdio transport: {}", e)))?;
let client = ().serve(transport).await.map_err(|e| {
McpError::ConnectionFailed(format!("initialize stdio client: {}", e))
})?;
tracing::info!("Connected to stdio server '{}'", config.name);
Ok(client)
}
McpTransport::Sse { url, token } => {
let transport = if let Some(tok) = token {
let client = reqwest::Client::builder()
.default_headers({
let mut headers = reqwest::header::HeaderMap::new();
headers.insert(
reqwest::header::AUTHORIZATION,
format!("Bearer {}", tok).parse().map_err(|e| {
McpError::Transport(format!("auth token: {}", e))
})?,
);
headers
})
.build()
.map_err(|e| McpError::Transport(format!("build HTTP client: {}", e)))?;
let cfg = SseClientConfig {
sse_endpoint: url.clone().into(),
..Default::default()
};
SseClientTransport::start_with_client(client, cfg)
.await
.map_err(|e| McpError::Transport(format!("create SSE transport: {}", e)))?
} else {
SseClientTransport::start(url.as_str())
.await
.map_err(|e| McpError::Transport(format!("create SSE transport: {}", e)))?
};
let client = ().serve(transport).await.map_err(|e| {
McpError::ConnectionFailed(format!("initialize SSE client: {}", e))
})?;
tracing::info!("Connected to SSE server '{}' at {}", config.name, url);
Ok(client)
}
McpTransport::Streamable { url, token } => {
let transport = if let Some(tok) = token {
let mut cfg = StreamableHttpClientTransportConfig::with_uri(url.as_str());
cfg.auth_header = Some(format!("Bearer {}", tok));
StreamableHttpClientTransport::from_config(cfg)
} else {
StreamableHttpClientTransport::from_uri(url.as_str())
};
let client = ().serve(transport).await.map_err(|e| {
McpError::ConnectionFailed(format!("initialize streamable client: {}", e))
})?;
tracing::info!(
"Connected to streamable HTTP server '{}' at {}",
config.name,
url
);
Ok(client)
}
}
}
// ===== Helpers =====
fn client_for(&self, server_name: &str) -> McpResult<&RunningService<RoleClient, ()>> {
self.clients
.get(server_name)
.ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))
}
fn tool_entry(&self, name: &str) -> McpResult<(String, McpTool)> {
self.tools
.get(name)
.map(|e| e.value().clone())
.ok_or_else(|| McpError::ToolNotFound(name.to_string()))
}
fn prompt_entry(&self, name: &str) -> McpResult<(String, Prompt)> {
self.prompts
.get(name)
.map(|e| e.value().clone())
.ok_or_else(|| McpError::PromptNotFound(name.to_string()))
}
fn resource_entry(&self, uri: &str) -> McpResult<(String, Resource)> {
self.resources
.get(uri)
.map(|e| e.value().clone())
.ok_or_else(|| McpError::ResourceNotFound(uri.to_string()))
}
// ===== Tool Methods =====
/// Call a tool by name
pub async fn call_tool(
&self,
tool_name: &str,
arguments: Option<serde_json::Map<String, serde_json::Value>>,
) -> McpResult<rmcp::model::CallToolResult> {
let (server_name, _tool) = self.tool_entry(tool_name)?;
let client = self.client_for(&server_name)?;
tracing::debug!("Calling tool '{}' on '{}'", tool_name, server_name);
client
.peer()
.call_tool(CallToolRequestParam {
name: Cow::Owned(tool_name.to_string()),
arguments,
})
.await
.map_err(|e| McpError::ToolExecution(format!("Tool call failed: {}", e)))
}
/// Get all available tools
pub fn list_tools(&self) -> Vec<ToolInfo> {
self.tools
.iter()
.map(|entry| {
let tool_name = entry.key().clone();
let (server_name, tool) = entry.value();
ToolInfo {
name: tool_name,
description: tool.description.as_deref().unwrap_or_default().to_string(),
server: server_name.clone(),
parameters: Some(serde_json::Value::Object((*tool.input_schema).clone())),
}
})
.collect()
}
/// Get a specific tool by name
pub fn get_tool(&self, name: &str) -> Option<ToolInfo> {
self.tools.get(name).map(|entry| {
let (server_name, tool) = entry.value();
ToolInfo {
name: name.to_string(),
description: tool.description.as_deref().unwrap_or_default().to_string(),
server: server_name.clone(),
parameters: Some(serde_json::Value::Object((*tool.input_schema).clone())),
}
})
}
/// Check if a tool exists
pub fn has_tool(&self, name: &str) -> bool {
self.tools.contains_key(name)
}
/// Get list of connected servers
pub fn list_servers(&self) -> Vec<String> {
self.clients.keys().cloned().collect()
}
// ===== Prompt Methods =====
/// Get a prompt by name with arguments
pub async fn get_prompt(
&self,
prompt_name: &str,
arguments: Option<serde_json::Map<String, serde_json::Value>>,
) -> McpResult<GetPromptResult> {
let (server_name, _prompt) = self.prompt_entry(prompt_name)?;
let client = self.client_for(&server_name)?;
tracing::debug!("Getting prompt '{}' from '{}'", prompt_name, server_name);
client
.peer()
.get_prompt(GetPromptRequestParam {
name: prompt_name.to_string(),
arguments,
})
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to get prompt: {}", e)))
}
/// List all available prompts
pub fn list_prompts(&self) -> Vec<PromptInfo> {
self.prompts
.iter()
.map(|entry| {
let name = entry.key().clone();
let (server_name, prompt) = entry.value();
PromptInfo {
name,
description: prompt.description.clone(),
server: server_name.clone(),
arguments: prompt
.arguments
.clone()
.map(|args| args.into_iter().map(|arg| serde_json::json!(arg)).collect()),
}
})
.collect()
}
/// Get a specific prompt info by name
pub fn get_prompt_info(&self, name: &str) -> Option<PromptInfo> {
self.prompts.get(name).map(|entry| {
let (server_name, prompt) = entry.value();
PromptInfo {
name: name.to_string(),
description: prompt.description.clone(),
server: server_name.clone(),
arguments: prompt
.arguments
.clone()
.map(|args| args.into_iter().map(|arg| serde_json::json!(arg)).collect()),
}
})
}
// ===== Resource Methods =====
/// Read a resource by URI
pub async fn read_resource(&self, uri: &str) -> McpResult<ReadResourceResult> {
let (server_name, _resource) = self.resource_entry(uri)?;
let client = self.client_for(&server_name)?;
tracing::debug!("Reading resource '{}' from '{}'", uri, server_name);
client
.peer()
.read_resource(ReadResourceRequestParam {
uri: uri.to_string(),
})
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to read resource: {}", e)))
}
/// List all available resources
pub fn list_resources(&self) -> Vec<ResourceInfo> {
self.resources
.iter()
.map(|entry| {
let uri = entry.key().clone();
let (server_name, resource) = entry.value();
ResourceInfo {
uri,
name: resource.name.clone(),
description: resource.description.clone(),
mime_type: resource.mime_type.clone(),
server: server_name.clone(),
}
})
.collect()
}
/// Get a specific resource info by URI
pub fn get_resource_info(&self, uri: &str) -> Option<ResourceInfo> {
self.resources.get(uri).map(|entry| {
let (server_name, resource) = entry.value();
ResourceInfo {
uri: uri.to_string(),
name: resource.name.clone(),
description: resource.description.clone(),
mime_type: resource.mime_type.clone(),
server: server_name.clone(),
}
})
}
/// Subscribe to resource changes
pub async fn subscribe_resource(&self, uri: &str) -> McpResult<()> {
let (server_name, _resource) = self.resource_entry(uri)?;
let client = self.client_for(&server_name)?;
tracing::debug!("Subscribing to '{}' on '{}'", uri, server_name);
client
.peer()
.subscribe(rmcp::model::SubscribeRequestParam {
uri: uri.to_string(),
})
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to subscribe: {}", e)))
}
/// Unsubscribe from resource changes
pub async fn unsubscribe_resource(&self, uri: &str) -> McpResult<()> {
let (server_name, _resource) = self.resource_entry(uri)?;
let client = self.client_for(&server_name)?;
tracing::debug!("Unsubscribing from '{}' on '{}'", uri, server_name);
client
.peer()
.unsubscribe(rmcp::model::UnsubscribeRequestParam {
uri: uri.to_string(),
})
.await
.map_err(|e| McpError::ToolExecution(format!("Failed to unsubscribe: {}", e)))
}
/// Disconnect from all servers (for cleanup)
pub async fn shutdown(&mut self) {
for (name, client) in self.clients.drain() {
if let Err(e) = client.cancel().await {
tracing::warn!("Error disconnecting from '{}': {}", name, e);
}
}
self.tools.clear();
self.prompts.clear();
self.resources.clear();
}
}

52
sgl-router/src/mcp/config.rs Normal file
View File

@@ -0,0 +1,52 @@
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpConfig {
pub servers: Vec<McpServerConfig>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct McpServerConfig {
pub name: String,
#[serde(flatten)]
pub transport: McpTransport,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "protocol", rename_all = "lowercase")]
pub enum McpTransport {
Stdio {
command: String,
#[serde(default)]
args: Vec<String>,
#[serde(default)]
envs: HashMap<String, String>,
},
Sse {
url: String,
#[serde(skip_serializing_if = "Option::is_none")]
token: Option<String>,
},
Streamable {
url: String,
#[serde(skip_serializing_if = "Option::is_none")]
token: Option<String>,
},
}
impl McpConfig {
/// Load configuration from a YAML file
pub async fn from_file(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
let content = tokio::fs::read_to_string(path).await?;
let config: Self = serde_yaml::from_str(&content)?;
Ok(config)
}
/// Load configuration from environment variables (optional)
pub fn from_env() -> Option<Self> {
// This could be expanded to read from env vars
// For now, return None to indicate env config not implemented
None
}
}

42
sgl-router/src/mcp/error.rs Normal file
View File

@@ -0,0 +1,42 @@
use thiserror::Error;
pub type McpResult<T> = Result<T, McpError>;
#[derive(Debug, Error)]
pub enum McpError {
#[error("Server not found: {0}")]
ServerNotFound(String),
#[error("Tool not found: {0}")]
ToolNotFound(String),
#[error("Transport error: {0}")]
Transport(String),
#[error("Tool execution failed: {0}")]
ToolExecution(String),
#[error("Connection failed: {0}")]
ConnectionFailed(String),
#[error("Configuration error: {0}")]
Config(String),
#[error("Authentication error: {0}")]
Auth(String),
#[error("Resource not found: {0}")]
ResourceNotFound(String),
#[error("Prompt not found: {0}")]
PromptNotFound(String),
#[error(transparent)]
Sdk(#[from] Box<rmcp::RmcpError>),
#[error(transparent)]
Io(#[from] std::io::Error),
#[error(transparent)]
Http(#[from] reqwest::Error),
}

18
sgl-router/src/mcp/mod.rs Normal file
View File

@@ -0,0 +1,18 @@
// MCP Client for SGLang Router
//
// This module provides a complete MCP (Model Context Protocol) client implementation
// supporting multiple transport types (stdio, SSE, HTTP) and all MCP features:
// - Tools: Discovery and execution
// - Prompts: Reusable templates for LLM interactions
// - Resources: File/data access with subscription support
// - OAuth: Secure authentication for remote servers
pub mod client_manager;
pub mod config;
pub mod error;
pub mod oauth;
// Re-export the main types for convenience
pub use client_manager::{McpClientManager, PromptInfo, ResourceInfo, ToolInfo};
pub use config::{McpConfig, McpServerConfig, McpTransport};
pub use error::{McpError, McpResult};

191
sgl-router/src/mcp/oauth.rs Normal file
View File

@@ -0,0 +1,191 @@
// OAuth authentication support for MCP servers
use axum::{
extract::{Query, State},
response::Html,
routing::get,
Router,
};
use rmcp::transport::auth::OAuthState;
use serde::Deserialize;
use std::{net::SocketAddr, sync::Arc};
use tokio::sync::{oneshot, Mutex};
use crate::mcp::error::{McpError, McpResult};
/// OAuth callback parameters
#[derive(Debug, Deserialize)]
struct CallbackParams {
code: String,
#[allow(dead_code)]
state: Option<String>,
}
/// State for the callback server
#[derive(Clone)]
struct CallbackState {
code_receiver: Arc<Mutex<Option<oneshot::Sender<String>>>>,
}
/// HTML page returned after successful OAuth callback
const CALLBACK_HTML: &str = r#"
<!DOCTYPE html>
<html>
<head>
<title>OAuth Success</title>
<style>
body {
font-family: Arial, sans-serif;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
margin: 0;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
}
.container {
background: white;
padding: 40px;
border-radius: 10px;
box-shadow: 0 10px 30px rgba(0,0,0,0.2);
text-align: center;
}
h1 { color: #333; }
p { color: #666; margin: 20px 0; }
.success { color: #4CAF50; font-size: 48px; }
</style>
</head>
<body>
<div class="container">
<div class="success">✓</div>
<h1>Authentication Successful!</h1>
<p>You can now close this window and return to your application.</p>
</div>
</body>
</html>
"#;
/// OAuth authentication helper for MCP servers
pub struct OAuthHelper {
server_url: String,
redirect_uri: String,
callback_port: u16,
}
impl OAuthHelper {
/// Create a new OAuth helper
pub fn new(server_url: String, redirect_uri: String, callback_port: u16) -> Self {
Self {
server_url,
redirect_uri,
callback_port,
}
}
/// Perform OAuth authentication flow
pub async fn authenticate(
&self,
scopes: &[&str],
) -> McpResult<rmcp::transport::auth::AuthorizationManager> {
// Initialize OAuth state machine
let mut oauth_state = OAuthState::new(&self.server_url, None)
.await
.map_err(|e| McpError::Auth(format!("Failed to initialize OAuth: {}", e)))?;
oauth_state
.start_authorization(scopes, &self.redirect_uri)
.await
.map_err(|e| McpError::Auth(format!("Failed to start authorization: {}", e)))?;
// Get authorization URL
let auth_url = oauth_state
.get_authorization_url()
.await
.map_err(|e| McpError::Auth(format!("Failed to get authorization URL: {}", e)))?;
tracing::info!("OAuth authorization URL: {}", auth_url);
// Start callback server and wait for code
let auth_code = self.start_callback_server().await?;
// Exchange code for token
oauth_state
.handle_callback(&auth_code)
.await
.map_err(|e| McpError::Auth(format!("Failed to handle OAuth callback: {}", e)))?;
// Get authorization manager
oauth_state
.into_authorization_manager()
.ok_or_else(|| McpError::Auth("Failed to get authorization manager".to_string()))
}
/// Start a local HTTP server to receive the OAuth callback
async fn start_callback_server(&self) -> McpResult<String> {
let (code_sender, code_receiver) = oneshot::channel::<String>();
let state = CallbackState {
code_receiver: Arc::new(Mutex::new(Some(code_sender))),
};
// Create router for callback
let app = Router::new()
.route("/callback", get(Self::callback_handler))
.with_state(state);
let addr = SocketAddr::from(([127, 0, 0, 1], self.callback_port));
// Start server in background
let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
McpError::Auth(format!(
"Failed to bind to callback port {}: {}",
self.callback_port, e
))
})?;
tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
tracing::info!(
"OAuth callback server started on port {}",
self.callback_port
);
// Wait for authorization code
code_receiver
.await
.map_err(|_| McpError::Auth("Failed to receive authorization code".to_string()))
}
/// Handle OAuth callback
async fn callback_handler(
Query(params): Query<CallbackParams>,
State(state): State<CallbackState>,
) -> Html<String> {
tracing::debug!("Received OAuth callback with code");
// Send code to waiting task
if let Some(sender) = state.code_receiver.lock().await.take() {
let _ = sender.send(params.code);
}
Html(CALLBACK_HTML.to_string())
}
}
/// Create an OAuth-authenticated client
pub async fn create_oauth_client(
server_url: String,
_sse_url: String,
redirect_uri: String,
callback_port: u16,
scopes: &[&str],
) -> McpResult<rmcp::transport::auth::AuthClient<reqwest::Client>> {
let helper = OAuthHelper::new(server_url, redirect_uri, callback_port);
let auth_manager = helper.authenticate(scopes).await?;
let client = rmcp::transport::auth::AuthClient::new(reqwest::Client::default(), auth_manager);
Ok(client)
}

1011
sgl-router/src/metrics.rs Normal file
View File

File diff suppressed because it is too large Load Diff

502
sgl-router/src/middleware.rs Normal file
View File

@@ -0,0 +1,502 @@
use axum::{
extract::Request, extract::State, http::HeaderValue, http::StatusCode, middleware::Next,
response::IntoResponse, response::Response,
};
use rand::Rng;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;
use tokio::sync::{mpsc, oneshot};
use tower::{Layer, Service};
use tower_http::trace::{MakeSpan, OnRequest, OnResponse, TraceLayer};
use tracing::{debug, error, field::Empty, info, info_span, warn, Span};
pub use crate::core::token_bucket::TokenBucket;
use crate::server::AppState;
/// Generate OpenAI-compatible request ID based on endpoint
fn generate_request_id(path: &str) -> String {
let prefix = if path.contains("/chat/completions") {
"chatcmpl-"
} else if path.contains("/completions") {
"cmpl-"
} else if path.contains("/generate") {
"gnt-"
} else {
"req-"
};
// Generate a random string similar to OpenAI's format
let chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
let mut rng = rand::rng();
let random_part: String = (0..24)
.map(|_| {
let idx = rng.random_range(0..chars.len());
chars.chars().nth(idx).unwrap()
})
.collect();
format!("{}{}", prefix, random_part)
}
/// Extension type for storing request ID
#[derive(Clone, Debug)]
pub struct RequestId(pub String);
/// Tower Layer for request ID middleware
#[derive(Clone)]
pub struct RequestIdLayer {
headers: Arc<Vec<String>>,
}
impl RequestIdLayer {
pub fn new(headers: Vec<String>) -> Self {
Self {
headers: Arc::new(headers),
}
}
}
impl<S> Layer<S> for RequestIdLayer {
type Service = RequestIdMiddleware<S>;
fn layer(&self, inner: S) -> Self::Service {
RequestIdMiddleware {
inner,
headers: self.headers.clone(),
}
}
}
/// Tower Service for request ID middleware
#[derive(Clone)]
pub struct RequestIdMiddleware<S> {
inner: S,
headers: Arc<Vec<String>>,
}
impl<S> Service<Request> for RequestIdMiddleware<S>
where
S: Service<Request, Response = Response> + Send + 'static,
S::Future: Send + 'static,
{
type Response = S::Response;
type Error = S::Error;
type Future = std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: Request) -> Self::Future {
let headers = self.headers.clone();
// Extract request ID from headers or generate new one
let mut request_id = None;
for header_name in headers.iter() {
if let Some(header_value) = req.headers().get(header_name) {
if let Ok(value) = header_value.to_str() {
request_id = Some(value.to_string());
break;
}
}
}
let request_id = request_id.unwrap_or_else(|| generate_request_id(req.uri().path()));
// Insert request ID into request extensions
req.extensions_mut().insert(RequestId(request_id.clone()));
// Create a span with the request ID for this request
let span = tracing::info_span!(
"http_request",
method = %req.method(),
uri = %req.uri(),
version = ?req.version(),
request_id = %request_id
);
// Log within the span
let _enter = span.enter();
tracing::info!(
target: "sglang_router_rs::request",
"started processing request"
);
drop(_enter);
// Capture values we need in the async block
let method = req.method().clone();
let uri = req.uri().clone();
let version = req.version();
// Call the inner service
let future = self.inner.call(req);
Box::pin(async move {
let start_time = Instant::now();
let mut response = future.await?;
let latency = start_time.elapsed();
// Add request ID to response headers
response.headers_mut().insert(
"x-request-id",
HeaderValue::from_str(&request_id)
.unwrap_or_else(|_| HeaderValue::from_static("invalid-request-id")),
);
// Log the response with proper request ID in span
let status = response.status();
let span = tracing::info_span!(
"http_request",
method = %method,
uri = %uri,
version = ?version,
request_id = %request_id,
status = %status,
latency = ?latency
);
let _enter = span.enter();
if status.is_server_error() {
tracing::error!(
target: "sglang_router_rs::response",
"request failed with server error"
);
} else if status.is_client_error() {
tracing::warn!(
target: "sglang_router_rs::response",
"request failed with client error"
);
} else {
tracing::info!(
target: "sglang_router_rs::response",
"finished processing request"
);
}
Ok(response)
})
}
}
// ============= Logging Middleware =============
/// Custom span maker that includes request ID
#[derive(Clone, Debug)]
pub struct RequestSpan;
impl<B> MakeSpan<B> for RequestSpan {
fn make_span(&mut self, request: &Request<B>) -> Span {
// Don't try to extract request ID here - it won't be available yet
// The RequestIdLayer runs after TraceLayer creates the span
info_span!(
"http_request",
method = %request.method(),
uri = %request.uri(),
version = ?request.version(),
request_id = Empty, // Will be set later
status_code = Empty,
latency = Empty,
error = Empty,
)
}
}
/// Custom on_request handler
#[derive(Clone, Debug)]
pub struct RequestLogger;
impl<B> OnRequest<B> for RequestLogger {
fn on_request(&mut self, request: &Request<B>, span: &Span) {
let _enter = span.enter();
// Try to get the request ID from extensions
// This will work if RequestIdLayer has already run
if let Some(request_id) = request.extensions().get::<RequestId>() {
span.record("request_id", request_id.0.as_str());
}
// Don't log here - we already log in RequestIdService with the proper request_id
}
}
/// Custom on_response handler
#[derive(Clone, Debug)]
pub struct ResponseLogger {
_start_time: Instant,
}
impl Default for ResponseLogger {
fn default() -> Self {
Self {
_start_time: Instant::now(),
}
}
}
impl<B> OnResponse<B> for ResponseLogger {
fn on_response(self, response: &Response<B>, latency: std::time::Duration, span: &Span) {
let status = response.status();
// Record these in the span for structured logging/observability tools
span.record("status_code", status.as_u16());
span.record("latency", format!("{:?}", latency));
// Don't log here - RequestIdService handles all logging with proper request IDs
}
}
/// Create a configured TraceLayer for HTTP logging
/// Note: Actual request/response logging with request IDs is done in RequestIdService
pub fn create_logging_layer() -> TraceLayer<
tower_http::classify::SharedClassifier<tower_http::classify::ServerErrorsAsFailures>,
RequestSpan,
RequestLogger,
ResponseLogger,
> {
TraceLayer::new_for_http()
.make_span_with(RequestSpan)
.on_request(RequestLogger)
.on_response(ResponseLogger::default())
}
/// Structured logging data for requests
#[derive(Debug, serde::Serialize)]
pub struct RequestLogEntry {
pub timestamp: String,
pub request_id: String,
pub method: String,
pub uri: String,
pub status: u16,
pub latency_ms: u64,
pub user_agent: Option<String>,
pub remote_addr: Option<String>,
pub error: Option<String>,
}
/// Log a request with structured data
pub fn log_request(entry: RequestLogEntry) {
if entry.status >= 500 {
tracing::error!(
target: "sglang_router_rs::http",
request_id = %entry.request_id,
method = %entry.method,
uri = %entry.uri,
status = entry.status,
latency_ms = entry.latency_ms,
user_agent = ?entry.user_agent,
remote_addr = ?entry.remote_addr,
error = ?entry.error,
"HTTP request failed"
);
} else if entry.status >= 400 {
tracing::warn!(
target: "sglang_router_rs::http",
request_id = %entry.request_id,
method = %entry.method,
uri = %entry.uri,
status = entry.status,
latency_ms = entry.latency_ms,
user_agent = ?entry.user_agent,
remote_addr = ?entry.remote_addr,
"HTTP request client error"
);
} else {
tracing::info!(
target: "sglang_router_rs::http",
request_id = %entry.request_id,
method = %entry.method,
uri = %entry.uri,
status = entry.status,
latency_ms = entry.latency_ms,
user_agent = ?entry.user_agent,
remote_addr = ?entry.remote_addr,
"HTTP request completed"
);
}
}
// ============ Concurrency Limiting with Queue Support ============
/// Request queue entry
pub struct QueuedRequest {
/// Time when the request was queued
queued_at: Instant,
/// Channel to send the permit back when acquired
permit_tx: oneshot::Sender<Result<(), StatusCode>>,
}
/// Queue metrics for monitoring
#[derive(Debug, Default)]
pub struct QueueMetrics {
pub total_queued: std::sync::atomic::AtomicU64,
pub current_queued: std::sync::atomic::AtomicU64,
pub total_timeout: std::sync::atomic::AtomicU64,
pub total_rejected: std::sync::atomic::AtomicU64,
}
/// Queue processor that handles queued requests
pub struct QueueProcessor {
token_bucket: Arc<TokenBucket>,
queue_rx: mpsc::Receiver<QueuedRequest>,
queue_timeout: Duration,
}
impl QueueProcessor {
pub fn new(
token_bucket: Arc<TokenBucket>,
queue_rx: mpsc::Receiver<QueuedRequest>,
queue_timeout: Duration,
) -> Self {
Self {
token_bucket,
queue_rx,
queue_timeout,
}
}
pub async fn run(mut self) {
info!("Starting concurrency queue processor");
// Process requests in a single task to reduce overhead
while let Some(queued) = self.queue_rx.recv().await {
// Check timeout immediately
let elapsed = queued.queued_at.elapsed();
if elapsed >= self.queue_timeout {
warn!("Request already timed out in queue");
let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
continue;
}
let remaining_timeout = self.queue_timeout - elapsed;
// Try to acquire token for this request
if self.token_bucket.try_acquire(1.0).await.is_ok() {
// Got token immediately
debug!("Queue: acquired token immediately for queued request");
let _ = queued.permit_tx.send(Ok(()));
} else {
// Need to wait for token
let token_bucket = self.token_bucket.clone();
// Spawn task only when we actually need to wait
tokio::spawn(async move {
if token_bucket
.acquire_timeout(1.0, remaining_timeout)
.await
.is_ok()
{
debug!("Queue: acquired token after waiting");
let _ = queued.permit_tx.send(Ok(()));
} else {
warn!("Queue: request timed out waiting for token");
let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
}
});
}
}
warn!("Concurrency queue processor shutting down");
}
}
/// State for the concurrency limiter
pub struct ConcurrencyLimiter {
pub queue_tx: Option<mpsc::Sender<QueuedRequest>>,
}
impl ConcurrencyLimiter {
/// Create new concurrency limiter with optional queue
pub fn new(
token_bucket: Arc<TokenBucket>,
queue_size: usize,
queue_timeout: Duration,
) -> (Self, Option<QueueProcessor>) {
if queue_size > 0 {
let (queue_tx, queue_rx) = mpsc::channel(queue_size);
let processor = QueueProcessor::new(token_bucket, queue_rx, queue_timeout);
(
Self {
queue_tx: Some(queue_tx),
},
Some(processor),
)
} else {
(Self { queue_tx: None }, None)
}
}
}
/// Middleware function for concurrency limiting with optional queuing
pub async fn concurrency_limit_middleware(
State(app_state): State<Arc<AppState>>,
request: Request<axum::body::Body>,
next: Next,
) -> Response {
let token_bucket = app_state.context.rate_limiter.clone();
// Try to acquire token immediately
if token_bucket.try_acquire(1.0).await.is_ok() {
debug!("Acquired token immediately");
let response = next.run(request).await;
// Return the token to the bucket
token_bucket.return_tokens(1.0).await;
response
} else {
// No tokens available, try to queue if enabled
if let Some(queue_tx) = &app_state.concurrency_queue_tx {
debug!("No tokens available, attempting to queue request");
// Create a channel for the token response
let (permit_tx, permit_rx) = oneshot::channel();
let queued = QueuedRequest {
queued_at: Instant::now(),
permit_tx,
};
// Try to send to queue
match queue_tx.try_send(queued) {
Ok(_) => {
// Wait for token from queue processor
match permit_rx.await {
Ok(Ok(())) => {
debug!("Acquired token from queue");
let response = next.run(request).await;
// Return the token to the bucket
token_bucket.return_tokens(1.0).await;
response
}
Ok(Err(status)) => {
warn!("Queue returned error status: {}", status);
status.into_response()
}
Err(_) => {
error!("Queue response channel closed");
StatusCode::INTERNAL_SERVER_ERROR.into_response()
}
}
}
Err(_) => {
warn!("Request queue is full, returning 429");
StatusCode::TOO_MANY_REQUESTS.into_response()
}
}
} else {
warn!("No tokens available and queuing is disabled, returning 429");
StatusCode::TOO_MANY_REQUESTS.into_response()
}
}
}

423
sgl-router/src/policies/cache_aware.rs Normal file
View File

@@ -0,0 +1,423 @@
/*
Cache-Aware Load Balancing Router
This router combines two strategies to optimize both cache utilization and request distribution:
1. Cache-Aware Routing (Approximate Tree)
2. Load Balancing (Shortest Queue with Balance Thresholds)
The router dynamically switches between these strategies based on load conditions:
- Uses load balancing when the system is imbalanced
- Uses cache-aware routing when the system is balanced
A system is considered imbalanced if both conditions are met:
1. (max - min) > abs_threshold
2. max > rel_threshold * min
Strategy Details:
1. Cache-Aware Routing (Approximate Tree)
-------------------------------------------
This strategy maintains an approximate radix tree for each worker based on request history,
eliminating the need for direct cache state queries. The tree stores raw text characters
instead of token IDs to avoid tokenization overhead.
Process:
a. For each request, find the worker with the highest prefix match
b. If match rate > cache_threshold:
Route to the worker with highest match (likely has relevant data cached)
c. If match rate ≤ cache_threshold:
Route to the worker with smallest tree size (most available cache capacity)
d. Background maintenance:
Periodically evict least recently used leaf nodes to prevent memory overflow
2. Load Balancing (Shortest Queue)
-------------------------------------------
This strategy tracks pending request counts per worker and routes new requests
to the least busy worker when the system is detected to be imbalanced.
Configuration Parameters:
------------------------
1. cache_threshold: (float, 0.0 to 1.0)
Minimum prefix match ratio to use highest-match routing.
Below this threshold, routes to worker with most available cache space.
2. balance_abs_threshold: (integer)
Absolute difference threshold for load imbalance detection.
System is potentially imbalanced if (max_load - min_load) > abs_threshold
3. balance_rel_threshold: (float)
Relative ratio threshold for load imbalance detection.
System is potentially imbalanced if max_load > min_load * rel_threshold
Used in conjunction with abs_threshold to determine final imbalance state.
4. eviction_interval_secs: (integer)
Interval between LRU eviction cycles for the approximate trees.
5. max_tree_size: (integer)
Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted
during the next eviction cycle.
*/
use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy};
use crate::core::Worker;
use crate::metrics::RouterMetrics;
use crate::tree::Tree;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use tracing::debug;
/// Cache-aware routing policy
///
/// Routes requests based on cache affinity when load is balanced,
/// switches to shortest-queue routing when load is imbalanced.
#[derive(Debug)]
pub struct CacheAwarePolicy {
config: CacheAwareConfig,
tree: Arc<Mutex<Tree>>,
eviction_handle: Option<thread::JoinHandle<()>>,
}
impl CacheAwarePolicy {
pub fn new() -> Self {
Self::with_config(CacheAwareConfig::default())
}
pub fn with_config(config: CacheAwareConfig) -> Self {
let tree = Arc::new(Mutex::new(Tree::new()));
// Start background eviction thread if configured
let eviction_handle = if config.eviction_interval_secs > 0 {
let tree_clone = Arc::clone(&tree);
let max_tree_size = config.max_tree_size;
let interval = config.eviction_interval_secs;
Some(thread::spawn(move || loop {
thread::sleep(Duration::from_secs(interval));
if let Ok(tree_guard) = tree_clone.lock() {
tree_guard.evict_tenant_by_size(max_tree_size);
debug!("Cache eviction completed, max_size: {}", max_tree_size);
}
}))
} else {
None
};
Self {
config,
tree,
eviction_handle,
}
}
/// Initialize the tree with worker URLs (used only during initial setup)
pub fn init_workers(&self, workers: &[Box<dyn Worker>]) {
if let Ok(tree) = self.tree.lock() {
for worker in workers {
tree.insert("", worker.url());
}
}
}
/// Add a single worker to the tree (incremental update)
pub fn add_worker(&self, url: &str) {
if let Ok(tree) = self.tree.lock() {
tree.insert("", url);
}
}
/// Remove a worker from the tree
pub fn remove_worker(&self, url: &str) {
if let Ok(tree) = self.tree.lock() {
tree.remove_tenant(url);
}
}
/// Run cache eviction to prevent unbounded growth
pub fn evict_cache(&self, max_size: usize) {
if let Ok(tree) = self.tree.lock() {
tree.evict_tenant_by_size(max_size);
}
}
}
impl LoadBalancingPolicy for CacheAwarePolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
// Get current load statistics
let loads: Vec<usize> = workers.iter().map(|w| w.load()).collect();
let max_load = *loads.iter().max().unwrap_or(&0);
let min_load = *loads.iter().min().unwrap_or(&0);
// Check if load is imbalanced
let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold
&& (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold);
if is_imbalanced {
// Log load balancing trigger
let worker_loads: Vec<(String, usize)> = workers
.iter()
.map(|w| (w.url().to_string(), w.load()))
.collect();
debug!(
"Load balancing triggered | max: {} | min: {} | workers: {:?}",
max_load, min_load, worker_loads
);
RouterMetrics::record_load_balancing_event();
RouterMetrics::set_load_range(max_load, min_load);
// Use shortest queue when imbalanced
let min_load_idx = healthy_indices
.iter()
.min_by_key(|&&idx| workers[idx].load())
.copied()?;
// Even in imbalanced mode, update the tree to maintain cache state
if let Some(text) = request_text {
if let Ok(tree) = self.tree.lock() {
tree.insert(text, workers[min_load_idx].url());
}
}
// Increment processed counter
workers[min_load_idx].increment_processed();
RouterMetrics::record_processed_request(workers[min_load_idx].url());
RouterMetrics::record_policy_decision(self.name(), workers[min_load_idx].url());
return Some(min_load_idx);
}
// Use cache-aware routing when balanced
let text = request_text.unwrap_or("");
if let Ok(tree) = self.tree.lock() {
let (matched_text, matched_worker) = tree.prefix_match(text);
let match_rate = if text.is_empty() {
0.0
} else {
matched_text.chars().count() as f32 / text.chars().count() as f32
};
let selected_url = if match_rate > self.config.cache_threshold {
RouterMetrics::record_cache_hit();
matched_worker.to_string()
} else {
RouterMetrics::record_cache_miss();
tree.get_smallest_tenant()
};
// Find the index of the selected worker
if let Some(selected_idx) = workers.iter().position(|w| w.url() == selected_url) {
// Only proceed if the worker is healthy
if workers[selected_idx].is_healthy() {
// Update the tree with this request
tree.insert(text, &selected_url);
// Increment processed counter
workers[selected_idx].increment_processed();
RouterMetrics::record_processed_request(&selected_url);
return Some(selected_idx);
}
} else {
// Selected worker no longer exists, remove it from tree
tree.remove_tenant(&selected_url);
debug!("Removed stale worker {} from cache tree", selected_url);
}
// Fallback to first healthy worker
return healthy_indices.first().copied();
}
// Fallback to first healthy worker if tree operations fail
healthy_indices.first().copied()
}
fn name(&self) -> &'static str {
"cache_aware"
}
fn needs_request_text(&self) -> bool {
true // Cache-aware policy needs request text for cache affinity
}
fn on_request_complete(&self, worker_url: &str, success: bool) {
// Could track success rates per worker for more intelligent routing
if !success {
// Optionally reduce affinity for failed requests
tracing::debug!(
"Request to {} completed with success={}",
worker_url,
success
);
}
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn select_worker_pair(
&self,
prefill_workers: &[Box<dyn Worker>],
decode_workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<(usize, usize)> {
// DEPRECATED: This method is no longer used when separate policies are configured.
// The PD router now uses separate policies for prefill and decode selection.
// This implementation remains for backward compatibility when a single policy is used.
// In PD mode with single policy:
// - Prefill: Use cache-aware routing for better cache utilization
// - Decode: Use least-load routing for better load distribution
// Select prefill worker using cache-aware logic
let prefill_idx = self.select_worker(prefill_workers, request_text)?;
// Select decode worker using least-load logic
let healthy_decode = get_healthy_worker_indices(decode_workers);
if healthy_decode.is_empty() {
return None;
}
let decode_idx = healthy_decode
.iter()
.min_by_key(|&&idx| decode_workers[idx].load())
.copied()?;
Some((prefill_idx, decode_idx))
}
}
impl Default for CacheAwarePolicy {
fn default() -> Self {
Self::new()
}
}
impl Drop for CacheAwarePolicy {
fn drop(&mut self) {
// Note: We can't properly stop the eviction thread since it's in an infinite loop
// In a production system, we'd use a channel or atomic flag to signal shutdown
if let Some(handle) = self.eviction_handle.take() {
// The thread will continue running until the program exits
// This is acceptable for now since the router typically runs for the lifetime of the program
drop(handle);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_cache_aware_with_balanced_load() {
// Create policy without eviction thread for testing
let config = CacheAwareConfig {
eviction_interval_secs: 0, // Disable eviction thread
..Default::default()
};
let policy = CacheAwarePolicy::with_config(config);
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Initialize the policy with workers
policy.init_workers(&workers);
// First request should be distributed
let idx1 = policy.select_worker(&workers, Some("hello world")).unwrap();
// Same request should go to same worker (cache hit)
let idx2 = policy.select_worker(&workers, Some("hello world")).unwrap();
assert_eq!(idx1, idx2);
// Similar request should also go to same worker
let idx3 = policy.select_worker(&workers, Some("hello")).unwrap();
assert_eq!(idx1, idx3);
}
#[test]
fn test_cache_aware_with_imbalanced_load() {
let policy = CacheAwarePolicy::with_config(CacheAwareConfig {
cache_threshold: 0.5,
balance_abs_threshold: 5,
balance_rel_threshold: 2.0,
eviction_interval_secs: 0, // Disable eviction thread
max_tree_size: 10000,
});
let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular);
let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular);
// Create significant load imbalance
for _ in 0..20 {
worker1.increment_load();
}
// worker2 has load 0
let workers: Vec<Box<dyn Worker>> = vec![Box::new(worker1), Box::new(worker2)];
policy.init_workers(&workers);
// Should select worker2 (lower load) despite cache affinity
for _ in 0..5 {
let idx = policy.select_worker(&workers, Some("test")).unwrap();
assert_eq!(idx, 1); // Should always pick worker2
}
}
#[test]
fn test_cache_aware_worker_removal() {
let config = CacheAwareConfig {
eviction_interval_secs: 0, // Disable eviction thread
..Default::default()
};
let policy = CacheAwarePolicy::with_config(config);
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
policy.init_workers(&workers);
// Route some requests
policy.select_worker(&workers, Some("test1"));
policy.select_worker(&workers, Some("test2"));
// Remove a worker
policy.remove_worker("http://w1:8000");
workers[0].set_healthy(false);
// All requests should now go to worker2
let idx = policy.select_worker(&workers, Some("test1")).unwrap();
assert_eq!(idx, 1);
}
}

94
sgl-router/src/policies/factory.rs Normal file
View File

@@ -0,0 +1,94 @@
//! Factory for creating load balancing policies
use super::{
CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy,
RoundRobinPolicy,
};
use crate::config::PolicyConfig;
use std::sync::Arc;
/// Factory for creating policy instances
pub struct PolicyFactory;
impl PolicyFactory {
/// Create a policy from configuration
pub fn create_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
match config {
PolicyConfig::Random => Arc::new(RandomPolicy::new()),
PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
PolicyConfig::CacheAware {
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
eviction_interval_secs,
max_tree_size,
} => {
let config = CacheAwareConfig {
cache_threshold: *cache_threshold,
balance_abs_threshold: *balance_abs_threshold,
balance_rel_threshold: *balance_rel_threshold,
eviction_interval_secs: *eviction_interval_secs,
max_tree_size: *max_tree_size,
};
Arc::new(CacheAwarePolicy::with_config(config))
}
}
}
/// Create a policy by name (for dynamic loading)
pub fn create_by_name(name: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
match name.to_lowercase().as_str() {
"random" => Some(Arc::new(RandomPolicy::new())),
"round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())),
"power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())),
"cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_from_config() {
// Test Random
let policy = PolicyFactory::create_from_config(&PolicyConfig::Random);
assert_eq!(policy.name(), "random");
// Test RoundRobin
let policy = PolicyFactory::create_from_config(&PolicyConfig::RoundRobin);
assert_eq!(policy.name(), "round_robin");
// Test PowerOfTwo
let policy = PolicyFactory::create_from_config(&PolicyConfig::PowerOfTwo {
load_check_interval_secs: 60,
});
assert_eq!(policy.name(), "power_of_two");
// Test CacheAware
let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAware {
cache_threshold: 0.7,
balance_abs_threshold: 10,
balance_rel_threshold: 1.5,
eviction_interval_secs: 30,
max_tree_size: 1000,
});
assert_eq!(policy.name(), "cache_aware");
}
#[test]
fn test_create_by_name() {
assert!(PolicyFactory::create_by_name("random").is_some());
assert!(PolicyFactory::create_by_name("RANDOM").is_some());
assert!(PolicyFactory::create_by_name("round_robin").is_some());
assert!(PolicyFactory::create_by_name("RoundRobin").is_some());
assert!(PolicyFactory::create_by_name("power_of_two").is_some());
assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some());
assert!(PolicyFactory::create_by_name("cache_aware").is_some());
assert!(PolicyFactory::create_by_name("CacheAware").is_some());
assert!(PolicyFactory::create_by_name("unknown").is_none());
}
}

148
sgl-router/src/policies/mod.rs Normal file
View File

@@ -0,0 +1,148 @@
//! Load balancing policies for SGLang router
//!
//! This module provides a unified abstraction for routing policies that work
//! across both regular and prefill-decode (PD) routing modes.
use crate::core::Worker;
use std::fmt::Debug;
mod cache_aware;
mod factory;
mod power_of_two;
mod random;
mod round_robin;
pub use cache_aware::CacheAwarePolicy;
pub use factory::PolicyFactory;
pub use power_of_two::PowerOfTwoPolicy;
pub use random::RandomPolicy;
pub use round_robin::RoundRobinPolicy;
/// Core trait for load balancing policies
///
/// This trait provides a unified interface for implementing routing algorithms
/// that can work with both regular single-worker selection and PD dual-worker selection.
pub trait LoadBalancingPolicy: Send + Sync + Debug {
/// Select a single worker from the available workers
///
/// This is used for regular routing mode where requests go to a single worker.
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize>;
/// Select a pair of workers (prefill and decode) for PD routing
///
/// Returns indices of (prefill_worker, decode_worker) from their respective arrays.
/// Default implementation uses select_worker for each array independently.
fn select_worker_pair(
&self,
prefill_workers: &[Box<dyn Worker>],
decode_workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<(usize, usize)> {
// Default implementation: independently select from each pool
let prefill_idx = self.select_worker(prefill_workers, request_text)?;
let decode_idx = self.select_worker(decode_workers, request_text)?;
Some((prefill_idx, decode_idx))
}
/// Update policy state after request completion
///
/// This is called when a request completes (successfully or not) to allow
/// policies to update their internal state.
fn on_request_complete(&self, _worker_url: &str, _success: bool) {
// Default: no-op for stateless policies
}
/// Get policy name for metrics and debugging
fn name(&self) -> &'static str;
/// Check if this policy needs request text for routing decisions
fn needs_request_text(&self) -> bool {
false // Default: most policies don't need request text
}
/// Update worker load information
///
/// This is called periodically with current load information for load-aware policies.
fn update_loads(&self, _loads: &std::collections::HashMap<String, isize>) {
// Default: no-op for policies that don't use load information
}
/// Reset any internal state
///
/// This is useful for policies that maintain state (e.g., round-robin counters).
fn reset(&self) {
// Default: no-op for stateless policies
}
/// Get as Any for downcasting
fn as_any(&self) -> &dyn std::any::Any;
}
/// Configuration for cache-aware policy
#[derive(Debug, Clone)]
pub struct CacheAwareConfig {
pub cache_threshold: f32,
pub balance_abs_threshold: usize,
pub balance_rel_threshold: f32,
pub eviction_interval_secs: u64,
pub max_tree_size: usize,
}
impl Default for CacheAwareConfig {
fn default() -> Self {
Self {
cache_threshold: 0.5,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
eviction_interval_secs: 30,
max_tree_size: 10000,
}
}
}
/// Helper function to filter healthy workers and return their indices
pub(crate) fn get_healthy_worker_indices(workers: &[Box<dyn Worker>]) -> Vec<usize> {
workers
.iter()
.enumerate()
.filter(|(_, w)| w.is_healthy() && w.circuit_breaker().can_execute())
.map(|(idx, _)| idx)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_get_healthy_worker_indices() {
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// All healthy initially
let indices = get_healthy_worker_indices(&workers);
assert_eq!(indices, vec![0, 1, 2]);
// Mark one unhealthy
workers[1].set_healthy(false);
let indices = get_healthy_worker_indices(&workers);
assert_eq!(indices, vec![0, 2]);
}
}

201
sgl-router/src/policies/power_of_two.rs Normal file
View File

@@ -0,0 +1,201 @@
//! Power-of-two choices load balancing policy
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
use crate::metrics::RouterMetrics;
use rand::Rng;
use std::collections::HashMap;
use std::sync::RwLock;
use tracing::info;
/// Power-of-two choices policy
///
/// Randomly selects two workers and routes to the one with lower load.
/// This provides good load distribution with minimal coordination overhead.
#[derive(Debug)]
pub struct PowerOfTwoPolicy {
/// Cached load information from external monitoring
cached_loads: RwLock<HashMap<String, isize>>,
}
impl PowerOfTwoPolicy {
pub fn new() -> Self {
Self {
cached_loads: RwLock::new(HashMap::new()),
}
}
fn get_worker_load(&self, worker: &dyn Worker) -> isize {
// First check cached loads (from external monitoring)
if let Ok(loads) = self.cached_loads.read() {
if let Some(&load) = loads.get(worker.url()) {
return load;
}
}
// Fall back to local load counter
worker.load() as isize
}
}
impl LoadBalancingPolicy for PowerOfTwoPolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
_request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
if healthy_indices.len() == 1 {
return Some(healthy_indices[0]);
}
// Select two random workers
let mut rng = rand::rng();
let idx1 = rng.random_range(0..healthy_indices.len());
let mut idx2 = rng.random_range(0..healthy_indices.len());
// Ensure we pick two different workers
while idx2 == idx1 {
idx2 = rng.random_range(0..healthy_indices.len());
}
let worker_idx1 = healthy_indices[idx1];
let worker_idx2 = healthy_indices[idx2];
// Compare loads and select the less loaded one
let load1 = self.get_worker_load(workers[worker_idx1].as_ref());
let load2 = self.get_worker_load(workers[worker_idx2].as_ref());
// Log selection for debugging
let selected_idx = if load1 <= load2 {
worker_idx1
} else {
worker_idx2
};
info!(
"Power-of-two selection: {}={} vs {}={} -> selected {}",
workers[worker_idx1].url(),
load1,
workers[worker_idx2].url(),
load2,
workers[selected_idx].url()
);
// Increment processed counter
workers[selected_idx].increment_processed();
RouterMetrics::record_processed_request(workers[selected_idx].url());
RouterMetrics::record_policy_decision(self.name(), workers[selected_idx].url());
Some(selected_idx)
}
fn name(&self) -> &'static str {
"power_of_two"
}
fn update_loads(&self, loads: &HashMap<String, isize>) {
if let Ok(mut cached) = self.cached_loads.write() {
*cached = loads.clone();
}
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
impl Default for PowerOfTwoPolicy {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_power_of_two_selection() {
let policy = PowerOfTwoPolicy::new();
let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular);
let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular);
let worker3 = BasicWorker::new("http://w3:8000".to_string(), WorkerType::Regular);
// Set different loads
for _ in 0..10 {
worker1.increment_load();
}
for _ in 0..5 {
worker2.increment_load();
}
// worker3 has load 0
let workers: Vec<Box<dyn Worker>> =
vec![Box::new(worker1), Box::new(worker2), Box::new(worker3)];
// Run multiple selections
let mut selected_counts = [0; 3];
for _ in 0..100 {
if let Some(idx) = policy.select_worker(&workers, None) {
selected_counts[idx] += 1;
}
}
// Worker with lowest load (worker3) should be selected most often
assert!(selected_counts[2] > selected_counts[1]);
assert!(selected_counts[1] > selected_counts[0]);
}
#[test]
fn test_power_of_two_with_cached_loads() {
let policy = PowerOfTwoPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Update cached loads
let mut loads = HashMap::new();
loads.insert("http://w1:8000".to_string(), 100);
loads.insert("http://w2:8000".to_string(), 10);
policy.update_loads(&loads);
// Should prefer worker2 with lower cached load
let mut w2_selected = 0;
for _ in 0..50 {
if let Some(idx) = policy.select_worker(&workers, None) {
if idx == 1 {
w2_selected += 1;
}
}
}
// Worker2 should be selected significantly more often
assert!(w2_selected > 35); // Should win most of the time
}
#[test]
fn test_power_of_two_single_worker() {
let policy = PowerOfTwoPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
))];
// With single worker, should always select it
assert_eq!(policy.select_worker(&workers, None), Some(0));
}
}

121
sgl-router/src/policies/random.rs Normal file
View File

@@ -0,0 +1,121 @@
//! Random load balancing policy
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
use crate::metrics::RouterMetrics;
use rand::Rng;
/// Random selection policy
///
/// Selects workers randomly with uniform distribution among healthy workers.
#[derive(Debug, Default)]
pub struct RandomPolicy;
impl RandomPolicy {
pub fn new() -> Self {
Self
}
}
impl LoadBalancingPolicy for RandomPolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
_request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
let mut rng = rand::rng();
let random_idx = rng.random_range(0..healthy_indices.len());
let worker = workers[healthy_indices[random_idx]].url();
RouterMetrics::record_processed_request(worker);
RouterMetrics::record_policy_decision(self.name(), worker);
Some(healthy_indices[random_idx])
}
fn name(&self) -> &'static str {
"random"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use std::collections::HashMap;
#[test]
fn test_random_selection() {
let policy = RandomPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// Test multiple selections to ensure randomness
let mut counts = HashMap::new();
for _ in 0..100 {
if let Some(idx) = policy.select_worker(&workers, None) {
*counts.entry(idx).or_insert(0) += 1;
}
}
// All workers should be selected at least once
assert_eq!(counts.len(), 3);
assert!(counts.values().all(|&count| count > 0));
}
#[test]
fn test_random_with_unhealthy_workers() {
let policy = RandomPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Mark first worker as unhealthy
workers[0].set_healthy(false);
// Should always select the healthy worker (index 1)
for _ in 0..10 {
assert_eq!(policy.select_worker(&workers, None), Some(1));
}
}
#[test]
fn test_random_no_healthy_workers() {
let policy = RandomPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
))];
workers[0].set_healthy(false);
assert_eq!(policy.select_worker(&workers, None), None);
}
}

140
sgl-router/src/policies/round_robin.rs Normal file
View File

@@ -0,0 +1,140 @@
//! Round-robin load balancing policy
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
use crate::metrics::RouterMetrics;
use std::sync::atomic::{AtomicUsize, Ordering};
/// Round-robin selection policy
///
/// Selects workers in sequential order, cycling through all healthy workers.
#[derive(Debug, Default)]
pub struct RoundRobinPolicy {
counter: AtomicUsize,
}
impl RoundRobinPolicy {
pub fn new() -> Self {
Self {
counter: AtomicUsize::new(0),
}
}
}
impl LoadBalancingPolicy for RoundRobinPolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
_request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
// Get and increment counter atomically
let count = self.counter.fetch_add(1, Ordering::Relaxed);
let selected_idx = count % healthy_indices.len();
let worker = workers[healthy_indices[selected_idx]].url();
RouterMetrics::record_processed_request(worker);
RouterMetrics::record_policy_decision(self.name(), worker);
Some(healthy_indices[selected_idx])
}
fn name(&self) -> &'static str {
"round_robin"
}
fn reset(&self) {
self.counter.store(0, Ordering::Relaxed);
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_round_robin_selection() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// Should select workers in order: 0, 1, 2, 0, 1, 2, ...
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
assert_eq!(policy.select_worker(&workers, None), Some(2));
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
}
#[test]
fn test_round_robin_with_unhealthy_workers() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// Mark middle worker as unhealthy
workers[1].set_healthy(false);
// Should skip unhealthy worker: 0, 2, 0, 2, ...
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(2));
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(2));
}
#[test]
fn test_round_robin_reset() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Advance the counter
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
// Reset should start from beginning
policy.reset();
assert_eq!(policy.select_worker(&workers, None), Some(0));
}
}

541
sgl-router/src/proto/sglang_scheduler.proto Normal file
View File

@@ -0,0 +1,541 @@
syntax = "proto3";
package sglang.grpc.scheduler;
import "google/protobuf/timestamp.proto";
import "google/protobuf/struct.proto";
// Service definition for SGLang scheduler communication
// This protocol bridges the Rust router and Python scheduler
service SglangScheduler {
// Initialize connection and get model info
rpc Initialize(InitializeRequest) returns (InitializeResponse);
// Submit a generation request (supports streaming)
rpc Generate(GenerateRequest) returns (stream GenerateResponse);
// Submit an embedding request
rpc Embed(EmbedRequest) returns (EmbedResponse);
// Health check and metrics
rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse);
// Abort a running request
rpc Abort(AbortRequest) returns (AbortResponse);
// Flush KV cache
rpc FlushCache(FlushCacheRequest) returns (FlushCacheResponse);
}
// =====================
// Common Types
// =====================
// Sampling parameters matching SGLang's SamplingParams
message SamplingParams {
float temperature = 1;
float top_p = 2;
int32 top_k = 3;
float min_p = 4;
float frequency_penalty = 5;
float presence_penalty = 6;
float repetition_penalty = 7;
int32 max_new_tokens = 8;
repeated string stop = 9;
repeated int32 stop_token_ids = 10;
bool skip_special_tokens = 11;
bool spaces_between_special_tokens = 12;
// Structured generation
oneof constraint {
string regex = 13;
string json_schema = 14;
string ebnf_grammar = 15;
}
// LoRA adapter
string lora_path = 16;
// Speculative decoding
int32 n = 17; // Number of samples
// Token healing
bool token_healing = 18;
// Additional parameters
int32 min_new_tokens = 19;
bool ignore_eos = 20;
bool no_stop_trim = 21;
int32 stream_interval = 22;
map<string, float> logit_bias = 23;
string structural_tag = 24;
// Custom parameters for extensibility
google.protobuf.Struct custom_params = 25;
}
// Session parameters for continual prompting
message SessionParams {
string session_id = 1;
string request_id = 2;
int32 offset = 3;
bool replace = 4;
bool drop_previous_output = 5;
}
// Disaggregated serving parameters
message DisaggregatedParams {
string bootstrap_host = 1;
int32 bootstrap_port = 2;
int32 bootstrap_room = 3;
}
// =====================
// Initialize
// =====================
message InitializeRequest {
string client_id = 1;
string client_version = 2;
// Operating mode
enum Mode {
REGULAR = 0; // Normal mode with local scheduler
PREFILL = 1; // Prefill-only mode for disaggregated serving
DECODE = 2; // Decode-only mode for disaggregated serving
}
Mode mode = 3;
}
message InitializeResponse {
bool success = 1;
string scheduler_version = 2;
// Model information
ModelInfo model_info = 3;
// Server capabilities
ServerCapabilities capabilities = 4;
// Error message if success is false
string error_message = 5;
}
message ModelInfo {
string model_name = 1;
int32 max_context_length = 2;
int32 vocab_size = 3;
bool supports_tool_calling = 4;
bool supports_vision = 5;
repeated string special_tokens = 6;
// Additional model metadata
string model_type = 7;
int32 num_layers = 8;
int32 hidden_size = 9;
int32 num_attention_heads = 10;
int32 num_key_value_heads = 11;
// Tokenizer info
string tokenizer_type = 12;
repeated int32 eos_token_ids = 13;
int32 pad_token_id = 14;
int32 bos_token_id = 15;
}
message ServerCapabilities {
bool continuous_batching = 1;
bool disaggregated_serving = 2;
bool speculative_decoding = 3;
int32 max_batch_size = 4;
int32 max_num_batched_tokens = 5;
int32 max_prefill_tokens = 6;
string attention_backend = 7; // "flashinfer", "triton", "torch"
// Additional capabilities
bool supports_lora = 8;
bool supports_grammar = 9;
bool supports_multimodal = 10;
repeated string supported_modalities = 11; // ["image", "video", "audio"]
bool supports_custom_logit_processor = 12;
bool supports_session = 13;
// Hardware info
int32 num_gpus = 14;
string gpu_type = 15;
int64 total_gpu_memory = 16;
// Parallelism info
int32 tensor_parallel_size = 17;
int32 pipeline_parallel_size = 18;
int32 data_parallel_size = 19;
}
// =====================
// Generate Request
// =====================
message GenerateRequest {
string request_id = 1;
// Input can be either text or tokenized
oneof input {
string text = 2;
TokenizedInput tokenized = 3;
}
// Multimodal inputs
MultimodalInputs mm_inputs = 4;
// Generation parameters
SamplingParams sampling_params = 5;
// Return options
bool return_logprob = 6;
int32 logprob_start_len = 7;
int32 top_logprobs_num = 8;
repeated int32 token_ids_logprob = 9;
bool return_hidden_states = 10;
// Session management
SessionParams session_params = 11;
// For disaggregated serving
DisaggregatedParams disaggregated_params = 12;
// Custom logit processor (serialized)
string custom_logit_processor = 13;
// Request metadata
google.protobuf.Timestamp timestamp = 14;
bool log_metrics = 15;
// Input embeddings (alternative to text/tokens)
repeated float input_embeds = 16;
// LoRA adapter ID (if pre-loaded)
string lora_id = 17;
// Data parallel routing
int32 data_parallel_rank = 18;
// For load balancing
int32 dp_balance_id = 19;
}
message TokenizedInput {
string original_text = 1; // For reference
repeated int32 input_ids = 2;
}
message MultimodalInputs {
// Simplified multimodal handling - actual data processed by tokenizer
repeated string image_urls = 1;
repeated string video_urls = 2;
repeated string audio_urls = 3;
// Pre-processed multimodal features (if available)
google.protobuf.Struct processed_features = 4;
// Raw data for direct processing
repeated bytes image_data = 5;
repeated bytes video_data = 6;
repeated bytes audio_data = 7;
// Modality metadata
repeated string modalities = 8;
}
// =====================
// Generate Response
// =====================
message GenerateResponse {
string request_id = 1;
// Response type
oneof response {
GenerateStreamChunk chunk = 2;
GenerateComplete complete = 3;
GenerateError error = 4;
}
}
message GenerateStreamChunk {
// Generated token
int32 token_id = 1;
string text = 2;
// Cumulative counts
int32 prompt_tokens = 3;
int32 completion_tokens = 4;
int32 cached_tokens = 5;
// Logprobs (if requested)
LogProbs logprobs = 6;
// Hidden states (if requested)
repeated float hidden_states = 7;
// Metadata
float generation_time = 8; // Time to generate this token
int32 queue_time = 9; // Time spent in queue
}
message GenerateComplete {
// Final output
repeated int32 output_ids = 1;
string output_text = 2;
// Finish reason
enum FinishReason {
// The model generated a stop sequence.
STOP = 0;
// The model reached the maximum generation length.
LENGTH = 1;
// The model generated an end-of-sequence (EOS) token.
EOS_TOKEN = 2;
// The model generated a user-provided stop string.
STOP_STR = 3;
// The request was aborted by the user or system.
ABORT = 4;
}
FinishReason finish_reason = 3;
// Final counts
int32 prompt_tokens = 4;
int32 completion_tokens = 5;
int32 cached_tokens = 6;
// Performance metrics
float total_generation_time = 7;
float time_to_first_token = 8;
float tokens_per_second = 9;
// Spec decode metrics
int32 spec_verify_count = 10;
// All logprobs if requested
repeated LogProbs all_logprobs = 11;
// All hidden states if requested
repeated HiddenStates all_hidden_states = 12;
}
message GenerateError {
string message = 1;
string http_status_code = 2;
string details = 3;
}
message LogProbs {
repeated float token_logprobs = 1;
repeated int32 token_ids = 2;
// Top logprobs at each position
repeated TopLogProbs top_logprobs = 3;
// Decoded text for tokens
repeated string token_texts = 4;
}
message TopLogProbs {
repeated float values = 1;
repeated int32 token_ids = 2;
repeated string token_texts = 3;
}
message HiddenStates {
repeated float values = 1;
int32 layer = 2;
int32 position = 3;
}
// =====================
// Embedding Request
// =====================
message EmbedRequest {
string request_id = 1;
oneof input {
string text = 2;
TokenizedInput tokenized = 3;
}
// Multimodal inputs
MultimodalInputs mm_inputs = 4;
// Dummy sampling params for compatibility
// EmbedRequest doesn't use sampling_params
SamplingParams sampling_params = 5;
bool log_metrics = 6;
// Token type IDs for models that require them
repeated int32 token_type_ids = 7;
// Data parallel routing
int32 data_parallel_rank = 8;
// For cross-encoder requests
bool is_cross_encoder = 9;
repeated string texts = 10; // For cross-encoder batch
}
message EmbedResponse {
string request_id = 1;
oneof response {
EmbedComplete complete = 2;
EmbedError error = 3;
}
}
message EmbedComplete {
repeated float embedding = 1;
int32 prompt_tokens = 2;
int32 cached_tokens = 3;
// Additional metadata
int32 embedding_dim = 4;
float generation_time = 5;
// For batch embeddings
repeated Embedding batch_embeddings = 6;
}
message Embedding {
repeated float values = 1;
int32 index = 2;
}
message EmbedError {
string message = 1;
string code = 2;
string details = 3;
}
// =====================
// Management Operations
// =====================
message HealthCheckRequest {
bool include_detailed_metrics = 1;
}
message HealthCheckResponse {
bool healthy = 1;
// Current load metrics
int32 num_requests_running = 2;
int32 num_requests_waiting = 3;
float gpu_cache_usage = 4;
float gpu_memory_usage = 5;
// KV cache metrics
int32 kv_cache_total_blocks = 6;
int32 kv_cache_used_blocks = 7;
float kv_cache_hit_rate = 8;
// Additional metrics
int32 num_grammar_queue_requests = 9;
float generation_throughput = 10; // tokens/sec
float average_queue_time = 11; // seconds
float average_generation_time = 12; // seconds
// System metrics
float cpu_usage = 13;
int64 memory_usage = 14;
// Disaggregation metrics
int32 num_prefill_requests = 15;
int32 num_decode_requests = 16;
// Detailed metrics (optional)
google.protobuf.Struct detailed_metrics = 17;
}
message AbortRequest {
string request_id = 1;
string reason = 2;
}
message AbortResponse {
bool success = 1;
string message = 2;
}
message FlushCacheRequest {
bool flush_all = 1;
repeated string session_ids = 2; // Flush specific sessions
}
message FlushCacheResponse {
bool success = 1;
int32 num_entries_flushed = 2;
int64 memory_freed = 3; // bytes
string message = 4;
}
// =====================
// Additional Operations (Future)
// =====================
// Load LoRA adapter
message LoadLoRARequest {
string adapter_id = 1;
string adapter_path = 2;
int32 rank = 3;
}
message LoadLoRAResponse {
bool success = 1;
string adapter_id = 2;
string message = 3;
}
// Unload LoRA adapter
message UnloadLoRARequest {
string adapter_id = 1;
}
message UnloadLoRAResponse {
bool success = 1;
string message = 2;
}
// Update weights
message UpdateWeightsRequest {
oneof source {
string disk_path = 1;
bytes tensor_data = 2;
string remote_url = 3;
}
string weight_name = 4;
}
message UpdateWeightsResponse {
bool success = 1;
string message = 2;
}
// Get internal state for debugging
message GetInternalStateRequest {
repeated string state_keys = 1;
}
message GetInternalStateResponse {
google.protobuf.Struct state = 1;
}
// Set internal state for testing
message SetInternalStateRequest {
google.protobuf.Struct state = 1;
}
message SetInternalStateResponse {
bool success = 1;
string message = 2;
}

5
sgl-router/src/protocols/mod.rs Normal file
View File

@@ -0,0 +1,5 @@
// Protocol definitions and validation for various LLM APIs
// This module provides a structured approach to handling different API protocols
pub mod spec;
pub mod validation;

2681
sgl-router/src/protocols/spec.rs Normal file
View File

File diff suppressed because it is too large Load Diff

1164
sgl-router/src/protocols/validation.rs Normal file
View File

File diff suppressed because it is too large Load Diff

474
sgl-router/src/reasoning_parser/README.md Normal file
View File

@@ -0,0 +1,474 @@
# Reasoning Parser Architecture
## 1. Executive Summary
### High-Level Overview
The reasoning parser layer provides a unified interface for detecting and extracting reasoning content from Large Language Model (LLM) outputs, particularly from models that support Chain-of-Thought (CoT) reasoning with explicit thinking blocks. The architecture follows a trait-based design pattern enabling pluggable parser implementations while maintaining consistent APIs across different model families that use various reasoning token formats.
**Key Components:**
- **Factory Pattern**: Registry-based creation and pooling of model-specific parsers
- **Trait System**: `ReasoningParser` trait for implementation flexibility
- **Parser Pooling**: Efficient reuse of parser instances across concurrent requests
- **Streaming Support**: Incremental parsing with partial token buffering
- **Model Detection**: Pattern-based matching for automatic parser selection
- **State Management**: Stateful parsing for streaming scenarios with buffer management
- **Thread Safety**: Arc<Mutex> based sharing for high-concurrency environments
- **Extensibility**: Easy addition of new model-specific parsers
**Data Flow:**
1. Request → Factory (model detection) → Pooled Parser Retrieval
2. One-Shot: Text → Parser → ParserResult (normal + reasoning text)
3. Streaming: Chunks → Parser (stateful) → Incremental ParserResult
4. Buffer Management: Partial Tokens → Buffer → Complete Token Detection
5. Reset: Parser State → Clear Buffers → Ready for Reuse
### Architecture Highlights
- **Model-Specific Parsers**: DeepSeek-R1, Qwen3, Kimi, GLM45, Step3 variants
- **Parser Pooling**: Singleton instances per model type for memory efficiency
- **High Concurrency**: Mutex-protected parsers handle 1000+ req/sec
- **Buffer Overflow Protection**: Configurable max buffer size (default 64KB)
- **Partial Token Detection**: Intelligent buffering for incomplete delimiters
- **Passthrough Mode**: Graceful fallback for unknown models
- **Zero-Copy Where Possible**: Efficient string handling in hot paths
## 2. Mermaid Diagrams
### Component Flow Diagram
```mermaid
graph TB
subgraph Input
R[Request] --> MID[Model ID]
end
subgraph Factory Layer
MID --> PF[ParserFactory]
PF --> REG[ParserRegistry]
REG --> PM[Pattern Matching]
PM --> PP[Parser Pool]
end
subgraph Parser Pool
PP --> DS[DeepSeek-R1]
PP --> QW[Qwen3]
PP --> QWT[Qwen3-Thinking]
PP --> KM[Kimi]
PP --> GL[GLM45]
PP --> S3[Step3]
PP --> PT[Passthrough]
end
subgraph Parser Instance
DS --> BP[BaseReasoningParser]
QW --> BP
KM --> BP
GL --> BP
S3 --> BP
end
subgraph Processing
BP --> DAP[detect_and_parse]
BP --> PSI[parse_streaming]
BP --> RST[reset]
end
subgraph State Management
BP --> BUF[Buffer]
BP --> IR[in_reasoning flag]
BP --> STS[stripped_think_start]
end
subgraph Output
DAP --> PR[ParserResult]
PSI --> PR
PR --> NT[normal_text]
PR --> RT[reasoning_text]
end
```
### Sequence Flow Diagram
```mermaid
sequenceDiagram
participant C as Client
participant F as ParserFactory
participant R as Registry
participant P as Parser Pool
participant BP as BaseParser
participant PR as ParserResult
C->>F: get_pooled("deepseek-r1-model")
F->>R: find_pooled_parser_for_model()
R->>R: pattern_match("deepseek-r1")
R->>P: get_pooled_parser("deepseek_r1")
alt Parser exists in pool
P-->>F: Arc<Mutex<Parser>>
else Create new parser
P->>BP: new DeepSeekR1Parser()
P->>P: insert into pool
P-->>F: Arc<Mutex<Parser>>
end
F-->>C: PooledParser
C->>BP: lock().parse_reasoning_streaming_incremental()
loop streaming chunks
C->>BP: parse_reasoning_streaming_incremental(chunk)
BP->>BP: buffer.push_str(chunk)
BP->>BP: check partial tokens
alt Complete token found
BP->>PR: create result
BP->>BP: clear buffer
BP-->>C: ParserResult
else Partial token
BP->>BP: keep buffering
BP-->>C: ParserResult::default()
end
end
C->>BP: reset()
BP->>BP: clear buffers & flags
C->>BP: unlock()
```
### Class/Type Diagram
```mermaid
classDiagram
class ReasoningParser {
<<trait>>
+detect_and_parse_reasoning(&mut self, text: &str) Result~ParserResult~
+parse_reasoning_streaming_incremental(&mut self, text: &str) Result~ParserResult~
+reset(&mut self)
+model_type(&self) &str
}
class ParserResult {
+normal_text: String
+reasoning_text: String
+new(normal: String, reasoning: String) Self
+normal(text: String) Self
+reasoning(text: String) Self
+is_empty() bool
}
class ParserConfig {
+think_start_token: String
+think_end_token: String
+stream_reasoning: bool
+max_buffer_size: usize
+initial_in_reasoning: bool
+default() Self
}
class BaseReasoningParser {
-config: ParserConfig
-in_reasoning: bool
-buffer: String
-stripped_think_start: bool
-model_type: String
+new(config: ParserConfig) Self
+with_model_type(model: String) Self
-is_partial_token(&self, text: &str) bool
}
class DeepSeekR1Parser {
-base: BaseReasoningParser
+new() Self
}
class Qwen3Parser {
-base: BaseReasoningParser
+new() Self
}
class QwenThinkingParser {
-base: BaseReasoningParser
+new() Self
}
class KimiParser {
-base: BaseReasoningParser
+new() Self
}
class Glm45Parser {
-base: BaseReasoningParser
+new() Self
}
class Step3Parser {
-base: BaseReasoningParser
+new() Self
}
class ParserFactory {
-registry: ParserRegistry
+new() Self
+get_pooled(model_id: &str) PooledParser
+create(model_id: &str) Result~Box~dyn ReasoningParser~~
+clear_pool()
}
class ParserRegistry {
-creators: Arc~RwLock~HashMap~~
-pool: Arc~RwLock~HashMap~~
-patterns: Arc~RwLock~Vec~~
+register_parser(name: &str, creator: F)
+register_pattern(pattern: &str, parser_name: &str)
+get_pooled_parser(name: &str) Option~PooledParser~
+find_pooled_parser_for_model(model: &str) Option~PooledParser~
}
ReasoningParser <|.. BaseReasoningParser
ReasoningParser <|.. DeepSeekR1Parser
ReasoningParser <|.. Qwen3Parser
ReasoningParser <|.. QwenThinkingParser
ReasoningParser <|.. KimiParser
ReasoningParser <|.. Glm45Parser
ReasoningParser <|.. Step3Parser
DeepSeekR1Parser o-- BaseReasoningParser
Qwen3Parser o-- BaseReasoningParser
QwenThinkingParser o-- BaseReasoningParser
KimiParser o-- BaseReasoningParser
Glm45Parser o-- BaseReasoningParser
Step3Parser o-- BaseReasoningParser
BaseReasoningParser o-- ParserConfig
ParserFactory o-- ParserRegistry
ParserRegistry o-- ReasoningParser
```
## 3. Module-by-Module Deep Dive
### 3.1 mod.rs (Main Module)
**Key Responsibilities:**
- Module organization and public API surface
- Re-exports for convenient access to core types
- Separation of concerns across submodules
**Module Structure:**
- `factory`: Parser creation and pooling logic
- `parsers`: Concrete parser implementations
- `traits`: Core trait definitions and types
### 3.2 traits.rs (Trait Definitions)
**ParserResult Methods**:
- `new()`: Create with both normal and reasoning text
- `normal()`: Create with only normal text (convenience)
- `reasoning()`: Create with only reasoning text (convenience)
- `is_empty()`: Check if result contains any text
**ReasoningParser Trait**:
- **`detect_and_parse_reasoning`**: One-shot parsing for complete text
- **`parse_reasoning_streaming_incremental`**: Stateful streaming parser
- **`reset`**: Clear state for parser reuse
- **`model_type`**: Identify parser variant for debugging
**ParserConfig Defaults**:
- Default tokens: `<think>` and `</think>`
- Stream reasoning: true (immediate output)
- Max buffer: 65536 bytes (64KB)
- Initial state: false (explicit reasoning blocks)
### 3.3 factory.rs (Parser Creation & Pooling)
**ParserRegistry Methods**:
1. **`register_parser`**:
- Register creator function for parser type
- Lazy instantiation when requested
- Thread-safe registration
2. **`register_pattern`**:
- Map model ID patterns to parser names
- First-match-wins ordering
- Case-insensitive matching
3. **`get_pooled_parser`**:
- Check pool for existing instance
- Create and pool if not present
- Return Arc<Mutex> for sharing
4. **`find_pooled_parser_for_model`**:
- Pattern match against model ID
- Delegate to get_pooled_parser
- Case-insensitive comparison
**ParserFactory Methods**:
1. **`new()`**:
- Register all built-in parsers
- Setup model pattern mappings
- Initialize empty pool
2. **`get_pooled`**:
- Primary API for getting parsers
- Automatic passthrough fallback
- Guaranteed non-null return
3. **`create`**:
- Create fresh parser instance
- No pooling (for testing/isolation)
- Returns Result for error handling
**Registered Parsers**:
- `base`: Generic configurable parser
- `deepseek_r1`: DeepSeek-R1 (initial_in_reasoning=true)
- `qwen3`: Qwen3 base model (initial_in_reasoning=false)
- `qwen3_thinking`: Qwen3 thinking variant (initial_in_reasoning=true)
- `kimi`: Kimi with Unicode tokens
- `glm45`: GLM-4.5 parser
- `step3`: Step3 parser
- `passthrough`: No-op fallback parser
**Model Pattern Mappings**:
```
"deepseek-r1" → "deepseek_r1"
"qwen3-thinking" → "qwen3_thinking"
"qwen-thinking" → "qwen3_thinking"
"qwen3" → "qwen3"
"qwen" → "qwen3"
"glm45" → "glm45"
"kimi" → "kimi"
"step3" → "step3"
```
### 3.4 parsers/base.rs (Base Implementation)
**Key Methods:**
**`detect_and_parse_reasoning`**:
```
Algorithm:
1. Check buffer overflow protection
2. Detect reasoning presence (in_reasoning OR contains start_token)
3. If no reasoning → return as normal text
4. Remove start token and trim
5. If no end token → assume truncated reasoning
6. Split on end token
7. Extract reasoning and normal portions
```
**`parse_reasoning_streaming_incremental`**:
```
Algorithm:
1. Check buffer capacity
2. Append text to buffer
3. Check if buffer is partial token prefix
4. If partial → buffer and return empty
5. Strip start token if present
6. Find end token position
7. Handle based on state:
- In reasoning + end found → split and return both
- In reasoning + streaming → return accumulated reasoning
- Not in reasoning → return as normal text
- In reasoning + no end → continue buffering
```
**Critical Features:**
1. **Partial Token Detection**:
- Prevents premature token matching
- Buffers incomplete delimiters
- Essential for streaming correctness
2. **Buffer Management**:
- Overflow protection
- Accumulation for partial content
- Clear on complete token detection
3. **State Tracking**:
- `in_reasoning`: Current parsing state
- `stripped_think_start`: Prevent double processing
- `buffer`: Accumulated partial content
## 4. Extensibility Guide
### Adding a New Parser
**Step 1: Create Parser Implementation**
```rust
// src/reasoning_parser/parsers/mymodel.rs
use crate::reasoning_parser::parsers::BaseReasoningParser;
use crate::reasoning_parser::traits::{ParserConfig, ReasoningParser};
pub struct MyModelParser {
base: BaseReasoningParser,
}
impl MyModelParser {
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "<reasoning>".to_string(),
think_end_token: "</reasoning>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: false, // or true for implicit
};
Self {
base: BaseReasoningParser::new(config)
.with_model_type("mymodel".to_string()),
}
}
}
impl ReasoningParser for MyModelParser {
// Delegate to base or implement custom logic
fn detect_and_parse_reasoning(&mut self, text: &str)
-> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
// ... other trait methods
}
```
**Step 2: Register in Factory**
```rust
// In factory.rs ParserFactory::new()
registry.register_parser("mymodel", || {
Box::new(MyModelParser::new())
});
// Register patterns
registry.register_pattern("my-model", "mymodel");
registry.register_pattern("mymodel", "mymodel");
```
**Step 3: Export from Module**
```rust
// In parsers/mod.rs
pub use self::mymodel::MyModelParser;
// In reasoning_parser/mod.rs
pub use parsers::MyModelParser;
```
### Custom Parsing Logic
For parsers requiring custom logic beyond configuration:
```rust
impl ReasoningParser for CustomParser {
fn parse_reasoning_streaming_incremental(&mut self, text: &str)
-> Result<ParserResult, ParseError> {
// Custom state machine
// Custom token detection
// Custom buffering strategy
// Return appropriate ParserResult
}
}
```

566
sgl-router/src/reasoning_parser/factory.rs Normal file
View File

@@ -0,0 +1,566 @@
// Factory and registry for creating model-specific reasoning parsers.
// Now with parser pooling support for efficient reuse across requests.
use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use crate::reasoning_parser::parsers::{
BaseReasoningParser, DeepSeekR1Parser, Glm45Parser, KimiParser, Qwen3Parser,
QwenThinkingParser, Step3Parser,
};
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ReasoningParser};
/// Type alias for pooled parser instances.
pub type PooledParser = Arc<Mutex<Box<dyn ReasoningParser>>>;
/// Type alias for parser creator functions.
type ParserCreator = Arc<dyn Fn() -> Box<dyn ReasoningParser> + Send + Sync>;
/// Registry for model-specific parsers with pooling support.
#[derive(Clone)]
pub struct ParserRegistry {
/// Creator functions for parsers (used when pool is empty)
creators: Arc<RwLock<HashMap<String, ParserCreator>>>,
/// Pooled parser instances for reuse
pool: Arc<RwLock<HashMap<String, PooledParser>>>,
/// Model pattern to parser name mappings
patterns: Arc<RwLock<Vec<(String, String)>>>, // (pattern, parser_name)
}
impl ParserRegistry {
/// Create a new empty registry.
pub fn new() -> Self {
Self {
creators: Arc::new(RwLock::new(HashMap::new())),
pool: Arc::new(RwLock::new(HashMap::new())),
patterns: Arc::new(RwLock::new(Vec::new())),
}
}
/// Register a parser creator for a given parser type.
pub fn register_parser<F>(&self, name: &str, creator: F)
where
F: Fn() -> Box<dyn ReasoningParser> + Send + Sync + 'static,
{
let mut creators = self.creators.write().unwrap();
creators.insert(name.to_string(), Arc::new(creator));
}
/// Register a model pattern to parser mapping.
/// Patterns are checked in order, first match wins.
pub fn register_pattern(&self, pattern: &str, parser_name: &str) {
let mut patterns = self.patterns.write().unwrap();
patterns.push((pattern.to_string(), parser_name.to_string()));
}
/// Get a pooled parser by exact name.
/// Returns a shared parser instance from the pool, creating one if needed.
pub fn get_pooled_parser(&self, name: &str) -> Option<PooledParser> {
// First check if we have a pooled instance
{
let pool = self.pool.read().unwrap();
if let Some(parser) = pool.get(name) {
return Some(Arc::clone(parser));
}
}
// If not in pool, create one and add to pool
let creators = self.creators.read().unwrap();
if let Some(creator) = creators.get(name) {
let parser = Arc::new(Mutex::new(creator()));
// Add to pool for future use
let mut pool = self.pool.write().unwrap();
pool.insert(name.to_string(), Arc::clone(&parser));
Some(parser)
} else {
None
}
}
/// Get a parser by exact name (creates new instance, not pooled).
/// Use this for compatibility or when you need a fresh instance.
pub fn get_parser(&self, name: &str) -> Option<Box<dyn ReasoningParser>> {
let creators = self.creators.read().unwrap();
creators.get(name).map(|creator| creator())
}
/// Find a pooled parser for a given model ID by pattern matching.
pub fn find_pooled_parser_for_model(&self, model_id: &str) -> Option<PooledParser> {
let patterns = self.patterns.read().unwrap();
let model_lower = model_id.to_lowercase();
for (pattern, parser_name) in patterns.iter() {
if model_lower.contains(&pattern.to_lowercase()) {
return self.get_pooled_parser(parser_name);
}
}
None
}
/// Find a parser for a given model ID by pattern matching (creates new instance).
pub fn find_parser_for_model(&self, model_id: &str) -> Option<Box<dyn ReasoningParser>> {
let patterns = self.patterns.read().unwrap();
let model_lower = model_id.to_lowercase();
for (pattern, parser_name) in patterns.iter() {
if model_lower.contains(&pattern.to_lowercase()) {
return self.get_parser(parser_name);
}
}
None
}
/// Clear the parser pool, forcing new instances to be created.
/// Useful for testing or when parsers need to be reset globally.
pub fn clear_pool(&self) {
let mut pool = self.pool.write().unwrap();
pool.clear();
}
}
impl Default for ParserRegistry {
fn default() -> Self {
Self::new()
}
}
/// Factory for creating reasoning parsers based on model type.
#[derive(Clone)]
pub struct ParserFactory {
registry: ParserRegistry,
}
impl ParserFactory {
/// Create a new factory with default parsers registered.
pub fn new() -> Self {
let registry = ParserRegistry::new();
// Register base parser
registry.register_parser("base", || {
Box::new(BaseReasoningParser::new(ParserConfig::default()))
});
// Register DeepSeek-R1 parser (starts with in_reasoning=true)
registry.register_parser("deepseek_r1", || Box::new(DeepSeekR1Parser::new()));
// Register Qwen3 parser (starts with in_reasoning=false)
registry.register_parser("qwen3", || Box::new(Qwen3Parser::new()));
// Register Qwen3-thinking parser (starts with in_reasoning=true)
registry.register_parser("qwen3_thinking", || Box::new(QwenThinkingParser::new()));
// Register Kimi parser with Unicode tokens (starts with in_reasoning=false)
registry.register_parser("kimi", || Box::new(KimiParser::new()));
// Register GLM45 parser (same format as Qwen3 but separate for debugging)
registry.register_parser("glm45", || Box::new(Glm45Parser::new()));
// Register Step3 parser (same format as DeepSeek-R1 but separate for debugging)
registry.register_parser("step3", || Box::new(Step3Parser::new()));
// Register model patterns
registry.register_pattern("deepseek-r1", "deepseek_r1");
registry.register_pattern("qwen3-thinking", "qwen3_thinking");
registry.register_pattern("qwen-thinking", "qwen3_thinking");
registry.register_pattern("qwen3", "qwen3");
registry.register_pattern("qwen", "qwen3");
registry.register_pattern("glm45", "glm45");
registry.register_pattern("kimi", "kimi");
registry.register_pattern("step3", "step3");
Self { registry }
}
/// Get a pooled parser for the given model ID.
/// Returns a shared instance that can be used concurrently.
/// Falls back to a passthrough parser if model is not recognized.
pub fn get_pooled(&self, model_id: &str) -> PooledParser {
// First try to find by pattern
if let Some(parser) = self.registry.find_pooled_parser_for_model(model_id) {
return parser;
}
// Fall back to no-op parser (get or create passthrough in pool)
self.registry
.get_pooled_parser("passthrough")
.unwrap_or_else(|| {
// Register passthrough if not already registered
self.registry.register_parser("passthrough", || {
let config = ParserConfig {
think_start_token: "".to_string(),
think_end_token: "".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: false,
};
Box::new(
BaseReasoningParser::new(config).with_model_type("passthrough".to_string()),
)
});
self.registry.get_pooled_parser("passthrough").unwrap()
})
}
/// Create a new parser instance for the given model ID.
/// Returns a fresh instance (not pooled).
/// Use this when you need an isolated parser instance.
pub fn create(&self, model_id: &str) -> Result<Box<dyn ReasoningParser>, ParseError> {
// First try to find by pattern
if let Some(parser) = self.registry.find_parser_for_model(model_id) {
return Ok(parser);
}
// Fall back to no-op parser (base parser without reasoning detection)
let config = ParserConfig {
think_start_token: "".to_string(),
think_end_token: "".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: false,
};
Ok(Box::new(
BaseReasoningParser::new(config).with_model_type("passthrough".to_string()),
))
}
/// Get the internal registry for custom registration.
pub fn registry(&self) -> &ParserRegistry {
&self.registry
}
/// Clear the parser pool.
/// Useful for testing or when parsers need to be reset globally.
pub fn clear_pool(&self) {
self.registry.clear_pool();
}
}
impl Default for ParserFactory {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_factory_creates_deepseek_r1() {
let factory = ParserFactory::new();
let parser = factory.create("deepseek-r1-distill").unwrap();
assert_eq!(parser.model_type(), "deepseek_r1");
}
#[test]
fn test_factory_creates_qwen3() {
let factory = ParserFactory::new();
let parser = factory.create("qwen3-7b").unwrap();
assert_eq!(parser.model_type(), "qwen3");
}
#[test]
fn test_factory_creates_kimi() {
let factory = ParserFactory::new();
let parser = factory.create("kimi-chat").unwrap();
assert_eq!(parser.model_type(), "kimi");
}
#[test]
fn test_factory_fallback_to_passthrough() {
let factory = ParserFactory::new();
let parser = factory.create("unknown-model").unwrap();
assert_eq!(parser.model_type(), "passthrough");
}
#[test]
fn test_case_insensitive_matching() {
let factory = ParserFactory::new();
let parser1 = factory.create("DeepSeek-R1").unwrap();
let parser2 = factory.create("QWEN3").unwrap();
let parser3 = factory.create("Kimi").unwrap();
assert_eq!(parser1.model_type(), "deepseek_r1");
assert_eq!(parser2.model_type(), "qwen3");
assert_eq!(parser3.model_type(), "kimi");
}
#[test]
fn test_step3_model() {
let factory = ParserFactory::new();
let step3 = factory.create("step3-model").unwrap();
assert_eq!(step3.model_type(), "step3");
}
#[test]
fn test_glm45_model() {
let factory = ParserFactory::new();
let glm45 = factory.create("glm45-v2").unwrap();
assert_eq!(glm45.model_type(), "glm45");
}
#[test]
fn test_pooled_parser_reuse() {
let factory = ParserFactory::new();
// Get the same parser twice - should be the same instance
let parser1 = factory.get_pooled("deepseek-r1");
let parser2 = factory.get_pooled("deepseek-r1");
// Both should point to the same Arc
assert!(Arc::ptr_eq(&parser1, &parser2));
// Different models should get different parsers
let parser3 = factory.get_pooled("qwen3");
assert!(!Arc::ptr_eq(&parser1, &parser3));
}
#[test]
fn test_pooled_parser_concurrent_access() {
use std::thread;
let factory = ParserFactory::new();
let parser = factory.get_pooled("deepseek-r1");
// Spawn multiple threads that use the same parser
let mut handles = vec![];
for i in 0..3 {
let parser_clone = Arc::clone(&parser);
let handle = thread::spawn(move || {
let mut parser = parser_clone.lock().unwrap();
let input = format!("thread {} reasoning</think>answer", i);
let result = parser.detect_and_parse_reasoning(&input).unwrap();
assert_eq!(result.normal_text, "answer");
assert!(result.reasoning_text.contains("reasoning"));
});
handles.push(handle);
}
// Wait for all threads to complete
for handle in handles {
handle.join().unwrap();
}
}
#[test]
fn test_pool_clearing() {
let factory = ParserFactory::new();
// Get a pooled parser
let parser1 = factory.get_pooled("deepseek-r1");
// Clear the pool
factory.clear_pool();
// Get another parser - should be a new instance
let parser2 = factory.get_pooled("deepseek-r1");
// They should be different instances (different Arc pointers)
assert!(!Arc::ptr_eq(&parser1, &parser2));
}
#[test]
fn test_passthrough_parser_pooling() {
let factory = ParserFactory::new();
// Unknown models should get passthrough parser
let parser1 = factory.get_pooled("unknown-model-1");
let parser2 = factory.get_pooled("unknown-model-2");
// Both should use the same passthrough parser instance
assert!(Arc::ptr_eq(&parser1, &parser2));
// Verify it's actually a passthrough parser
let parser = parser1.lock().unwrap();
assert_eq!(parser.model_type(), "passthrough");
}
#[test]
fn test_high_concurrency_parser_access() {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;
use std::time::Instant;
let factory = ParserFactory::new();
let num_threads = 100;
let requests_per_thread = 50;
let models = vec!["deepseek-r1", "qwen3", "kimi", "qwen3-thinking"];
// Track successful operations
let success_count = Arc::new(AtomicUsize::new(0));
let error_count = Arc::new(AtomicUsize::new(0));
let start = Instant::now();
let mut handles = vec![];
for thread_id in 0..num_threads {
let factory = factory.clone();
let models = models.clone();
let success_count = Arc::clone(&success_count);
let error_count = Arc::clone(&error_count);
let handle = thread::spawn(move || {
for request_id in 0..requests_per_thread {
// Rotate through different models
let model = &models[(thread_id + request_id) % models.len()];
let parser = factory.get_pooled(model);
// Use blocking lock - this is the realistic scenario
// In production, requests would wait for the parser to be available
// Handle poisoned locks gracefully
let mut p = match parser.lock() {
Ok(guard) => guard,
Err(_poisoned) => {
// Lock was poisoned by a panicking thread
// In production, we might want to recreate the parser
// For testing, we'll just skip this iteration
error_count.fetch_add(1, Ordering::Relaxed);
continue;
}
};
// Simulate realistic parsing work with substantial text
// Typical reasoning can be 500-5000 tokens
let reasoning_text = format!(
"Thread {} is processing request {}. Let me think through this step by step. \
First, I need to understand the problem. The problem involves analyzing data \
and making calculations. Let me break this down: \n\
1. Initial analysis shows that we have multiple variables to consider. \
2. The data suggests a pattern that needs further investigation. \
3. Computing the values: {} * {} = {}. \
4. Cross-referencing with previous results indicates consistency. \
5. The mathematical proof follows from the axioms... \
6. Considering edge cases and boundary conditions... \
7. Validating against known constraints... \
8. The conclusion follows logically from premises A, B, and C. \
This reasoning chain demonstrates the validity of our approach.",
thread_id, request_id, thread_id, request_id, thread_id * request_id
);
let answer_text = format!(
"Based on my analysis, the answer for thread {} request {} is: \
The solution involves multiple steps as outlined in the reasoning. \
The final result is {} with confidence level high. \
This conclusion is supported by rigorous mathematical analysis \
and has been validated against multiple test cases. \
The implementation should handle edge cases appropriately.",
thread_id,
request_id,
thread_id * request_id
);
let input = format!("<think>{}</think>{}", reasoning_text, answer_text);
match p.detect_and_parse_reasoning(&input) {
Ok(result) => {
// Verify parsing worked correctly with substantial content
// Note: Some parsers with stream_reasoning=true won't accumulate reasoning text
assert!(result
.normal_text
.contains(&format!("thread {}", thread_id)));
// For parsers that accumulate reasoning (stream_reasoning=false)
// the reasoning_text should be populated
if !result.reasoning_text.is_empty() {
assert!(result
.reasoning_text
.contains(&format!("Thread {}", thread_id)));
assert!(result.reasoning_text.len() > 500); // Ensure substantial reasoning
}
// Normal text should always be present
assert!(result.normal_text.len() > 100); // Ensure substantial answer
success_count.fetch_add(1, Ordering::Relaxed);
}
Err(e) => {
eprintln!("Parse error: {:?}", e);
error_count.fetch_add(1, Ordering::Relaxed);
}
}
// Explicitly drop the lock to release it quickly
drop(p);
}
});
handles.push(handle);
}
// Wait for all threads
for handle in handles {
handle.join().unwrap();
}
let duration = start.elapsed();
let total_requests = num_threads * requests_per_thread;
let successes = success_count.load(Ordering::Relaxed);
let errors = error_count.load(Ordering::Relaxed);
// Print stats for debugging
println!(
"High concurrency test: {} threads, {} requests each",
num_threads, requests_per_thread
);
println!(
"Completed in {:?}, {} successes, {} errors",
duration, successes, errors
);
println!(
"Throughput: {:.0} requests/sec",
(total_requests as f64) / duration.as_secs_f64()
);
// All requests should succeed
assert_eq!(successes, total_requests);
assert_eq!(errors, 0);
// Performance check: should handle at least 1000 req/sec
let throughput = (total_requests as f64) / duration.as_secs_f64();
assert!(
throughput > 1000.0,
"Throughput too low: {:.0} req/sec",
throughput
);
}
#[test]
fn test_concurrent_pool_modifications() {
use std::thread;
let factory = ParserFactory::new();
let mut handles = vec![];
// Thread 1: Continuously get parsers
let factory1 = factory.clone();
handles.push(thread::spawn(move || {
for _ in 0..100 {
let _parser = factory1.get_pooled("deepseek-r1");
}
}));
// Thread 2: Continuously clear pool
let factory2 = factory.clone();
handles.push(thread::spawn(move || {
for _ in 0..10 {
factory2.clear_pool();
thread::sleep(std::time::Duration::from_micros(100));
}
}));
// Thread 3: Get different parsers
let factory3 = factory.clone();
handles.push(thread::spawn(move || {
for i in 0..100 {
let models = ["qwen3", "kimi", "unknown"];
let _parser = factory3.get_pooled(models[i % 3]);
}
}));
// Wait for all threads - should not deadlock or panic
for handle in handles {
handle.join().unwrap();
}
}
}

10
sgl-router/src/reasoning_parser/mod.rs Normal file
View File

@@ -0,0 +1,10 @@
pub mod factory;
pub mod parsers;
pub mod traits;
pub use factory::{ParserFactory, ParserRegistry, PooledParser};
pub use parsers::{
BaseReasoningParser, DeepSeekR1Parser, Glm45Parser, KimiParser, Qwen3Parser,
QwenThinkingParser, Step3Parser,
};
pub use traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};

386
sgl-router/src/reasoning_parser/parsers/base.rs Normal file
View File

@@ -0,0 +1,386 @@
// Base implementation of reasoning parser that handles common logic
// for detecting and extracting reasoning blocks from text.
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};
use tracing as log;
/// Base reasoning parser implementation.
///
/// This parser handles the common logic for detecting reasoning blocks
/// delimited by start and end tokens (e.g., <think> and </think>).
#[derive(Debug, Clone)]
pub struct BaseReasoningParser {
config: ParserConfig,
in_reasoning: bool,
buffer: String,
stripped_think_start: bool,
model_type: String,
}
impl BaseReasoningParser {
/// Create a new BaseReasoningParser with the given configuration.
pub fn new(config: ParserConfig) -> Self {
let in_reasoning = config.initial_in_reasoning;
Self {
config,
in_reasoning,
buffer: String::new(),
stripped_think_start: false,
model_type: "base".to_string(),
}
}
/// Create with custom model type identifier.
pub fn with_model_type(mut self, model_type: String) -> Self {
self.model_type = model_type;
self
}
/// Check if the current buffer is a prefix of one of the tokens.
fn is_partial_token(&self, text: &str) -> bool {
(self.config.think_start_token.starts_with(text) && self.config.think_start_token != text)
|| (self.config.think_end_token.starts_with(text)
&& self.config.think_end_token != text)
}
}
impl ReasoningParser for BaseReasoningParser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
log::debug!("detect_and_parse_reasoning called with text: {:?}", text);
// Check input size against buffer limit
if text.len() > self.config.max_buffer_size {
return Err(ParseError::BufferOverflow(text.len()));
}
let in_reasoning = self.in_reasoning || text.contains(&self.config.think_start_token);
log::debug!("in_reasoning: {}", in_reasoning);
if !in_reasoning {
log::debug!("No reasoning detected, returning normal text.");
return Ok(ParserResult::normal(text.to_string()));
}
// The text is considered to be in a reasoning block.
let processed_text = text
.replace(&self.config.think_start_token, "")
.trim()
.to_string();
log::debug!(
"Processed text after removing think_start_token: {:?}",
processed_text
);
if !processed_text.contains(&self.config.think_end_token) {
log::debug!(
"Reasoning truncated, think_end_token not found. Returning reasoning text."
);
// Assume reasoning was truncated before end token
return Ok(ParserResult::reasoning(processed_text));
}
// Extract reasoning content
let splits: Vec<&str> = processed_text
.splitn(2, &self.config.think_end_token)
.collect();
let reasoning_text = splits.first().unwrap_or(&"").to_string();
let normal_text = splits
.get(1)
.map(|s| s.trim().to_string())
.unwrap_or_default();
log::debug!("Extracted reasoning_text: {:?}", reasoning_text);
log::debug!("Extracted normal_text: {:?}", normal_text);
Ok(ParserResult::new(normal_text, reasoning_text))
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
// Check if adding this text would exceed buffer limit
if self.buffer.len() + text.len() > self.config.max_buffer_size {
return Err(ParseError::BufferOverflow(self.buffer.len() + text.len()));
}
// Incrementally parse the streaming text
self.buffer.push_str(text);
let mut current_text = self.buffer.clone();
log::debug!(
"parse_reasoning_streaming_incremental called with text: {:?}",
text
);
log::debug!("current buffer: {:?}", self.buffer);
log::debug!("current_text: {:?}", current_text);
log::debug!(
"in_reasoning: {}, stripped_think_start: {}, stream_reasoning: {}",
self.in_reasoning,
self.stripped_think_start,
self.config.stream_reasoning
);
// If the current text is a prefix of a token, keep buffering
if self.is_partial_token(&current_text) {
return Ok(ParserResult::default());
}
// Strip start token if present
if !self.stripped_think_start && current_text.contains(&self.config.think_start_token) {
current_text = current_text.replace(&self.config.think_start_token, "");
self.buffer = current_text.clone();
self.stripped_think_start = true;
self.in_reasoning = true;
}
// Handle end of reasoning block
let think_end_idx = if self.in_reasoning {
current_text
.find(&self.config.think_end_token)
.unwrap_or(current_text.len())
} else {
current_text.len()
};
if self.in_reasoning && think_end_idx < current_text.len() {
let reasoning_text = &current_text[..think_end_idx];
self.buffer.clear();
self.in_reasoning = false;
let start_idx = think_end_idx + self.config.think_end_token.len();
let normal_text = if start_idx < current_text.len() {
&current_text[start_idx..]
} else {
""
};
return Ok(ParserResult::new(
normal_text.to_string(),
reasoning_text.trim().to_string(),
));
}
// Continue with reasoning content
if self.in_reasoning && self.config.stream_reasoning {
// Stream the content immediately
let reasoning_text = current_text;
self.buffer.clear();
Ok(ParserResult::reasoning(reasoning_text))
} else if !self.in_reasoning {
// If we're not in a reasoning block, return as normal text
// CRITICAL FIX: Return current_text (with buffer) not just text
// This prevents buffer loss when partial tokens are followed by normal text
let normal_text = current_text;
self.buffer.clear();
Ok(ParserResult::normal(normal_text))
} else {
// If we are in a reasoning block but no end token is found, buffer it
Ok(ParserResult::default())
}
}
fn reset(&mut self) {
self.in_reasoning = self.config.initial_in_reasoning;
self.buffer.clear();
self.stripped_think_start = false;
}
fn model_type(&self) -> &str {
&self.model_type
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_parser(
initial_in_reasoning: bool,
stream_reasoning: bool,
) -> BaseReasoningParser {
let config = ParserConfig {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning,
max_buffer_size: 65536,
initial_in_reasoning,
};
BaseReasoningParser::new(config)
}
#[test]
fn test_detect_and_parse_reasoning() {
let mut parser = create_test_parser(false, true);
let result = parser
.detect_and_parse_reasoning("<think>with reasoning</think> and more text.")
.unwrap();
assert_eq!(result.normal_text, "and more text.");
assert_eq!(result.reasoning_text, "with reasoning");
}
#[test]
fn test_detect_and_parse_no_reasoning() {
let mut parser = create_test_parser(false, true);
let result = parser
.detect_and_parse_reasoning("This is a test without reasoning.")
.unwrap();
assert_eq!(result.normal_text, "This is a test without reasoning.");
assert_eq!(result.reasoning_text, "");
}
#[test]
fn test_detect_and_parse_truncated_reasoning() {
let mut parser = create_test_parser(false, true);
let result = parser
.detect_and_parse_reasoning("<think>with truncated reasoning")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "with truncated reasoning");
}
#[test]
fn test_parse_streaming_partial_token() {
let mut parser = create_test_parser(false, true);
let result = parser
.parse_reasoning_streaming_incremental("<thi")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "");
}
#[test]
fn test_parse_streaming_complete() {
let mut parser = create_test_parser(false, true);
let result = parser
.parse_reasoning_streaming_incremental("<think>with reasoning</think> and more text.")
.unwrap();
assert_eq!(result.normal_text, " and more text.");
assert_eq!(result.reasoning_text, "with reasoning");
}
#[test]
fn test_parse_streaming_no_end_token() {
let mut parser = create_test_parser(true, true);
let result = parser
.parse_reasoning_streaming_incremental("<think>with reasoning")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "with reasoning");
}
#[test]
fn test_initial_in_reasoning_true() {
// Parser starts with in_reasoning=true (like DeepSeek-R1)
let mut parser = create_test_parser(true, true);
let result = parser
.detect_and_parse_reasoning("no think tags here")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "no think tags here");
}
#[test]
fn test_buffer_loss_bug_fix() {
// Critical test for buffer preservation
let mut parser = create_test_parser(false, true);
// Step 1: Send partial end tag when not in reasoning mode
let result1 = parser.parse_reasoning_streaming_incremental("</").unwrap();
assert_eq!(result1.normal_text, "");
assert_eq!(result1.reasoning_text, "");
// Step 2: Send normal text that doesn't complete the end tag
// Must return "</answer" not just "answer"
let result2 = parser
.parse_reasoning_streaming_incremental("answer")
.unwrap();
assert_eq!(result2.normal_text, "</answer");
assert_eq!(result2.reasoning_text, "");
}
#[test]
fn test_streaming_with_stream_reasoning_enabled() {
let mut parser = create_test_parser(false, true);
// Start reasoning block
let result1 = parser
.parse_reasoning_streaming_incremental("<think>reasoning ")
.unwrap();
assert_eq!(result1.normal_text, "");
assert_eq!(result1.reasoning_text, "reasoning ");
// Continue streaming reasoning
let result2 = parser
.parse_reasoning_streaming_incremental("content ")
.unwrap();
assert_eq!(result2.normal_text, "");
assert_eq!(result2.reasoning_text, "content ");
// End reasoning block
let result3 = parser
.parse_reasoning_streaming_incremental("more</think> normal")
.unwrap();
assert_eq!(result3.normal_text, " normal");
assert_eq!(result3.reasoning_text, "more");
}
#[test]
fn test_reset_state() {
let mut parser = create_test_parser(false, true);
// Process some text
parser
.parse_reasoning_streaming_incremental("<think>reasoning</think> normal")
.unwrap();
// Reset and verify state
parser.reset();
assert!(!parser.in_reasoning);
assert!(parser.buffer.is_empty());
assert!(!parser.stripped_think_start);
}
#[test]
fn test_buffer_overflow_detect_and_parse() {
let config = ParserConfig {
max_buffer_size: 10, // Set a very small buffer
..Default::default()
};
let mut parser = BaseReasoningParser::new(config);
let large_text = "a".repeat(20);
let result = parser.detect_and_parse_reasoning(&large_text);
assert!(result.is_err());
match result {
Err(ParseError::BufferOverflow(size)) => {
assert_eq!(size, 20);
}
_ => panic!("Expected BufferOverflow error"),
}
}
#[test]
fn test_buffer_overflow_streaming() {
let config = ParserConfig {
max_buffer_size: 10, // Set a very small buffer
..Default::default()
};
let mut parser = BaseReasoningParser::new(config);
// Send a partial token that will be buffered
let result1 = parser.parse_reasoning_streaming_incremental("<thi");
assert!(result1.is_ok());
assert_eq!(result1.unwrap().normal_text, "");
// Second chunk would exceed buffer
// Buffer has "<thi" (4 chars) + "this_is_too_large" (17 chars) = 21 total
let result2 = parser.parse_reasoning_streaming_incremental("this_is_too_large");
assert!(result2.is_err());
match result2 {
Err(ParseError::BufferOverflow(size)) => {
assert_eq!(size, 21); // 4 + 17
}
_ => panic!("Expected BufferOverflow error"),
}
}
}

112
sgl-router/src/reasoning_parser/parsers/deepseek_r1.rs Normal file
View File

@@ -0,0 +1,112 @@
// DeepSeek-R1 specific reasoning parser.
// This parser starts with in_reasoning=true, assuming all text is reasoning
// until an end token is encountered.
use crate::reasoning_parser::parsers::BaseReasoningParser;
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};
/// DeepSeek-R1 reasoning parser.
///
/// This parser assumes reasoning from the start of text (in_reasoning=true)
/// and uses <think> and </think> tokens.
pub struct DeepSeekR1Parser {
base: BaseReasoningParser,
}
impl DeepSeekR1Parser {
/// Create a new DeepSeek-R1 parser.
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: true, // Always starts with reasoning
};
Self {
base: BaseReasoningParser::new(config).with_model_type("deepseek_r1".to_string()),
}
}
}
impl Default for DeepSeekR1Parser {
fn default() -> Self {
Self::new()
}
}
impl ReasoningParser for DeepSeekR1Parser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
self.base.parse_reasoning_streaming_incremental(text)
}
fn reset(&mut self) {
self.base.reset()
}
fn model_type(&self) -> &str {
self.base.model_type()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_deepseek_r1_initial_state() {
let mut parser = DeepSeekR1Parser::new();
// Should treat text as reasoning even without start token
let result = parser
.detect_and_parse_reasoning("This is reasoning content")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "This is reasoning content");
}
#[test]
fn test_deepseek_r1_with_end_token() {
let mut parser = DeepSeekR1Parser::new();
// Should extract reasoning until end token
let result = parser
.detect_and_parse_reasoning("reasoning content</think>normal content")
.unwrap();
assert_eq!(result.normal_text, "normal content");
assert_eq!(result.reasoning_text, "reasoning content");
}
#[test]
fn test_deepseek_r1_streaming() {
let mut parser = DeepSeekR1Parser::new();
// First chunk - all reasoning
let result1 = parser
.parse_reasoning_streaming_incremental("thinking about")
.unwrap();
assert_eq!(result1.reasoning_text, "thinking about");
assert_eq!(result1.normal_text, "");
// Second chunk - ends reasoning
let result2 = parser
.parse_reasoning_streaming_incremental(" the problem</think>answer")
.unwrap();
assert_eq!(result2.reasoning_text, "the problem"); // Text is trimmed
assert_eq!(result2.normal_text, "answer");
}
#[test]
fn test_model_type() {
let parser = DeepSeekR1Parser::new();
assert_eq!(parser.model_type(), "deepseek_r1");
}
}

118
sgl-router/src/reasoning_parser/parsers/glm45.rs Normal file
View File

@@ -0,0 +1,118 @@
// GLM45 specific reasoning parser.
// Uses the same format as Qwen3 but has its own implementation for debugging.
use crate::reasoning_parser::parsers::BaseReasoningParser;
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};
/// GLM45 reasoning parser.
///
/// This parser uses the same format as Qwen3 (<think>...</think>) but has
/// its own implementation for better debugging and potential future customization.
pub struct Glm45Parser {
base: BaseReasoningParser,
}
impl Glm45Parser {
/// Create a new GLM45 parser.
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: false, // Requires explicit start token like Qwen3
};
Self {
base: BaseReasoningParser::new(config).with_model_type("glm45".to_string()),
}
}
}
impl Default for Glm45Parser {
fn default() -> Self {
Self::new()
}
}
impl ReasoningParser for Glm45Parser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
self.base.parse_reasoning_streaming_incremental(text)
}
fn reset(&mut self) {
self.base.reset()
}
fn model_type(&self) -> &str {
self.base.model_type()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_glm45_initial_state() {
let mut parser = Glm45Parser::new();
// Should NOT treat text as reasoning without start token
let result = parser
.detect_and_parse_reasoning("This is normal content")
.unwrap();
assert_eq!(result.normal_text, "This is normal content");
assert_eq!(result.reasoning_text, "");
}
#[test]
fn test_glm45_with_tokens() {
let mut parser = Glm45Parser::new();
// Should extract reasoning with proper tokens
let result = parser
.detect_and_parse_reasoning("<think>reasoning content</think>answer")
.unwrap();
assert_eq!(result.normal_text, "answer");
assert_eq!(result.reasoning_text, "reasoning content");
}
#[test]
fn test_glm45_streaming() {
let mut parser = Glm45Parser::new();
// First chunk - normal text
let result1 = parser
.parse_reasoning_streaming_incremental("normal text ")
.unwrap();
assert_eq!(result1.normal_text, "normal text ");
assert_eq!(result1.reasoning_text, "");
// Second chunk - enters reasoning
let result2 = parser
.parse_reasoning_streaming_incremental("<think>reasoning")
.unwrap();
assert_eq!(result2.normal_text, "");
assert_eq!(result2.reasoning_text, "reasoning");
// Third chunk - exits reasoning
let result3 = parser
.parse_reasoning_streaming_incremental("</think>answer")
.unwrap();
assert_eq!(result3.normal_text, "answer");
assert_eq!(result3.reasoning_text, "");
}
#[test]
fn test_model_type() {
let parser = Glm45Parser::new();
assert_eq!(parser.model_type(), "glm45");
}
}

137
sgl-router/src/reasoning_parser/parsers/kimi.rs Normal file
View File

@@ -0,0 +1,137 @@
// Kimi specific reasoning parser.
// This parser uses Unicode tokens and starts with in_reasoning=false.
use crate::reasoning_parser::parsers::BaseReasoningParser;
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};
/// Kimi reasoning parser.
///
/// This parser uses Unicode tokens (◁think▷ and ◁/think▷) and requires
/// explicit start tokens to enter reasoning mode.
pub struct KimiParser {
base: BaseReasoningParser,
}
impl KimiParser {
/// Create a new Kimi parser.
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "◁think▷".to_string(),
think_end_token: "◁/think▷".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: false, // Requires explicit start token
};
Self {
base: BaseReasoningParser::new(config).with_model_type("kimi".to_string()),
}
}
}
impl Default for KimiParser {
fn default() -> Self {
Self::new()
}
}
impl ReasoningParser for KimiParser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
self.base.parse_reasoning_streaming_incremental(text)
}
fn reset(&mut self) {
self.base.reset()
}
fn model_type(&self) -> &str {
self.base.model_type()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_kimi_initial_state() {
let mut parser = KimiParser::new();
// Should NOT treat text as reasoning without start token
let result = parser
.detect_and_parse_reasoning("This is normal content")
.unwrap();
assert_eq!(result.normal_text, "This is normal content");
assert_eq!(result.reasoning_text, "");
}
#[test]
fn test_kimi_with_unicode_tokens() {
let mut parser = KimiParser::new();
// Should extract reasoning with Unicode tokens
let result = parser
.detect_and_parse_reasoning("◁think▷reasoning content◁/think▷answer")
.unwrap();
assert_eq!(result.normal_text, "answer");
assert_eq!(result.reasoning_text, "reasoning content");
}
#[test]
fn test_kimi_partial_unicode() {
let mut parser = KimiParser::new();
// Test partial Unicode token buffering
let result1 = parser
.parse_reasoning_streaming_incremental("◁thi")
.unwrap();
assert_eq!(result1.normal_text, "");
assert_eq!(result1.reasoning_text, "");
// Complete the token
let result2 = parser
.parse_reasoning_streaming_incremental("nk▷reasoning")
.unwrap();
assert_eq!(result2.normal_text, "");
assert_eq!(result2.reasoning_text, "reasoning");
}
#[test]
fn test_kimi_streaming() {
let mut parser = KimiParser::new();
// Normal text first
let result1 = parser
.parse_reasoning_streaming_incremental("normal ")
.unwrap();
assert_eq!(result1.normal_text, "normal ");
assert_eq!(result1.reasoning_text, "");
// Enter reasoning with Unicode token
let result2 = parser
.parse_reasoning_streaming_incremental("◁think▷thinking")
.unwrap();
assert_eq!(result2.normal_text, "");
assert_eq!(result2.reasoning_text, "thinking");
// Exit reasoning
let result3 = parser
.parse_reasoning_streaming_incremental("◁/think▷answer")
.unwrap();
assert_eq!(result3.normal_text, "answer");
assert_eq!(result3.reasoning_text, ""); // Already returned in stream mode
}
#[test]
fn test_model_type() {
let parser = KimiParser::new();
assert_eq!(parser.model_type(), "kimi");
}
}

13
sgl-router/src/reasoning_parser/parsers/mod.rs Normal file
View File

@@ -0,0 +1,13 @@
pub mod base;
pub mod deepseek_r1;
pub mod glm45;
pub mod kimi;
pub mod qwen3;
pub mod step3;
pub use base::BaseReasoningParser;
pub use deepseek_r1::DeepSeekR1Parser;
pub use glm45::Glm45Parser;
pub use kimi::KimiParser;
pub use qwen3::{Qwen3Parser, QwenThinkingParser};
pub use step3::Step3Parser;

178
sgl-router/src/reasoning_parser/parsers/qwen3.rs Normal file
View File

@@ -0,0 +1,178 @@
// Qwen3 specific reasoning parser.
// This parser starts with in_reasoning=false, requiring an explicit
// start token to enter reasoning mode.
use crate::reasoning_parser::parsers::BaseReasoningParser;
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};
/// Qwen3 reasoning parser.
///
/// This parser requires explicit <think> tokens to enter reasoning mode
/// (in_reasoning=false initially).
pub struct Qwen3Parser {
base: BaseReasoningParser,
}
impl Qwen3Parser {
/// Create a new Qwen3 parser.
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: false, // Requires explicit start token
};
Self {
base: BaseReasoningParser::new(config).with_model_type("qwen3".to_string()),
}
}
}
impl Default for Qwen3Parser {
fn default() -> Self {
Self::new()
}
}
impl ReasoningParser for Qwen3Parser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
self.base.parse_reasoning_streaming_incremental(text)
}
fn reset(&mut self) {
self.base.reset()
}
fn model_type(&self) -> &str {
self.base.model_type()
}
}
/// QwenThinking parser - variant that assumes reasoning from start.
///
/// This is for qwen*thinking models that behave like DeepSeek-R1.
pub struct QwenThinkingParser {
base: BaseReasoningParser,
}
impl QwenThinkingParser {
/// Create a new QwenThinking parser.
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: true, // Assumes reasoning from start
};
Self {
base: BaseReasoningParser::new(config).with_model_type("qwen_thinking".to_string()),
}
}
}
impl Default for QwenThinkingParser {
fn default() -> Self {
Self::new()
}
}
impl ReasoningParser for QwenThinkingParser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
self.base.parse_reasoning_streaming_incremental(text)
}
fn reset(&mut self) {
self.base.reset()
}
fn model_type(&self) -> &str {
self.base.model_type()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_qwen3_initial_state() {
let mut parser = Qwen3Parser::new();
// Should NOT treat text as reasoning without start token
let result = parser
.detect_and_parse_reasoning("This is normal content")
.unwrap();
assert_eq!(result.normal_text, "This is normal content");
assert_eq!(result.reasoning_text, "");
}
#[test]
fn test_qwen3_with_tokens() {
let mut parser = Qwen3Parser::new();
// Should extract reasoning with proper tokens
let result = parser
.detect_and_parse_reasoning("<think>reasoning</think>answer")
.unwrap();
assert_eq!(result.normal_text, "answer");
assert_eq!(result.reasoning_text, "reasoning");
}
#[test]
fn test_qwen_thinking_initial_state() {
let mut parser = QwenThinkingParser::new();
// Should treat text as reasoning even without start token
let result = parser
.detect_and_parse_reasoning("This is reasoning content")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "This is reasoning content");
}
#[test]
fn test_qwen3_streaming() {
let mut parser = Qwen3Parser::new();
// First chunk - normal text (no start token yet)
let result1 = parser
.parse_reasoning_streaming_incremental("normal text ")
.unwrap();
assert_eq!(result1.normal_text, "normal text ");
assert_eq!(result1.reasoning_text, "");
// Second chunk - enters reasoning
let result2 = parser
.parse_reasoning_streaming_incremental("<think>reasoning")
.unwrap();
assert_eq!(result2.normal_text, "");
assert_eq!(result2.reasoning_text, "reasoning");
}
#[test]
fn test_model_types() {
let qwen3 = Qwen3Parser::new();
assert_eq!(qwen3.model_type(), "qwen3");
let qwen_thinking = QwenThinkingParser::new();
assert_eq!(qwen_thinking.model_type(), "qwen_thinking");
}
}

123
sgl-router/src/reasoning_parser/parsers/step3.rs Normal file
View File

@@ -0,0 +1,123 @@
// Step3 specific reasoning parser.
// Uses the same format as DeepSeek-R1 but has its own implementation for debugging.
use crate::reasoning_parser::parsers::BaseReasoningParser;
use crate::reasoning_parser::traits::{ParseError, ParserConfig, ParserResult, ReasoningParser};
/// Step3 reasoning parser.
///
/// This parser uses the same format as DeepSeek-R1 (<think>...</think>) but has
/// its own implementation for better debugging and potential future customization.
pub struct Step3Parser {
base: BaseReasoningParser,
}
impl Step3Parser {
/// Create a new Step3 parser.
pub fn new() -> Self {
let config = ParserConfig {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536,
initial_in_reasoning: true, // Assumes reasoning from start like DeepSeek-R1
};
Self {
base: BaseReasoningParser::new(config).with_model_type("step3".to_string()),
}
}
}
impl Default for Step3Parser {
fn default() -> Self {
Self::new()
}
}
impl ReasoningParser for Step3Parser {
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError> {
self.base.detect_and_parse_reasoning(text)
}
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError> {
self.base.parse_reasoning_streaming_incremental(text)
}
fn reset(&mut self) {
self.base.reset()
}
fn model_type(&self) -> &str {
self.base.model_type()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_step3_initial_state() {
let mut parser = Step3Parser::new();
// Should treat text as reasoning even without start token
let result = parser
.detect_and_parse_reasoning("This is reasoning content")
.unwrap();
assert_eq!(result.normal_text, "");
assert_eq!(result.reasoning_text, "This is reasoning content");
}
#[test]
fn test_step3_with_end_token() {
let mut parser = Step3Parser::new();
// Should handle text with end token
let result = parser
.detect_and_parse_reasoning("reasoning content</think>answer")
.unwrap();
assert_eq!(result.normal_text, "answer");
assert_eq!(result.reasoning_text, "reasoning content");
}
#[test]
fn test_step3_with_both_tokens() {
let mut parser = Step3Parser::new();
// Should handle both start and end tokens
let result = parser
.detect_and_parse_reasoning("<think>reasoning content</think>answer")
.unwrap();
assert_eq!(result.normal_text, "answer");
assert_eq!(result.reasoning_text, "reasoning content");
}
#[test]
fn test_step3_streaming() {
let mut parser = Step3Parser::new();
// First chunk - treated as reasoning (initial_in_reasoning=true)
let result1 = parser
.parse_reasoning_streaming_incremental("reasoning text ")
.unwrap();
assert_eq!(result1.normal_text, "");
assert_eq!(result1.reasoning_text, "reasoning text ");
// Second chunk - continues reasoning until end token
let result2 = parser
.parse_reasoning_streaming_incremental("more reasoning</think>answer")
.unwrap();
assert_eq!(result2.normal_text, "answer");
assert_eq!(result2.reasoning_text, "more reasoning");
}
#[test]
fn test_model_type() {
let parser = Step3Parser::new();
assert_eq!(parser.model_type(), "step3");
}
}

130
sgl-router/src/reasoning_parser/traits.rs Normal file
View File

@@ -0,0 +1,130 @@
use std::fmt;
/// Result of parsing text for reasoning content.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ParserResult {
/// The normal text outside of reasoning blocks.
pub normal_text: String,
/// The extracted reasoning text from within reasoning blocks.
pub reasoning_text: String,
}
impl ParserResult {
/// Create a new ParserResult with the given normal and reasoning text.
pub fn new(normal_text: String, reasoning_text: String) -> Self {
Self {
normal_text,
reasoning_text,
}
}
/// Create a result with only normal text.
pub fn normal(text: String) -> Self {
Self {
normal_text: text,
reasoning_text: String::new(),
}
}
/// Create a result with only reasoning text.
pub fn reasoning(text: String) -> Self {
Self {
normal_text: String::new(),
reasoning_text: text,
}
}
/// Check if this result contains any text.
pub fn is_empty(&self) -> bool {
self.normal_text.is_empty() && self.reasoning_text.is_empty()
}
}
/// Trait for parsing reasoning content from LLM outputs.
pub trait ReasoningParser: Send + Sync {
/// Detects and parses reasoning from the input text (one-time parsing).
///
/// This method is used for non-streaming scenarios where the complete
/// text is available at once.
///
/// Returns an error if the text exceeds buffer limits or contains invalid UTF-8.
fn detect_and_parse_reasoning(&mut self, text: &str) -> Result<ParserResult, ParseError>;
/// Parses reasoning incrementally from streaming input.
///
/// This method maintains internal state across calls to handle partial
/// tokens and chunk boundaries correctly.
///
/// Returns an error if the buffer exceeds max_buffer_size.
fn parse_reasoning_streaming_incremental(
&mut self,
text: &str,
) -> Result<ParserResult, ParseError>;
/// Reset the parser state for reuse.
///
/// This should clear any buffers and reset flags to initial state.
fn reset(&mut self);
/// Get the model type this parser is designed for.
fn model_type(&self) -> &str;
}
/// Error types for reasoning parsing operations.
#[derive(Debug, thiserror::Error)]
pub enum ParseError {
#[error("Invalid UTF-8 in stream: {0}")]
Utf8Error(#[from] std::str::Utf8Error),
#[error("Buffer overflow: {0} bytes exceeds maximum")]
BufferOverflow(usize),
#[error("Unknown model type: {0}")]
UnknownModel(String),
#[error("Parser configuration error: {0}")]
ConfigError(String),
}
/// Configuration for parser behavior.
#[derive(Debug, Clone)]
pub struct ParserConfig {
/// The token that marks the start of reasoning content.
pub think_start_token: String,
/// The token that marks the end of reasoning content.
pub think_end_token: String,
/// Whether to stream reasoning content as it arrives.
pub stream_reasoning: bool,
/// Maximum buffer size in bytes.
pub max_buffer_size: usize,
/// Initial state for in_reasoning flag (fixed per parser type).
pub initial_in_reasoning: bool,
}
impl Default for ParserConfig {
fn default() -> Self {
Self {
think_start_token: "<think>".to_string(),
think_end_token: "</think>".to_string(),
stream_reasoning: true,
max_buffer_size: 65536, // 64KB default
initial_in_reasoning: false, // Default to false (explicit reasoning)
}
}
}
impl fmt::Display for ParserResult {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"ParserResult {{ normal: {} chars, reasoning: {} chars }}",
self.normal_text.len(),
self.reasoning_text.len()
)
}
}

195
sgl-router/src/routers/factory.rs Normal file
View File

@@ -0,0 +1,195 @@
//! Factory for creating router instances
use super::{
http::{openai_router::OpenAIRouter, pd_router::PDRouter, router::Router},
RouterTrait,
};
use crate::config::{ConnectionMode, PolicyConfig, RoutingMode};
use crate::policies::PolicyFactory;
use crate::server::AppContext;
use std::sync::Arc;
/// Factory for creating router instances based on configuration
pub struct RouterFactory;
impl RouterFactory {
/// Create a router instance from application context
pub async fn create_router(ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
// Check if IGW mode is enabled
if ctx.router_config.enable_igw {
return Self::create_igw_router(ctx).await;
}
// Check connection mode and route to appropriate implementation
match ctx.router_config.connection_mode {
ConnectionMode::Grpc => {
// Route to gRPC implementation based on routing mode
match &ctx.router_config.mode {
RoutingMode::Regular { worker_urls } => {
Self::create_grpc_router(worker_urls, &ctx.router_config.policy, ctx).await
}
RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
prefill_policy,
decode_policy,
} => {
Self::create_grpc_pd_router(
prefill_urls,
decode_urls,
prefill_policy.as_ref(),
decode_policy.as_ref(),
&ctx.router_config.policy,
ctx,
)
.await
}
RoutingMode::OpenAI { .. } => {
Err("OpenAI mode requires HTTP connection_mode".to_string())
}
}
}
ConnectionMode::Http => {
// Route to HTTP implementation based on routing mode
match &ctx.router_config.mode {
RoutingMode::Regular { worker_urls } => {
Self::create_regular_router(worker_urls, &ctx.router_config.policy, ctx)
.await
}
RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
prefill_policy,
decode_policy,
} => {
Self::create_pd_router(
prefill_urls,
decode_urls,
prefill_policy.as_ref(),
decode_policy.as_ref(),
&ctx.router_config.policy,
ctx,
)
.await
}
RoutingMode::OpenAI { worker_urls, .. } => {
Self::create_openai_router(worker_urls.clone(), ctx).await
}
}
}
}
}
/// Create a regular router with injected policy
async fn create_regular_router(
worker_urls: &[String],
policy_config: &PolicyConfig,
ctx: &Arc<AppContext>,
) -> Result<Box<dyn RouterTrait>, String> {
// Create policy
let policy = PolicyFactory::create_from_config(policy_config);
// Create regular router with injected policy and context
let router = Router::new(worker_urls.to_vec(), policy, ctx).await?;
Ok(Box::new(router))
}
/// Create a PD router with injected policy
async fn create_pd_router(
prefill_urls: &[(String, Option<u16>)],
decode_urls: &[String],
prefill_policy_config: Option<&PolicyConfig>,
decode_policy_config: Option<&PolicyConfig>,
main_policy_config: &PolicyConfig,
ctx: &Arc<AppContext>,
) -> Result<Box<dyn RouterTrait>, String> {
// Create policies - use specific policies if provided, otherwise fall back to main policy
let prefill_policy =
PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config));
let decode_policy =
PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config));
// Create PD router with separate policies and context
let router = PDRouter::new(
prefill_urls.to_vec(),
decode_urls.to_vec(),
prefill_policy,
decode_policy,
ctx,
)
.await?;
Ok(Box::new(router))
}
/// Create a gRPC router with injected policy
pub async fn create_grpc_router(
worker_urls: &[String],
policy_config: &PolicyConfig,
ctx: &Arc<AppContext>,
) -> Result<Box<dyn RouterTrait>, String> {
use super::grpc::router::GrpcRouter;
// Create policy
let policy = PolicyFactory::create_from_config(policy_config);
// Create gRPC router with context
let router = GrpcRouter::new(worker_urls.to_vec(), policy, ctx).await?;
Ok(Box::new(router))
}
/// Create a gRPC PD router with tokenizer and worker configuration
pub async fn create_grpc_pd_router(
prefill_urls: &[(String, Option<u16>)],
decode_urls: &[String],
prefill_policy_config: Option<&PolicyConfig>,
decode_policy_config: Option<&PolicyConfig>,
main_policy_config: &PolicyConfig,
ctx: &Arc<AppContext>,
) -> Result<Box<dyn RouterTrait>, String> {
use super::grpc::pd_router::GrpcPDRouter;
// Create policies - use specific policies if provided, otherwise fall back to main policy
let prefill_policy =
PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config));
let decode_policy =
PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config));
// Create gRPC PD router with context
let router = GrpcPDRouter::new(
prefill_urls.to_vec(),
decode_urls.to_vec(),
prefill_policy,
decode_policy,
ctx,
)
.await?;
Ok(Box::new(router))
}
/// Create an OpenAI router
async fn create_openai_router(
worker_urls: Vec<String>,
ctx: &Arc<AppContext>,
) -> Result<Box<dyn RouterTrait>, String> {
// Use the first worker URL as the OpenAI-compatible base
let base_url = worker_urls
.first()
.cloned()
.ok_or_else(|| "OpenAI mode requires at least one worker URL".to_string())?;
let router =
OpenAIRouter::new(base_url, Some(ctx.router_config.circuit_breaker.clone())).await?;
Ok(Box::new(router))
}
/// Create an IGW router (placeholder for future implementation)
async fn create_igw_router(_ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
// For now, return an error indicating IGW is not yet implemented
Err("IGW mode is not yet implemented".to_string())
}
}

4
sgl-router/src/routers/grpc/mod.rs Normal file
View File

@@ -0,0 +1,4 @@
//! gRPC router implementations
pub mod pd_router;
pub mod router;

328
sgl-router/src/routers/grpc/pd_router.rs Normal file
View File

@@ -0,0 +1,328 @@
// PD (Prefill-Decode) gRPC Router Implementation
use crate::config::types::RetryConfig;
use crate::core::{
BasicWorker, CircuitBreakerConfig, HealthChecker, HealthConfig, Worker, WorkerType,
};
use crate::grpc::SglangSchedulerClient;
use crate::metrics::RouterMetrics;
use crate::policies::LoadBalancingPolicy;
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterTrait, WorkerManagement};
use crate::tokenizer::traits::Tokenizer;
use crate::tool_parser::ParserRegistry;
use async_trait::async_trait;
use axum::{
body::Body,
extract::Request,
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tracing::{info, warn};
/// gRPC PD (Prefill-Decode) router implementation for SGLang
#[allow(dead_code)] // Fields will be used once implementation is complete
pub struct GrpcPDRouter {
/// Prefill worker connections
prefill_workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
/// Decode worker connections
decode_workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
/// gRPC clients for prefill workers
prefill_grpc_clients: Arc<RwLock<HashMap<String, SglangSchedulerClient>>>,
/// gRPC clients for decode workers
decode_grpc_clients: Arc<RwLock<HashMap<String, SglangSchedulerClient>>>,
/// Load balancing policy for prefill
prefill_policy: Arc<dyn LoadBalancingPolicy>,
/// Load balancing policy for decode
decode_policy: Arc<dyn LoadBalancingPolicy>,
/// Tokenizer for handling text encoding/decoding
tokenizer: Arc<dyn Tokenizer>,
/// Reasoning parser factory for structured reasoning outputs
reasoning_parser_factory: ParserFactory,
/// Tool parser registry for function/tool calls
tool_parser_registry: &'static ParserRegistry,
/// Worker health checkers
_prefill_health_checker: Option<HealthChecker>,
_decode_health_checker: Option<HealthChecker>,
/// Configuration
timeout_secs: u64,
interval_secs: u64,
dp_aware: bool,
api_key: Option<String>,
retry_config: RetryConfig,
circuit_breaker_config: CircuitBreakerConfig,
}
impl GrpcPDRouter {
/// Create a new gRPC PD router
pub async fn new(
prefill_urls: Vec<(String, Option<u16>)>,
decode_urls: Vec<String>,
prefill_policy: Arc<dyn LoadBalancingPolicy>,
decode_policy: Arc<dyn LoadBalancingPolicy>,
ctx: &Arc<crate::server::AppContext>,
) -> Result<Self, String> {
// Update metrics
RouterMetrics::set_active_workers(prefill_urls.len() + decode_urls.len());
// Extract necessary components from context
let tokenizer = ctx
.tokenizer
.as_ref()
.ok_or_else(|| "gRPC PD router requires tokenizer".to_string())?
.clone();
let reasoning_parser_factory = ctx
.reasoning_parser_factory
.as_ref()
.ok_or_else(|| "gRPC PD router requires reasoning parser factory".to_string())?
.clone();
let tool_parser_registry = ctx
.tool_parser_registry
.ok_or_else(|| "gRPC PD router requires tool parser registry".to_string())?;
// Convert config CircuitBreakerConfig to core CircuitBreakerConfig
let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config();
let core_cb_config = CircuitBreakerConfig {
failure_threshold: circuit_breaker_config.failure_threshold,
success_threshold: circuit_breaker_config.success_threshold,
timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs),
window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs),
};
// Create gRPC clients for prefill workers
let mut prefill_grpc_clients = HashMap::new();
for (url, _bootstrap_port) in &prefill_urls {
match SglangSchedulerClient::connect(url).await {
Ok(client) => {
prefill_grpc_clients.insert(url.clone(), client);
info!("Connected to gRPC prefill worker at {}", url);
}
Err(e) => {
warn!("Failed to connect to gRPC prefill worker at {}: {}", url, e);
// Continue with other workers
}
}
}
// Create gRPC clients for decode workers
let mut decode_grpc_clients = HashMap::new();
for url in &decode_urls {
match SglangSchedulerClient::connect(url).await {
Ok(client) => {
decode_grpc_clients.insert(url.clone(), client);
info!("Connected to gRPC decode worker at {}", url);
}
Err(e) => {
warn!("Failed to connect to gRPC decode worker at {}: {}", url, e);
// Continue with other workers
}
}
}
if prefill_grpc_clients.is_empty() && decode_grpc_clients.is_empty() {
return Err("Failed to connect to any gRPC workers".to_string());
}
// Create Prefill Worker trait objects with gRPC connection mode
let prefill_workers: Vec<Box<dyn Worker>> = prefill_urls
.iter()
.map(|(url, bootstrap_port)| {
let worker = BasicWorker::with_connection_mode(
url.clone(),
WorkerType::Prefill {
bootstrap_port: *bootstrap_port,
},
crate::core::ConnectionMode::Grpc {
port: *bootstrap_port,
},
)
.with_circuit_breaker_config(core_cb_config.clone())
.with_health_config(HealthConfig {
timeout_secs: ctx.router_config.health_check.timeout_secs,
check_interval_secs: ctx.router_config.health_check.check_interval_secs,
endpoint: ctx.router_config.health_check.endpoint.clone(),
failure_threshold: ctx.router_config.health_check.failure_threshold,
success_threshold: ctx.router_config.health_check.success_threshold,
});
Box::new(worker) as Box<dyn Worker>
})
.collect();
// Create Decode Worker trait objects with gRPC connection mode
let decode_workers: Vec<Box<dyn Worker>> = decode_urls
.iter()
.map(|url| {
let worker = BasicWorker::with_connection_mode(
url.clone(),
WorkerType::Decode,
crate::core::ConnectionMode::Grpc { port: None },
)
.with_circuit_breaker_config(core_cb_config.clone())
.with_health_config(HealthConfig {
timeout_secs: ctx.router_config.health_check.timeout_secs,
check_interval_secs: ctx.router_config.health_check.check_interval_secs,
endpoint: ctx.router_config.health_check.endpoint.clone(),
failure_threshold: ctx.router_config.health_check.failure_threshold,
success_threshold: ctx.router_config.health_check.success_threshold,
});
Box::new(worker) as Box<dyn Worker>
})
.collect();
// Initialize policies with workers if needed
if let Some(cache_aware) = prefill_policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_aware.init_workers(&prefill_workers);
}
if let Some(cache_aware) = decode_policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_aware.init_workers(&decode_workers);
}
let prefill_workers = Arc::new(RwLock::new(prefill_workers));
let decode_workers = Arc::new(RwLock::new(decode_workers));
let prefill_health_checker = crate::core::start_health_checker(
Arc::clone(&prefill_workers),
ctx.router_config.worker_startup_check_interval_secs,
);
let decode_health_checker = crate::core::start_health_checker(
Arc::clone(&decode_workers),
ctx.router_config.worker_startup_check_interval_secs,
);
Ok(GrpcPDRouter {
prefill_workers,
decode_workers,
prefill_grpc_clients: Arc::new(RwLock::new(prefill_grpc_clients)),
decode_grpc_clients: Arc::new(RwLock::new(decode_grpc_clients)),
prefill_policy,
decode_policy,
tokenizer,
reasoning_parser_factory,
tool_parser_registry,
_prefill_health_checker: Some(prefill_health_checker),
_decode_health_checker: Some(decode_health_checker),
timeout_secs: ctx.router_config.worker_startup_timeout_secs,
interval_secs: ctx.router_config.worker_startup_check_interval_secs,
dp_aware: ctx.router_config.dp_aware,
api_key: ctx.router_config.api_key.clone(),
retry_config: ctx.router_config.effective_retry_config(),
circuit_breaker_config: core_cb_config,
})
}
}
impl std::fmt::Debug for GrpcPDRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GrpcPDRouter")
.field(
"prefill_workers_count",
&self.prefill_workers.read().unwrap().len(),
)
.field(
"decode_workers_count",
&self.decode_workers.read().unwrap().len(),
)
.field("timeout_secs", &self.timeout_secs)
.field("interval_secs", &self.interval_secs)
.field("dp_aware", &self.dp_aware)
.finish()
}
}
#[async_trait]
impl RouterTrait for GrpcPDRouter {
fn as_any(&self) -> &dyn std::any::Any {
self
}
async fn health(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn health_generate(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_server_info(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_models(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_model_info(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_generate(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::GenerateRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_chat(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::ChatCompletionRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_completion(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::CompletionRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_embeddings(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn flush_cache(&self) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_worker_loads(&self) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
fn router_type(&self) -> &'static str {
"grpc_pd"
}
fn readiness(&self) -> Response {
(StatusCode::SERVICE_UNAVAILABLE).into_response()
}
}
#[async_trait]
impl WorkerManagement for GrpcPDRouter {
async fn add_worker(&self, _worker_url: &str) -> Result<String, String> {
Err("Not implemented".to_string())
}
fn remove_worker(&self, _worker_url: &str) {}
fn get_worker_urls(&self) -> Vec<String> {
vec![]
}
}

266
sgl-router/src/routers/grpc/router.rs Normal file
View File

@@ -0,0 +1,266 @@
// gRPC Router Implementation
use crate::config::types::RetryConfig;
use crate::core::{
BasicWorker, CircuitBreakerConfig, HealthChecker, HealthConfig, Worker, WorkerType,
};
use crate::grpc::SglangSchedulerClient;
use crate::metrics::RouterMetrics;
use crate::policies::LoadBalancingPolicy;
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterTrait, WorkerManagement};
use crate::tokenizer::traits::Tokenizer;
use crate::tool_parser::ParserRegistry;
use async_trait::async_trait;
use axum::{
body::Body,
extract::Request,
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tracing::{info, warn};
/// gRPC router implementation for SGLang
#[allow(dead_code)] // Fields will be used once implementation is complete
pub struct GrpcRouter {
/// Worker connections
workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
/// gRPC clients for each worker
grpc_clients: Arc<RwLock<HashMap<String, SglangSchedulerClient>>>,
/// Load balancing policy
policy: Arc<dyn LoadBalancingPolicy>,
/// Tokenizer for handling text encoding/decoding
tokenizer: Arc<dyn Tokenizer>,
/// Reasoning parser factory for structured reasoning outputs
reasoning_parser_factory: ParserFactory,
/// Tool parser registry for function/tool calls
tool_parser_registry: &'static ParserRegistry,
/// Worker health checker
_health_checker: Option<HealthChecker>,
/// Configuration
timeout_secs: u64,
interval_secs: u64,
dp_aware: bool,
api_key: Option<String>,
retry_config: RetryConfig,
circuit_breaker_config: CircuitBreakerConfig,
}
impl GrpcRouter {
/// Create a new gRPC router
pub async fn new(
worker_urls: Vec<String>,
policy: Arc<dyn LoadBalancingPolicy>,
ctx: &Arc<crate::server::AppContext>,
) -> Result<Self, String> {
// Update metrics
RouterMetrics::set_active_workers(worker_urls.len());
// Extract necessary components from context
let tokenizer = ctx
.tokenizer
.as_ref()
.ok_or_else(|| "gRPC router requires tokenizer".to_string())?
.clone();
let reasoning_parser_factory = ctx
.reasoning_parser_factory
.as_ref()
.ok_or_else(|| "gRPC router requires reasoning parser factory".to_string())?
.clone();
let tool_parser_registry = ctx
.tool_parser_registry
.ok_or_else(|| "gRPC router requires tool parser registry".to_string())?;
// Convert config CircuitBreakerConfig to core CircuitBreakerConfig
let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config();
let core_cb_config = CircuitBreakerConfig {
failure_threshold: circuit_breaker_config.failure_threshold,
success_threshold: circuit_breaker_config.success_threshold,
timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs),
window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs),
};
// Create gRPC clients for each worker
let mut grpc_clients = HashMap::new();
for url in &worker_urls {
match SglangSchedulerClient::connect(url).await {
Ok(client) => {
grpc_clients.insert(url.clone(), client);
info!("Connected to gRPC worker at {}", url);
}
Err(e) => {
warn!("Failed to connect to gRPC worker at {}: {}", url, e);
// Continue with other workers
}
}
}
if grpc_clients.is_empty() {
return Err("Failed to connect to any gRPC workers".to_string());
}
// Create Worker trait objects with gRPC connection mode
let mut workers: Vec<Box<dyn Worker>> = Vec::new();
// Move clients from the HashMap to the workers
for url in &worker_urls {
if let Some(client) = grpc_clients.remove(url) {
let worker = BasicWorker::with_connection_mode(
url.clone(),
WorkerType::Regular,
crate::core::ConnectionMode::Grpc { port: None },
)
.with_circuit_breaker_config(core_cb_config.clone())
.with_health_config(HealthConfig {
timeout_secs: ctx.router_config.health_check.timeout_secs,
check_interval_secs: ctx.router_config.health_check.check_interval_secs,
endpoint: ctx.router_config.health_check.endpoint.clone(),
failure_threshold: ctx.router_config.health_check.failure_threshold,
success_threshold: ctx.router_config.health_check.success_threshold,
})
.with_grpc_client(client);
workers.push(Box::new(worker) as Box<dyn Worker>);
} else {
warn!("No gRPC client for worker {}, skipping", url);
}
}
// Initialize policy with workers if needed
if let Some(cache_aware) = policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_aware.init_workers(&workers);
}
let workers = Arc::new(RwLock::new(workers));
let health_checker = crate::core::start_health_checker(
Arc::clone(&workers),
ctx.router_config.worker_startup_check_interval_secs,
);
Ok(GrpcRouter {
workers,
grpc_clients: Arc::new(RwLock::new(grpc_clients)),
policy,
tokenizer,
reasoning_parser_factory,
tool_parser_registry,
_health_checker: Some(health_checker),
timeout_secs: ctx.router_config.worker_startup_timeout_secs,
interval_secs: ctx.router_config.worker_startup_check_interval_secs,
dp_aware: ctx.router_config.dp_aware,
api_key: ctx.router_config.api_key.clone(),
retry_config: ctx.router_config.effective_retry_config(),
circuit_breaker_config: core_cb_config,
})
}
}
impl std::fmt::Debug for GrpcRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GrpcRouter")
.field("workers_count", &self.workers.read().unwrap().len())
.field("timeout_secs", &self.timeout_secs)
.field("interval_secs", &self.interval_secs)
.field("dp_aware", &self.dp_aware)
.finish()
}
}
#[async_trait]
impl RouterTrait for GrpcRouter {
fn as_any(&self) -> &dyn std::any::Any {
self
}
async fn health(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn health_generate(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_server_info(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_models(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_model_info(&self, _req: Request<Body>) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_generate(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::GenerateRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_chat(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::ChatCompletionRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_completion(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::CompletionRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_embeddings(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn flush_cache(&self) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn get_worker_loads(&self) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
fn router_type(&self) -> &'static str {
"grpc"
}
fn readiness(&self) -> Response {
(StatusCode::SERVICE_UNAVAILABLE).into_response()
}
}
#[async_trait]
impl WorkerManagement for GrpcRouter {
async fn add_worker(&self, _worker_url: &str) -> Result<String, String> {
Err("Not implemented".to_string())
}
fn remove_worker(&self, _worker_url: &str) {}
fn get_worker_urls(&self) -> Vec<String> {
self.workers
.read()
.unwrap()
.iter()
.map(|w| w.url().to_string())
.collect()
}
}

53
sgl-router/src/routers/header_utils.rs Normal file
View File

@@ -0,0 +1,53 @@
use axum::body::Body;
use axum::extract::Request;
use axum::http::HeaderMap;
/// Copy request headers to a Vec of name-value string pairs
/// Used for forwarding headers to backend workers
pub fn copy_request_headers(req: &Request<Body>) -> Vec<(String, String)> {
req.headers()
.iter()
.filter_map(|(name, value)| {
// Convert header value to string, skipping non-UTF8 headers
value
.to_str()
.ok()
.map(|v| (name.to_string(), v.to_string()))
})
.collect()
}
/// Convert headers from reqwest Response to axum HeaderMap
/// Filters out hop-by-hop headers that shouldn't be forwarded
pub fn preserve_response_headers(reqwest_headers: &HeaderMap) -> HeaderMap {
let mut headers = HeaderMap::new();
for (name, value) in reqwest_headers.iter() {
// Skip hop-by-hop headers that shouldn't be forwarded
let name_str = name.as_str().to_lowercase();
if should_forward_header(&name_str) {
// The original name and value are already valid, so we can just clone them
headers.insert(name.clone(), value.clone());
}
}
headers
}
/// Determine if a header should be forwarded from backend to client
fn should_forward_header(name: &str) -> bool {
// List of headers that should NOT be forwarded (hop-by-hop headers)
!matches!(
name,
"connection" |
"keep-alive" |
"proxy-authenticate" |
"proxy-authorization" |
"te" |
"trailers" |
"transfer-encoding" |
"upgrade" |
"content-encoding" | // Let axum/hyper handle encoding
"host" // Should not forward the backend's host header
)
}

6
sgl-router/src/routers/http/mod.rs Normal file
View File

@@ -0,0 +1,6 @@
//! HTTP router implementations
pub mod openai_router;
pub mod pd_router;
pub mod pd_types;
pub mod router;

379
sgl-router/src/routers/http/openai_router.rs Normal file
View File

@@ -0,0 +1,379 @@
//! OpenAI router implementation (reqwest-based)
use crate::config::CircuitBreakerConfig;
use crate::core::{CircuitBreaker, CircuitBreakerConfig as CoreCircuitBreakerConfig};
use crate::protocols::spec::{ChatCompletionRequest, CompletionRequest, GenerateRequest};
use async_trait::async_trait;
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
};
use futures_util::StreamExt;
use std::{
any::Any,
sync::atomic::{AtomicBool, Ordering},
};
/// Router for OpenAI backend
#[derive(Debug)]
pub struct OpenAIRouter {
/// HTTP client for upstream OpenAI-compatible API
client: reqwest::Client,
/// Base URL for identification (no trailing slash)
base_url: String,
/// Circuit breaker
circuit_breaker: CircuitBreaker,
/// Health status
healthy: AtomicBool,
}
impl OpenAIRouter {
/// Create a new OpenAI router
pub async fn new(
base_url: String,
circuit_breaker_config: Option<CircuitBreakerConfig>,
) -> Result<Self, String> {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(300))
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
let base_url = base_url.trim_end_matches('/').to_string();
// Convert circuit breaker config
let core_cb_config = circuit_breaker_config
.map(|cb| CoreCircuitBreakerConfig {
failure_threshold: cb.failure_threshold,
success_threshold: cb.success_threshold,
timeout_duration: std::time::Duration::from_secs(cb.timeout_duration_secs),
window_duration: std::time::Duration::from_secs(cb.window_duration_secs),
})
.unwrap_or_default();
let circuit_breaker = CircuitBreaker::with_config(core_cb_config);
Ok(Self {
client,
base_url,
circuit_breaker,
healthy: AtomicBool::new(true),
})
}
}
#[async_trait]
impl super::super::WorkerManagement for OpenAIRouter {
async fn add_worker(&self, _worker_url: &str) -> Result<String, String> {
Err("Cannot add workers to OpenAI router".to_string())
}
fn remove_worker(&self, _worker_url: &str) {
// No-op for OpenAI router
}
fn get_worker_urls(&self) -> Vec<String> {
vec![self.base_url.clone()]
}
}
#[async_trait]
impl super::super::RouterTrait for OpenAIRouter {
fn as_any(&self) -> &dyn Any {
self
}
async fn health(&self, _req: Request<Body>) -> Response {
// Simple upstream probe: GET {base}/v1/models without auth
let url = format!("{}/v1/models", self.base_url);
match self
.client
.get(&url)
.timeout(std::time::Duration::from_secs(2))
.send()
.await
{
Ok(resp) => {
let code = resp.status();
// Treat success and auth-required as healthy (endpoint reachable)
if code.is_success() || code.as_u16() == 401 || code.as_u16() == 403 {
(StatusCode::OK, "OK").into_response()
} else {
(
StatusCode::SERVICE_UNAVAILABLE,
format!("Upstream status: {}", code),
)
.into_response()
}
}
Err(e) => (
StatusCode::SERVICE_UNAVAILABLE,
format!("Upstream error: {}", e),
)
.into_response(),
}
}
async fn health_generate(&self, _req: Request<Body>) -> Response {
// For OpenAI, health_generate is the same as health
self.health(_req).await
}
async fn get_server_info(&self, _req: Request<Body>) -> Response {
let info = serde_json::json!({
"router_type": "openai",
"workers": 1,
"base_url": &self.base_url
});
(StatusCode::OK, info.to_string()).into_response()
}
async fn get_models(&self, req: Request<Body>) -> Response {
// Proxy to upstream /v1/models; forward Authorization header if provided
let headers = req.headers();
let mut upstream = self.client.get(format!("{}/v1/models", self.base_url));
if let Some(auth) = headers
.get("authorization")
.or_else(|| headers.get("Authorization"))
{
upstream = upstream.header("Authorization", auth);
}
match upstream.send().await {
Ok(res) => {
let status = StatusCode::from_u16(res.status().as_u16())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
let content_type = res.headers().get(CONTENT_TYPE).cloned();
match res.bytes().await {
Ok(body) => {
let mut response = Response::new(axum::body::Body::from(body));
*response.status_mut() = status;
if let Some(ct) = content_type {
response.headers_mut().insert(CONTENT_TYPE, ct);
}
response
}
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to read upstream response: {}", e),
)
.into_response(),
}
}
Err(e) => (
StatusCode::BAD_GATEWAY,
format!("Failed to contact upstream: {}", e),
)
.into_response(),
}
}
async fn get_model_info(&self, _req: Request<Body>) -> Response {
// Not directly supported without model param; return 501
(
StatusCode::NOT_IMPLEMENTED,
"get_model_info not implemented for OpenAI router",
)
.into_response()
}
async fn route_generate(
&self,
_headers: Option<&HeaderMap>,
_body: &GenerateRequest,
) -> Response {
// Generate endpoint is SGLang-specific, not supported for OpenAI backend
(
StatusCode::NOT_IMPLEMENTED,
"Generate endpoint not supported for OpenAI backend",
)
.into_response()
}
async fn route_chat(
&self,
headers: Option<&HeaderMap>,
body: &ChatCompletionRequest,
) -> Response {
if !self.circuit_breaker.can_execute() {
return (StatusCode::SERVICE_UNAVAILABLE, "Circuit breaker open").into_response();
}
// Serialize request body, removing SGLang-only fields
let mut payload = match serde_json::to_value(body) {
Ok(v) => v,
Err(e) => {
return (
StatusCode::BAD_REQUEST,
format!("Failed to serialize request: {}", e),
)
.into_response();
}
};
if let Some(obj) = payload.as_object_mut() {
for key in [
"top_k",
"min_p",
"min_tokens",
"regex",
"ebnf",
"stop_token_ids",
"no_stop_trim",
"ignore_eos",
"continue_final_message",
"skip_special_tokens",
"lora_path",
"session_params",
"separate_reasoning",
"stream_reasoning",
"chat_template_kwargs",
"return_hidden_states",
"repetition_penalty",
] {
obj.remove(key);
}
}
let url = format!("{}/v1/chat/completions", self.base_url);
let mut req = self.client.post(&url).json(&payload);
// Forward Authorization header if provided
if let Some(h) = headers {
if let Some(auth) = h.get("authorization").or_else(|| h.get("Authorization")) {
req = req.header("Authorization", auth);
}
}
// Accept SSE when stream=true
if body.stream {
req = req.header("Accept", "text/event-stream");
}
let resp = match req.send().await {
Ok(r) => r,
Err(e) => {
self.circuit_breaker.record_failure();
return (
StatusCode::SERVICE_UNAVAILABLE,
format!("Failed to contact upstream: {}", e),
)
.into_response();
}
};
let status = StatusCode::from_u16(resp.status().as_u16())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
if !body.stream {
// Capture Content-Type before consuming response body
let content_type = resp.headers().get(CONTENT_TYPE).cloned();
match resp.bytes().await {
Ok(body) => {
self.circuit_breaker.record_success();
let mut response = Response::new(axum::body::Body::from(body));
*response.status_mut() = status;
if let Some(ct) = content_type {
response.headers_mut().insert(CONTENT_TYPE, ct);
}
response
}
Err(e) => {
self.circuit_breaker.record_failure();
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to read response: {}", e),
)
.into_response()
}
}
} else {
// Stream SSE bytes to client
let stream = resp.bytes_stream();
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
tokio::spawn(async move {
let mut s = stream;
while let Some(chunk) = s.next().await {
match chunk {
Ok(bytes) => {
if tx.send(Ok(bytes)).is_err() {
break;
}
}
Err(e) => {
let _ = tx.send(Err(format!("Stream error: {}", e)));
break;
}
}
}
});
let mut response = Response::new(Body::from_stream(
tokio_stream::wrappers::UnboundedReceiverStream::new(rx),
));
*response.status_mut() = status;
response
.headers_mut()
.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
response
}
}
async fn route_completion(
&self,
_headers: Option<&HeaderMap>,
_body: &CompletionRequest,
) -> Response {
// Completion endpoint not implemented for OpenAI backend
(
StatusCode::NOT_IMPLEMENTED,
"Completion endpoint not implemented for OpenAI backend",
)
.into_response()
}
async fn flush_cache(&self) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"flush_cache not supported for OpenAI router",
)
.into_response()
}
async fn get_worker_loads(&self) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"get_worker_loads not supported for OpenAI router",
)
.into_response()
}
fn router_type(&self) -> &'static str {
"openai"
}
fn readiness(&self) -> Response {
if self.healthy.load(Ordering::Acquire) && self.circuit_breaker.can_execute() {
(StatusCode::OK, "Ready").into_response()
} else {
(StatusCode::SERVICE_UNAVAILABLE, "Not ready").into_response()
}
}
async fn route_embeddings(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Embeddings endpoint not implemented for OpenAI backend",
)
.into_response()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Rerank endpoint not implemented for OpenAI backend",
)
.into_response()
}
}

Some files were not shown because too many files have changed in this diff Show More