[model-gateway] change sgl-router to sgl-model-gateway (#14312)
This commit is contained in:
Simo Lin
@@ -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",
|
||||
]
|
||||
@@ -0,0 +1,181 @@
|
||||
[package]
|
||||
name = "sgl-model-gateway"
|
||||
version = "0.2.3"
|
||||
edition = "2021"
|
||||
|
||||
[features]
|
||||
default = ["grpc-client"]
|
||||
grpc-client = []
|
||||
grpc-server = []
|
||||
|
||||
vendored-openssl = ["openssl/vendored"]
|
||||
|
||||
[lints.rust]
|
||||
unused_qualifications = "warn"
|
||||
|
||||
[lib]
|
||||
name = "sgl_model_gateway"
|
||||
crate-type = ["rlib"]
|
||||
|
||||
[[bin]]
|
||||
name = "sgl-model-gateway"
|
||||
path = "src/main.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "smg"
|
||||
path = "src/main.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "amg"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
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 = { version = "1.0", default-features = false, features = [
|
||||
"std",
|
||||
"preserve_order",
|
||||
] }
|
||||
serde_bytes = "0.11"
|
||||
bytes = "1.8.0"
|
||||
rand = "0.9.2"
|
||||
reqwest = { version = "0.12.8", features = ["stream", "blocking", "json", "rustls-tls"], default-features = false }
|
||||
futures-util = "0.3"
|
||||
futures = "0.3"
|
||||
dashmap = "6.1.0"
|
||||
lru = "0.16.2"
|
||||
blake3 = "1.5"
|
||||
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"] }
|
||||
ulid = "1.2.1"
|
||||
parking_lot = "0.12.4"
|
||||
rayon = "1.10"
|
||||
thiserror = "2.0.12"
|
||||
regex = "1.10"
|
||||
url = "2.5.4"
|
||||
validator = { version = "0.20.0", features = ["derive"] }
|
||||
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", features = ["unstable_machinery", "json", "builtins"] }
|
||||
minijinja-contrib = { version = "2.0", features = ["pycompat"] }
|
||||
rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
|
||||
openssl = "0.10.73"
|
||||
hf-hub = { version = "0.4.3", features = ["tokio"] }
|
||||
rmcp = { version = "0.8.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"
|
||||
oracle = { version = "0.6.3", features = ["chrono"] }
|
||||
subtle = "2.6"
|
||||
rustpython-parser = "0.4.0"
|
||||
num-traits = "0.2"
|
||||
image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] }
|
||||
ndarray = "0.16"
|
||||
base64 = "0.22"
|
||||
openai-harmony = { git = "https://github.com/openai/harmony", tag = "v0.0.4" }
|
||||
openmetrics-parser = "0.4.4"
|
||||
|
||||
# gRPC and Protobuf dependencies
|
||||
tonic = { version = "0.14.2", features = ["gzip", "transport"] }
|
||||
prost = "0.14.1"
|
||||
prost-types = "0.14.1"
|
||||
tonic-prost = "0.14.2"
|
||||
deadpool = { version = "0.12", features = ["managed", "rt_tokio_1"] }
|
||||
backoff = { version = "0.4", features = ["tokio"] }
|
||||
strum = { version = "0.26", features = ["derive"] }
|
||||
bitflags = "2.10.0"
|
||||
once_cell = "1.21.3"
|
||||
tokio-postgres = { version = "0.7.15", features = ["runtime","with-chrono-0_4","with-serde_json-1","array-impls"] }
|
||||
deadpool-postgres = "0.14.1"
|
||||
|
||||
|
||||
# wasm dependencies
|
||||
sha2 = "0.10"
|
||||
wasmtime = { version = "38.0", features = ["component-model", "async"] }
|
||||
wasmtime-wasi = "38.0"
|
||||
async-channel = "2.5"
|
||||
|
||||
[build-dependencies]
|
||||
tonic-prost-build = "0.14.2"
|
||||
prost-build = "0.14.1"
|
||||
chrono = { version = "0.4", features = ["clock"] }
|
||||
toml = "0.9"
|
||||
|
||||
[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"
|
||||
wasm-encoder = "0.242"
|
||||
npyz = { version = "0.8", features = ["npz"] } # For reading numpy .npz files in golden tests
|
||||
|
||||
[[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]
|
||||
opt-level = "z" # Optimize for size
|
||||
lto = "fat" # Full LTO for smaller binaries
|
||||
codegen-units = 1 # Better optimization, slower compile
|
||||
strip = true # Strip debug symbols
|
||||
|
||||
[profile.ci]
|
||||
inherits = "release"
|
||||
opt-level = 2 # Lighter optimization (still fast runtime, much faster compile)
|
||||
lto = "thin" # Thin LTO - good balance
|
||||
codegen-units = 16 # More parallelization for faster builds
|
||||
strip = true
|
||||
|
||||
[profile.dev]
|
||||
opt-level = 0
|
||||
debug = 1
|
||||
split-debuginfo = "unpacked"
|
||||
incremental = true
|
||||
codegen-units = 256
|
||||
|
||||
[profile.dev.package."*"]
|
||||
opt-level = 2
|
||||
debug = false
|
||||
|
||||
|
||||
[profile.dev.build-override]
|
||||
opt-level = 3
|
||||
codegen-units = 1
|
||||
|
||||
[profile.dev-opt]
|
||||
inherits = "dev"
|
||||
opt-level = 1
|
||||
Symlink
+1
@@ -0,0 +1 @@
|
||||
../LICENSE
|
||||
@@ -0,0 +1,159 @@
|
||||
# Model Gateway Makefile
|
||||
# Provides convenient shortcuts for common development tasks
|
||||
|
||||
# Python bindings directory
|
||||
PYTHON_DIR := bindings/python
|
||||
|
||||
# Auto-detect CPU cores and cap at reasonable limit to avoid thread exhaustion
|
||||
# Can be overridden: make python-dev JOBS=4
|
||||
NPROC := $(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 8)
|
||||
JOBS ?= $(shell echo $$(($(NPROC) > 16 ? 16 : $(NPROC))))
|
||||
|
||||
# Check if sccache is available and set RUSTC_WRAPPER accordingly
|
||||
SCCACHE := $(shell which sccache 2>/dev/null)
|
||||
ifdef SCCACHE
|
||||
export RUSTC_WRAPPER := $(SCCACHE)
|
||||
$(info Using sccache for compilation caching)
|
||||
else
|
||||
$(info sccache not found. Install it for faster builds: cargo install sccache)
|
||||
endif
|
||||
|
||||
.PHONY: help build test clean docs check fmt dev-setup pre-commit setup-sccache sccache-stats sccache-clean sccache-stop \
|
||||
python-dev python-build python-build-release python-install python-clean python-test python-check \
|
||||
release-notes
|
||||
|
||||
help: ## Show this help message
|
||||
@echo "Model Gateway 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 Model Gateway..."
|
||||
@cargo build --release
|
||||
|
||||
test: ## Run all tests
|
||||
@echo "Running tests..."
|
||||
@cargo test
|
||||
|
||||
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 --all-targets --all-features -- -D warnings
|
||||
|
||||
fmt: ## Format code with rustfmt
|
||||
@echo "Formatting code..."
|
||||
@rustup run nightly cargo fmt
|
||||
|
||||
# Development workflow shortcuts
|
||||
dev-setup: build test ## Set up development environment
|
||||
@echo "Development environment ready!"
|
||||
|
||||
pre-commit: fmt check test ## Run pre-commit checks
|
||||
@echo "Pre-commit checks passed!"
|
||||
|
||||
# sccache management targets
|
||||
setup-sccache: ## Install and configure sccache
|
||||
@echo "Setting up sccache..."
|
||||
@./scripts/setup-sccache.sh
|
||||
|
||||
sccache-stats: ## Show sccache statistics
|
||||
@if [ -n "$(SCCACHE)" ]; then \
|
||||
echo "sccache statistics:"; \
|
||||
sccache -s; \
|
||||
else \
|
||||
echo "sccache not installed. Run 'make setup-sccache' to install it."; \
|
||||
fi
|
||||
|
||||
sccache-clean: ## Clear sccache cache
|
||||
@if [ -n "$(SCCACHE)" ]; then \
|
||||
echo "Clearing sccache cache..."; \
|
||||
sccache -C; \
|
||||
echo "sccache cache cleared"; \
|
||||
else \
|
||||
echo "sccache not installed"; \
|
||||
fi
|
||||
|
||||
sccache-stop: ## Stop the sccache server
|
||||
@if [ -n "$(SCCACHE)" ]; then \
|
||||
echo "Stopping sccache server..."; \
|
||||
sccache --stop-server || true; \
|
||||
else \
|
||||
echo "sccache not installed"; \
|
||||
fi
|
||||
|
||||
# Python bindings (maturin) targets
|
||||
python-dev: ## Build Python bindings in development mode (fast, debug build)
|
||||
@echo "Building Python bindings in development mode (using $(JOBS) parallel jobs with sccache)..."
|
||||
@cd $(PYTHON_DIR) && CARGO_BUILD_JOBS=$(JOBS) maturin develop
|
||||
|
||||
python-build: ## Build Python wheel (release mode with vendored OpenSSL)
|
||||
@echo "Building Python wheel (release, vendored OpenSSL, using $(JOBS) parallel jobs with sccache)..."
|
||||
@cd $(PYTHON_DIR) && CARGO_BUILD_JOBS=$(JOBS) maturin build --release --out dist --features vendored-openssl
|
||||
|
||||
python-build-release: python-build ## Alias for python-build
|
||||
|
||||
python-install: python-build ## Build and install Python wheel
|
||||
@echo "Installing Python wheel..."
|
||||
@pip install --force-reinstall $(PYTHON_DIR)/dist/*.whl
|
||||
@echo "Python package installed!"
|
||||
|
||||
python-clean: ## Clean Python build artifacts
|
||||
@echo "Cleaning Python build artifacts..."
|
||||
@rm -rf $(PYTHON_DIR)/dist/
|
||||
@rm -rf $(PYTHON_DIR)/target/
|
||||
@rm -rf $(PYTHON_DIR)/sglang_router.egg-info/
|
||||
@rm -rf $(PYTHON_DIR)/sglang_router/__pycache__/
|
||||
@find $(PYTHON_DIR) -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true
|
||||
@find $(PYTHON_DIR) -name "*.pyc" -delete 2>/dev/null || true
|
||||
@echo "Python build artifacts cleaned!"
|
||||
|
||||
python-test: ## Run Python tests
|
||||
@echo "Running Python tests..."
|
||||
@pytest py_test/ -v
|
||||
|
||||
python-check: ## Check Python package with twine
|
||||
@echo "Checking Python package..."
|
||||
@cd $(PYTHON_DIR) && CARGO_BUILD_JOBS=$(JOBS) maturin build --release --out dist --features vendored-openssl
|
||||
@pip install twine 2>/dev/null || true
|
||||
@twine check $(PYTHON_DIR)/dist/*
|
||||
@echo "Python package check passed!"
|
||||
|
||||
# Combined shortcuts
|
||||
dev: python-dev ## Quick development setup (build Python bindings in dev mode)
|
||||
|
||||
install: python-install ## Build and install everything
|
||||
|
||||
# Release management
|
||||
release-notes: ## Generate release notes for gateway (usage: make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0)
|
||||
@if [ -z "$(PREV)" ] || [ -z "$(CURR)" ]; then \
|
||||
echo "Usage: make release-notes PREV=<previous-tag> CURR=<current-tag>"; \
|
||||
echo "Example: make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0"; \
|
||||
echo ""; \
|
||||
echo "Options:"; \
|
||||
echo " OUTPUT=<file> Save to file (default: stdout)"; \
|
||||
echo " CREATE_RELEASE=1 Create GitHub draft release via gh CLI (default: draft)"; \
|
||||
echo " DRAFT=0 Publish release immediately (skip draft)"; \
|
||||
exit 1; \
|
||||
fi
|
||||
@ARGS="$(PREV) $(CURR)"; \
|
||||
if [ -n "$(OUTPUT)" ]; then \
|
||||
ARGS="$$ARGS --output $(OUTPUT)"; \
|
||||
fi; \
|
||||
if [ "$(CREATE_RELEASE)" = "1" ]; then \
|
||||
ARGS="$$ARGS --create-release"; \
|
||||
if [ "$(DRAFT)" = "0" ]; then \
|
||||
ARGS="$$ARGS --no-draft"; \
|
||||
fi; \
|
||||
fi; \
|
||||
./scripts/generate_gateway_release_notes.sh $$ARGS
|
||||
@@ -0,0 +1,676 @@
|
||||
# SGLang Model Gateway
|
||||
|
||||
High-performance model routing control and data plane for large-scale LLM deployments. The gateway orchestrates fleets of workers, balances traffic across HTTP and gRPC backends, and exposes OpenAI-compatible APIs with pluggable history storage and tool integrations—while remaining deeply optimized for the SGLang serving runtime.
|
||||
|
||||
## Overview
|
||||
- Unified control plane for registering, monitoring, and orchestrating prefill, decode, and regular workers across heterogeneous model fleets.
|
||||
- Data plane that routes requests across HTTP, PD (prefill/decode), gRPC, and OpenAI-compatible backends with shared reliability features.
|
||||
- Industry-first gRPC pipeline with native Rust tokenization, reasoning, and tool-call execution for high-throughput OpenAI-compatible serving.
|
||||
- Multi-model inference gateway mode (`--enable-igw`) that runs several routers at once and applies per-model policies.
|
||||
- Conversation, response, and chat-history connectors that centralize state at the router, enabling compliant sharing across models/MCP loops with in-memory, no-op, or Oracle ATP storage options.
|
||||
- Built-in reliability primitives: retries with exponential backoff, circuit breakers, token-bucket rate limiting, and queuing.
|
||||
- First-class observability with structured logging and Prometheus metrics.
|
||||
|
||||
### Architecture at a Glance
|
||||
**Control Plane**
|
||||
- Worker Manager validates workers, discovers capabilities, and keeps the registry in sync.
|
||||
- Job Queue serializes background operations (add/remove) and exposes status via `/workers/{url}`.
|
||||
- Background health checker and load monitor keep circuit breakers and policies informed.
|
||||
- Optional Kubernetes service discovery keeps the registry aligned with pods.
|
||||
|
||||
**Data Plane**
|
||||
- SGLang HTTP routers for regular and PD (prefill/decode) traffic with policy-aware selection.
|
||||
- SGLang gRPC router and pipeline that stream tokenized requests through SRT gRPC workers with fully Rust tokenizer, reasoning parser, and tool parser implementations for maximal OpenAI API performance, supporting both single-stage and PD serving topologies.
|
||||
- OpenAI router that proxies OpenAI-style requests, responses, and conversations to remote vendors (OpenAI, xAI, Gemini, and other OpenAI-compatible providers) while preserving streaming/SSE semantics.
|
||||
- Router Manager coordinates multiple router implementations when IGW is enabled.
|
||||
- Resilience layer delivers token-bucket rate limiting, request queuing, retry executor, and per-worker circuit breakers to keep traffic flowing through failures.
|
||||
- Advanced load balancing with cache-aware request reuse, load-aware (power-of-two) selection, and per-model policy overrides.
|
||||
|
||||
## Feature Highlights
|
||||
- Multiple load balancing strategies (`random`, `round_robin`, `cache_aware`, `power_of_two`) with DP-aware scheduling.
|
||||
- Multi-model HTTP serving and inference gateway routing with model-specific policies.
|
||||
- Prefill/decode disaggregation, including bootstrap port handling and cache-aware merging.
|
||||
- gRPC routing with fully Rust tokenizer loading, reasoning parser selection, and tool parser integration for OpenAI-compatible endpoints—supporting streaming and non-streaming modes across DeepSeek, Llama, Kimi K2, Qwen, GPT-OSS, Mistral, Step-3, GLM4, and other reasoning-capable models.
|
||||
- OpenAI-compatible `/v1/chat/completions`, `/v1/responses`, `/v1/conversations`, `/v1/embeddings`, and `/v1/rerank` endpoints.
|
||||
- Native MCP client integration supporting all MCP transport protocols (STDIO, HTTP, SSE, and Streamable) for tool execution loops.
|
||||
- Pluggable history connectors: in-memory, disabled, or Oracle ATP (with pooling and credential support).
|
||||
- Reliability controls: retry with jitter, worker-scoped circuit breakers, token bucket limiter with optional queue, and cache flush APIs.
|
||||
- Service discovery for regular and PD workloads with independent selectors.
|
||||
- Prometheus metrics and structured tracing for every stage of routing.
|
||||
|
||||
## Documentation
|
||||
- **User Guide**: [docs.sglang.io/advanced_features/router.html](https://docs.sglang.io/advanced_features/router.html)
|
||||
- Additional guides, API references, and deployment patterns are continuously updated alongside SGLang releases.
|
||||
|
||||
## Installation
|
||||
### Prerequisites
|
||||
- **Rust and Cargo**
|
||||
```bash
|
||||
# Install rustup (Rust installer and version manager)
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
|
||||
|
||||
# Reload shell environment
|
||||
source "$HOME/.cargo/env"
|
||||
|
||||
# Verify installation
|
||||
rustc --version
|
||||
cargo --version
|
||||
```
|
||||
- **Python** with `pip` and virtualenv tooling available.
|
||||
|
||||
### Rust Binary
|
||||
```bash
|
||||
# Build release binary
|
||||
cargo build --release
|
||||
```
|
||||
|
||||
### Python Package
|
||||
```bash
|
||||
pip install maturin
|
||||
|
||||
# Fast development mode (debug build, no wheel, instant)
|
||||
# Uses system OpenSSL (requires libssl-dev/openssl-devel)
|
||||
cd bindings/python
|
||||
maturin develop
|
||||
|
||||
# Production build (optimized, creates wheel)
|
||||
# Uses vendored OpenSSL (cross-platform compatibility)
|
||||
cd bindings/python
|
||||
maturin build --release --out dist --features vendored-openssl
|
||||
pip install --force-reinstall dist/*.whl
|
||||
|
||||
# Development build with system OpenSSL (faster)
|
||||
# Requires: apt install libssl-dev pkg-config (Ubuntu/Debian)
|
||||
# or: yum install openssl-devel (RHEL/CentOS)
|
||||
cd bindings/python
|
||||
maturin build --release --out dist
|
||||
pip install --force-reinstall dist/*.whl
|
||||
```
|
||||
> **Note:** Python bindings are located in `bindings/python/` with their own Cargo.toml. Use `maturin develop` for fast iteration during development (builds in debug mode and installs directly). Use `maturin build --release --features vendored-openssl` for production wheels with full optimizations (opt-level="z", lto="fat") and cross-platform compatibility. The package uses abi3 support for Python 3.8+ compatibility.
|
||||
|
||||
## Checking Version
|
||||
|
||||
After installation, verify the installation and check version information:
|
||||
|
||||
```bash
|
||||
# Simple version (Rust binary)
|
||||
./target/release/sgl-model-gateway --version
|
||||
# or use aliases
|
||||
./target/release/smg --version
|
||||
./target/release/amg --version
|
||||
|
||||
# Full version info with build details
|
||||
./target/release/sgl-model-gateway --version-verbose
|
||||
|
||||
# Python CLI
|
||||
amg --version
|
||||
amg --version-verbose
|
||||
python3 -m sglang_router --version
|
||||
```
|
||||
|
||||
The `--version` (or `-V`) flag displays the version string. Use `--version-verbose` for comprehensive build information including Git commit, build time, compiler versions, and platform details.
|
||||
|
||||
## Quick Start
|
||||
### Regular HTTP Routing
|
||||
- **Rust binary**
|
||||
```bash
|
||||
./target/release/sgl-model-gateway \
|
||||
--worker-urls http://worker1:8000 http://worker2:8000 \
|
||||
--policy cache_aware
|
||||
```
|
||||
`cargo run --release -- …` provides the same behavior during development.
|
||||
- **Python launcher**
|
||||
```bash
|
||||
python3 -m sglang_router.launch_router \
|
||||
--worker-urls http://worker1:8000 http://worker2:8000 \
|
||||
--policy cache_aware
|
||||
```
|
||||
|
||||
### Prefill/Decode Disaggregation (PD)
|
||||
- **Rust binary**
|
||||
```bash
|
||||
./target/release/sgl-model-gateway \
|
||||
--pd-disaggregation \
|
||||
--prefill http://prefill1:30001 9001 \
|
||||
--prefill http://prefill2:30002 \
|
||||
--decode http://decode1:30011 \
|
||||
--decode http://decode2:30012 \
|
||||
--policy cache_aware \
|
||||
--prefill-policy cache_aware \
|
||||
--decode-policy power_of_two
|
||||
```
|
||||
- **Python launcher**
|
||||
```bash
|
||||
python3 -m sglang_router.launch_router \
|
||||
--pd-disaggregation \
|
||||
--prefill http://prefill1:30001 9001 \
|
||||
--prefill http://prefill2:30002 \
|
||||
--decode http://decode1:30011 \
|
||||
--decode http://decode2:30012 \
|
||||
--policy cache_aware
|
||||
```
|
||||
Prefill entries accept an optional bootstrap port. PD mode merges prefill metadata with decode outputs and streams results back to the client.
|
||||
|
||||
### Multi-Model Inference Gateway
|
||||
Enable IGW mode to route multiple models through a single router while applying per-model policies:
|
||||
```bash
|
||||
./target/release/sgl-model-gateway \
|
||||
--enable-igw \
|
||||
--policy cache_aware \
|
||||
--max-concurrent-requests 512
|
||||
|
||||
# Register workers dynamically
|
||||
curl -X POST http://localhost:30000/workers \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"url": "http://worker-a:8000",
|
||||
"model_id": "mistral",
|
||||
"priority": 10,
|
||||
"labels": {"tier": "gold"}
|
||||
}'
|
||||
|
||||
# Add another worker with a different model/policy hint
|
||||
curl -X POST http://localhost:30000/workers \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"url": "http://worker-b:8000",
|
||||
"model_id": "llama3",
|
||||
"priority": 20,
|
||||
"labels": {"policy": "power_of_two", "tier": "silver"}
|
||||
}'
|
||||
|
||||
# Inspect registered workers
|
||||
curl http://localhost:30000/workers
|
||||
```
|
||||
Sample response (http workers):
|
||||
```json
|
||||
{
|
||||
"workers": [
|
||||
{"id":"http://0.0.0.0:31378","url":"http://0.0.0.0:31378","model_id":"mistral","priority":50,"cost":1.0,"worker_type":"regular","is_healthy":true,"load":0,"connection_mode":"Http"},
|
||||
{"id":"http://0.0.0.0:34881","url":"http://0.0.0.0:34881","model_id":"llama3","priority":50,"cost":1.0,"worker_type":"regular","is_healthy":true,"load":0,"connection_mode":"Http"}
|
||||
],
|
||||
"total": 2,
|
||||
"stats": {
|
||||
"prefill_count": 0,
|
||||
"decode_count": 0,
|
||||
"regular_count": 2
|
||||
}
|
||||
}
|
||||
```
|
||||
Add more workers with the same API; include optional `labels` (for per-model policies) or `tokenizer_path` / `reasoning_parser` / `tool_parser` fields as needed. `/workers/{url}` exposes queued job status while background jobs finalize registration.
|
||||
|
||||
### gRPC Routing
|
||||
- **Rust binary**
|
||||
```bash
|
||||
./target/release/sgl-model-gateway \
|
||||
--worker-urls grpc://worker-grpc-0:31001 grpc://worker-grpc-1:31002 \
|
||||
--tokenizer-path /path/to/tokenizer.json \
|
||||
--reasoning-parser deepseek-r1 \
|
||||
--tool-call-parser json
|
||||
```
|
||||
- **Python router**
|
||||
```bash
|
||||
python3 -m sglang_router.launch_router \
|
||||
--worker-urls grpc://127.0.0.1:20000 \
|
||||
--model-path meta-llama/Llama-3.1-8B-Instruct \
|
||||
--host 0.0.0.0 \
|
||||
--port 8080
|
||||
```
|
||||
The gRPC router tokenizes inputs locally, supports tool-call parsing, and streams responses. It supports both regular HTTP-equivalent serving and PD (prefill/decode) serving when the worker registry contains PD workers. Provide `--model-path` or `--tokenizer-path` (HuggingFace ID or local directory) whenever connection mode resolves to gRPC.
|
||||
Use `--reasoning-parser` to select built-in reasoning pipelines (DeepSeek-R1, Qwen3, Step-3, GLM4, etc.) and `--tool-call-parser` for JSON/Pythonic/XML tool contracts in streaming or non-streaming modes.
|
||||
|
||||
### OpenAI Backend Mode
|
||||
Route requests to OpenAI or OpenAI-compatible endpoints:
|
||||
|
||||
```bash
|
||||
# Route to OpenAI API
|
||||
python3 -m sglang_router.launch_router \
|
||||
--backend openai \
|
||||
--worker-urls https://api.openai.com \
|
||||
|
||||
# Route to custom OpenAI-compatible endpoint (Gemini, xAI, etc.)
|
||||
python3 -m sglang_router.launch_router \
|
||||
--backend openai \
|
||||
--worker-urls http://my-openai-compatible-service:8000 \
|
||||
```
|
||||
|
||||
**Notes**
|
||||
- OpenAI backend mode acts as a proxy to a single remote endpoint; load balancing is not applied.
|
||||
- Provide exactly one `--worker-urls` entry per router instance.
|
||||
- The Rust binary supports the same flags (`./target/release/sgl-model-gateway --backend openai ...`).
|
||||
|
||||
### MCP Integration
|
||||
The SGL Model Gateway provides native Model Context Protocol (MCP) client integration, enabling tool calling across STDIO, SSE, and Streamable transports. MCP servers are configured via a YAML configuration file and registered at startup through the workflow engine.
|
||||
|
||||
#### Basic Usage
|
||||
```bash
|
||||
# Rust binary
|
||||
./target/release/sgl-model-gateway \
|
||||
--mcp-config-path /path/to/mcp-config.yaml \
|
||||
--worker-urls http://worker1:8000
|
||||
|
||||
# Python launcher
|
||||
python3 -m sglang_router.launch_router \
|
||||
--mcp-config-path /path/to/mcp-config.yaml \
|
||||
--worker-urls http://worker1:8000
|
||||
```
|
||||
|
||||
#### MCP Configuration File
|
||||
Create an MCP configuration file to define servers, transports, and connection settings:
|
||||
|
||||
```yaml
|
||||
servers:
|
||||
- name: "filesystem"
|
||||
command: "npx"
|
||||
args: ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
|
||||
protocol: "stdio"
|
||||
required: false
|
||||
|
||||
- name: "github"
|
||||
url: "https://api.github.com/mcp"
|
||||
token: "ghp_xxxxx"
|
||||
protocol: "sse"
|
||||
required: false
|
||||
|
||||
- name: "custom-tools"
|
||||
url: "https://tools.example.com/mcp"
|
||||
protocol: "streamable"
|
||||
required: true
|
||||
|
||||
pool:
|
||||
max_connections: 100
|
||||
idle_timeout: 300 # seconds
|
||||
|
||||
proxy:
|
||||
http: "http://proxy.internal:8080"
|
||||
https: "https://proxy.internal:8443"
|
||||
no_proxy: "localhost,127.0.0.1,*.internal"
|
||||
|
||||
inventory:
|
||||
enable_refresh: true
|
||||
tool_ttl: 300 # seconds - how long tools are considered fresh
|
||||
refresh_interval: 300 # seconds - background refresh interval
|
||||
```
|
||||
|
||||
#### Configuration Options
|
||||
|
||||
**Server Configuration** (`servers` array):
|
||||
- `name`: Unique identifier for the MCP server
|
||||
- `command` + `args`: For STDIO transport (local process execution)
|
||||
- `url`: For SSE or Streamable transports (HTTP/HTTPS endpoints)
|
||||
- `token`: Optional authentication token for HTTP-based transports
|
||||
- `protocol`: Protocol type (`"sse"`, `"streamable"`, or `"stdio"`)
|
||||
- `required`: If `true`, router fails to start if server is unreachable (default: `false`)
|
||||
- `envs`: Environment variables for STDIO processes (optional)
|
||||
- `proxy`: Per-server proxy override (set to `null` to bypass global proxy)
|
||||
|
||||
**Connection Pool** (`pool`):
|
||||
- `max_connections`: Maximum pooled connections for dynamic servers (default: 100)
|
||||
- `idle_timeout`: Idle connection timeout in seconds before cleanup (default: 300)
|
||||
|
||||
**Proxy Configuration** (`proxy`):
|
||||
- `http`/`https`: Proxy URLs for MCP server connections (not LLM traffic)
|
||||
- `no_proxy`: Comma-separated hosts to exclude from proxying (supports wildcards)
|
||||
- **Note**: Proxy settings are currently ignored for `streamable` transport. Use STDIO or SSE transports if proxy support is required.
|
||||
|
||||
**Inventory Settings** (`inventory`):
|
||||
- `enable_refresh`: Enable automatic background refresh of tool inventory (default: true)
|
||||
- `tool_ttl`: Tool cache TTL in seconds - how long tools are considered fresh (default: 300)
|
||||
- `refresh_interval`: Background refresh interval in seconds - proactive inventory refresh (default: 300)
|
||||
|
||||
#### Transport Types
|
||||
|
||||
**STDIO** (Local Process):
|
||||
```yaml
|
||||
name: "local-tools"
|
||||
command: "python"
|
||||
args: ["-m", "my_mcp_server"]
|
||||
envs:
|
||||
API_KEY: "secret"
|
||||
DEBUG: "true"
|
||||
```
|
||||
|
||||
**SSE** (Server-Sent Events):
|
||||
```yaml
|
||||
name: "remote-sse"
|
||||
url: "https://mcp.example.com/events"
|
||||
token: "bearer-token"
|
||||
protocol: "sse"
|
||||
```
|
||||
|
||||
**Streamable** (Bidirectional Streaming):
|
||||
```yaml
|
||||
name: "streaming-tools"
|
||||
url: "https://mcp.example.com/stream"
|
||||
protocol: "streamable"
|
||||
required: true
|
||||
```
|
||||
|
||||
#### Server Lifecycle
|
||||
- MCP servers are registered via the workflow engine with retry logic (100 attempts, 2-hour timeout for STDIO servers)
|
||||
- Discovery phase identifies tools, prompts, and resources
|
||||
- Tool inventory is cached with configurable TTL and periodic refresh
|
||||
- Failed optional servers log warnings; required servers halt startup
|
||||
- Static servers (from config) are permanent; dynamic servers (per-request) use connection pooling
|
||||
|
||||
Check Prometheus metrics for MCP activity (`mcp_*` metrics) and workflow job status via the admin API.
|
||||
|
||||
### Python Launcher (Router + Workers)
|
||||
Launch router and SGLang worker processes together; `launch_server` spins up workers (HTTP or gRPC) and the router in one shot.
|
||||
```bash
|
||||
python3 -m sglang_router.launch_server --host 0.0.0.0
|
||||
```
|
||||
Add flags as needed for production deployments:
|
||||
```bash
|
||||
python3 -m sglang_router.launch_server \
|
||||
--host 0.0.0.0 \
|
||||
--port 8080 \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--tp-size 1 \
|
||||
--dp-size 8 \
|
||||
--grpc-mode
|
||||
```
|
||||
Omit `--grpc-mode` to start HTTP workers; the router automatically configures worker URLs and schedules them based on the provided DP size.
|
||||
|
||||
### Mini Load Balancer (Debug)
|
||||
```bash
|
||||
python3 -m sglang_router.launch_router \
|
||||
--mini-lb \
|
||||
--pd-disaggregation \
|
||||
--prefill http://localhost:30001 \
|
||||
--decode http://localhost:30011
|
||||
```
|
||||
MiniLB forwards PD requests using simple random routing and is intended for local debugging only.
|
||||
|
||||
### Running Worker Servers
|
||||
Use upstream SGLang binaries to start dedicated worker processes.
|
||||
- **Prefill worker server (gRPC mode)**:
|
||||
```bash
|
||||
python3 -m sglang.launch_server \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--port 20000 \
|
||||
--tp-size 1 \
|
||||
--grpc-mode
|
||||
```
|
||||
Remove `--grpc-mode` for HTTP workers. Combine with the router commands above to register the worker via CLI flags or the control-plane API.
|
||||
|
||||
## Control Plane
|
||||
|
||||
### Worker Lifecycle & Job Queue
|
||||
- `JobQueue` handles asynchronous add/remove operations to avoid blocking clients.
|
||||
- `WorkerManager` inspects worker metadata (`/get_server_info`, `/get_model_info`), tracks load, and exposes `flush_cache` and `get_loads`.
|
||||
- Per-worker circuit breakers and health probes keep the registry healthy; load monitor feeds metrics to cache-aware and power-of-two policies.
|
||||
|
||||
### Administrative & Worker APIs
|
||||
| Method | Path | Description |
|
||||
|----------|------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `POST` | `/workers` | Queue worker registration (prefill/decode/regular). Body matches `WorkerConfigRequest`. Returns `202 Accepted` while the job queue processes the request. |
|
||||
| `GET` | `/workers` | List workers with health, load, policy metadata, and queued job status. |
|
||||
| `GET` | `/workers/{url}` | Inspect a specific worker or job queue entry. |
|
||||
| `DELETE` | `/workers/{url}` | Queue worker removal. |
|
||||
| `POST` | `/flush_cache` | Trigger cache flush across HTTP workers with success/failure breakdown. |
|
||||
| `GET` | `/get_loads` | Sample current load reported by each worker. |
|
||||
|
||||
All administrative routes inherit router API-key protection when `--api-key` is supplied. Job status includes `pending`, `processing`, and `failed` phases with timestamps.
|
||||
|
||||
### Service Discovery
|
||||
Enable Kubernetes discovery to reconcile workers automatically:
|
||||
```bash
|
||||
./target/release/sgl-model-gateway \
|
||||
--service-discovery \
|
||||
--selector app=sglang-worker role=inference \
|
||||
--service-discovery-namespace sglang-system \
|
||||
--service-discovery-port 8000
|
||||
```
|
||||
PD mode accepts dedicated selectors:
|
||||
```bash
|
||||
--pd-disaggregation \
|
||||
--prefill-selector app=sglang component=prefill \
|
||||
--decode-selector app=sglang component=decode \
|
||||
--service-discovery
|
||||
```
|
||||
Prefill pods can expose bootstrap ports via the `sglang.ai/bootstrap-port` annotation. RBAC must allow `get`, `list`, and `watch` on pods.
|
||||
|
||||
## Data Plane
|
||||
|
||||
### Router Capabilities (HTTP & gRPC)
|
||||
Both router stacks:
|
||||
- Share load-balancing policies (random, round-robin, cache-aware, power-of-two) with DP-aware scheduling, retries, circuit breakers, and rate limiting.
|
||||
- Record metrics per request, track running load, and integrate with the router-wide policy registry.
|
||||
|
||||
The HTTP router exposes the full OpenAI-compatible surface area (`/generate`, `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/responses`, `/v1/rerank`, etc.). The gRPC router delivers blazing-fast `/generate` and `/v1/chat/completions` today, with the remaining endpoints returning `501 Not Implemented` until their pipelines are finalised.
|
||||
|
||||
#### HTTP Router specifics
|
||||
- **Regular router** handles classic single-stage workers with per-model policy overrides.
|
||||
- **Prefill/Decode router** coordinates disaggregated prefill and decode workers, merges metadata, and manages streaming fan-in.
|
||||
|
||||
#### gRPC Router specifics
|
||||
- Industry-first fully Rust implementation of an OpenAI-compatible gRPC inference gateway, including tokenizer, reasoning parser, and tool parser execution in-process for maximum throughput.
|
||||
- Supports both single-stage and PD (prefill/decode) worker topologies; the router automatically selects the appropriate pipeline per model.
|
||||
- Provides the same `/v1/*` APIs as the HTTP router while streaming tokenized requests/responses directly to SRT gRPC workers.
|
||||
- Built-in reasoning parsers for DeepSeek, Qwen, Llama, Mistral, GPT-OSS, Step-3, GLM4, Kimi K2, and other structured-thought models.
|
||||
- Tool-call parsers for JSON, Pythonic, XML, and custom schemas with streaming and non-streaming execution loops.
|
||||
- Tokenizer factory supporting HuggingFace models, local tokenizer.json files, and chat template overrides (see `src/tokenizer`).
|
||||
- Explore the code paths in `src/reasoning_parser`, `src/tool_parser`, and `src/tokenizer` for the end-to-end Rust implementations that power gRPC mode.
|
||||
|
||||
### OpenAI Router
|
||||
- Proxies OpenAI-compatible chat completions and responses APIs, preserving headers and SSE streams end-to-end.
|
||||
- Supports `/v1/responses` background jobs with cancellation, deletion, and listing input items—enabling agentic, multi-turn orchestration without persisting data at remote vendor endpoints.
|
||||
- Conversation APIs (`/v1/conversations` and `/v1/conversations/{id}/items`) interact with the configured conversation storage backend for compliant chat-history management. Conversation state lives at the router tier, so the same history can drive different models or MCP loops without leaking data to upstream vendors.
|
||||
- Chat history, agentic multi-turn `/v1/responses`, and the native MCP client (STDIO/HTTP/SSE/Streamable transports) are designed to satisfy enterprise data-privacy requirements by keeping sensitive state within the router.
|
||||
|
||||
### Request Endpoints
|
||||
| Endpoint | Notes |
|
||||
|----------------------------------------------------------------------------------|------------------------------------------------------------|
|
||||
| `POST /generate` | SGLang generate API. |
|
||||
| `POST /v1/chat/completions` | OpenAI-compatible chat. Supports streaming and tool calls. |
|
||||
| `POST /v1/completions` | OpenAI-compatible text completions. |
|
||||
| `POST /v1/responses` | Create background responses, returns response IDs. |
|
||||
| `GET /v1/responses/{id}` | Retrieve stored responses. |
|
||||
| Conversation endpoints (`/v1/conversations`, `/v1/conversations/{id}`, `/v1/conversations/{id}/items`) | Manage chat history. |
|
||||
| `POST /v1/embeddings` | Forward embedding requests. |
|
||||
| `POST /v1/rerank`, `POST /rerank` | Ranking APIs. |
|
||||
|
||||
Public health endpoints (`/liveness`, `/readiness`, `/health`, `/health_generate`) reflect registry state; readiness ensures PD workers are paired and IGW has at least one healthy route.
|
||||
|
||||
## Conversations, Responses, and Data Connectors
|
||||
- `--history-backend memory` (default) stores responses and conversations in-process.
|
||||
- `--history-backend none` disables persistence while keeping APIs.
|
||||
- `--history-backend oracle` uses Oracle Autonomous Database; provide credentials via flags or environment variables.
|
||||
- `--history-backend postgres` uses PostgreSQL Database.
|
||||
- Conversation item storage mirrors the history backend (Oracle or memory). The same storage powers OpenAI `/responses` and conversation APIs.
|
||||
|
||||
### History Backend (OpenAI Router Mode)
|
||||
Store conversation and response data for tracking, debugging, or analytics.
|
||||
|
||||
> **Note:** History backends are currently supported only when running with `--backend openai`. gRPC mode support for the `/v1/responses` API is planned.
|
||||
|
||||
#### Available storage options
|
||||
- **Memory** (default): In-memory storage, fast but ephemeral.
|
||||
- **None**: No storage, minimal overhead.
|
||||
- **Oracle**: Persistent storage backed by Oracle Autonomous Database.
|
||||
- **Postgres**: Persistent storage backed by PostgreSQL Database.
|
||||
|
||||
```bash
|
||||
# Memory backend (default)
|
||||
python3 -m sglang_router.launch_router \
|
||||
--backend openai \
|
||||
--worker-urls https://api.openai.com \
|
||||
--history-backend memory
|
||||
|
||||
# No storage for maximum performance
|
||||
python3 -m sglang_router.launch_router \
|
||||
--backend openai \
|
||||
--worker-urls https://api.openai.com \
|
||||
--history-backend none
|
||||
|
||||
# Oracle ATP backend (see configuration below)
|
||||
python3 -m sglang_router.launch_router \
|
||||
--backend openai \
|
||||
--worker-urls https://api.openai.com \
|
||||
--history-backend oracle
|
||||
|
||||
# PostgreSQL backend
|
||||
python3 -m sglang_router.launch_router \
|
||||
--backend openai \
|
||||
--worker-urls https://api.openai.com \
|
||||
--history-backend postgres
|
||||
```
|
||||
|
||||
#### Oracle configuration
|
||||
Install the Oracle Instant Client and set `LD_LIBRARY_PATH` accordingly. Choose **one** connection method:
|
||||
```bash
|
||||
# Option 1: Full connection descriptor
|
||||
export ATP_DSN="(description=(address=(protocol=tcps)(port=1522)(host=adb.region.oraclecloud.com))(connect_data=(service_name=service_name)))"
|
||||
|
||||
# Option 2: TNS alias (requires wallet)
|
||||
export ATP_TNS_ALIAS="sglroutertestatp_high"
|
||||
export ATP_WALLET_PATH="/path/to/wallet"
|
||||
```
|
||||
Provide database credentials and optional pool sizing:
|
||||
```bash
|
||||
export ATP_USER="admin"
|
||||
export ATP_PASSWORD="YourPassword123"
|
||||
export ATP_POOL_MIN=4
|
||||
export ATP_POOL_MAX=32
|
||||
```
|
||||
|
||||
Router flags map to these values:
|
||||
- `--oracle-dsn` (env: `ATP_DSN`) or `--oracle-tns-alias` with `--oracle-wallet-path`.
|
||||
- `--oracle-user` / `--oracle-password` (`ATP_USER` / `ATP_PASSWORD`).
|
||||
- `--oracle-wallet-path` (`ATP_WALLET_PATH`) when using TNS alias.
|
||||
- `--oracle-pool-min`, `--oracle-pool-max`, `--oracle-pool-timeout-secs`.
|
||||
|
||||
Only one of `--oracle-dsn` or `--oracle-tns-alias` should be supplied.
|
||||
|
||||
## Reliability & Flow Control
|
||||
- **Retries**: Default max retries = 5 with exponential backoff (`--retry-max-retries`, `--retry-initial-backoff-ms`, `--retry-max-backoff-ms`, `--retry-backoff-multiplier`, `--retry-jitter-factor`). Retries trigger on 408/429/500/502/503/504.
|
||||
- **Circuit Breakers**: Per worker thresholds (`--cb-failure-threshold`, `--cb-success-threshold`, `--cb-timeout-duration-secs`, `--cb-window-duration-secs`). Disable via `--disable-circuit-breaker`.
|
||||
- **Rate Limiting**: Token bucket driven by `--max-concurrent-requests`. Set `--rate-limit-tokens-per-second` to override refill rate. Configure request queue via `--queue-size` and `--queue-timeout-secs`; queued requests observe FIFO order and respect cancellation.
|
||||
- **Health Checks**: Runtime probes via `--health-check-interval-secs`, `--health-check-timeout-secs`, failure/success thresholds, and `--health-check-endpoint`.
|
||||
- **Cache Management**: `/flush_cache` ensures LRU eviction when redeploying PD workers.
|
||||
|
||||
## Load Balancing Policies
|
||||
- `random`: uniform random worker selection.
|
||||
- `round_robin`: sequential rotation with atomic counters.
|
||||
- `cache_aware`: maintains a prefix tree of prompts to route repeat traffic and evens load with configurable thresholds (`--cache-threshold`, `--balance-abs-threshold`, `--balance-rel-threshold`, `--eviction-interval`, `--max-tree-size`).
|
||||
- `power_of_two`: chooses the lighter worker among two random candidates; integrates with `LoadMonitor`.
|
||||
Per-model overrides are available in PD mode (`--prefill-policy`, `--decode-policy`) and IGW mode via the worker registry.
|
||||
|
||||
## Observability
|
||||
- **Logging**: Structured tracing through `tracing` with optional file sink (`--log-dir`) and `--log-level` (`debug`, `info`, `warn`, `error`).
|
||||
- **Prometheus Metrics**: Enable with `--prometheus-host`/`--prometheus-port` (defaults to `0.0.0.0:29000`). Metrics cover request latency, retry behavior, circuit breaker states, worker health/load, queue depth, PD pipeline stats, tokenizer timings, and MCP activity.
|
||||
- **Request IDs**: Configurable headers via `--request-id-headers`; responses include `x-request-id`.
|
||||
- **CORS**: Set `--cors-allowed-origins` for browser access.
|
||||
|
||||
## Security
|
||||
|
||||
### Router and Worker API Keys
|
||||
- **Router API key (`--api-key`)** protects client access to router endpoints; all protected routes expect `Authorization: Bearer <key>`.
|
||||
- Workers listed in `--worker-urls` inherit the router API key automatically.
|
||||
- When adding workers dynamically, provide explicit API keys via payload or query string; they do **not** inherit automatically.
|
||||
|
||||
```bash
|
||||
# Router and initial workers share the same key
|
||||
python3 -m sglang_router.launch_router \
|
||||
--api-key "shared-api-key" \
|
||||
--worker-urls http://worker1:8000 http://worker2:8000
|
||||
|
||||
# Adding a worker without key while router has one triggers a warning and leaves the worker unprotected
|
||||
curl -X POST http://localhost:8080/add_worker?url=http://worker3:8000
|
||||
|
||||
# Add worker with explicit key
|
||||
curl -X POST "http://localhost:8080/add_worker?url=http://worker3:8000&api_key=worker3-specific-key"
|
||||
```
|
||||
|
||||
### Security Configurations
|
||||
1. **No Authentication** (default): Router and workers accept requests without keys—use only in trusted environments.
|
||||
2. **Router-only Authentication**: Provide `--api-key`; clients must present the key, router accesses workers without credentials.
|
||||
3. **Worker-only Authentication**: Router open to clients; each worker requires its own key. Supply keys when calling `/workers` or `/add_worker`.
|
||||
4. **Full Authentication**: Set router API key and provide per-worker keys. Example:
|
||||
```bash
|
||||
python3 -m sglang_router.launch_router --api-key "router-key"
|
||||
curl -H "Authorization: Bearer router-key" \
|
||||
-X POST http://localhost:8080/add_worker?url=http://worker:8000&api_key=worker-key
|
||||
```
|
||||
|
||||
### Important Notes
|
||||
- Initial workers declared via CLI inherit the router key; dynamic workers must supply keys explicitly.
|
||||
- Router logs a warning when a worker is registered without a key while the router expects authentication.
|
||||
- When router and workers share the same key, still include the key when invoking dynamic registration APIs.
|
||||
|
||||
## Development & Testing
|
||||
```bash
|
||||
# Build Rust components (debug mode, fast)
|
||||
cargo build
|
||||
|
||||
# Run Rust tests
|
||||
cargo test
|
||||
|
||||
# Fast Python development (rebuilds and installs in debug mode)
|
||||
cd bindings/python && maturin develop
|
||||
|
||||
# Run Python tests
|
||||
cd ../.. # Back to sgl-model-gateway root
|
||||
pytest py_test/
|
||||
```
|
||||
For production builds, use `maturin build --release --out dist` from the `bindings/python/` directory to create optimized wheels. During development, `maturin develop` rebuilds and installs instantly without creating wheel files. Use `python -m sglang_router.launch_server` to co-launch router and SGLang workers in small clusters for local validation.
|
||||
|
||||
---
|
||||
|
||||
## Release Management
|
||||
|
||||
### Creating Gateway Releases
|
||||
|
||||
Create releases for the Gateway/Router component with filtered commits:
|
||||
|
||||
```bash
|
||||
# Using make
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0
|
||||
|
||||
# Save to file
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0 OUTPUT=RELEASE_NOTES.md
|
||||
|
||||
# Create draft release (requires gh CLI, DEFAULT behavior)
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0 CREATE_RELEASE=1
|
||||
|
||||
# Publish release immediately (requires gh CLI)
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0 CREATE_RELEASE=1 DRAFT=0
|
||||
```
|
||||
|
||||
**Tag Naming**: Use `gateway-*` or `router-*` prefixes to avoid triggering unrelated CI workflows.
|
||||
|
||||
### Release Workflow
|
||||
|
||||
1. **Create and push tag**:
|
||||
```bash
|
||||
git tag -a gateway-v1.0.0 <commit-hash> -m "Gateway release v1.0.0"
|
||||
git push origin gateway-v1.0.0
|
||||
```
|
||||
|
||||
2. **Generate release notes** (automatically filters gateway-related commits):
|
||||
```bash
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0
|
||||
```
|
||||
|
||||
3. **Create GitHub release**:
|
||||
```bash
|
||||
# Create draft (DEFAULT - review before publishing)
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0 CREATE_RELEASE=1
|
||||
|
||||
# Or publish immediately (skip draft)
|
||||
make release-notes PREV=gateway-v0.2.2 CURR=gateway-v1.0.0 CREATE_RELEASE=1 DRAFT=0
|
||||
```
|
||||
|
||||
### Filtered Paths
|
||||
|
||||
Release notes only include commits touching:
|
||||
- `sgl-model-gateway/` - Router codebase
|
||||
- `python/sglang/srt/grpc/` - gRPC protocol
|
||||
- `python/sglang/srt/entrypoints/grpc_server.py` - gRPC server
|
||||
|
||||
The script automatically extracts author attribution, PR links, and identifies new contributors.
|
||||
|
||||
---
|
||||
|
||||
SGLang Model Gateway continues to evolve alongside the core SGLang runtime. Contributions should keep CLI flags, documentation, and Python bindings in sync with the Rust implementation.
|
||||
@@ -0,0 +1,670 @@
|
||||
use std::time::Instant;
|
||||
|
||||
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
|
||||
use serde_json::{from_str, to_string, to_value, to_vec};
|
||||
use sgl_model_gateway::{
|
||||
core::{BasicWorker, BasicWorkerBuilder, Worker, WorkerType},
|
||||
protocols::{
|
||||
chat::{ChatCompletionRequest, ChatMessage, MessageContent},
|
||||
common::StringOrArray,
|
||||
completion::CompletionRequest,
|
||||
generate::GenerateRequest,
|
||||
sampling_params::SamplingParams,
|
||||
},
|
||||
routers::http::pd_types::{generate_room_id, RequestWithBootstrap},
|
||||
};
|
||||
|
||||
fn create_test_worker() -> BasicWorker {
|
||||
BasicWorkerBuilder::new("http://test-server:8000")
|
||||
.worker_type(WorkerType::Prefill {
|
||||
bootstrap_port: Some(5678),
|
||||
})
|
||||
.build()
|
||||
}
|
||||
|
||||
// Helper function to get bootstrap info from worker
|
||||
fn get_bootstrap_info(worker: &BasicWorker) -> (String, Option<u16>) {
|
||||
let hostname = worker.bootstrap_host().to_string();
|
||||
let bootstrap_port = worker.bootstrap_port();
|
||||
(hostname, bootstrap_port)
|
||||
}
|
||||
|
||||
/// Create a default GenerateRequest for benchmarks with minimal fields set
|
||||
fn default_generate_request() -> GenerateRequest {
|
||||
GenerateRequest {
|
||||
text: None,
|
||||
model: None,
|
||||
input_ids: None,
|
||||
input_embeds: None,
|
||||
image_data: None,
|
||||
video_data: None,
|
||||
audio_data: None,
|
||||
sampling_params: None,
|
||||
return_logprob: None,
|
||||
logprob_start_len: None,
|
||||
top_logprobs_num: None,
|
||||
token_ids_logprob: None,
|
||||
return_text_in_logprobs: false,
|
||||
stream: false,
|
||||
log_metrics: true,
|
||||
return_hidden_states: false,
|
||||
modalities: None,
|
||||
session_params: None,
|
||||
lora_path: None,
|
||||
lora_id: None,
|
||||
custom_logit_processor: None,
|
||||
bootstrap_host: None,
|
||||
bootstrap_port: None,
|
||||
bootstrap_room: None,
|
||||
bootstrap_pair_key: None,
|
||||
data_parallel_rank: None,
|
||||
background: false,
|
||||
conversation_id: None,
|
||||
priority: None,
|
||||
extra_key: None,
|
||||
no_logs: false,
|
||||
custom_labels: None,
|
||||
return_bytes: false,
|
||||
return_entropy: false,
|
||||
rid: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a default ChatCompletionRequest for benchmarks with minimal fields set
|
||||
#[allow(deprecated)]
|
||||
fn default_chat_completion_request() -> ChatCompletionRequest {
|
||||
ChatCompletionRequest {
|
||||
// Required fields in OpenAI order
|
||||
messages: vec![],
|
||||
model: String::new(),
|
||||
|
||||
// Use default for all other fields
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
sampling_seed: None,
|
||||
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()),
|
||||
sampling_params: Some(SamplingParams {
|
||||
max_new_tokens: Some(100),
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(deprecated)]
|
||||
fn create_sample_chat_completion_request() -> ChatCompletionRequest {
|
||||
ChatCompletionRequest {
|
||||
model: "gpt-3.5-turbo".to_string(),
|
||||
messages: vec![
|
||||
ChatMessage::System {
|
||||
content: MessageContent::Text("You are a helpful assistant".to_string()),
|
||||
name: None,
|
||||
},
|
||||
ChatMessage::User {
|
||||
content: MessageContent::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()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(deprecated)]
|
||||
fn create_large_chat_completion_request() -> ChatCompletionRequest {
|
||||
let mut messages = vec![ChatMessage::System {
|
||||
content: MessageContent::Text(
|
||||
"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 {
|
||||
content: MessageContent::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 {
|
||||
content: Some(MessageContent::Text(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,
|
||||
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),
|
||||
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);
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,887 @@
|
||||
//! 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 std::{
|
||||
collections::BTreeMap,
|
||||
sync::{
|
||||
atomic::{AtomicBool, AtomicU64, Ordering},
|
||||
Arc, Mutex,
|
||||
},
|
||||
thread,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use criterion::{black_box, criterion_group, BenchmarkId, Criterion, Throughput};
|
||||
use serde_json::json;
|
||||
use sgl_model_gateway::{
|
||||
protocols::common::{Function, Tool},
|
||||
tool_parser::{JsonParser, ParserFactory as ToolParserFactory, ToolParser},
|
||||
};
|
||||
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>"#;
|
||||
|
||||
// Create test tools for parsers that need them
|
||||
fn create_test_tools() -> Vec<Tool> {
|
||||
vec![
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "search".to_string(),
|
||||
description: Some("Search for information".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string"},
|
||||
"limit": {"type": "number"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "code_interpreter".to_string(),
|
||||
description: Some("Execute code".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"language": {"type": "string"},
|
||||
"code": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
// 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(ToolParserFactory::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(ToolParserFactory::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(ToolParserFactory::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();
|
||||
|
||||
// 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 rt = rt.handle().clone();
|
||||
|
||||
b.iter_custom(|iters| {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let start = Instant::now();
|
||||
for _ in 0..iters {
|
||||
let mut parser = JsonParser::new();
|
||||
let mut complete_tools = Vec::new();
|
||||
|
||||
rt.block_on(async {
|
||||
for chunk in &chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
if !result.calls.is_empty() {
|
||||
complete_tools.extend(result.calls);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
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(ToolParserFactory::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((_normal_text, 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(ToolParserFactory::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 = ToolParserFactory::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(ToolParserFactory::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(ToolParserFactory::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);
|
||||
@@ -0,0 +1,24 @@
|
||||
# Build artifacts
|
||||
target/
|
||||
lib/
|
||||
|
||||
# Compiled binaries
|
||||
examples/simple/simple
|
||||
examples/streaming/streaming
|
||||
|
||||
# Go build artifacts
|
||||
*.o
|
||||
*.a
|
||||
*.so
|
||||
*.dylib
|
||||
|
||||
# IDE and editor files
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Environment files
|
||||
.env
|
||||
.env.local
|
||||
@@ -0,0 +1,47 @@
|
||||
[package]
|
||||
name = "sgl-model-gateway-golang"
|
||||
version = "0.2.3"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
name = "sgl_model_gateway_go"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1.42.0", features = ["full"] }
|
||||
serde_json = { version = "1.0", default-features = false, features = [
|
||||
"std",
|
||||
"preserve_order",
|
||||
] }
|
||||
uuid = { version = "1.10", features = ["v4", "serde"] }
|
||||
once_cell = "1.21.3"
|
||||
futures-util = "0.3"
|
||||
tracing = "0.1"
|
||||
|
||||
[dependencies.sgl-model-gateway]
|
||||
path = "../.."
|
||||
default-features = true
|
||||
|
||||
[features]
|
||||
default = []
|
||||
vendored-openssl = ["sgl-model-gateway/vendored-openssl"]
|
||||
|
||||
[profile.release]
|
||||
opt-level = "z" # Optimize for size
|
||||
lto = "fat" # Full LTO for smaller binaries
|
||||
codegen-units = 1 # Better optimization, slower compile
|
||||
strip = true # Strip debug symbols
|
||||
|
||||
[profile.ci]
|
||||
inherits = "release"
|
||||
opt-level = 2 # Lighter optimization (still fast runtime, much faster compile)
|
||||
lto = "thin" # Thin LTO - good balance
|
||||
codegen-units = 16 # More parallelization for faster builds
|
||||
strip = true
|
||||
|
||||
[profile.dev]
|
||||
opt-level = 0
|
||||
debug = 1
|
||||
split-debuginfo = "unpacked"
|
||||
incremental = true
|
||||
codegen-units = 256
|
||||
@@ -0,0 +1,103 @@
|
||||
# Makefile for sglang-router golang bindings
|
||||
# This builds the Rust FFI library and provides convenience targets for Go development
|
||||
|
||||
# Configuration
|
||||
CARGO_BUILD_DIR ?= $(shell pwd)/target
|
||||
BUILD_MODE ?= release
|
||||
LIB_NAME = libsglang_router_rs
|
||||
|
||||
# Detect OS
|
||||
UNAME_S := $(shell uname -s)
|
||||
ifeq ($(UNAME_S),Linux)
|
||||
LIB_EXT = .so
|
||||
LD_LIBRARY_PATH_VAR = LD_LIBRARY_PATH
|
||||
endif
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
LIB_EXT = .dylib
|
||||
LD_LIBRARY_PATH_VAR = DYLD_LIBRARY_PATH
|
||||
endif
|
||||
|
||||
# Paths
|
||||
ROOT_DIR := $(shell pwd)
|
||||
RUST_SRC_DIR := $(ROOT_DIR)/src
|
||||
LIB_BUILD_DIR := $(CARGO_BUILD_DIR)/$(BUILD_MODE)
|
||||
LIB_BUILD_PATH := $(LIB_BUILD_DIR)/$(LIB_NAME)$(LIB_EXT)
|
||||
LIB_EXPORT_DIR := $(ROOT_DIR)/lib
|
||||
LIB_EXPORT_PATH := $(LIB_EXPORT_DIR)/$(LIB_NAME)$(LIB_EXT)
|
||||
|
||||
# Python LDFLAGS (needed for Rust FFI that depends on Python)
|
||||
PYTHON_LDFLAGS := $(shell python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
|
||||
|
||||
# CGO flags - use exported lib directory if available, otherwise build directory
|
||||
LIB_DIR := $(if $(wildcard $(LIB_EXPORT_PATH)),$(LIB_EXPORT_DIR),$(LIB_BUILD_DIR))
|
||||
export CGO_LDFLAGS = -L$(LIB_DIR) -lsglang_router_rs $(PYTHON_LDFLAGS) -ldl
|
||||
export $(LD_LIBRARY_PATH_VAR) := $(LIB_DIR):$($(LD_LIBRARY_PATH_VAR))
|
||||
|
||||
.PHONY: all build build-dev lib lib-clean clean test examples help run-simple run-streaming check-lib
|
||||
|
||||
help:
|
||||
@echo "Available targets:"
|
||||
@echo " build - Build release version of Rust FFI library"
|
||||
@echo " build-dev - Build debug version of Rust FFI library"
|
||||
@echo " lib - Copy built library to ./lib directory"
|
||||
@echo " lib-clean - Clean ./lib directory"
|
||||
@echo " clean - Clean build artifacts"
|
||||
@echo " test - Run Go tests"
|
||||
@echo " examples - Build example programs"
|
||||
@echo " run-simple - Run simple example"
|
||||
@echo " run-streaming - Run streaming example"
|
||||
|
||||
all: build
|
||||
|
||||
build:
|
||||
@echo "Building Rust FFI library (release mode)..."
|
||||
@CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo build --release --manifest-path Cargo.toml
|
||||
@echo "Library built at: $(LIB_BUILD_PATH)"
|
||||
|
||||
build-dev:
|
||||
@echo "Building Rust FFI library (debug mode)..."
|
||||
@CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo build --manifest-path Cargo.toml
|
||||
@echo "Library built at: $(LIB_BUILD_DIR)/debug/$(LIB_NAME)$(LIB_EXT)"
|
||||
|
||||
lib: build
|
||||
@echo "Copying library to ./lib directory..."
|
||||
@mkdir -p $(LIB_EXPORT_DIR)
|
||||
@cp $(LIB_BUILD_PATH) $(LIB_EXPORT_PATH)
|
||||
@echo "Library exported at: $(LIB_EXPORT_PATH)"
|
||||
|
||||
lib-clean:
|
||||
@echo "Cleaning ./lib directory..."
|
||||
@rm -rf $(LIB_EXPORT_DIR)
|
||||
@echo "Lib directory cleaned"
|
||||
|
||||
clean: lib-clean
|
||||
@echo "Cleaning build artifacts..."
|
||||
@CARGO_TARGET_DIR=$(CARGO_BUILD_DIR) cargo clean --manifest-path Cargo.toml
|
||||
@echo "Clean complete"
|
||||
|
||||
test: build
|
||||
@echo "Running Go tests..."
|
||||
@go test ./...
|
||||
|
||||
examples: build
|
||||
@echo "Building example programs..."
|
||||
@cd examples/simple && go build -o simple main.go
|
||||
@cd examples/streaming && go build -o streaming main.go
|
||||
@echo "Examples built"
|
||||
|
||||
run-simple: build
|
||||
@echo "Running simple example..."
|
||||
@cd examples/simple && bash run.sh
|
||||
|
||||
run-streaming: build
|
||||
@echo "Running streaming example..."
|
||||
@cd examples/streaming && bash run.sh
|
||||
|
||||
# Check if library exists (either in lib dir or build dir)
|
||||
check-lib:
|
||||
@if [ ! -f "$(LIB_EXPORT_PATH)" ] && [ ! -f "$(LIB_BUILD_PATH)" ]; then \
|
||||
echo "Error: Library not found at $(LIB_EXPORT_PATH) or $(LIB_BUILD_PATH)"; \
|
||||
echo "Run 'make build' or 'make lib' first"; \
|
||||
exit 1; \
|
||||
fi
|
||||
@echo "Library found at: $(LIB_DIR)/$(LIB_NAME)$(LIB_EXT)"
|
||||
@@ -0,0 +1,552 @@
|
||||
# SGLang Go gRPC SDK
|
||||
|
||||
A high-level Go SDK for interacting with SGLang gRPC API, designed with an OpenAI-style API for familiarity and ease of use.
|
||||
|
||||
**Location**: `sgl-model-gateway/bindings/golang/`
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Features](#features)
|
||||
- [Installation](#installation)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Basic Usage](#basic-usage)
|
||||
- [Streaming Usage](#streaming-usage)
|
||||
- [Examples](#examples)
|
||||
- [Configuration](#configuration)
|
||||
- [API Reference](#api-reference)
|
||||
- [Testing](#testing)
|
||||
- [Unit Tests](#unit-tests)
|
||||
- [Integration Tests](#integration-tests)
|
||||
- [Benchmarks](#benchmarks)
|
||||
- [Documentation](#documentation)
|
||||
- [Development](#development)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [License](#license)
|
||||
|
||||
## Features
|
||||
|
||||
- **OpenAI-style API**: Familiar interface similar to OpenAI Go SDK
|
||||
- **Streaming Support**: Real-time streaming chat completions
|
||||
- **Non-streaming Support**: Simple request/response API
|
||||
- **Tool Calling**: Support for function calling and tool use
|
||||
- **Type-safe**: Full Go type definitions for requests and responses
|
||||
- **Comprehensive Testing**: 18+ unit and integration tests
|
||||
- **Thread-safe**: All public methods are safe for concurrent use
|
||||
- **Well-documented**: Full API documentation with examples
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
go get github.com/sglang/sglang-go-grpc-sdk
|
||||
```
|
||||
|
||||
### Build Requirements
|
||||
|
||||
- Go 1.21 or later
|
||||
- Rust toolchain (for building the FFI library)
|
||||
- Python 3.x (for Python bindings in Rust FFI)
|
||||
- Tokio runtime for async operations
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Basic Usage (Non-streaming)
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
|
||||
"github.com/sglang/sglang-go-grpc-sdk"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Create client
|
||||
client, err := sglang.NewClient(sglang.ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Create completion
|
||||
resp, err := client.CreateChatCompletion(context.Background(), sglang.ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []sglang.ChatMessage{
|
||||
{Role: "user", Content: "Hello!"},
|
||||
},
|
||||
Stream: false,
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
fmt.Println(resp.Choices[0].Message.Content)
|
||||
fmt.Printf("Usage: Prompt=%d, Completion=%d, Total=%d\n",
|
||||
resp.Usage.PromptTokens,
|
||||
resp.Usage.CompletionTokens,
|
||||
resp.Usage.TotalTokens)
|
||||
}
|
||||
```
|
||||
|
||||
### Streaming Usage
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
|
||||
"github.com/sglang/sglang-go-grpc-sdk"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Create client
|
||||
client, err := sglang.NewClient(sglang.ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Create streaming completion
|
||||
ctx := context.Background()
|
||||
stream, err := client.CreateChatCompletionStream(ctx, sglang.ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []sglang.ChatMessage{
|
||||
{Role: "user", Content: "Tell me a story"},
|
||||
},
|
||||
Stream: true,
|
||||
MaxCompletionTokens: intPtr(500),
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
// Read streaming response
|
||||
for {
|
||||
chunk, err := stream.Recv()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Delta.Content != "" {
|
||||
fmt.Print(choice.Delta.Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
fmt.Println() // newline
|
||||
}
|
||||
|
||||
// Helper functions for optional pointer fields
|
||||
func intPtr(i int) *int {
|
||||
return &i
|
||||
}
|
||||
|
||||
func float32Ptr(f float32) *float32 {
|
||||
return &f
|
||||
}
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
The SDK includes several examples in the `examples/` directory:
|
||||
|
||||
- **simple**: Basic non-streaming chat completion example
|
||||
- **streaming**: Real-time streaming with performance metrics
|
||||
|
||||
### Running Examples
|
||||
|
||||
```bash
|
||||
# Run simple example
|
||||
cd bindings/golang/examples/simple
|
||||
bash run.sh
|
||||
|
||||
# Run streaming example
|
||||
cd bindings/golang/examples/streaming
|
||||
bash run.sh
|
||||
|
||||
# Or use Makefile from bindings/golang directory
|
||||
cd bindings/golang
|
||||
make run-simple
|
||||
make run-streaming
|
||||
```
|
||||
|
||||
Examples automatically detect the server endpoint and tokenizer path via environment variables or defaults.
|
||||
|
||||
## Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
- `SGL_GRPC_ENDPOINT`: gRPC server endpoint (default: `grpc://localhost:20000`)
|
||||
- `SGL_TOKENIZER_PATH`: Path to tokenizer directory (required)
|
||||
- `CARGO_BUILD_DIR`: Rust build output directory (auto-detected if not set)
|
||||
|
||||
### ClientConfig
|
||||
|
||||
```go
|
||||
type ClientConfig struct {
|
||||
// Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000")
|
||||
// Required field. Must include the scheme (grpc://) and port number.
|
||||
Endpoint string
|
||||
|
||||
// TokenizerPath is the path to the tokenizer directory containing
|
||||
// tokenizer configuration files (e.g., tokenizer.json, vocab.json)
|
||||
// Required field.
|
||||
TokenizerPath string
|
||||
}
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Client Methods
|
||||
|
||||
```go
|
||||
type Client struct {
|
||||
// Thread-safe client for SGLang gRPC API
|
||||
}
|
||||
|
||||
// Creates a new client with the given configuration
|
||||
func NewClient(config ClientConfig) (*Client, error)
|
||||
|
||||
// Closes the client and releases all resources
|
||||
func (c *Client) Close() error
|
||||
|
||||
// Creates a non-streaming chat completion
|
||||
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error)
|
||||
|
||||
// Creates a streaming chat completion
|
||||
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error)
|
||||
```
|
||||
|
||||
### Request Types
|
||||
|
||||
- `ChatCompletionRequest`: Main request type for chat completions
|
||||
- Model, Messages, Stream, Temperature, TopP, MaxCompletionTokens, Tools, etc.
|
||||
- `ChatMessage`: Individual message in a conversation
|
||||
- Role, Content
|
||||
- `Tool`: Tool/function definition for function calling
|
||||
- Type, Function (name, description, parameters)
|
||||
|
||||
### Response Types
|
||||
|
||||
- `ChatCompletionResponse`: Non-streaming response
|
||||
- ID, Model, Created, Choices, Usage
|
||||
- `ChatCompletionStreamResponse`: Streaming response chunk
|
||||
- Same structure as above but for incremental updates
|
||||
- `Message`: Complete message with content and tool calls
|
||||
- `ToolCall`: Tool call information with function and arguments
|
||||
- `Usage`: Token usage statistics
|
||||
- PromptTokens, CompletionTokens, TotalTokens
|
||||
|
||||
## Testing
|
||||
|
||||
The SDK includes comprehensive testing infrastructure with both unit and integration tests.
|
||||
|
||||
### Unit Tests
|
||||
|
||||
Unit tests are located in `client_test.go` and test individual components without requiring a server.
|
||||
|
||||
#### Running Unit Tests
|
||||
|
||||
```bash
|
||||
# Run all unit tests
|
||||
go test ./...
|
||||
|
||||
# Run with verbose output
|
||||
go test -v ./...
|
||||
|
||||
# Run specific test
|
||||
go test -run TestClientConfig
|
||||
|
||||
# Run tests with race detector (detects concurrency issues)
|
||||
go test -race ./...
|
||||
|
||||
# Run with coverage analysis
|
||||
go test -cover ./...
|
||||
|
||||
# Generate detailed coverage report
|
||||
go test -coverprofile=coverage.out ./...
|
||||
go tool cover -html=coverage.out -o coverage.html
|
||||
```
|
||||
|
||||
#### Unit Test Coverage
|
||||
|
||||
- **Configuration validation** (`TestClientConfig`) - Validates ClientConfig requirements
|
||||
- **Type structures** - Verifying all struct types work correctly
|
||||
- **Response handling** - Testing response parsing and validation
|
||||
- **Concurrent operations** (`TestConcurrentClientOperations`) - Thread-safety verification
|
||||
- **Benchmarks** (`BenchmarkChatCompletionRequest`) - Performance measurement
|
||||
|
||||
**Test Files**:
|
||||
- `client_test.go` - 10 unit tests covering core functionality
|
||||
- Tests cover: config validation, message types, request validation, close operations, response types, streaming, tools, concurrency, and context cancellation
|
||||
|
||||
### Integration Tests
|
||||
|
||||
Integration tests require a running SGLang server and test the full client-server interaction.
|
||||
|
||||
#### Prerequisites
|
||||
|
||||
1. Start an SGLang server:
|
||||
|
||||
```bash
|
||||
# Using Python (requires sglang package installed)
|
||||
python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-hf
|
||||
|
||||
# Or using pre-built Docker image
|
||||
docker run -p 20000:20000 lmsys/sglang:latest
|
||||
|
||||
# Or build your own
|
||||
sglang launch_server --model-path <model_path>
|
||||
```
|
||||
|
||||
2. Set required environment variables:
|
||||
|
||||
```bash
|
||||
# Set the gRPC endpoint (default: grpc://localhost:20000)
|
||||
export SGL_GRPC_ENDPOINT=grpc://localhost:20000
|
||||
|
||||
# Set the tokenizer path (required)
|
||||
export SGL_TOKENIZER_PATH=/path/to/tokenizer
|
||||
```
|
||||
|
||||
#### Running Integration Tests
|
||||
|
||||
```bash
|
||||
# Run all integration tests
|
||||
go test -tags=integration ./...
|
||||
|
||||
# Run specific integration test
|
||||
go test -tags=integration -run TestIntegrationNonStreamingCompletion
|
||||
|
||||
# Run with verbose output
|
||||
go test -tags=integration -v ./...
|
||||
|
||||
# Run with race detector
|
||||
go test -tags=integration -race ./...
|
||||
```
|
||||
|
||||
#### Integration Test Coverage
|
||||
|
||||
**Test File**: `integration_test.go` - 4 integration tests
|
||||
|
||||
- `TestIntegrationNonStreamingCompletion` - Basic non-streaming request/response
|
||||
- `TestIntegrationStreamingCompletion` - Streaming response handling
|
||||
- `TestIntegrationConcurrentRequests` - Multiple simultaneous requests
|
||||
- `TestIntegrationContextCancellation` - Context timeout and cancellation
|
||||
|
||||
### Benchmarks
|
||||
|
||||
Measure performance of SDK operations:
|
||||
|
||||
```bash
|
||||
# Run all benchmarks
|
||||
go test -bench=. -benchmem ./...
|
||||
|
||||
# Run specific benchmark
|
||||
go test -bench=BenchmarkChatCompletionRequest -benchmem
|
||||
|
||||
# Run for longer duration
|
||||
go test -bench=. -benchtime=10s ./...
|
||||
```
|
||||
|
||||
Current benchmarks:
|
||||
- `BenchmarkChatCompletionRequest` - Measures request creation performance
|
||||
|
||||
### CI/CD Integration
|
||||
|
||||
Add to your GitHub Actions workflow:
|
||||
|
||||
```yaml
|
||||
- name: Run Go tests
|
||||
run: |
|
||||
go test -race -cover ./...
|
||||
|
||||
- name: Run integration tests (on main branch)
|
||||
if: github.ref == 'refs/heads/main'
|
||||
env:
|
||||
SGL_GRPC_ENDPOINT: grpc://localhost:20000
|
||||
SGL_TOKENIZER_PATH: /path/to/tokenizer
|
||||
run: go test -tags=integration ./...
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
### Code Documentation
|
||||
|
||||
All public types and functions include comprehensive documentation:
|
||||
|
||||
1. **Package-level documentation** in `client.go` with usage examples
|
||||
2. **Type documentation** for all structs with field descriptions
|
||||
3. **Function documentation** with:
|
||||
- Purpose and behavior description
|
||||
- Parameter documentation with types and constraints
|
||||
- Return value documentation
|
||||
- Error cases and handling
|
||||
- Safety notes (for FFI functions)
|
||||
- Usage examples
|
||||
|
||||
### Key Documented Components
|
||||
|
||||
- `Client` - Main client with thread-safety notes
|
||||
- `ClientConfig` - Configuration requirements and validation rules
|
||||
- `ChatCompletionRequest` - Request structure with field descriptions
|
||||
- `ChatCompletionResponse` - Response structure and usage
|
||||
- `ChatCompletionStreamResponse` - Streaming response format
|
||||
- `Usage` - Token usage information structure
|
||||
- `Tool`, `Function`, `ToolCall` - Tool call structures
|
||||
|
||||
### Viewing Documentation
|
||||
|
||||
Generate and view HTML documentation:
|
||||
|
||||
```bash
|
||||
# Install godoc (if not already installed)
|
||||
go install golang.org/x/tools/cmd/godoc@latest
|
||||
|
||||
# Generate and serve documentation
|
||||
godoc -http=:6060
|
||||
|
||||
# Visit: http://localhost:6060/pkg/github.com/sglang/sglang-go-grpc-sdk/
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
### Building
|
||||
|
||||
```bash
|
||||
cd bindings/golang
|
||||
|
||||
# Build the Go bindings (compiles Rust FFI library)
|
||||
make build
|
||||
|
||||
# Clean build
|
||||
make clean && make build
|
||||
```
|
||||
|
||||
### Code Quality
|
||||
|
||||
Ensure code quality before committing:
|
||||
|
||||
```bash
|
||||
# Run Go vet (check for potential bugs)
|
||||
go vet ./...
|
||||
|
||||
# Format code
|
||||
go fmt ./...
|
||||
|
||||
# Run all tests with race detection
|
||||
go test -race ./...
|
||||
```
|
||||
|
||||
### Project Structure
|
||||
|
||||
```
|
||||
bindings/golang/
|
||||
├── client.go # Main client implementation
|
||||
├── client_test.go # Unit tests
|
||||
├── integration_test.go # Integration tests
|
||||
├── README.md # This file
|
||||
├── Makefile # Build automation
|
||||
├── Cargo.toml # Rust FFI dependencies
|
||||
├── examples/ # Example programs
|
||||
│ ├── simple/ # Non-streaming example
|
||||
│ └── streaming/ # Streaming example
|
||||
├── src/ # Rust FFI source
|
||||
│ ├── client.rs # Client FFI
|
||||
│ ├── stream.rs # Stream handling
|
||||
│ ├── grpc_converter.rs # Response conversion
|
||||
│ └── ...
|
||||
└── internal/ # Internal packages
|
||||
└── ffi/ # FFI bindings
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Connection Errors
|
||||
|
||||
**Error**: `connection refused` or `failed to dial`
|
||||
|
||||
**Solution**:
|
||||
1. Ensure SGLang server is running: `python -m sglang.launch_server`
|
||||
2. Check endpoint: `echo $SGL_GRPC_ENDPOINT`
|
||||
3. Verify port is not blocked: `nc -zv localhost 20000`
|
||||
|
||||
### Tokenizer Not Found
|
||||
|
||||
**Error**: `tokenizer path not found` or `tokenizer configuration missing`
|
||||
|
||||
**Solution**:
|
||||
1. Set `SGL_TOKENIZER_PATH` environment variable
|
||||
2. Verify path contains required files: `ls $SGL_TOKENIZER_PATH`
|
||||
3. Files should include: `tokenizer.json`, `vocab.json`, `config.json`
|
||||
|
||||
### Build Failures
|
||||
|
||||
**Error**: `library 'sglang_router_rs' not found`
|
||||
|
||||
**Solution**:
|
||||
1. Rebuild Rust library: `cd sgl-model-gateway/bindings/golang && make build`
|
||||
2. Or manually with cargo: `cd sgl-model-gateway/bindings/golang && cargo build --release`
|
||||
3. Set `CARGO_BUILD_DIR` if using non-standard build location
|
||||
4. Ensure Rust toolchain is installed: `rustup toolchain list`
|
||||
|
||||
### Tests Hanging
|
||||
|
||||
**Error**: Tests seem to hang indefinitely
|
||||
|
||||
**Solution**:
|
||||
1. Use timeout for hanging tests: `timeout 30s go test ./...`
|
||||
2. Run with verbose output to see which test hangs: `go test -v ./...`
|
||||
3. Ensure server is responsive: `grpcurl -plaintext localhost:20000 list`
|
||||
|
||||
### Memory Issues
|
||||
|
||||
**Error**: Out of memory during tests
|
||||
|
||||
**Solution**:
|
||||
```bash
|
||||
# Run with memory limit for long-running tests
|
||||
GODEBUG=madvdontneed=1 go test -timeout 5m ./...
|
||||
|
||||
# Monitor memory during tests
|
||||
watch -n1 'ps aux | grep test'
|
||||
```
|
||||
|
||||
## Contributing
|
||||
|
||||
When adding new features:
|
||||
|
||||
1. Add comprehensive documentation to public types/functions
|
||||
2. Include usage examples for complex APIs
|
||||
3. Add unit tests covering happy path and error cases
|
||||
4. Add integration tests if server interaction required
|
||||
5. Ensure code passes `go vet` and `go test -race`
|
||||
6. Update this README if adding new features
|
||||
|
||||
## License
|
||||
|
||||
See LICENSE file for details.
|
||||
|
||||
---
|
||||
|
||||
**Need Help?**
|
||||
- Check examples in `examples/` directory
|
||||
- Run tests to see working code: `go test -v ./...`
|
||||
- Review function documentation: `godoc` or inline comments
|
||||
- Check troubleshooting section above
|
||||
@@ -0,0 +1,510 @@
|
||||
// Package sglang provides a Go SDK for SGLang gRPC API.
|
||||
//
|
||||
// SGLang is a fast language model serving framework. This package provides a Go client
|
||||
// library for interacting with SGLang's gRPC API, following the style of OpenAI's Go SDK.
|
||||
//
|
||||
// Basic usage:
|
||||
//
|
||||
// client, err := sglang.NewClient(sglang.ClientConfig{
|
||||
// Endpoint: "grpc://localhost:20000",
|
||||
// TokenizerPath: "/path/to/tokenizer",
|
||||
// })
|
||||
// if err != nil {
|
||||
// log.Fatal(err)
|
||||
// }
|
||||
// defer client.Close()
|
||||
//
|
||||
// resp, err := client.CreateChatCompletion(ctx, sglang.ChatCompletionRequest{
|
||||
// Model: "default",
|
||||
// Messages: []sglang.ChatMessage{
|
||||
// {Role: "user", Content: "Hello"},
|
||||
// },
|
||||
// })
|
||||
//
|
||||
// For streaming responses, use CreateChatCompletionStream instead.
|
||||
package sglang
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/sglang/sglang-go-grpc-sdk/internal/ffi"
|
||||
)
|
||||
|
||||
// Client is the main client for interacting with SGLang gRPC API.
|
||||
// It manages the connection to the SGLang server and handles both streaming
|
||||
// and non-streaming chat completions.
|
||||
//
|
||||
// Thread-safe: All public methods are safe for concurrent use.
|
||||
type Client struct {
|
||||
endpoint string
|
||||
tokenizerPath string
|
||||
clientHandle *ffi.SglangClientHandle
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// ClientConfig holds configuration for creating a new client.
|
||||
type ClientConfig struct {
|
||||
// Endpoint is the gRPC endpoint URL (e.g., "grpc://localhost:20000").
|
||||
// Required field. Must include the scheme (grpc://) and port number.
|
||||
Endpoint string
|
||||
|
||||
// TokenizerPath is the path to the tokenizer directory containing
|
||||
// tokenizer configuration files (e.g., tokenizer.json, vocab.json).
|
||||
// Required field.
|
||||
TokenizerPath string
|
||||
}
|
||||
|
||||
// NewClient creates a new SGLang client with the given configuration.
|
||||
//
|
||||
// The client maintains a long-lived connection to the SGLang server and should
|
||||
// be reused for multiple requests. Call Close() to release resources.
|
||||
//
|
||||
// Returns an error if:
|
||||
// - Endpoint is empty
|
||||
// - TokenizerPath is empty
|
||||
// - Connection to the server fails
|
||||
func NewClient(config ClientConfig) (*Client, error) {
|
||||
if config.Endpoint == "" {
|
||||
return nil, errors.New("endpoint is required")
|
||||
}
|
||||
if config.TokenizerPath == "" {
|
||||
return nil, errors.New("tokenizer path is required")
|
||||
}
|
||||
|
||||
clientHandle, err := ffi.NewClient(config.Endpoint, config.TokenizerPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create client: %w", err)
|
||||
}
|
||||
|
||||
return &Client{
|
||||
endpoint: config.Endpoint,
|
||||
tokenizerPath: config.TokenizerPath,
|
||||
clientHandle: clientHandle,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Close closes the client and releases all resources.
|
||||
//
|
||||
// After Close() is called, the client cannot be used for further requests.
|
||||
// Calling Close() multiple times is safe and idempotent.
|
||||
func (c *Client) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.clientHandle != nil {
|
||||
c.clientHandle.Free()
|
||||
c.clientHandle = nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ChatCompletionRequest represents a request for chat completion.
|
||||
// It follows the OpenAI API style for familiar usage.
|
||||
type ChatCompletionRequest struct {
|
||||
// Model specifies the model to use for completion (e.g., "default")
|
||||
Model string `json:"model"`
|
||||
// Messages is the list of messages in the conversation
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Temperature *float32 `json:"temperature,omitempty"`
|
||||
TopP *float32 `json:"top_p,omitempty"`
|
||||
TopK *int `json:"top_k,omitempty"`
|
||||
MaxCompletionTokens *int `json:"max_completion_tokens,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
Stop interface{} `json:"stop,omitempty"`
|
||||
StopTokenIDs []int `json:"stop_token_ids,omitempty"`
|
||||
SkipSpecialTokens bool `json:"skip_special_tokens,omitempty"`
|
||||
FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
|
||||
PresencePenalty *float32 `json:"presence_penalty,omitempty"`
|
||||
ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
|
||||
Seed *int `json:"seed,omitempty"`
|
||||
Logprobs bool `json:"logprobs,omitempty"`
|
||||
TopLogprobs *int `json:"top_logprobs,omitempty"`
|
||||
User string `json:"user,omitempty"`
|
||||
}
|
||||
|
||||
// ChatMessage represents a single message in a chat conversation
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content interface{} `json:"content"`
|
||||
Name string `json:"name,omitempty"`
|
||||
}
|
||||
|
||||
// Tool represents a tool/function that can be called
|
||||
type Tool struct {
|
||||
Type string `json:"type"`
|
||||
Function Function `json:"function"`
|
||||
}
|
||||
|
||||
// Function represents a function definition
|
||||
type Function struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters map[string]interface{} `json:"parameters"`
|
||||
}
|
||||
|
||||
// ResponseFormat represents the response format
|
||||
type ResponseFormat struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
// ChatCompletionResponse represents a non-streaming chat completion response
|
||||
type ChatCompletionResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
Choices []Choice `json:"choices"`
|
||||
Usage Usage `json:"usage"`
|
||||
}
|
||||
|
||||
// Choice represents a choice in the completion response
|
||||
type Choice struct {
|
||||
Index int `json:"index"`
|
||||
Message Message `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
// Message represents a message in the response
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
}
|
||||
|
||||
// ToolCall represents a tool call in the response
|
||||
type ToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function FunctionCall `json:"function"`
|
||||
}
|
||||
|
||||
// FunctionCall represents a function call
|
||||
type FunctionCall struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}
|
||||
|
||||
// Usage represents token usage information
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
}
|
||||
|
||||
// ChatCompletionStreamResponse represents a streaming chat completion response
|
||||
type ChatCompletionStreamResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
Choices []StreamChoice `json:"choices"`
|
||||
Usage *Usage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
// StreamChoice represents a choice in a streaming response
|
||||
type StreamChoice struct {
|
||||
Index int `json:"index"`
|
||||
Delta MessageDelta `json:"delta"`
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
}
|
||||
|
||||
// MessageDelta represents incremental message updates
|
||||
type MessageDelta struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
}
|
||||
|
||||
// CreateChatCompletion creates a non-streaming chat completion with context support.
|
||||
//
|
||||
// Context Support:
|
||||
// The ctx parameter is fully supported for cancellation and timeouts:
|
||||
// - If ctx is cancelled, the request will be interrupted on the next stream.Recv() call
|
||||
// - If ctx times out, the request will return context.DeadlineExceeded
|
||||
//
|
||||
// Example with timeout:
|
||||
//
|
||||
// ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
// defer cancel()
|
||||
// resp, err := client.CreateChatCompletion(ctx, req)
|
||||
//
|
||||
// Note: Internally, this creates a stream and collects all chunks,
|
||||
// so context monitoring happens at the chunk level.
|
||||
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) {
|
||||
// For non-streaming, we'll collect all chunks and return the final response
|
||||
req.Stream = true // We still use streaming internally, but collect all chunks
|
||||
|
||||
// Prepare request: if Tools is empty, set to nil for proper JSON serialization
|
||||
if len(req.Tools) == 0 {
|
||||
req.Tools = nil
|
||||
}
|
||||
|
||||
stream, err := c.CreateChatCompletionStream(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
var fullContent strings.Builder
|
||||
var fullToolCalls []ToolCall
|
||||
var finishReason string
|
||||
var usage Usage
|
||||
var responseID string
|
||||
var created int64
|
||||
var model string
|
||||
var systemFingerprint string
|
||||
|
||||
for {
|
||||
chunk, err := stream.Recv()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if chunk.ID != "" {
|
||||
responseID = chunk.ID
|
||||
}
|
||||
if chunk.Created > 0 {
|
||||
created = chunk.Created
|
||||
}
|
||||
if chunk.Model != "" {
|
||||
model = chunk.Model
|
||||
}
|
||||
if chunk.SystemFingerprint != "" {
|
||||
systemFingerprint = chunk.SystemFingerprint
|
||||
}
|
||||
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Delta.Content != "" {
|
||||
fullContent.WriteString(choice.Delta.Content)
|
||||
}
|
||||
if len(choice.Delta.ToolCalls) > 0 {
|
||||
fullToolCalls = append(fullToolCalls, choice.Delta.ToolCalls...)
|
||||
}
|
||||
// Always update finish_reason if present (even if empty string, but should not be empty)
|
||||
// The last chunk (Complete message) should have finish_reason set
|
||||
if choice.FinishReason != "" {
|
||||
finishReason = choice.FinishReason
|
||||
}
|
||||
}
|
||||
|
||||
// Extract usage from chunk if available (usually in the last chunk)
|
||||
// Always update usage if present, as the last chunk should have the final usage
|
||||
if chunk.Usage != nil {
|
||||
usage = *chunk.Usage
|
||||
}
|
||||
}
|
||||
|
||||
// Build final response
|
||||
message := Message{
|
||||
Role: "assistant",
|
||||
Content: fullContent.String(),
|
||||
}
|
||||
if len(fullToolCalls) > 0 {
|
||||
message.ToolCalls = fullToolCalls
|
||||
}
|
||||
|
||||
// Ensure finish_reason is set (defensive check)
|
||||
// If finish_reason is still empty, default to "stop"
|
||||
if finishReason == "" {
|
||||
finishReason = "stop"
|
||||
}
|
||||
|
||||
return &ChatCompletionResponse{
|
||||
ID: responseID,
|
||||
Object: "chat.completion",
|
||||
Created: created,
|
||||
Model: model,
|
||||
SystemFingerprint: systemFingerprint,
|
||||
Choices: []Choice{
|
||||
{
|
||||
Index: 0,
|
||||
Message: message,
|
||||
FinishReason: finishReason,
|
||||
},
|
||||
},
|
||||
Usage: usage,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ChatCompletionStream represents a streaming chat completion
|
||||
type ChatCompletionStream struct {
|
||||
stream *ffi.SglangStreamHandle
|
||||
mu sync.Mutex
|
||||
done bool // Track if stream has been marked as done
|
||||
ctx context.Context // Context for cancellation support
|
||||
cancel context.CancelFunc // Cancel function to stop monitoring goroutine
|
||||
closed chan struct{} // Signal when stream is closed
|
||||
}
|
||||
|
||||
// Recv receives the next chunk from the stream.
|
||||
//
|
||||
// Supports context cancellation: if the context passed to CreateChatCompletionStream
|
||||
// is cancelled, Recv will return context.Canceled error on the next call.
|
||||
func (s *ChatCompletionStream) Recv() (*ChatCompletionStreamResponse, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// Check if context was cancelled
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
return nil, s.ctx.Err() // Returns context.Canceled or context.DeadlineExceeded
|
||||
default:
|
||||
}
|
||||
|
||||
if s.stream == nil {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
// If stream was already marked as done, immediately return EOF
|
||||
// This prevents calling ReadNext() again after isDone=1
|
||||
if s.done {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
// Loop to handle empty responses (Ok(None) from Rust)
|
||||
// Keep reading until we get actual data or stream ends
|
||||
for {
|
||||
responseJSON, isDone, err := s.stream.ReadNext()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Mark stream as done if ReadNext indicates completion
|
||||
if isDone {
|
||||
s.done = true
|
||||
}
|
||||
|
||||
// If we have a response, parse and return it
|
||||
if responseJSON != "" {
|
||||
var response ChatCompletionStreamResponse
|
||||
if err := json.Unmarshal([]byte(responseJSON), &response); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
return &response, nil
|
||||
}
|
||||
|
||||
// If stream is done but no response, return EOF
|
||||
if isDone {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
// Empty response and stream not done - loop to read next chunk
|
||||
// This handles Ok(None) cases where Rust returns no data but stream continues
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the stream and cancels any pending operations.
|
||||
func (s *ChatCompletionStream) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// Cancel the context to signal the monitoring goroutine to stop
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
|
||||
// Signal that stream is closed
|
||||
select {
|
||||
case <-s.closed:
|
||||
// Already closed
|
||||
default:
|
||||
close(s.closed)
|
||||
}
|
||||
|
||||
// Free the stream to mark it as completed
|
||||
// This prevents AbortOnDropStream from sending abort when dropped
|
||||
if s.stream != nil {
|
||||
s.stream.Free()
|
||||
s.stream = nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateChatCompletionStream creates a streaming chat completion with context cancellation support.
|
||||
//
|
||||
// Context Support:
|
||||
// The ctx parameter is now fully supported for cancellation and timeouts:
|
||||
// - If ctx is cancelled, stream.Recv() will return context.Canceled on the next call
|
||||
// - If ctx times out (WithTimeout), stream.Recv() will return context.DeadlineExceeded
|
||||
// - Calling stream.Close() also cancels the context
|
||||
//
|
||||
// Example with timeout:
|
||||
//
|
||||
// ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
// defer cancel()
|
||||
// stream, err := client.CreateChatCompletionStream(ctx, req)
|
||||
// // Stream will auto-close if 30 seconds elapse
|
||||
//
|
||||
// Example with cancellation:
|
||||
//
|
||||
// ctx, cancel := context.WithCancel(context.Background())
|
||||
// stream, err := client.CreateChatCompletionStream(ctx, req)
|
||||
// go func() {
|
||||
// time.Sleep(5*time.Second)
|
||||
// cancel() // Cancel after 5 seconds
|
||||
// }()
|
||||
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionStream, error) {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
|
||||
if c.clientHandle == nil {
|
||||
return nil, errors.New("client is closed")
|
||||
}
|
||||
|
||||
// Marshal request to JSON, then ensure tools field is always present.
|
||||
// Due to omitempty tag, empty Tools slice will be omitted from JSON.
|
||||
// We need to ensure tools field is always present as [] when empty (not omitted),
|
||||
// matching the behavior of complete_sdk example.
|
||||
reqJSON, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
// Unmarshal into map and ensure tools field is present
|
||||
var reqMap map[string]interface{}
|
||||
if err := json.Unmarshal(reqJSON, &reqMap); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal request to map: %w", err)
|
||||
}
|
||||
|
||||
// Add empty tools array if not present
|
||||
if _, exists := reqMap["tools"]; !exists {
|
||||
reqMap["tools"] = []interface{}{}
|
||||
}
|
||||
|
||||
// Marshal back to JSON
|
||||
reqJSON, err = json.Marshal(reqMap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request map to JSON: %w", err)
|
||||
}
|
||||
|
||||
// Create stream
|
||||
streamHandle, err := c.clientHandle.ChatCompletionStream(string(reqJSON))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create stream: %w", err)
|
||||
}
|
||||
|
||||
// Create a child context from the provided context for cancellation support
|
||||
streamCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
stream := &ChatCompletionStream{
|
||||
stream: streamHandle,
|
||||
ctx: streamCtx,
|
||||
cancel: cancel,
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
|
||||
return stream, nil
|
||||
}
|
||||
@@ -0,0 +1,325 @@
|
||||
package sglang
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestClientConfig tests ClientConfig validation
|
||||
func TestClientConfig(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
config ClientConfig
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid config",
|
||||
config: ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "missing endpoint",
|
||||
config: ClientConfig{
|
||||
Endpoint: "",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "missing tokenizer path",
|
||||
config: ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "both missing",
|
||||
config: ClientConfig{
|
||||
Endpoint: "",
|
||||
TokenizerPath: "",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, err := NewClient(tt.config)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("NewClient() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestChatMessageTypes tests ChatMessage struct and its variants
|
||||
func TestChatMessageTypes(t *testing.T) {
|
||||
msg := ChatMessage{
|
||||
Role: "user",
|
||||
Content: "Hello",
|
||||
}
|
||||
|
||||
if msg.Role != "user" {
|
||||
t.Errorf("Expected role 'user', got '%s'", msg.Role)
|
||||
}
|
||||
if msg.Content != "Hello" {
|
||||
t.Errorf("Expected content 'Hello', got '%s'", msg.Content)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChatCompletionRequestValidation tests ChatCompletionRequest validation
|
||||
func TestChatCompletionRequestValidation(t *testing.T) {
|
||||
// Test valid request
|
||||
req := ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "test"},
|
||||
},
|
||||
Stream: false,
|
||||
}
|
||||
|
||||
if req.Model == "" {
|
||||
t.Error("Expected model to be set")
|
||||
}
|
||||
|
||||
if len(req.Messages) == 0 {
|
||||
t.Error("Expected messages to be non-empty")
|
||||
}
|
||||
|
||||
if req.Messages[0].Role != "user" {
|
||||
t.Errorf("Expected first message role 'user', got '%s'", req.Messages[0].Role)
|
||||
}
|
||||
}
|
||||
|
||||
// TestClientClose tests that Close can be called multiple times safely
|
||||
func TestClientClose(t *testing.T) {
|
||||
// Create a mock client (note: in real tests, you might want to skip this
|
||||
// if it requires actual server connection)
|
||||
config := ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
}
|
||||
|
||||
// Skip if connection fails (expected in unit test environment)
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Skip("Skipping client close test: server not available")
|
||||
}
|
||||
|
||||
// First close should succeed
|
||||
if err := client.Close(); err != nil {
|
||||
t.Errorf("First Close() failed: %v", err)
|
||||
}
|
||||
|
||||
// Second close should also succeed (idempotent)
|
||||
if err := client.Close(); err != nil {
|
||||
t.Errorf("Second Close() failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChatCompletionResponseTypes tests response type structures
|
||||
func TestChatCompletionResponseTypes(t *testing.T) {
|
||||
resp := ChatCompletionResponse{
|
||||
ID: "test-id",
|
||||
Model: "default",
|
||||
Created: 1234567890,
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
Content: "Hello",
|
||||
},
|
||||
FinishReason: "stop",
|
||||
},
|
||||
},
|
||||
Usage: Usage{
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 20,
|
||||
TotalTokens: 30,
|
||||
},
|
||||
}
|
||||
|
||||
if resp.ID != "test-id" {
|
||||
t.Errorf("Expected ID 'test-id', got '%s'", resp.ID)
|
||||
}
|
||||
|
||||
if len(resp.Choices) != 1 {
|
||||
t.Errorf("Expected 1 choice, got %d", len(resp.Choices))
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.Content != "Hello" {
|
||||
t.Errorf("Expected content 'Hello', got '%s'", resp.Choices[0].Message.Content)
|
||||
}
|
||||
|
||||
if resp.Usage.TotalTokens != 30 {
|
||||
t.Errorf("Expected total tokens 30, got %d", resp.Usage.TotalTokens)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamingResponseTypes tests streaming response structures
|
||||
func TestStreamingResponseTypes(t *testing.T) {
|
||||
chunk := ChatCompletionStreamResponse{
|
||||
ID: "stream-id",
|
||||
Created: 1234567890,
|
||||
Choices: []StreamChoice{
|
||||
{
|
||||
Index: 0,
|
||||
Delta: MessageDelta{
|
||||
Content: "Hello",
|
||||
},
|
||||
FinishReason: "",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
if chunk.ID != "stream-id" {
|
||||
t.Errorf("Expected ID 'stream-id', got '%s'", chunk.ID)
|
||||
}
|
||||
|
||||
if len(chunk.Choices) == 0 {
|
||||
t.Error("Expected at least one choice")
|
||||
}
|
||||
|
||||
if chunk.Choices[0].Delta.Content != "Hello" {
|
||||
t.Errorf("Expected delta content 'Hello', got '%s'", chunk.Choices[0].Delta.Content)
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolCallStructure tests Tool and ToolCall structures
|
||||
func TestToolCallStructure(t *testing.T) {
|
||||
tool := Tool{
|
||||
Type: "function",
|
||||
Function: Function{
|
||||
Name: "get_weather",
|
||||
Description: "Get the weather",
|
||||
Parameters: map[string]interface{}{
|
||||
"location": "string",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
if tool.Type != "function" {
|
||||
t.Errorf("Expected tool type 'function', got '%s'", tool.Type)
|
||||
}
|
||||
|
||||
if tool.Function.Name != "get_weather" {
|
||||
t.Errorf("Expected function name 'get_weather', got '%s'", tool.Function.Name)
|
||||
}
|
||||
|
||||
toolCall := ToolCall{
|
||||
ID: "call-123",
|
||||
Type: "function",
|
||||
Function: FunctionCall{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location": "San Francisco"}`,
|
||||
},
|
||||
}
|
||||
|
||||
if toolCall.ID != "call-123" {
|
||||
t.Errorf("Expected tool call ID 'call-123', got '%s'", toolCall.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentClientOperations tests thread safety
|
||||
// This is a basic test that just verifies concurrent calls don't panic
|
||||
func TestConcurrentClientOperations(t *testing.T) {
|
||||
config := ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
}
|
||||
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Skip("Skipping concurrent operations test: server not available")
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Try concurrent Close calls (should not panic or race)
|
||||
done := make(chan bool, 2)
|
||||
|
||||
go func() {
|
||||
client.Close()
|
||||
done <- true
|
||||
}()
|
||||
|
||||
go func() {
|
||||
client.Close()
|
||||
done <- true
|
||||
}()
|
||||
|
||||
<-done
|
||||
<-done
|
||||
}
|
||||
|
||||
// BenchmarkChatCompletionRequest benchmarks request creation
|
||||
func BenchmarkChatCompletionRequest(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
_ = ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "test message"},
|
||||
},
|
||||
Stream: false,
|
||||
Temperature: floatPtr(0.7),
|
||||
MaxCompletionTokens: intPtr(100),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions for benchmarks
|
||||
func floatPtr(f float32) *float32 {
|
||||
return &f
|
||||
}
|
||||
|
||||
func intPtr(i int) *int {
|
||||
return &i
|
||||
}
|
||||
|
||||
// TestContextCancellation tests that cancelled context is handled gracefully.
|
||||
//
|
||||
// NOTE: Currently, the FFI layer is blocking and doesn't actively monitor context cancellation.
|
||||
// This test verifies that the client at least returns an error rather than panicking or
|
||||
// hanging indefinitely when a pre-cancelled context is passed.
|
||||
//
|
||||
// Future: When FFI supports context cancellation (via signals or async operations),
|
||||
// this test should be updated to assert that the error is context.Canceled or wrapped
|
||||
// context cancellation error.
|
||||
func TestContextCancellation(t *testing.T) {
|
||||
config := ClientConfig{
|
||||
Endpoint: "grpc://localhost:20000",
|
||||
TokenizerPath: "/path/to/tokenizer",
|
||||
}
|
||||
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Skip("Skipping context cancellation test: server not available")
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Create a pre-cancelled context
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
req := ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
// Attempt request with cancelled context
|
||||
// Since FFI is blocking, we expect either:
|
||||
// 1. An error from the server/network
|
||||
// 2. The call to complete normally (FFI doesn't check context)
|
||||
// What we DON'T expect is a panic or indefinite hang
|
||||
_, err = client.CreateChatCompletion(ctx, req)
|
||||
if err != nil {
|
||||
t.Logf("Request with cancelled context returned error: %v", err)
|
||||
} else {
|
||||
t.Logf("Request with cancelled context completed (FFI may not support context cancellation)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
// Simple example demonstrating basic usage of SGLang Go SDK
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
|
||||
"github.com/sglang/sglang-go-grpc-sdk"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Get configuration from environment or command line
|
||||
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
|
||||
if endpoint == "" {
|
||||
endpoint = "grpc://localhost:20000"
|
||||
}
|
||||
|
||||
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
|
||||
if tokenizerPath == "" {
|
||||
tokenizerPath = "./examples/tokenizer"
|
||||
}
|
||||
|
||||
// Create client
|
||||
client, err := sglang.NewClient(sglang.ClientConfig{
|
||||
Endpoint: endpoint,
|
||||
TokenizerPath: tokenizerPath,
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Create chat completion request
|
||||
req := sglang.ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []sglang.ChatMessage{
|
||||
{
|
||||
Role: "system",
|
||||
Content: "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: "写一首歌关于夏天",
|
||||
},
|
||||
},
|
||||
Stream: false,
|
||||
Temperature: float32Ptr(0.7),
|
||||
MaxCompletionTokens: intPtr(200),
|
||||
SkipSpecialTokens: true,
|
||||
Tools: nil, // Use nil instead of empty slice to avoid template errors
|
||||
}
|
||||
|
||||
// Create completion
|
||||
ctx := context.Background()
|
||||
resp, err := client.CreateChatCompletion(ctx, req)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to create completion: %v", err)
|
||||
}
|
||||
|
||||
// Print response
|
||||
fmt.Println("=== Response ===")
|
||||
fmt.Printf("ID: %s\n", resp.ID)
|
||||
fmt.Printf("Model: %s\n", resp.Model)
|
||||
fmt.Printf("Created: %d\n", resp.Created)
|
||||
fmt.Println("\nContent:")
|
||||
for _, choice := range resp.Choices {
|
||||
fmt.Println(choice.Message.Content)
|
||||
}
|
||||
fmt.Printf("\nFinish Reason: %s\n", resp.Choices[0].FinishReason)
|
||||
fmt.Printf("\nUsage: Prompt=%d, Completion=%d, Total=%d\n",
|
||||
resp.Usage.PromptTokens,
|
||||
resp.Usage.CompletionTokens,
|
||||
resp.Usage.TotalTokens,
|
||||
)
|
||||
}
|
||||
|
||||
func float32Ptr(f float32) *float32 {
|
||||
return &f
|
||||
}
|
||||
|
||||
func intPtr(i int) *int {
|
||||
return &i
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Simple example runner
|
||||
# Usage: ./run.sh [tokenizer_path] [endpoint]
|
||||
|
||||
# Set library path for Rust FFI library
|
||||
# The library should be in ./lib directory (created by 'make lib')
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
LIB_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)/lib"
|
||||
|
||||
# Check if lib directory exists
|
||||
if [ ! -d "$LIB_DIR" ]; then
|
||||
echo "Error: Library directory not found at $LIB_DIR"
|
||||
echo "Please run 'make lib' first to build and export the library"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Get Python LDFLAGS (needed for Rust FFI that depends on Python)
|
||||
PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
|
||||
|
||||
# Set CGO_LDFLAGS to link with the Rust library
|
||||
export CGO_LDFLAGS="-L${LIB_DIR} -lsglang_router_rs ${PYTHON_LDFLAGS} -ldl"
|
||||
|
||||
# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH
|
||||
if [[ "$OSTYPE" == "darwin"* ]]; then
|
||||
export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}"
|
||||
else
|
||||
export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}"
|
||||
fi
|
||||
|
||||
# Default configuration (can be overridden by environment variables or command line arguments)
|
||||
# Tokenizer path: ../tokenizer (relative to this script)
|
||||
DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}"
|
||||
DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}"
|
||||
|
||||
TOKENIZER_PATH="${1:-${DEFAULT_TOKENIZER_PATH}}"
|
||||
ENDPOINT="${2:-${DEFAULT_ENDPOINT}}"
|
||||
|
||||
echo "Running simple example..."
|
||||
echo "Library path: ${LIB_DIR}"
|
||||
echo "Tokenizer: $TOKENIZER_PATH"
|
||||
echo "Endpoint: $ENDPOINT"
|
||||
echo ""
|
||||
|
||||
cd "$(dirname "${BASH_SOURCE[0]}")"
|
||||
SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" go run main.go
|
||||
@@ -0,0 +1,125 @@
|
||||
// Streaming example demonstrating real-time streaming with SGLang Go SDK
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sglang/sglang-go-grpc-sdk"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Get configuration from environment or command line
|
||||
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
|
||||
if endpoint == "" {
|
||||
endpoint = "grpc://localhost:20000"
|
||||
}
|
||||
|
||||
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
|
||||
if tokenizerPath == "" {
|
||||
tokenizerPath = "./examples/tokenizer"
|
||||
}
|
||||
|
||||
// Create client
|
||||
client, err := sglang.NewClient(sglang.ClientConfig{
|
||||
Endpoint: endpoint,
|
||||
TokenizerPath: tokenizerPath,
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Create streaming chat completion request
|
||||
req := sglang.ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []sglang.ChatMessage{
|
||||
{
|
||||
Role: "system",
|
||||
Content: "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: "写一首春天的诗歌",
|
||||
},
|
||||
},
|
||||
Stream: true,
|
||||
Temperature: float32Ptr(0.7),
|
||||
MaxCompletionTokens: intPtr(500),
|
||||
SkipSpecialTokens: true,
|
||||
Tools: nil, // Use nil instead of empty slice to avoid template errors
|
||||
}
|
||||
|
||||
// Create streaming completion
|
||||
ctx := context.Background()
|
||||
stream, err := client.CreateChatCompletionStream(ctx, req)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to create stream: %v", err)
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
fmt.Println("=== Streaming Response ===")
|
||||
fmt.Println()
|
||||
|
||||
var fullContent strings.Builder
|
||||
chunkCount := 0
|
||||
startTime := time.Now()
|
||||
var firstTokenTime time.Time
|
||||
firstTokenReceived := false
|
||||
|
||||
for {
|
||||
chunk, err := stream.Recv()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
log.Fatalf("Stream error: %v", err)
|
||||
}
|
||||
|
||||
chunkCount++
|
||||
|
||||
// Extract content from delta
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Delta.Content != "" {
|
||||
fmt.Print(choice.Delta.Content)
|
||||
fullContent.WriteString(choice.Delta.Content)
|
||||
|
||||
// Track first token time (TTFT)
|
||||
if !firstTokenReceived {
|
||||
firstTokenTime = time.Now()
|
||||
firstTokenReceived = true
|
||||
ttft := firstTokenTime.Sub(startTime)
|
||||
fmt.Printf("\n[TTFT: %v]\n", ttft)
|
||||
}
|
||||
}
|
||||
|
||||
if choice.FinishReason != "" {
|
||||
fmt.Printf("\n\n[Finished: %s]\n", choice.FinishReason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate metrics
|
||||
if firstTokenReceived {
|
||||
elapsed := time.Since(startTime)
|
||||
tokensPerSecond := float64(fullContent.Len()) / elapsed.Seconds()
|
||||
fmt.Printf("\n=== Metrics ===\n")
|
||||
fmt.Printf("Total chunks: %d\n", chunkCount)
|
||||
fmt.Printf("Total content length: %d characters\n", fullContent.Len())
|
||||
fmt.Printf("Time elapsed: %v\n", elapsed)
|
||||
fmt.Printf("Tokens per second: %.2f\n", tokensPerSecond)
|
||||
}
|
||||
}
|
||||
|
||||
func float32Ptr(f float32) *float32 {
|
||||
return &f
|
||||
}
|
||||
|
||||
func intPtr(i int) *int {
|
||||
return &i
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Streaming example runner
|
||||
# Usage: ./run.sh [tokenizer_path] [endpoint]
|
||||
|
||||
# Set library path for Rust FFI library
|
||||
# The library should be in ./lib directory (created by 'make lib')
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
LIB_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)/lib"
|
||||
|
||||
# Check if lib directory exists
|
||||
if [ ! -d "$LIB_DIR" ]; then
|
||||
echo "Error: Library directory not found at $LIB_DIR"
|
||||
echo "Please run 'make lib' first to build and export the library"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Get Python LDFLAGS (needed for Rust FFI that depends on Python)
|
||||
PYTHON_LDFLAGS=$(python3-config --ldflags --embed 2>/dev/null || python3-config --ldflags 2>/dev/null || echo "")
|
||||
|
||||
# Set CGO_LDFLAGS to link with the Rust library
|
||||
export CGO_LDFLAGS="-L${LIB_DIR} -lsglang_router_rs ${PYTHON_LDFLAGS} -ldl"
|
||||
|
||||
# macOS uses DYLD_LIBRARY_PATH, Linux uses LD_LIBRARY_PATH
|
||||
if [[ "$OSTYPE" == "darwin"* ]]; then
|
||||
export DYLD_LIBRARY_PATH="${LIB_DIR}:${DYLD_LIBRARY_PATH}"
|
||||
else
|
||||
export LD_LIBRARY_PATH="${LIB_DIR}:${LD_LIBRARY_PATH}"
|
||||
fi
|
||||
|
||||
# Default configuration (can be overridden by environment variables or command line arguments)
|
||||
# Tokenizer path: ../tokenizer (relative to this script)
|
||||
DEFAULT_TOKENIZER_PATH="${SGL_TOKENIZER_PATH:-../tokenizer}"
|
||||
DEFAULT_ENDPOINT="${SGL_GRPC_ENDPOINT:-grpc://localhost:20000}"
|
||||
|
||||
TOKENIZER_PATH="${1:-${DEFAULT_TOKENIZER_PATH}}"
|
||||
ENDPOINT="${2:-${DEFAULT_ENDPOINT}}"
|
||||
|
||||
echo "Running streaming example..."
|
||||
echo "Library path: ${LIB_DIR}"
|
||||
echo "Tokenizer: $TOKENIZER_PATH"
|
||||
echo "Endpoint: $ENDPOINT"
|
||||
echo ""
|
||||
|
||||
cd "$(dirname "${BASH_SOURCE[0]}")"
|
||||
SGL_TOKENIZER_PATH="$TOKENIZER_PATH" SGL_GRPC_ENDPOINT="$ENDPOINT" go run main.go
|
||||
@@ -0,0 +1,228 @@
|
||||
//go:build integration
|
||||
// +build integration
|
||||
|
||||
// integration_test.go contains integration tests that require a running SGLang server
|
||||
//
|
||||
// To run these tests:
|
||||
// 1. Start an SGLang server: python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-hf
|
||||
// 2. Run: go test -tags=integration -run TestIntegration
|
||||
|
||||
package sglang
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// getTestConfig returns test configuration from environment or defaults
|
||||
func getTestConfig(t *testing.T) ClientConfig {
|
||||
endpoint := os.Getenv("SGL_GRPC_ENDPOINT")
|
||||
if endpoint == "" {
|
||||
endpoint = "grpc://localhost:20000"
|
||||
}
|
||||
|
||||
tokenizerPath := os.Getenv("SGL_TOKENIZER_PATH")
|
||||
if tokenizerPath == "" {
|
||||
t.Skip("SGL_TOKENIZER_PATH not set")
|
||||
}
|
||||
|
||||
return ClientConfig{
|
||||
Endpoint: endpoint,
|
||||
TokenizerPath: tokenizerPath,
|
||||
}
|
||||
}
|
||||
|
||||
// TestIntegrationNonStreamingCompletion tests non-streaming chat completion
|
||||
func TestIntegrationNonStreamingCompletion(t *testing.T) {
|
||||
config := getTestConfig(t)
|
||||
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
req := ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Say 'Hello, World!' only"},
|
||||
},
|
||||
Stream: false,
|
||||
Temperature: float32Ptr(0.0),
|
||||
MaxCompletionTokens: intPtr(50),
|
||||
}
|
||||
|
||||
resp, err := client.CreateChatCompletion(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateChatCompletion failed: %v", err)
|
||||
}
|
||||
|
||||
if resp.ID == "" {
|
||||
t.Error("Response ID is empty")
|
||||
}
|
||||
|
||||
if len(resp.Choices) == 0 {
|
||||
t.Error("Response has no choices")
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.Content == "" {
|
||||
t.Error("Response content is empty")
|
||||
}
|
||||
|
||||
if resp.Usage == nil || resp.Usage.TotalTokens == 0 {
|
||||
t.Error("Usage information is missing or invalid")
|
||||
}
|
||||
|
||||
t.Logf("Response: %s", resp.Choices[0].Message.Content)
|
||||
t.Logf("Usage: %+v", resp.Usage)
|
||||
}
|
||||
|
||||
// TestIntegrationStreamingCompletion tests streaming chat completion
|
||||
func TestIntegrationStreamingCompletion(t *testing.T) {
|
||||
config := getTestConfig(t)
|
||||
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
req := ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Count from 1 to 5"},
|
||||
},
|
||||
Stream: true,
|
||||
Temperature: float32Ptr(0.0),
|
||||
MaxCompletionTokens: intPtr(100),
|
||||
}
|
||||
|
||||
stream, err := client.CreateChatCompletionStream(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateChatCompletionStream failed: %v", err)
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
chunkCount := 0
|
||||
totalContent := ""
|
||||
|
||||
for {
|
||||
chunk, err := stream.Recv()
|
||||
if err == io.EOF {
|
||||
// io.EOF is expected at end of stream
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("Stream error: %v", err)
|
||||
}
|
||||
|
||||
chunkCount++
|
||||
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Delta.Content != "" {
|
||||
totalContent += choice.Delta.Content
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if chunkCount == 0 {
|
||||
t.Error("Received no chunks from stream")
|
||||
}
|
||||
|
||||
if totalContent == "" {
|
||||
t.Error("Received no content from stream")
|
||||
}
|
||||
|
||||
t.Logf("Received %d chunks with content: %s", chunkCount, totalContent)
|
||||
}
|
||||
|
||||
// TestIntegrationConcurrentRequests tests multiple concurrent requests
|
||||
func TestIntegrationConcurrentRequests(t *testing.T) {
|
||||
config := getTestConfig(t)
|
||||
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
numRequests := 3
|
||||
done := make(chan error, numRequests)
|
||||
|
||||
for i := 0; i < numRequests; i++ {
|
||||
go func(idx int) {
|
||||
req := ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "Say 'test'"},
|
||||
},
|
||||
Stream: false,
|
||||
MaxCompletionTokens: intPtr(50),
|
||||
}
|
||||
|
||||
_, err := client.CreateChatCompletion(ctx, req)
|
||||
done <- err
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Collect results
|
||||
for i := 0; i < numRequests; i++ {
|
||||
if err := <-done; err != nil {
|
||||
t.Errorf("Request %d failed: %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
t.Logf("All %d concurrent requests completed successfully", numRequests)
|
||||
}
|
||||
|
||||
// TestIntegrationContextCancellation tests that context cancellation is handled
|
||||
func TestIntegrationContextCancellation(t *testing.T) {
|
||||
config := getTestConfig(t)
|
||||
|
||||
client, err := NewClient(config)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Create a context that cancels immediately
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
req := ChatCompletionRequest{
|
||||
Model: "default",
|
||||
Messages: []ChatMessage{
|
||||
{Role: "user", Content: "test"},
|
||||
},
|
||||
Stream: false,
|
||||
}
|
||||
|
||||
// Should handle cancelled context gracefully
|
||||
_, err = client.CreateChatCompletion(ctx, req)
|
||||
if err == nil {
|
||||
t.Error("Expected error from cancelled context")
|
||||
}
|
||||
|
||||
t.Logf("Cancelled context handled: %v", err)
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
func float32Ptr(f float32) *float32 {
|
||||
return &f
|
||||
}
|
||||
|
||||
func intPtr(i int) *int {
|
||||
return &i
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
// Package ffi provides Go bindings for SGLang's Rust FFI (Foreign Function Interface).
|
||||
//
|
||||
// This package wraps the Rust FFI layer of SGLang, providing low-level access to:
|
||||
// - Client creation and connection management
|
||||
// - Chat completion streaming
|
||||
// - Stream reading and response conversion
|
||||
// - Memory management for C strings
|
||||
//
|
||||
// Internal use only: This package is intended for internal use by the sglang package.
|
||||
// End users should use the public sglang package instead.
|
||||
package ffi
|
||||
|
||||
/*
|
||||
#cgo LDFLAGS: -lsglang_router_rs -ldl
|
||||
#include <stdlib.h>
|
||||
#include <stdint.h>
|
||||
|
||||
// Error codes
|
||||
typedef enum {
|
||||
SGL_ERROR_SUCCESS = 0,
|
||||
SGL_ERROR_INVALID_ARGUMENT = 1,
|
||||
SGL_ERROR_TOKENIZATION_ERROR = 2,
|
||||
SGL_ERROR_PARSING_ERROR = 3,
|
||||
SGL_ERROR_MEMORY_ERROR = 4,
|
||||
SGL_ERROR_UNKNOWN = 99
|
||||
} SglErrorCode;
|
||||
|
||||
// Opaque handles
|
||||
typedef void* SglangClientHandle;
|
||||
typedef void* SglangStreamHandle;
|
||||
|
||||
// Client SDK functions
|
||||
SglangClientHandle* sgl_client_create(const char* endpoint, const char* tokenizer_path, char** error_out);
|
||||
void sgl_client_free(SglangClientHandle* handle);
|
||||
SglErrorCode sgl_client_chat_completion_stream(SglangClientHandle* client_handle, const char* request_json, SglangStreamHandle** stream_handle_out, char** error_out);
|
||||
SglErrorCode sgl_stream_read_next(SglangStreamHandle* stream_handle, char** response_json_out, int* is_done_out, char** error_out);
|
||||
void sgl_stream_free(SglangStreamHandle* handle);
|
||||
void sgl_free_string(char* s);
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// ErrorCode represents FFI error codes returned by Rust functions.
|
||||
//
|
||||
// These codes indicate the result of FFI operations. Use Error() to get a human-readable
|
||||
// error message.
|
||||
type ErrorCode int
|
||||
|
||||
const (
|
||||
// ErrorSuccess indicates the operation completed successfully
|
||||
ErrorSuccess ErrorCode = 0
|
||||
// ErrorInvalidArgument indicates invalid arguments were passed to the FFI function
|
||||
ErrorInvalidArgument ErrorCode = 1
|
||||
// ErrorTokenizationError indicates an error during tokenization
|
||||
ErrorTokenizationError ErrorCode = 2
|
||||
// ErrorParsingError indicates an error parsing the response or request
|
||||
ErrorParsingError ErrorCode = 3
|
||||
// ErrorMemoryError indicates a memory allocation error
|
||||
ErrorMemoryError ErrorCode = 4
|
||||
// ErrorUnknown indicates an unclassified error
|
||||
ErrorUnknown ErrorCode = 99
|
||||
)
|
||||
|
||||
// Error implements the error interface for ErrorCode.
|
||||
func (e ErrorCode) Error() string {
|
||||
switch e {
|
||||
case ErrorSuccess:
|
||||
return "success"
|
||||
case ErrorInvalidArgument:
|
||||
return "invalid argument"
|
||||
case ErrorTokenizationError:
|
||||
return "tokenization error"
|
||||
case ErrorParsingError:
|
||||
return "parsing error"
|
||||
case ErrorMemoryError:
|
||||
return "memory error"
|
||||
case ErrorUnknown:
|
||||
return "unknown error"
|
||||
default:
|
||||
return fmt.Sprintf("unknown error code: %d", e)
|
||||
}
|
||||
}
|
||||
|
||||
// SglangClientHandle wraps the Rust client SDK FFI handle.
|
||||
//
|
||||
// This struct maintains a connection to the SGLang gRPC server and is used
|
||||
// to create streams and manage the underlying Rust client resources.
|
||||
type SglangClientHandle struct {
|
||||
handle *C.SglangClientHandle
|
||||
}
|
||||
|
||||
// NewClient creates a new SGLang client handle via FFI.
|
||||
//
|
||||
// This function initializes the Rust client with the given endpoint and tokenizer path.
|
||||
//
|
||||
// Parameters:
|
||||
// - endpoint: gRPC endpoint URL (e.g., "grpc://localhost:20000")
|
||||
// - tokenizerPath: Path to tokenizer directory
|
||||
//
|
||||
// Returns:
|
||||
// - *SglangClientHandle: A new client handle
|
||||
// - error: An error if client creation failed
|
||||
func NewClient(endpoint, tokenizerPath string) (*SglangClientHandle, error) {
|
||||
cEndpoint := C.CString(endpoint)
|
||||
defer C.free(unsafe.Pointer(cEndpoint))
|
||||
|
||||
cTokenizerPath := C.CString(tokenizerPath)
|
||||
defer C.free(unsafe.Pointer(cTokenizerPath))
|
||||
|
||||
var errorPtr *C.char
|
||||
handle := C.sgl_client_create(cEndpoint, cTokenizerPath, &errorPtr)
|
||||
|
||||
if handle == nil {
|
||||
errorMsg := ""
|
||||
if errorPtr != nil {
|
||||
errorMsg = C.GoString(errorPtr)
|
||||
C.sgl_free_string(errorPtr)
|
||||
}
|
||||
if errorMsg == "" {
|
||||
errorMsg = "failed to create client"
|
||||
}
|
||||
return nil, fmt.Errorf("%s", errorMsg)
|
||||
}
|
||||
|
||||
return &SglangClientHandle{handle: handle}, nil
|
||||
}
|
||||
|
||||
// Free releases the client handle
|
||||
func (h *SglangClientHandle) Free() {
|
||||
if h.handle != nil {
|
||||
C.sgl_client_free(h.handle)
|
||||
h.handle = nil
|
||||
}
|
||||
}
|
||||
|
||||
// ChatCompletionStream creates a streaming chat completion request
|
||||
func (h *SglangClientHandle) ChatCompletionStream(requestJSON string) (*SglangStreamHandle, error) {
|
||||
if h.handle == nil {
|
||||
return nil, fmt.Errorf("client handle is nil")
|
||||
}
|
||||
|
||||
cRequestJSON := C.CString(requestJSON)
|
||||
defer C.free(unsafe.Pointer(cRequestJSON))
|
||||
|
||||
var streamHandle *C.SglangStreamHandle
|
||||
var errorPtr *C.char
|
||||
|
||||
result := C.sgl_client_chat_completion_stream(
|
||||
h.handle,
|
||||
cRequestJSON,
|
||||
&streamHandle,
|
||||
&errorPtr,
|
||||
)
|
||||
|
||||
if ErrorCode(result) != ErrorSuccess {
|
||||
errorMsg := ""
|
||||
if errorPtr != nil {
|
||||
errorMsg = C.GoString(errorPtr)
|
||||
C.sgl_free_string(errorPtr)
|
||||
}
|
||||
if errorMsg == "" {
|
||||
errorMsg = fmt.Sprintf("error code %d", result)
|
||||
}
|
||||
return nil, fmt.Errorf("%s", errorMsg)
|
||||
}
|
||||
|
||||
if streamHandle == nil {
|
||||
return nil, fmt.Errorf("stream handle is nil")
|
||||
}
|
||||
|
||||
return &SglangStreamHandle{handle: streamHandle}, nil
|
||||
}
|
||||
|
||||
// SglangStreamHandle wraps the Rust stream FFI handle
|
||||
type SglangStreamHandle struct {
|
||||
handle *C.SglangStreamHandle
|
||||
}
|
||||
|
||||
// ReadNext reads the next chunk from the stream
|
||||
// Returns: (responseJSON, isDone, error)
|
||||
func (h *SglangStreamHandle) ReadNext() (string, bool, error) {
|
||||
if h.handle == nil {
|
||||
return "", true, fmt.Errorf("stream handle is nil")
|
||||
}
|
||||
|
||||
var responseJSON *C.char
|
||||
var isDone C.int
|
||||
var errorPtr *C.char
|
||||
|
||||
result := C.sgl_stream_read_next(
|
||||
h.handle,
|
||||
&responseJSON,
|
||||
&isDone,
|
||||
&errorPtr,
|
||||
)
|
||||
|
||||
if ErrorCode(result) != ErrorSuccess {
|
||||
errorMsg := ""
|
||||
if errorPtr != nil {
|
||||
errorMsg = C.GoString(errorPtr)
|
||||
C.sgl_free_string(errorPtr)
|
||||
}
|
||||
if errorMsg == "" {
|
||||
errorMsg = fmt.Sprintf("error code %d", result)
|
||||
}
|
||||
return "", isDone == 1, fmt.Errorf("%s", errorMsg)
|
||||
}
|
||||
|
||||
responseStr := ""
|
||||
if responseJSON != nil {
|
||||
responseStr = C.GoString(responseJSON)
|
||||
C.sgl_free_string(responseJSON)
|
||||
}
|
||||
|
||||
return responseStr, isDone == 1, nil
|
||||
}
|
||||
|
||||
// Free releases the stream handle
|
||||
func (h *SglangStreamHandle) Free() {
|
||||
if h.handle != nil {
|
||||
C.sgl_stream_free(h.handle)
|
||||
h.handle = nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
//! Client SDK FFI functions
|
||||
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::{c_char};
|
||||
use std::ptr;
|
||||
use std::sync::Arc;
|
||||
use tokio::runtime::Runtime;
|
||||
use once_cell::sync::Lazy;
|
||||
use uuid::Uuid;
|
||||
|
||||
use sgl_model_gateway::tokenizer::create_tokenizer_from_file;
|
||||
use sgl_model_gateway::tokenizer::traits::Tokenizer;
|
||||
use sgl_model_gateway::grpc_client::sglang_scheduler::SglangSchedulerClient;
|
||||
use sgl_model_gateway::protocols::chat::ChatCompletionRequest;
|
||||
use sgl_model_gateway::routers::grpc::utils::{process_chat_messages, generate_tool_constraints};
|
||||
|
||||
use super::error::{SglErrorCode, set_error_message};
|
||||
use super::grpc_converter::sgl_grpc_response_converter_create;
|
||||
use super::tokenizer::TokenizerHandle;
|
||||
use super::stream::SglangStreamHandle;
|
||||
|
||||
/// Global tokio runtime for async operations
|
||||
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
|
||||
Runtime::new().expect("Failed to create tokio runtime for client FFI")
|
||||
});
|
||||
|
||||
/// Handle for complete client SDK (gRPC client + tokenizer)
|
||||
/// This handle manages the connection to sglang and provides a complete SDK interface
|
||||
pub struct SglangClientHandle {
|
||||
pub(crate) client: Arc<SglangSchedulerClient>,
|
||||
pub(crate) tokenizer: Arc<dyn Tokenizer>,
|
||||
}
|
||||
|
||||
/// Handle for streaming request (includes prompt token count)
|
||||
#[allow(dead_code)]
|
||||
pub struct StreamRequestState {
|
||||
pub(crate) prompt_tokens: i32, // Number of prompt tokens for this request
|
||||
}
|
||||
|
||||
/// Create a new SGLang client handle
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `endpoint` - gRPC endpoint (e.g., "grpc://localhost:20000")
|
||||
/// * `tokenizer_path` - Path to tokenizer directory
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * Pointer to SglangClientHandle on success, null on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_client_create(
|
||||
endpoint: *const c_char,
|
||||
tokenizer_path: *const c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> *mut SglangClientHandle {
|
||||
if endpoint.is_null() || tokenizer_path.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
|
||||
let endpoint_str = match CStr::from_ptr(endpoint).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in endpoint");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
let tokenizer_path_str = match CStr::from_ptr(tokenizer_path).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in tokenizer_path");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
// Create tokenizer
|
||||
let tokenizer = match create_tokenizer_from_file(tokenizer_path_str) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create tokenizer: {}", e));
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
// Create gRPC client
|
||||
let client = match RUNTIME.block_on(async {
|
||||
SglangSchedulerClient::connect(endpoint_str).await
|
||||
}) {
|
||||
Ok(c) => Arc::new(c),
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to connect to endpoint: {}", e));
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
Box::into_raw(Box::new(SglangClientHandle {
|
||||
client,
|
||||
tokenizer,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Free a client handle
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_client_free(handle: *mut SglangClientHandle) {
|
||||
if !handle.is_null() {
|
||||
let _ = Box::from_raw(handle);
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a chat completion request and start streaming
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `client_handle` - Client handle
|
||||
/// * `request_json` - OpenAI ChatCompletionRequest as JSON string
|
||||
/// * `stream_handle_out` - Pointer to receive stream handle
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_client_chat_completion_stream(
|
||||
client_handle: *mut SglangClientHandle,
|
||||
request_json: *const c_char,
|
||||
stream_handle_out: *mut *mut SglangStreamHandle,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if client_handle.is_null() || request_json.is_null() || stream_handle_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let request_str = match CStr::from_ptr(request_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in request_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
let client_ref = &*client_handle;
|
||||
let client = Arc::clone(&client_ref.client);
|
||||
let tokenizer = Arc::clone(&client_ref.tokenizer);
|
||||
|
||||
// Parse OpenAI ChatCompletionRequest
|
||||
let chat_request: ChatCompletionRequest = match serde_json::from_str(request_str) {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse request JSON: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Process messages and apply chat template
|
||||
let processed_messages = match process_chat_messages(&chat_request, tokenizer.as_ref()) {
|
||||
Ok(msgs) => msgs,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to process messages: {}", e));
|
||||
return SglErrorCode::TokenizationError;
|
||||
}
|
||||
};
|
||||
|
||||
// Tokenize
|
||||
let token_ids = match tokenizer.encode(&processed_messages.text) {
|
||||
Ok(encoding) => encoding.token_ids().to_vec(),
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to tokenize: {}", e));
|
||||
return SglErrorCode::TokenizationError;
|
||||
}
|
||||
};
|
||||
let prompt_tokens = token_ids.len() as i32; // Save prompt token count
|
||||
|
||||
// Generate tool constraints if needed
|
||||
let tool_constraint = if let Some(tools) = chat_request.tools.as_ref() {
|
||||
match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) {
|
||||
Ok(Some((constraint_type, constraint_value))) => Some((constraint_type, constraint_value)),
|
||||
Ok(None) => None,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to generate tool constraints: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Build GenerateRequest
|
||||
let request_id = format!("chatcmpl-{}", Uuid::new_v4());
|
||||
let proto_request = match client.build_generate_request_from_chat(
|
||||
request_id.clone(),
|
||||
&chat_request,
|
||||
processed_messages.text,
|
||||
token_ids,
|
||||
processed_messages.multimodal_inputs,
|
||||
tool_constraint,
|
||||
) {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to build generate request: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Send request and get stream
|
||||
let stream = match RUNTIME.block_on(async {
|
||||
client.generate(proto_request).await
|
||||
}) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to send request: {}", e));
|
||||
return SglErrorCode::UnknownError;
|
||||
}
|
||||
};
|
||||
|
||||
// Create response converter
|
||||
let tools_json = chat_request.tools.as_ref()
|
||||
.and_then(|t| serde_json::to_string(t).ok())
|
||||
.map(|s| CString::new(s).unwrap().into_raw());
|
||||
let tool_choice_json = chat_request.tool_choice.as_ref()
|
||||
.and_then(|tc| serde_json::to_string(tc).ok())
|
||||
.map(|s| CString::new(s).unwrap().into_raw());
|
||||
let stop_json = chat_request.stop.as_ref()
|
||||
.and_then(|s| serde_json::to_string(s).ok())
|
||||
.map(|s| CString::new(s).unwrap().into_raw());
|
||||
let stop_token_ids_json = chat_request.stop_token_ids.as_ref()
|
||||
.and_then(|ids| serde_json::to_string(ids).ok())
|
||||
.map(|s| CString::new(s).unwrap().into_raw());
|
||||
|
||||
// Create tokenizer handle for converter (we'll create a temporary one)
|
||||
let tokenizer_handle = Box::into_raw(Box::new(TokenizerHandle {
|
||||
tokenizer: Arc::clone(&tokenizer),
|
||||
}));
|
||||
|
||||
let converter = sgl_grpc_response_converter_create(
|
||||
tokenizer_handle,
|
||||
CString::new(chat_request.model.clone()).unwrap().as_ptr(),
|
||||
CString::new(request_id.clone()).unwrap().as_ptr(),
|
||||
tools_json.unwrap_or(ptr::null_mut()),
|
||||
tool_choice_json.unwrap_or(ptr::null_mut()),
|
||||
stop_json.unwrap_or(ptr::null_mut()),
|
||||
stop_token_ids_json.unwrap_or(ptr::null_mut()),
|
||||
if chat_request.skip_special_tokens { 1 } else { 0 },
|
||||
error_out,
|
||||
);
|
||||
|
||||
// Free temporary tokenizer handle (converter now owns the tokenizer)
|
||||
let _ = Box::from_raw(tokenizer_handle);
|
||||
|
||||
if converter.is_null() {
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
|
||||
// Clean up temporary CStrings
|
||||
if let Some(ptr) = tools_json {
|
||||
let _ = CString::from_raw(ptr);
|
||||
}
|
||||
if let Some(ptr) = tool_choice_json {
|
||||
let _ = CString::from_raw(ptr);
|
||||
}
|
||||
if let Some(ptr) = stop_json {
|
||||
let _ = CString::from_raw(ptr);
|
||||
}
|
||||
if let Some(ptr) = stop_token_ids_json {
|
||||
let _ = CString::from_raw(ptr);
|
||||
}
|
||||
|
||||
// Create converter handle and set initial_prompt_tokens immediately
|
||||
let mut converter_handle = *Box::from_raw(converter);
|
||||
converter_handle.initial_prompt_tokens = Some(prompt_tokens);
|
||||
|
||||
// Create stream handle with prompt_tokens
|
||||
*stream_handle_out = Box::into_raw(Box::new(SglangStreamHandle {
|
||||
stream: Arc::new(tokio::sync::Mutex::new(stream)),
|
||||
converter: Arc::new(tokio::sync::Mutex::new(converter_handle)),
|
||||
client: Arc::clone(&client),
|
||||
prompt_tokens,
|
||||
}));
|
||||
|
||||
SglErrorCode::Success
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Error handling for FFI functions
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::raw::c_char;
|
||||
use std::ptr;
|
||||
|
||||
/// Error codes returned by FFI functions
|
||||
#[repr(C)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SglErrorCode {
|
||||
Success = 0,
|
||||
InvalidArgument = 1,
|
||||
TokenizationError = 2,
|
||||
ParsingError = 3,
|
||||
MemoryError = 4,
|
||||
UnknownError = 99,
|
||||
}
|
||||
|
||||
/// Helper to set error message in FFI output parameter
|
||||
pub fn set_error_message(error_out: *mut *mut c_char, message: &str) {
|
||||
unsafe {
|
||||
if !error_out.is_null() {
|
||||
if let Ok(cstr) = CString::new(message) {
|
||||
*error_out = cstr.into_raw();
|
||||
} else {
|
||||
*error_out = ptr::null_mut();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to set error message from format string
|
||||
pub fn set_error_message_fmt(error_out: *mut *mut c_char, fmt: std::fmt::Arguments) {
|
||||
if !error_out.is_null() {
|
||||
let msg = format!("{}", fmt);
|
||||
set_error_message(error_out, &msg);
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to clear error message
|
||||
pub fn clear_error_message(error_out: *mut *mut c_char) {
|
||||
unsafe {
|
||||
if !error_out.is_null() {
|
||||
*error_out = ptr::null_mut();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions for error handling
|
||||
// Note: Some helper functions are kept for potential future use
|
||||
@@ -0,0 +1,758 @@
|
||||
//! gRPC response converter FFI functions
|
||||
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::{c_char, c_int};
|
||||
use std::ptr;
|
||||
use std::sync::Arc;
|
||||
use std::collections::HashMap;
|
||||
use serde_json::Value;
|
||||
use tokio::runtime::Runtime;
|
||||
use once_cell::sync::Lazy;
|
||||
|
||||
use sgl_model_gateway::tokenizer::traits::Tokenizer;
|
||||
use sgl_model_gateway::tokenizer::stream::DecodeStream;
|
||||
use sgl_model_gateway::tool_parser::ToolParser;
|
||||
use sgl_model_gateway::protocols::common::{Tool, ToolChoice, ToolChoiceValue, ToolCallDelta, FunctionCallDelta, Usage, StringOrArray};
|
||||
use sgl_model_gateway::tokenizer::stop::StopSequenceDecoder;
|
||||
use sgl_model_gateway::grpc_client::sglang_proto as proto;
|
||||
|
||||
use super::error::{SglErrorCode, set_error_message, clear_error_message};
|
||||
use super::tokenizer::TokenizerHandle;
|
||||
use super::utils::generate_tool_call_id;
|
||||
|
||||
/// Global parser factory (initialized once)
|
||||
// Use the re-exported ParserFactory from tool_parser module
|
||||
static PARSER_FACTORY: Lazy<sgl_model_gateway::tool_parser::ParserFactory> = Lazy::new(|| {
|
||||
// ParserFactory is re-exported from tool_parser::factory, so we can use it directly
|
||||
sgl_model_gateway::tool_parser::ParserFactory::default()
|
||||
});
|
||||
|
||||
/// Global tokio runtime for async operations
|
||||
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
|
||||
Runtime::new().expect("Failed to create tokio runtime for gRPC converter FFI")
|
||||
});
|
||||
|
||||
/// Handle for gRPC response converter (maintains state for streaming)
|
||||
#[repr(C)]
|
||||
pub struct GrpcResponseConverterHandle {
|
||||
pub(crate) tokenizer: Arc<dyn Tokenizer>,
|
||||
pub(crate) tool_parser: Option<Arc<tokio::sync::Mutex<Box<dyn ToolParser>>>>,
|
||||
pub(crate) stop_decoder: Option<Arc<tokio::sync::Mutex<StopSequenceDecoder>>>,
|
||||
pub(crate) model: String,
|
||||
pub(crate) request_id: String,
|
||||
pub(crate) created: u64,
|
||||
pub(crate) system_fingerprint: Option<String>,
|
||||
pub(crate) tools: Option<Vec<Tool>>,
|
||||
pub(crate) tool_choice: Option<ToolChoice>,
|
||||
pub(crate) history_tool_calls_count: usize,
|
||||
pub(crate) stream_buffers: HashMap<u32, String>, // Per-index text buffers
|
||||
pub(crate) decode_streams: HashMap<u32, DecodeStream>, // Per-index incremental decoders
|
||||
pub(crate) has_tool_calls: HashMap<u32, bool>, // Track if tool calls were emitted
|
||||
pub(crate) is_first_chunk: HashMap<u32, bool>, // Track first chunk per index
|
||||
pub(crate) prompt_tokens: HashMap<u32, i32>, // Track prompt tokens per index (from chunks)
|
||||
pub(crate) completion_tokens: HashMap<u32, i32>, // Track completion tokens per index (cumulative)
|
||||
pub(crate) initial_prompt_tokens: Option<i32>, // Initial prompt tokens from request (if available)
|
||||
pub(crate) skip_special_tokens: bool, // Whether to skip special tokens when decoding
|
||||
}
|
||||
|
||||
/// Create a gRPC response converter handle
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `tokenizer_handle` - Tokenizer handle (must be valid)
|
||||
/// * `model` - Model name
|
||||
/// * `request_id` - Request ID
|
||||
/// * `tools_json` - Optional JSON array of tools
|
||||
/// * `tool_choice_json` - Optional JSON object for tool_choice
|
||||
/// * `stop` - Optional stop sequences (JSON array)
|
||||
/// * `stop_token_ids` - Optional stop token IDs (JSON array)
|
||||
/// * `skip_special_tokens` - Whether to skip special tokens
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * Pointer to GrpcResponseConverterHandle on success, null on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_grpc_response_converter_create(
|
||||
tokenizer_handle: *mut TokenizerHandle,
|
||||
model: *const c_char,
|
||||
request_id: *const c_char,
|
||||
tools_json: *const c_char,
|
||||
tool_choice_json: *const c_char,
|
||||
stop: *const c_char,
|
||||
stop_token_ids: *const c_char,
|
||||
skip_special_tokens: c_int,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> *mut GrpcResponseConverterHandle {
|
||||
if tokenizer_handle.is_null() || model.is_null() || request_id.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
|
||||
let model_str = match CStr::from_ptr(model).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in model");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
let request_id_str = match CStr::from_ptr(request_id).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in request_id");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
let handle_ref = &*tokenizer_handle;
|
||||
let tokenizer = Arc::clone(&handle_ref.tokenizer);
|
||||
|
||||
// Parse tools if provided
|
||||
let tools: Option<Vec<Tool>> = if !tools_json.is_null() {
|
||||
match CStr::from_ptr(tools_json).to_str() {
|
||||
Ok(s) => serde_json::from_str::<Vec<Tool>>(s).ok(),
|
||||
Err(_) => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Parse tool_choice if provided
|
||||
let tool_choice: Option<ToolChoice> = if !tool_choice_json.is_null() {
|
||||
match CStr::from_ptr(tool_choice_json).to_str() {
|
||||
Ok(s) => serde_json::from_str::<ToolChoice>(s).ok(),
|
||||
Err(_) => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Parse stop sequences
|
||||
let stop: Option<StringOrArray> = if !stop.is_null() {
|
||||
let stop_str = match CStr::from_ptr(stop).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => return ptr::null_mut(),
|
||||
};
|
||||
serde_json::from_str::<StringOrArray>(stop_str).ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Parse stop token IDs
|
||||
let stop_token_ids: Option<Vec<u32>> = if !stop_token_ids.is_null() {
|
||||
let ids_str = match CStr::from_ptr(stop_token_ids).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => return ptr::null_mut(),
|
||||
};
|
||||
serde_json::from_str::<Vec<u32>>(ids_str).ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Create stop decoder if needed
|
||||
let stop_decoder = if stop.is_some() || stop_token_ids.is_some() {
|
||||
Some(Arc::new(tokio::sync::Mutex::new(
|
||||
sgl_model_gateway::routers::grpc::utils::create_stop_decoder(
|
||||
&tokenizer,
|
||||
stop.as_ref(),
|
||||
stop_token_ids.as_ref(),
|
||||
skip_special_tokens != 0,
|
||||
false, // no_stop_trim
|
||||
),
|
||||
)))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Create tool parser if tools are provided
|
||||
let tool_parser = if tools.is_some() {
|
||||
PARSER_FACTORY.registry().create_for_model(model_str)
|
||||
.map(|p| Arc::new(tokio::sync::Mutex::new(p)))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Get system fingerprint from model (simplified)
|
||||
let system_fingerprint = Some("fp_placeholder".to_string()); // TODO: Get actual fingerprint
|
||||
|
||||
Box::into_raw(Box::new(GrpcResponseConverterHandle {
|
||||
tokenizer,
|
||||
tool_parser,
|
||||
stop_decoder,
|
||||
model: model_str.to_string(),
|
||||
request_id: request_id_str.to_string(),
|
||||
created: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs(),
|
||||
system_fingerprint,
|
||||
tools,
|
||||
tool_choice,
|
||||
history_tool_calls_count: 0,
|
||||
stream_buffers: HashMap::new(),
|
||||
decode_streams: HashMap::new(),
|
||||
has_tool_calls: HashMap::new(),
|
||||
is_first_chunk: HashMap::new(),
|
||||
prompt_tokens: HashMap::new(),
|
||||
completion_tokens: HashMap::new(),
|
||||
initial_prompt_tokens: None, // Will be set from stream handle
|
||||
skip_special_tokens: skip_special_tokens != 0,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Convert a gRPC GenerateResponse chunk to OpenAI format
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Converter handle
|
||||
/// * `response_json` - JSON string of proto.GenerateResponse
|
||||
/// * `result_json_out` - Pointer to receive OpenAI format JSON (must be freed with sgl_free_string)
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_grpc_response_converter_convert_chunk(
|
||||
handle: *mut GrpcResponseConverterHandle,
|
||||
response_json: *const c_char,
|
||||
result_json_out: *mut *mut c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || response_json.is_null() || result_json_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let response_str = match CStr::from_ptr(response_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in response_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse proto.GenerateResponse from JSON
|
||||
let json_value: Value = match serde_json::from_str(response_str) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse response JSON: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
// Build proto::GenerateResponse from JSON value
|
||||
let mut proto_response = proto::GenerateResponse {
|
||||
request_id: json_value.get("request_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
response: None,
|
||||
};
|
||||
|
||||
// Parse the response oneof field
|
||||
if let Some(chunk_json) = json_value.get("chunk") {
|
||||
let chunk = proto::GenerateStreamChunk {
|
||||
token_ids: chunk_json.get("token_ids")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as u32)).collect())
|
||||
.unwrap_or_default(),
|
||||
prompt_tokens: chunk_json.get("prompt_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|n| n as i32)
|
||||
.unwrap_or(0),
|
||||
completion_tokens: chunk_json.get("completion_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|n| n as i32)
|
||||
.unwrap_or(0),
|
||||
cached_tokens: chunk_json.get("cached_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|n| n as i32)
|
||||
.unwrap_or(0),
|
||||
output_logprobs: None,
|
||||
hidden_states: vec![],
|
||||
input_logprobs: None,
|
||||
index: 0,
|
||||
};
|
||||
proto_response.response = Some(proto::generate_response::Response::Chunk(chunk));
|
||||
} else if let Some(complete_json) = json_value.get("complete") {
|
||||
let complete = proto::GenerateComplete {
|
||||
output_ids: complete_json.get("output_ids")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as u32)).collect())
|
||||
.unwrap_or_default(),
|
||||
finish_reason: complete_json.get("finish_reason")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
prompt_tokens: complete_json.get("prompt_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|n| n as i32)
|
||||
.unwrap_or(0),
|
||||
completion_tokens: complete_json.get("completion_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|n| n as i32)
|
||||
.unwrap_or(0),
|
||||
cached_tokens: complete_json.get("cached_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|n| n as i32)
|
||||
.unwrap_or(0),
|
||||
output_logprobs: None,
|
||||
all_hidden_states: vec![],
|
||||
input_logprobs: None,
|
||||
matched_stop: None,
|
||||
index: 0,
|
||||
};
|
||||
proto_response.response = Some(proto::generate_response::Response::Complete(complete));
|
||||
} else if let Some(error_json) = json_value.get("error") {
|
||||
let error = proto::GenerateError {
|
||||
message: error_json.get("message")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
http_status_code: error_json.get("http_status_code")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("500")
|
||||
.to_string(),
|
||||
details: error_json.get("details")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
};
|
||||
proto_response.response = Some(proto::generate_response::Response::Error(error));
|
||||
} else {
|
||||
set_error_message(error_out, "Response JSON must contain 'chunk', 'complete', or 'error' field");
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
|
||||
let handle_ref = &mut *handle;
|
||||
let tokenizer = Arc::clone(&handle_ref.tokenizer);
|
||||
let model = handle_ref.model.clone();
|
||||
let request_id = handle_ref.request_id.clone();
|
||||
let created = handle_ref.created;
|
||||
let system_fingerprint = handle_ref.system_fingerprint.clone();
|
||||
|
||||
// Use tokio runtime to run async code
|
||||
let result = RUNTIME.block_on(async {
|
||||
convert_proto_chunk_to_openai(
|
||||
proto_response,
|
||||
handle_ref,
|
||||
&tokenizer,
|
||||
&model,
|
||||
&request_id,
|
||||
created,
|
||||
system_fingerprint.as_deref(),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(Some(openai_response)) => {
|
||||
// Serialize to JSON
|
||||
let result_str = match serde_json::to_string(&openai_response) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to serialize response: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
let result_cstr = match CString::new(result_str) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
|
||||
*result_json_out = result_cstr.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Ok(None) => {
|
||||
// No response to send (e.g., empty chunk)
|
||||
let empty = CString::new("").unwrap();
|
||||
*result_json_out = empty.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Conversion error: {}", e));
|
||||
SglErrorCode::ParsingError
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper function to convert proto chunk to OpenAI format
|
||||
pub(crate) async fn convert_proto_chunk_to_openai(
|
||||
proto_response: proto::GenerateResponse,
|
||||
handle: &mut GrpcResponseConverterHandle,
|
||||
tokenizer: &Arc<dyn Tokenizer>,
|
||||
model: &str,
|
||||
request_id: &str,
|
||||
created: u64,
|
||||
system_fingerprint: Option<&str>,
|
||||
) -> Result<Option<sgl_model_gateway::protocols::chat::ChatCompletionStreamResponse>, String> {
|
||||
use sgl_model_gateway::grpc_client::sglang_proto::generate_response::Response::*;
|
||||
use sgl_model_gateway::protocols::chat::{ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice};
|
||||
|
||||
match proto_response.response {
|
||||
Some(Chunk(chunk)) => {
|
||||
let index = chunk.index;
|
||||
|
||||
// Mark as not first chunk if we've seen this index before
|
||||
let is_first = handle.is_first_chunk.entry(index).or_insert(true);
|
||||
let first_chunk = *is_first;
|
||||
*is_first = false;
|
||||
|
||||
// Track token counts from chunks (cumulative values from proto)
|
||||
// These are cumulative values, so we always use the latest value
|
||||
// For prompt_tokens, if chunk value is 0, preserve existing value or use initial_prompt_tokens
|
||||
// This prevents overwriting valid prompt_tokens with 0
|
||||
if chunk.prompt_tokens > 0 {
|
||||
handle.prompt_tokens.insert(index, chunk.prompt_tokens);
|
||||
} else {
|
||||
// If chunk.prompt_tokens is 0, try to preserve existing value or use initial_prompt_tokens
|
||||
if !handle.prompt_tokens.contains_key(&index) {
|
||||
// No existing value, try to use initial_prompt_tokens
|
||||
if let Some(initial_prompt) = handle.initial_prompt_tokens {
|
||||
handle.prompt_tokens.insert(index, initial_prompt);
|
||||
}
|
||||
}
|
||||
// If existing value exists, keep it (don't overwrite with 0)
|
||||
}
|
||||
// For completion_tokens, always update (even if 0) as it's cumulative
|
||||
handle.completion_tokens.insert(index, chunk.completion_tokens);
|
||||
|
||||
// Process tokens through stop decoder if available, otherwise use incremental decoder
|
||||
let chunk_text = if let Some(ref stop_decoder) = handle.stop_decoder {
|
||||
let mut decoder_guard = stop_decoder.lock().await;
|
||||
let mut text = String::new();
|
||||
for &token_id in &chunk.token_ids {
|
||||
match decoder_guard.process_token(token_id).unwrap_or_else(|_| {
|
||||
sgl_model_gateway::tokenizer::stop::SequenceDecoderOutput::Held
|
||||
}) {
|
||||
sgl_model_gateway::tokenizer::stop::SequenceDecoderOutput::Text(t) => {
|
||||
text.push_str(&t);
|
||||
}
|
||||
sgl_model_gateway::tokenizer::stop::SequenceDecoderOutput::StoppedWithText(t) => {
|
||||
text.push_str(&t);
|
||||
break;
|
||||
}
|
||||
sgl_model_gateway::tokenizer::stop::SequenceDecoderOutput::Stopped => {
|
||||
break;
|
||||
}
|
||||
sgl_model_gateway::tokenizer::stop::SequenceDecoderOutput::Held => {}
|
||||
}
|
||||
}
|
||||
text
|
||||
} else {
|
||||
// Use incremental decoder to handle multi-byte character boundaries
|
||||
let decode_stream = handle.decode_streams.entry(index).or_insert_with(|| {
|
||||
DecodeStream::new(
|
||||
Arc::clone(&tokenizer),
|
||||
&[], // No prompt tokens for completion
|
||||
handle.skip_special_tokens,
|
||||
)
|
||||
});
|
||||
|
||||
// Process tokens incrementally
|
||||
let mut text_parts = Vec::new();
|
||||
for &token_id in &chunk.token_ids {
|
||||
if let Ok(Some(text)) = decode_stream.step(token_id) {
|
||||
text_parts.push(text);
|
||||
}
|
||||
}
|
||||
text_parts.join("")
|
||||
};
|
||||
|
||||
if chunk_text.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Send first chunk with role
|
||||
if first_chunk {
|
||||
let first_response = ChatCompletionStreamResponse {
|
||||
id: request_id.to_string(),
|
||||
object: "chat.completion.chunk".to_string(),
|
||||
created,
|
||||
model: model.to_string(),
|
||||
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
|
||||
choices: vec![ChatStreamChoice {
|
||||
index,
|
||||
delta: ChatMessageDelta {
|
||||
role: Some("assistant".to_string()),
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
logprobs: None,
|
||||
finish_reason: None,
|
||||
matched_stop: None,
|
||||
}],
|
||||
usage: None,
|
||||
};
|
||||
return Ok(Some(first_response));
|
||||
}
|
||||
|
||||
// Update stream buffer
|
||||
let stream_buffer = handle.stream_buffers.entry(index).or_default();
|
||||
stream_buffer.push_str(&chunk_text);
|
||||
|
||||
// Handle tool calls if tools are provided
|
||||
if let (Some(ref tools), Some(ref tool_parser)) = (handle.tools.as_ref(), handle.tool_parser.as_ref()) {
|
||||
let tool_choice_enabled = !matches!(
|
||||
handle.tool_choice,
|
||||
Some(ToolChoice::Value(ToolChoiceValue::None))
|
||||
);
|
||||
|
||||
if tool_choice_enabled {
|
||||
let mut parser_guard = tool_parser.lock().await;
|
||||
match parser_guard.parse_incremental(&chunk_text, tools).await {
|
||||
Ok(streaming_result) => {
|
||||
if !streaming_result.calls.is_empty() {
|
||||
handle.has_tool_calls.insert(index, true);
|
||||
// Convert tool call items to OpenAI format
|
||||
let tool_call_deltas: Vec<_> = streaming_result
|
||||
.calls
|
||||
.into_iter()
|
||||
.map(|item| {
|
||||
let id = if let Some(ref name) = item.name {
|
||||
generate_tool_call_id(
|
||||
model,
|
||||
name,
|
||||
item.tool_index,
|
||||
handle.history_tool_calls_count,
|
||||
)
|
||||
} else {
|
||||
format!("call_{}", item.tool_index)
|
||||
};
|
||||
|
||||
ToolCallDelta {
|
||||
index: item.tool_index as u32,
|
||||
id: Some(id),
|
||||
tool_type: if item.name.is_some() {
|
||||
Some("function".to_string())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
function: Some(FunctionCallDelta {
|
||||
name: item.name,
|
||||
arguments: if !item.parameters.is_empty() {
|
||||
Some(item.parameters)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
}),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let tool_response = ChatCompletionStreamResponse {
|
||||
id: request_id.to_string(),
|
||||
object: "chat.completion.chunk".to_string(),
|
||||
created,
|
||||
model: model.to_string(),
|
||||
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
|
||||
choices: vec![ChatStreamChoice {
|
||||
index,
|
||||
delta: ChatMessageDelta {
|
||||
role: Some("assistant".to_string()),
|
||||
content: None,
|
||||
tool_calls: Some(tool_call_deltas),
|
||||
reasoning_content: None,
|
||||
},
|
||||
logprobs: None,
|
||||
finish_reason: None,
|
||||
matched_stop: None,
|
||||
}],
|
||||
usage: None,
|
||||
};
|
||||
return Ok(Some(tool_response));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Log error but continue with regular content
|
||||
tracing::warn!("Tool parser error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Regular content emission
|
||||
let content_response = ChatCompletionStreamResponse {
|
||||
id: request_id.to_string(),
|
||||
object: "chat.completion.chunk".to_string(),
|
||||
created,
|
||||
model: model.to_string(),
|
||||
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
|
||||
choices: vec![ChatStreamChoice {
|
||||
index,
|
||||
delta: ChatMessageDelta {
|
||||
role: Some("assistant".to_string()),
|
||||
content: Some(chunk_text),
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
logprobs: None,
|
||||
finish_reason: None,
|
||||
matched_stop: None,
|
||||
}],
|
||||
usage: None,
|
||||
};
|
||||
|
||||
Ok(Some(content_response))
|
||||
}
|
||||
Some(Complete(complete)) => {
|
||||
let index = complete.index;
|
||||
|
||||
// Flush any remaining text
|
||||
// Flush any remaining text from decode stream
|
||||
let mut final_text = handle.stream_buffers.remove(&index).unwrap_or_default();
|
||||
if let Some(ref mut decode_stream) = handle.decode_streams.get_mut(&index) {
|
||||
if let Ok(Some(remaining)) = decode_stream.flush() {
|
||||
final_text.push_str(&remaining);
|
||||
}
|
||||
}
|
||||
handle.decode_streams.remove(&index);
|
||||
|
||||
// Determine finish reason - ensure it's never empty
|
||||
// If finish_reason is empty, try to infer from other fields or use default
|
||||
let finish_reason = if handle.has_tool_calls.get(&index).copied().unwrap_or(false)
|
||||
&& (complete.finish_reason == "stop" || complete.finish_reason.is_empty())
|
||||
{
|
||||
"tool_calls".to_string()
|
||||
} else if complete.finish_reason.is_empty() || complete.finish_reason.trim().is_empty() {
|
||||
// If finish_reason is empty, try to infer from completion_tokens or use default
|
||||
if complete.completion_tokens > 0 {
|
||||
// If we have completion tokens, likely stopped normally
|
||||
"stop".to_string()
|
||||
} else if !complete.output_ids.is_empty() {
|
||||
// If we have output_ids, likely stopped normally
|
||||
"stop".to_string()
|
||||
} else {
|
||||
// Default fallback - always ensure we have a value
|
||||
"stop".to_string()
|
||||
}
|
||||
} else {
|
||||
complete.finish_reason.clone()
|
||||
};
|
||||
|
||||
// Ensure finish_reason is never empty (defensive check)
|
||||
let finish_reason = if finish_reason.is_empty() || finish_reason.trim().is_empty() {
|
||||
"stop".to_string()
|
||||
} else {
|
||||
finish_reason
|
||||
};
|
||||
|
||||
// Extract matched_stop
|
||||
let matched_stop = match &complete.matched_stop {
|
||||
Some(proto::generate_complete::MatchedStop::MatchedTokenId(token_id)) => {
|
||||
Some(Value::Number(serde_json::Number::from(*token_id)))
|
||||
}
|
||||
Some(proto::generate_complete::MatchedStop::MatchedStopStr(stop_str)) => {
|
||||
Some(Value::String(stop_str.clone()))
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
||||
// Build usage - prefer values from complete message, but fallback to accumulated values from chunks
|
||||
// Complete message should have the final values, but sometimes they might be 0 or missing
|
||||
// Always use the latest cumulative value from chunks if available, otherwise use complete message value
|
||||
let mut prompt_tokens = handle.prompt_tokens.get(&index)
|
||||
.copied()
|
||||
.filter(|&v| v > 0)
|
||||
.unwrap_or(complete.prompt_tokens);
|
||||
let mut completion_tokens = handle.completion_tokens.get(&index)
|
||||
.copied()
|
||||
.filter(|&v| v > 0)
|
||||
.unwrap_or(complete.completion_tokens);
|
||||
|
||||
// Always try to use initial_prompt_tokens if prompt_tokens is 0 or missing
|
||||
// This is the most reliable source for prompt tokens since we calculate it from the request
|
||||
if prompt_tokens == 0 {
|
||||
if let Some(initial_prompt) = handle.initial_prompt_tokens {
|
||||
prompt_tokens = initial_prompt;
|
||||
}
|
||||
}
|
||||
|
||||
// If completion_tokens is 0, try to infer from output_ids or accumulated chunks
|
||||
if completion_tokens == 0 {
|
||||
// Try to use completion_tokens from complete message even if 0
|
||||
// Or calculate from output_ids
|
||||
if complete.completion_tokens > 0 {
|
||||
completion_tokens = complete.completion_tokens;
|
||||
} else if !complete.output_ids.is_empty() {
|
||||
completion_tokens = complete.output_ids.len() as i32;
|
||||
} else if let Some(&last_completion) = handle.completion_tokens.get(&index) {
|
||||
completion_tokens = last_completion;
|
||||
}
|
||||
}
|
||||
|
||||
// Final fallback: if both are still 0, try to use initial_prompt_tokens for prompt
|
||||
// and calculate completion from output_ids
|
||||
if prompt_tokens == 0 && completion_tokens == 0 {
|
||||
// Try to infer from output_ids if available
|
||||
let output_ids_len = complete.output_ids.len() as i32;
|
||||
if output_ids_len > 0 {
|
||||
completion_tokens = output_ids_len;
|
||||
// Always try to use initial_prompt_tokens for prompt
|
||||
if let Some(initial_prompt) = handle.initial_prompt_tokens {
|
||||
prompt_tokens = initial_prompt;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Final defensive check: ensure prompt_tokens is set if we have initial_prompt_tokens
|
||||
if prompt_tokens == 0 {
|
||||
if let Some(initial_prompt) = handle.initial_prompt_tokens {
|
||||
prompt_tokens = initial_prompt;
|
||||
}
|
||||
}
|
||||
|
||||
// Always create usage, even if values are 0 (defensive)
|
||||
let usage = Some(Usage {
|
||||
prompt_tokens: prompt_tokens.max(0) as u32,
|
||||
completion_tokens: completion_tokens.max(0) as u32,
|
||||
total_tokens: (prompt_tokens.max(0) + completion_tokens.max(0)) as u32,
|
||||
completion_tokens_details: None,
|
||||
});
|
||||
|
||||
let finish_response = ChatCompletionStreamResponse {
|
||||
id: request_id.to_string(),
|
||||
object: "chat.completion.chunk".to_string(),
|
||||
created,
|
||||
model: model.to_string(),
|
||||
system_fingerprint: system_fingerprint.map(|s| s.to_string()),
|
||||
choices: vec![ChatStreamChoice {
|
||||
index,
|
||||
delta: ChatMessageDelta {
|
||||
role: Some("assistant".to_string()),
|
||||
content: if !final_text.is_empty() {
|
||||
Some(final_text)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
tool_calls: None,
|
||||
reasoning_content: None,
|
||||
},
|
||||
logprobs: None,
|
||||
finish_reason: Some(finish_reason),
|
||||
matched_stop,
|
||||
}],
|
||||
usage,
|
||||
};
|
||||
|
||||
Ok(Some(finish_response))
|
||||
}
|
||||
Some(Error(error)) => {
|
||||
Err(format!("Server error: {} (status: {})", error.message, error.http_status_code))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Free a gRPC response converter handle
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_grpc_response_converter_free(handle: *mut GrpcResponseConverterHandle) {
|
||||
if !handle.is_null() {
|
||||
let _ = Box::from_raw(handle);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
//! FFI module for exposing sgl-model-gateway preprocessing and postprocessing functions
|
||||
//! to C-compatible languages (e.g., Golang via cgo)
|
||||
//!
|
||||
//! This module provides C-compatible function signatures for:
|
||||
//! - Tokenizer operations (encode, decode, chat template)
|
||||
//! - Tool parser operations (parse tool calls)
|
||||
//! - Tool constraint generation
|
||||
//! - gRPC client SDK (complete request-response flow)
|
||||
//!
|
||||
//! # Safety
|
||||
//! All functions marked with `#[no_mangle]` and `extern "C"` must be called
|
||||
//! with valid pointers and follow the documented memory management rules.
|
||||
|
||||
// Re-export error types
|
||||
pub use error::{SglErrorCode, set_error_message, set_error_message_fmt, clear_error_message};
|
||||
|
||||
// Re-export memory management functions
|
||||
pub use memory::{sgl_free_string, sgl_free_token_ids};
|
||||
|
||||
// Re-export tokenizer functions
|
||||
pub use tokenizer::{
|
||||
TokenizerHandle,
|
||||
sgl_tokenizer_create_from_file,
|
||||
sgl_tokenizer_encode,
|
||||
sgl_tokenizer_apply_chat_template,
|
||||
sgl_tokenizer_apply_chat_template_with_tools,
|
||||
sgl_tokenizer_decode,
|
||||
sgl_tokenizer_free,
|
||||
};
|
||||
|
||||
// Re-export tool parser functions
|
||||
pub use tool_parser::{
|
||||
ToolParserHandle,
|
||||
sgl_tool_parser_create,
|
||||
sgl_tool_parser_parse_complete,
|
||||
sgl_tool_parser_parse_incremental,
|
||||
sgl_tool_parser_reset,
|
||||
sgl_tool_parser_free,
|
||||
};
|
||||
|
||||
// Re-export gRPC converter functions
|
||||
pub use grpc_converter::{
|
||||
GrpcResponseConverterHandle,
|
||||
sgl_grpc_response_converter_create,
|
||||
sgl_grpc_response_converter_convert_chunk,
|
||||
sgl_grpc_response_converter_free,
|
||||
};
|
||||
|
||||
// Re-export client SDK functions
|
||||
pub use client::{
|
||||
SglangClientHandle,
|
||||
sgl_client_create,
|
||||
sgl_client_free,
|
||||
};
|
||||
|
||||
// Re-export stream functions
|
||||
pub use stream::{
|
||||
SglangStreamHandle,
|
||||
sgl_stream_read_next,
|
||||
sgl_stream_free,
|
||||
};
|
||||
|
||||
// Re-export client stream function (defined in client.rs but used by stream)
|
||||
pub use client::sgl_client_chat_completion_stream;
|
||||
|
||||
// Re-export utility functions
|
||||
pub use utils::sgl_generate_tool_constraints;
|
||||
|
||||
// Sub-modules
|
||||
mod error;
|
||||
mod memory;
|
||||
mod tokenizer;
|
||||
mod tool_parser;
|
||||
mod grpc_converter;
|
||||
mod client;
|
||||
mod stream;
|
||||
mod utils;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_error_codes() {
|
||||
assert_eq!(SglErrorCode::Success as i32, 0);
|
||||
assert_eq!(SglErrorCode::InvalidArgument as i32, 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
//! Memory management for FFI functions
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::raw::c_char;
|
||||
|
||||
/// Free a C string allocated by Rust
|
||||
///
|
||||
/// # Safety
|
||||
/// This function must only be called with pointers returned by other FFI functions.
|
||||
/// Calling with arbitrary pointers or multiple times on the same pointer is undefined behavior.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_free_string(s: *mut c_char) {
|
||||
if !s.is_null() {
|
||||
let _ = CString::from_raw(s);
|
||||
}
|
||||
}
|
||||
|
||||
/// Free token IDs array allocated by Rust
|
||||
///
|
||||
/// # Safety
|
||||
/// This function must only be called with pointers returned by `sgl_tokenizer_encode`.
|
||||
/// The `count` parameter must match the length of the array.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_free_token_ids(ptr: *mut u32, count: usize) {
|
||||
if !ptr.is_null() && count > 0 {
|
||||
let _ = Vec::from_raw_parts(ptr, count, count);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,288 @@
|
||||
//! Stream handling FFI functions
|
||||
//!
|
||||
//! This module provides FFI (Foreign Function Interface) functions for managing
|
||||
//! streaming responses from the SGLang gRPC API. It handles:
|
||||
//!
|
||||
//! - Creating and managing stream handles
|
||||
//! - Reading chunks from streams and converting them to OpenAI format
|
||||
//! - Managing automatic abort on stream drop (via AbortOnDropStream)
|
||||
//! - Thread-safe access to streams and response converters
|
||||
//!
|
||||
//! # Safety
|
||||
//!
|
||||
//! All FFI functions are marked `unsafe` as per Rust FFI conventions. Callers must:
|
||||
//! - Pass valid pointers
|
||||
//! - Ensure proper pointer lifetime management
|
||||
//! - Call corresponding free functions for cleanup
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::raw::{c_char, c_int};
|
||||
use std::ptr;
|
||||
use std::sync::Arc;
|
||||
use tokio::runtime::Runtime;
|
||||
use once_cell::sync::Lazy;
|
||||
use futures_util::StreamExt;
|
||||
|
||||
use sgl_model_gateway::grpc_client::{sglang_proto as proto, sglang_scheduler::{SglangSchedulerClient, AbortOnDropStream}};
|
||||
|
||||
use super::error::{SglErrorCode, set_error_message};
|
||||
use super::grpc_converter::{GrpcResponseConverterHandle, convert_proto_chunk_to_openai};
|
||||
|
||||
/// Global tokio runtime for async operations
|
||||
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
|
||||
Runtime::new().expect("Failed to create tokio runtime for stream FFI")
|
||||
});
|
||||
|
||||
/// Handle for an active streaming request.
|
||||
///
|
||||
/// This struct manages the stream and response converter for a single request.
|
||||
/// It is wrapped in Arc and Mutex for thread-safe concurrent access.
|
||||
///
|
||||
/// # Fields
|
||||
///
|
||||
/// * `stream` - The gRPC stream wrapped in AbortOnDropStream for automatic cleanup
|
||||
/// * `converter` - Response converter that transforms proto messages to OpenAI format
|
||||
/// * `client` - The underlying gRPC client connection
|
||||
/// * `prompt_tokens` - Number of prompt tokens from the original request
|
||||
pub struct SglangStreamHandle {
|
||||
pub(crate) stream: Arc<tokio::sync::Mutex<AbortOnDropStream>>,
|
||||
pub(crate) converter: Arc<tokio::sync::Mutex<GrpcResponseConverterHandle>>,
|
||||
#[allow(dead_code)]
|
||||
pub(crate) client: Arc<SglangSchedulerClient>,
|
||||
#[allow(dead_code)]
|
||||
pub(crate) prompt_tokens: i32, // Number of prompt tokens for this request
|
||||
}
|
||||
|
||||
/// Read next chunk from stream and convert to OpenAI format.
|
||||
///
|
||||
/// This function reads the next chunk from the gRPC stream, converts it from the
|
||||
/// internal protocol format to OpenAI-compatible JSON format, and returns it via
|
||||
/// the output parameters.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `stream_handle` - Mutable pointer to the stream handle
|
||||
/// * `response_json_out` - Pointer to receive OpenAI format JSON string
|
||||
/// - Caller must free this with `sgl_free_string`
|
||||
/// - May be NULL if no data available
|
||||
/// * `is_done_out` - Pointer to receive completion status
|
||||
/// - 0 = stream has more data
|
||||
/// - 1 = stream is complete
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
/// - Only set if function returns an error code
|
||||
/// - Must be freed with `sgl_free_string` if not NULL
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// * `SglErrorCode::Success` - Successfully read a chunk or reached end of stream
|
||||
/// * Other error codes - See `SglErrorCode` for details
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// - All pointers must be valid and properly aligned
|
||||
/// - `stream_handle` must point to a valid `SglangStreamHandle`
|
||||
/// - Output pointers must be writable
|
||||
///
|
||||
/// # Notes
|
||||
///
|
||||
/// - Complete messages are identified by the presence of `proto::GenerateResponse::Complete`
|
||||
/// - When is_done=1, this may be the last readable chunk or the stream may be ending
|
||||
/// - Subsequent calls after is_done=1 will mark the stream as complete internally
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_stream_read_next(
|
||||
stream_handle: *mut SglangStreamHandle,
|
||||
response_json_out: *mut *mut c_char,
|
||||
is_done_out: *mut c_int,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if stream_handle.is_null() || response_json_out.is_null() || is_done_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let handle_ref = &*stream_handle;
|
||||
let stream = Arc::clone(&handle_ref.stream);
|
||||
let converter = Arc::clone(&handle_ref.converter);
|
||||
|
||||
// Read next chunk from stream
|
||||
let chunk_result = RUNTIME.block_on(async {
|
||||
let mut stream_guard = stream.lock().await;
|
||||
stream_guard.next().await
|
||||
});
|
||||
|
||||
match chunk_result {
|
||||
Some(Ok(proto_response)) => {
|
||||
// Convert proto response to OpenAI format
|
||||
// We need to get the converter lock first
|
||||
let conversion_result = RUNTIME.block_on(async {
|
||||
let mut converter_guard = converter.lock().await;
|
||||
|
||||
// Clone necessary fields for conversion
|
||||
let tokenizer = Arc::clone(&converter_guard.tokenizer);
|
||||
let model = converter_guard.model.clone();
|
||||
let request_id = converter_guard.request_id.clone();
|
||||
let created = converter_guard.created;
|
||||
let system_fingerprint = converter_guard.system_fingerprint.clone();
|
||||
|
||||
// Call the conversion function
|
||||
convert_proto_chunk_to_openai(
|
||||
proto_response.clone(),
|
||||
&mut *converter_guard,
|
||||
&tokenizer,
|
||||
&model,
|
||||
&request_id,
|
||||
created,
|
||||
system_fingerprint.as_deref(),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
match conversion_result {
|
||||
Ok(Some(openai_response)) => {
|
||||
// Serialize to JSON
|
||||
let result_str = match serde_json::to_string(&openai_response) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to serialize response: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
let result_cstr = match CString::new(result_str) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
|
||||
// Check if this is a complete response (stream done)
|
||||
let is_complete = matches!(proto_response.response, Some(proto::generate_response::Response::Complete(_)) | Some(proto::generate_response::Response::Error(_)));
|
||||
|
||||
*response_json_out = result_cstr.into_raw();
|
||||
*is_done_out = if is_complete { 1 } else { 0 };
|
||||
|
||||
if is_complete {
|
||||
// Mark stream as completed
|
||||
// Ensure mark_completed() completes and is visible before returning
|
||||
// Use yield_now to ensure Release ordering is fully propagated
|
||||
RUNTIME.block_on(async {
|
||||
let stream_guard = stream.lock().await;
|
||||
stream_guard.mark_completed();
|
||||
// Keep the guard until mark_completed() is fully executed
|
||||
drop(stream_guard);
|
||||
// Yield to ensure Release ordering is propagated before returning
|
||||
// This prevents race condition where Free() is called immediately
|
||||
// and Drop might not see the mark_completed() write
|
||||
tokio::task::yield_now().await;
|
||||
});
|
||||
}
|
||||
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Ok(None) => {
|
||||
// No response to send (e.g., empty chunk)
|
||||
// Don't mark as completed - stream might continue
|
||||
// Just return null and let caller read more
|
||||
*response_json_out = ptr::null_mut();
|
||||
*is_done_out = 0; // Keep stream open, not done yet
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
// Conversion error - don't mark as completed
|
||||
// Let the stream end naturally or return error without stopping stream
|
||||
set_error_message(error_out, &format!("Conversion error: {}", e));
|
||||
*response_json_out = ptr::null_mut();
|
||||
*is_done_out = 0; // Don't mark as done - let caller decide
|
||||
SglErrorCode::ParsingError
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
// Stream error - mark as completed to prevent abort
|
||||
RUNTIME.block_on(async {
|
||||
let stream_guard = stream.lock().await;
|
||||
stream_guard.mark_completed();
|
||||
drop(stream_guard);
|
||||
// Yield to ensure Release ordering is propagated
|
||||
tokio::task::yield_now().await;
|
||||
});
|
||||
|
||||
set_error_message(error_out, &format!("Stream error: {}", e));
|
||||
*is_done_out = 1;
|
||||
SglErrorCode::UnknownError
|
||||
}
|
||||
None => {
|
||||
// Stream ended naturally (no more chunks)
|
||||
// Mark stream as completed before returning to prevent abort
|
||||
RUNTIME.block_on(async {
|
||||
let stream_guard = stream.lock().await;
|
||||
stream_guard.mark_completed();
|
||||
drop(stream_guard);
|
||||
// Yield to ensure Release ordering is propagated
|
||||
tokio::task::yield_now().await;
|
||||
});
|
||||
|
||||
*response_json_out = ptr::null_mut();
|
||||
*is_done_out = 1;
|
||||
SglErrorCode::Success
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Free a stream handle and release all associated resources.
|
||||
///
|
||||
/// This function must be called exactly once for each stream handle returned by
|
||||
/// `sgl_client_chat_completion_stream`. It marks the stream as completed internally
|
||||
/// to prevent abort signals from being sent when resources are cleaned up.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `handle` - Mutable pointer to the stream handle to free
|
||||
/// - If NULL, this function does nothing
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// - Must be called only once per handle
|
||||
/// - Handle must not be used after calling this function
|
||||
/// - After this call, the stream is no longer valid
|
||||
///
|
||||
/// # Notes
|
||||
///
|
||||
/// - This function internally calls `mark_completed()` before freeing to ensure
|
||||
/// the stream cleanup doesn't trigger an abort RPC to the server
|
||||
/// - Memory fences are used to ensure visibility across threads
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_stream_free(handle: *mut SglangStreamHandle) {
|
||||
if !handle.is_null() {
|
||||
let handle_ref = Box::from_raw(handle);
|
||||
|
||||
// Mark stream as completed to prevent abort on drop
|
||||
// By this point, the stream should already be completed by ReadNext()
|
||||
// but we call it again to be safe
|
||||
RUNTIME.block_on(async {
|
||||
let stream_guard = handle_ref.stream.lock().await;
|
||||
stream_guard.mark_completed();
|
||||
// Keep guard alive to ensure mark_completed() write completes
|
||||
drop(stream_guard);
|
||||
// Yield to ensure the atomic write is visible
|
||||
tokio::task::yield_now().await;
|
||||
});
|
||||
|
||||
// Use a strong memory fence to ensure mark_completed()'s Release write
|
||||
// is visible before we drop the last Arc reference
|
||||
std::sync::atomic::fence(std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
// Now drop all references - if mark_completed() was called successfully,
|
||||
// the drop won't send an abort
|
||||
drop(handle_ref.stream);
|
||||
|
||||
// Free converter
|
||||
let converter = Arc::try_unwrap(handle_ref.converter)
|
||||
.ok()
|
||||
.map(|m| m.into_inner());
|
||||
if let Some(conv) = converter {
|
||||
super::grpc_converter::sgl_grpc_response_converter_free(Box::into_raw(Box::new(conv)));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,379 @@
|
||||
//! Tokenizer FFI functions
|
||||
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::{c_char, c_int};
|
||||
use std::ptr;
|
||||
use std::sync::Arc;
|
||||
use serde_json::Value;
|
||||
|
||||
use sgl_model_gateway::tokenizer::{
|
||||
create_tokenizer_from_file,
|
||||
traits::Tokenizer as TokenizerTrait,
|
||||
chat_template::ChatTemplateParams,
|
||||
huggingface::HuggingFaceTokenizer,
|
||||
};
|
||||
|
||||
use super::error::{SglErrorCode, set_error_message, clear_error_message};
|
||||
|
||||
/// Opaque handle for a tokenizer instance
|
||||
#[repr(C)]
|
||||
pub struct TokenizerHandle {
|
||||
pub(crate) tokenizer: Arc<dyn TokenizerTrait>,
|
||||
}
|
||||
|
||||
/// Create a tokenizer from a file path
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `path` - Path to tokenizer.json file (null-terminated C string)
|
||||
/// * `error_out` - Optional pointer to receive error message (must be freed with sgl_free_string)
|
||||
///
|
||||
/// # Returns
|
||||
/// * Pointer to TokenizerHandle on success, null on failure
|
||||
///
|
||||
/// # Safety
|
||||
/// The returned handle must be freed with `sgl_tokenizer_free`.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tokenizer_create_from_file(
|
||||
path: *const c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> *mut TokenizerHandle {
|
||||
if path.is_null() {
|
||||
set_error_message(error_out, "path cannot be null");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
|
||||
let path_str = match CStr::from_ptr(path).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Invalid UTF-8 in path: {}", e));
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
match create_tokenizer_from_file(path_str) {
|
||||
Ok(tokenizer) => {
|
||||
clear_error_message(error_out);
|
||||
Box::into_raw(Box::new(TokenizerHandle {
|
||||
tokenizer,
|
||||
}))
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &e.to_string());
|
||||
ptr::null_mut()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode text to token IDs
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Tokenizer handle (must not be null)
|
||||
/// * `text` - Input text (null-terminated C string)
|
||||
/// * `token_ids_out` - Pointer to receive array of token IDs (must be freed with sgl_free_token_ids)
|
||||
/// * `token_count_out` - Pointer to receive token count
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
///
|
||||
/// # Safety
|
||||
/// The token_ids_out array must be freed with sgl_free_token_ids() after use.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tokenizer_encode(
|
||||
handle: *mut TokenizerHandle,
|
||||
text: *const c_char,
|
||||
token_ids_out: *mut *mut u32,
|
||||
token_count_out: *mut usize,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || text.is_null() || token_ids_out.is_null() || token_count_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let text_str = match CStr::from_ptr(text).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in text");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
let tokenizer = &(*handle).tokenizer;
|
||||
match tokenizer.encode(text_str) {
|
||||
Ok(encoding) => {
|
||||
let token_ids = encoding.token_ids();
|
||||
let count = token_ids.len();
|
||||
|
||||
// Allocate memory for token IDs using Vec, then leak to give ownership to C
|
||||
let vec = token_ids.to_vec();
|
||||
let ptr = vec.as_ptr() as *mut u32;
|
||||
let _ = std::mem::ManuallyDrop::new(vec);
|
||||
|
||||
*token_ids_out = ptr;
|
||||
*token_count_out = count;
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &e.to_string());
|
||||
SglErrorCode::TokenizationError
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply chat template to messages with tools support
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Tokenizer handle
|
||||
/// * `messages_json` - JSON string of messages array
|
||||
/// * `tools_json` - Optional JSON string of tools array (null or empty string for no tools)
|
||||
/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string)
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tokenizer_apply_chat_template_with_tools(
|
||||
handle: *mut TokenizerHandle,
|
||||
messages_json: *const c_char,
|
||||
tools_json: *const c_char,
|
||||
result_out: *mut *mut c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || messages_json.is_null() || result_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let messages_str = match CStr::from_ptr(messages_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in messages_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse JSON messages
|
||||
let messages: Vec<Value> = match serde_json::from_str(messages_str) {
|
||||
Ok(msgs) => msgs,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse messages JSON: {}", e));
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse tools JSON if provided
|
||||
let tools: Option<Vec<Value>> = if tools_json.is_null() {
|
||||
None
|
||||
} else {
|
||||
let tools_str = match CStr::from_ptr(tools_json).to_str() {
|
||||
Ok(s) => {
|
||||
if s.is_empty() {
|
||||
None
|
||||
} else {
|
||||
match serde_json::from_str::<Vec<Value>>(s) {
|
||||
Ok(t) => Some(t),
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse tools JSON: {}", e));
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in tools_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
tools_str
|
||||
};
|
||||
|
||||
// Get the tokenizer from handle
|
||||
let handle_ref = &*handle;
|
||||
let tokenizer = &handle_ref.tokenizer;
|
||||
|
||||
// Try to downcast to HuggingFaceTokenizer
|
||||
if let Some(hf_tokenizer) = tokenizer.as_any().downcast_ref::<HuggingFaceTokenizer>() {
|
||||
// Apply chat template with tools
|
||||
let empty_docs: [Value; 0] = [];
|
||||
let tools_slice = tools.as_ref().map(|t| t.as_slice());
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true,
|
||||
tools: tools_slice,
|
||||
documents: Some(&empty_docs),
|
||||
template_kwargs: None,
|
||||
};
|
||||
|
||||
match hf_tokenizer.apply_chat_template(&messages, params) {
|
||||
Ok(result) => {
|
||||
let result_cstr = match CString::new(result) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
*result_out = result_cstr.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to apply chat template: {}", e));
|
||||
SglErrorCode::TokenizationError
|
||||
}
|
||||
}
|
||||
} else {
|
||||
set_error_message(error_out, "Chat template is only supported for HuggingFace tokenizers");
|
||||
SglErrorCode::TokenizationError
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply chat template to messages
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Tokenizer handle
|
||||
/// * `messages_json` - JSON string of messages array
|
||||
/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string)
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tokenizer_apply_chat_template(
|
||||
handle: *mut TokenizerHandle,
|
||||
messages_json: *const c_char,
|
||||
result_out: *mut *mut c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || messages_json.is_null() || result_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let messages_str = match CStr::from_ptr(messages_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in messages_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse JSON messages
|
||||
let messages: Vec<Value> = match serde_json::from_str(messages_str) {
|
||||
Ok(msgs) => msgs,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to parse messages JSON: {}", e));
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Get the tokenizer from handle
|
||||
let handle_ref = &*handle;
|
||||
let tokenizer = &handle_ref.tokenizer;
|
||||
|
||||
// Try to downcast to HuggingFaceTokenizer
|
||||
if let Some(hf_tokenizer) = tokenizer.as_any().downcast_ref::<HuggingFaceTokenizer>() {
|
||||
// Apply chat template with default parameters
|
||||
// Use empty arrays instead of None to avoid template errors
|
||||
// Set add_generation_prompt to true so the model knows to start generating
|
||||
let empty_tools: [Value; 0] = [];
|
||||
let empty_docs: [Value; 0] = [];
|
||||
let params = ChatTemplateParams {
|
||||
add_generation_prompt: true, // Important: tells the model to start generating
|
||||
tools: Some(&empty_tools),
|
||||
documents: Some(&empty_docs),
|
||||
template_kwargs: None,
|
||||
};
|
||||
|
||||
match hf_tokenizer.apply_chat_template(&messages, params) {
|
||||
Ok(result) => {
|
||||
let result_cstr = match CString::new(result) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
*result_out = result_cstr.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to apply chat template: {}", e));
|
||||
SglErrorCode::TokenizationError
|
||||
}
|
||||
}
|
||||
} else {
|
||||
set_error_message(error_out, "Chat template is only supported for HuggingFace tokenizers");
|
||||
SglErrorCode::TokenizationError
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode token IDs to text
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Tokenizer handle
|
||||
/// * `token_ids` - Array of token IDs
|
||||
/// * `token_count` - Number of tokens
|
||||
/// * `skip_special_tokens` - Whether to skip special tokens
|
||||
/// * `result_out` - Pointer to receive result string (must be freed with sgl_free_string)
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tokenizer_decode(
|
||||
handle: *mut TokenizerHandle,
|
||||
token_ids: *const u32,
|
||||
token_count: usize,
|
||||
skip_special_tokens: c_int,
|
||||
result_out: *mut *mut c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || token_ids.is_null() || result_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
if token_count == 0 {
|
||||
let empty = CString::new("").unwrap();
|
||||
*result_out = empty.into_raw();
|
||||
clear_error_message(error_out);
|
||||
return SglErrorCode::Success;
|
||||
}
|
||||
|
||||
// Convert C array to Rust slice
|
||||
let token_slice = std::slice::from_raw_parts(token_ids, token_count);
|
||||
|
||||
let tokenizer = &(*handle).tokenizer;
|
||||
match tokenizer.decode(token_slice, skip_special_tokens != 0) {
|
||||
Ok(text) => {
|
||||
let result_cstr = match CString::new(text) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
*result_out = result_cstr.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &e.to_string());
|
||||
SglErrorCode::TokenizationError
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Free a tokenizer handle
|
||||
///
|
||||
/// # Safety
|
||||
/// This function must only be called once per handle, and the handle must not be used after calling.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tokenizer_free(handle: *mut TokenizerHandle) {
|
||||
if !handle.is_null() {
|
||||
let _ = Box::from_raw(handle);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
//! Tool parser FFI functions
|
||||
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::{c_char};
|
||||
use std::ptr;
|
||||
use std::sync::Arc;
|
||||
use std::collections::HashMap;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::runtime::Runtime;
|
||||
use once_cell::sync::Lazy;
|
||||
|
||||
use sgl_model_gateway::tool_parser::{ParserFactory, ToolParser};
|
||||
use sgl_model_gateway::protocols::common::Tool;
|
||||
|
||||
use super::error::{SglErrorCode, set_error_message, clear_error_message};
|
||||
use super::utils::generate_tool_call_id;
|
||||
|
||||
/// Global parser factory (initialized once)
|
||||
static PARSER_FACTORY: Lazy<ParserFactory> = Lazy::new(|| ParserFactory::new());
|
||||
|
||||
/// Global tokio runtime for async operations
|
||||
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
|
||||
Runtime::new().expect("Failed to create tokio runtime for tool parser FFI")
|
||||
});
|
||||
|
||||
/// Opaque handle for a tool parser instance
|
||||
/// Note: For streaming, we need mutable access, so we use Arc<Mutex<>> internally
|
||||
/// Note: This is an opaque handle, C code doesn't access fields directly
|
||||
pub struct ToolParserHandle {
|
||||
parser: Arc<tokio::sync::Mutex<Box<dyn ToolParser>>>,
|
||||
model: String, // Store model name for ID generation
|
||||
history_tool_calls_count: usize, // Track tool call count for ID generation
|
||||
tool_index_to_id: HashMap<usize, String>, // Map tool_index to ID for incremental updates
|
||||
}
|
||||
|
||||
/// Create a tool parser
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `parser_type` - Parser type name (e.g., "json", "llama", "mistral") or model name (e.g., "gpt-4")
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * Pointer to ToolParserHandle on success, null on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tool_parser_create(
|
||||
parser_type: *const c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> *mut ToolParserHandle {
|
||||
if parser_type.is_null() {
|
||||
set_error_message(error_out, "parser_type cannot be null");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
|
||||
let type_str = match CStr::from_ptr(parser_type).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in parser_type");
|
||||
return ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
// Create parser using factory
|
||||
// The factory will determine the parser type based on model name or use the provided type
|
||||
let parser = if let Some(parser_box) = PARSER_FACTORY.registry().create_for_model(type_str) {
|
||||
parser_box
|
||||
} else if let Some(parser_box) = PARSER_FACTORY.registry().create_parser(type_str) {
|
||||
parser_box
|
||||
} else {
|
||||
set_error_message(error_out, &format!("Unknown parser type: {}", type_str));
|
||||
return ptr::null_mut();
|
||||
};
|
||||
|
||||
Box::into_raw(Box::new(ToolParserHandle {
|
||||
parser: Arc::new(tokio::sync::Mutex::new(parser)),
|
||||
model: type_str.to_string(),
|
||||
history_tool_calls_count: 0,
|
||||
tool_index_to_id: HashMap::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
/// Parse complete tool calls from text
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Tool parser handle
|
||||
/// * `text` - Input text to parse
|
||||
/// * `result_json_out` - Pointer to receive JSON result (must be freed with sgl_free_string)
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tool_parser_parse_complete(
|
||||
handle: *mut ToolParserHandle,
|
||||
text: *const c_char,
|
||||
result_json_out: *mut *mut c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || text.is_null() || result_json_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let text_str = match CStr::from_ptr(text).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in text");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
let handle_ref = &*handle;
|
||||
let parser = Arc::clone(&handle_ref.parser);
|
||||
let model = handle_ref.model.clone();
|
||||
let history_count = handle_ref.history_tool_calls_count;
|
||||
|
||||
// Use tokio runtime to run async code
|
||||
let result = RUNTIME.block_on(async {
|
||||
let parser_guard = parser.lock().await;
|
||||
parser_guard.parse_complete(text_str).await
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok((normal_text, tool_calls)) => {
|
||||
// Convert Rust ToolCall to OpenAI format
|
||||
let openai_tool_calls: Vec<Value> = tool_calls
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, tc)| {
|
||||
// Generate ID for this tool call
|
||||
let id = generate_tool_call_id(&model, &tc.function.name, index, history_count);
|
||||
json!({
|
||||
"id": id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Build result JSON
|
||||
let result_json = json!({
|
||||
"normal_text": normal_text,
|
||||
"tool_calls": openai_tool_calls
|
||||
});
|
||||
|
||||
let result_str = match serde_json::to_string(&result_json) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to serialize JSON: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
let result_cstr = match CString::new(result_str) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
|
||||
*result_json_out = result_cstr.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Parse error: {}", e));
|
||||
SglErrorCode::ParsingError
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse tool calls incrementally from streaming chunks
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `handle` - Tool parser handle
|
||||
/// * `chunk` - New text chunk from stream
|
||||
/// * `tools_json` - JSON array of available tools (for validation, can be null/empty)
|
||||
/// * `result_json_out` - Pointer to receive JSON result (must be freed with sgl_free_string)
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tool_parser_parse_incremental(
|
||||
handle: *mut ToolParserHandle,
|
||||
chunk: *const c_char,
|
||||
tools_json: *const c_char,
|
||||
result_json_out: *mut *mut c_char,
|
||||
error_out: *mut *mut c_char,
|
||||
) -> SglErrorCode {
|
||||
if handle.is_null() || chunk.is_null() || result_json_out.is_null() {
|
||||
set_error_message(error_out, "Invalid arguments: null pointer");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
|
||||
let chunk_str = match CStr::from_ptr(chunk).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in chunk");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse tools JSON if provided
|
||||
let tools: Vec<Tool> = if !tools_json.is_null() {
|
||||
let tools_str = match CStr::from_ptr(tools_json).to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
set_error_message(error_out, "Invalid UTF-8 in tools_json");
|
||||
return SglErrorCode::InvalidArgument;
|
||||
}
|
||||
};
|
||||
match serde_json::from_str::<Vec<Tool>>(tools_str) {
|
||||
Ok(t) => t,
|
||||
Err(_) => vec![], // If parsing fails, use empty tools
|
||||
}
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
|
||||
let handle_ref = &*handle;
|
||||
let parser = Arc::clone(&handle_ref.parser);
|
||||
let model = handle_ref.model.clone();
|
||||
let history_count = handle_ref.history_tool_calls_count;
|
||||
|
||||
// Use tokio runtime to run async code
|
||||
let result = RUNTIME.block_on(async {
|
||||
let mut parser_guard = parser.lock().await;
|
||||
parser_guard.parse_incremental(chunk_str, &tools).await
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(streaming_result) => {
|
||||
// Convert StreamingParseResult to OpenAI format
|
||||
let handle_mut = &mut *handle;
|
||||
let openai_tool_calls: Vec<Value> = streaming_result
|
||||
.calls
|
||||
.into_iter()
|
||||
.map(|item| {
|
||||
// For incremental parsing, we may not have complete tool calls yet
|
||||
// Generate or reuse ID based on tool_index
|
||||
let id = if let Some(ref name) = item.name {
|
||||
// New tool call with name - generate ID and store it
|
||||
let id = generate_tool_call_id(&model, name, item.tool_index, history_count);
|
||||
handle_mut.tool_index_to_id.insert(item.tool_index, id.clone());
|
||||
id
|
||||
} else {
|
||||
// Parameter update - reuse existing ID for this tool_index
|
||||
handle_mut.tool_index_to_id
|
||||
.get(&item.tool_index)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("call_{}", item.tool_index))
|
||||
};
|
||||
|
||||
json!({
|
||||
"id": id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": item.name.unwrap_or_default(),
|
||||
"arguments": item.parameters
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Build result JSON
|
||||
let result_json = json!({
|
||||
"normal_text": streaming_result.normal_text,
|
||||
"tool_calls": openai_tool_calls
|
||||
});
|
||||
|
||||
let result_str = match serde_json::to_string(&result_json) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to serialize JSON: {}", e));
|
||||
return SglErrorCode::ParsingError;
|
||||
}
|
||||
};
|
||||
|
||||
let result_cstr = match CString::new(result_str) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Failed to create result string: {}", e));
|
||||
return SglErrorCode::MemoryError;
|
||||
}
|
||||
};
|
||||
|
||||
*result_json_out = result_cstr.into_raw();
|
||||
clear_error_message(error_out);
|
||||
SglErrorCode::Success
|
||||
}
|
||||
Err(e) => {
|
||||
set_error_message(error_out, &format!("Parse incremental error: {}", e));
|
||||
SglErrorCode::ParsingError
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the parser state for reuse
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tool_parser_reset(handle: *mut ToolParserHandle) {
|
||||
if handle.is_null() {
|
||||
return;
|
||||
}
|
||||
|
||||
let handle_ref = &mut *handle;
|
||||
let parser = Arc::clone(&handle_ref.parser);
|
||||
|
||||
// Reset parser state
|
||||
RUNTIME.block_on(async {
|
||||
let mut parser_guard = parser.lock().await;
|
||||
parser_guard.reset();
|
||||
});
|
||||
|
||||
// Reset history count and tool index mapping
|
||||
handle_ref.history_tool_calls_count = 0;
|
||||
handle_ref.tool_index_to_id.clear();
|
||||
}
|
||||
|
||||
/// Free a tool parser handle
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_tool_parser_free(handle: *mut ToolParserHandle) {
|
||||
if !handle.is_null() {
|
||||
let _ = Box::from_raw(handle);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
//! Utility functions for FFI
|
||||
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Helper function to generate tool call ID (matches router implementation)
|
||||
pub fn generate_tool_call_id(
|
||||
model: &str,
|
||||
function_name: &str,
|
||||
index: usize,
|
||||
history_tool_calls_count: usize,
|
||||
) -> String {
|
||||
if model.to_lowercase().contains("kimi") {
|
||||
// KimiK2 format: functions.{name}:{global_index}
|
||||
format!("functions.{}:{}", function_name, history_tool_calls_count + index)
|
||||
} else {
|
||||
// Standard OpenAI format: call_{24-char-uuid}
|
||||
format!("call_{}", &Uuid::new_v4().simple().to_string()[..24])
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate tool constraints (placeholder implementation)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `tools_json` - JSON array of tools
|
||||
/// * `tool_choice_json` - JSON object representing tool_choice
|
||||
/// * `constraint_type_out` - Pointer to receive constraint type (e.g., "json_schema")
|
||||
/// * `constraint_schema_out` - Pointer to receive constraint schema JSON
|
||||
/// * `error_out` - Optional pointer to receive error message
|
||||
///
|
||||
/// # Returns
|
||||
/// * SglErrorCode::Success on success, error code on failure
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn sgl_generate_tool_constraints(
|
||||
_tools_json: *const std::os::raw::c_char,
|
||||
_tool_choice_json: *const std::os::raw::c_char,
|
||||
_constraint_type_out: *mut *mut std::os::raw::c_char,
|
||||
_constraint_schema_out: *mut *mut std::os::raw::c_char,
|
||||
error_out: *mut *mut std::os::raw::c_char,
|
||||
) -> super::error::SglErrorCode {
|
||||
// Implementation would parse JSON and call generate_tool_constraints
|
||||
// This is a placeholder
|
||||
super::error::set_error_message(error_out, "Tool constraint generation not yet implemented in FFI");
|
||||
super::error::SglErrorCode::UnknownError
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
[run]
|
||||
source = sglang_router
|
||||
omit =
|
||||
*/mini_lb.py
|
||||
*/cli.py
|
||||
*/__main__.py
|
||||
|
||||
[report]
|
||||
fail_under = 80
|
||||
@@ -0,0 +1,27 @@
|
||||
[package]
|
||||
name = "sgl-model-gateway-python"
|
||||
version = "0.2.3"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
name = "sglang_router_rs"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
pyo3 = { version = "0.27.1", features = ["extension-module", "abi3-py38"] }
|
||||
tokio = { version = "1.42.0", features = ["full"] }
|
||||
|
||||
[dependencies.sgl-model-gateway]
|
||||
path = "../.."
|
||||
default-features = true
|
||||
|
||||
[features]
|
||||
default = ["pyo3/extension-module"]
|
||||
vendored-openssl = ["sgl-model-gateway/vendored-openssl"]
|
||||
|
||||
[profile.ci]
|
||||
inherits = "release"
|
||||
opt-level = 2 # Lighter optimization (still fast runtime, much faster compile)
|
||||
lto = "thin" # Thin LTO - good balance
|
||||
codegen-units = 16 # More parallelization for faster builds
|
||||
strip = true
|
||||
@@ -0,0 +1,9 @@
|
||||
# Must include:
|
||||
include Cargo.toml # Python bindings Cargo configuration
|
||||
include ../../Cargo.toml # Main Rust project configuration
|
||||
include ../../build.rs # Build script for protobuf generation
|
||||
include ../../LICENSE
|
||||
recursive-include src *.rs # Python bindings wrapper
|
||||
recursive-include ../../src *.rs # Main Rust source files
|
||||
recursive-include ../../src/proto *.proto # Protobuf definitions
|
||||
recursive-include sglang_router *.py # Python source files
|
||||
@@ -0,0 +1,71 @@
|
||||
# SGLang Model Gateway Python Bindings
|
||||
|
||||
This directory contains the Python bindings for the SGLang Router, built using [maturin](https://github.com/PyO3/maturin) and [PyO3](https://github.com/PyO3/pyo3).
|
||||
|
||||
## Directory Structure
|
||||
|
||||
```
|
||||
bindings/python/
|
||||
├── src/ # Rust source code for Python bindings
|
||||
│ └── lib.rs # PyO3 bindings implementation
|
||||
├── sglang_router/ # Python source code
|
||||
│ ├── __init__.py
|
||||
│ ├── version.py
|
||||
│ ├── launch_server.py
|
||||
│ ├── launch_router.py
|
||||
│ ├── router.py
|
||||
│ ├── router_args.py
|
||||
│ └── mini_lb.py
|
||||
├── Cargo.toml # Rust package configuration for bindings
|
||||
├── pyproject.toml # Python package configuration
|
||||
├── setup.py # Setup configuration
|
||||
├── MANIFEST.in # Package manifest
|
||||
├── .coveragerc # Test coverage configuration
|
||||
└── README.md # This file
|
||||
```
|
||||
|
||||
## Building
|
||||
|
||||
### Development Build
|
||||
|
||||
```bash
|
||||
# Install maturin
|
||||
pip install maturin
|
||||
|
||||
# Build and install in development mode
|
||||
cd sgl-model-gateway/bindings/python
|
||||
maturin develop --features vendored-openssl
|
||||
```
|
||||
|
||||
### Production Build
|
||||
|
||||
```bash
|
||||
# Build wheel
|
||||
cd sgl-model-gateway/bindings/python
|
||||
maturin build --release --out dist --features vendored-openssl
|
||||
|
||||
# Install the built wheel
|
||||
pip install dist/sglang_router-*.whl
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Run Python tests
|
||||
cd sgl-model-gateway
|
||||
pytest py_test/
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
- **pyproject.toml**: Defines package metadata, dependencies, and build configuration
|
||||
- **python-source**: Set to "." to indicate Python source is in the same directory as pyproject.toml
|
||||
- **module-name**: `sglang_router.sglang_router_rs` - the Rust extension module name
|
||||
|
||||
## Notes
|
||||
|
||||
- The Rust bindings source code is located in `src/lib.rs`
|
||||
- The bindings have their own `Cargo.toml` in this directory
|
||||
- The main sglang-router library is located in `../../` and is used as a dependency
|
||||
- The package includes both Python code and Rust extensions built with PyO3
|
||||
- PyO3 types are prefixed with `Py` in Rust but exposed to Python without the prefix using the `name` attribute
|
||||
@@ -0,0 +1,54 @@
|
||||
[build-system]
|
||||
requires = ["maturin>=1.0,<2.0"]
|
||||
build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "sglang-router"
|
||||
version = "0.2.3"
|
||||
description = "High-performance Rust-based load balancer for SGLang with multiple routing algorithms and prefill-decode disaggregation support"
|
||||
authors = [
|
||||
{name = "Simo Lin", email = "linsimo.mark@gmail.com"},
|
||||
{name = "Chang Su", email = "mckvtl@gmail.com"},
|
||||
{name = "Keyang Ru", email = "rukeyang@gmail.com"},
|
||||
{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",
|
||||
"Programming Language :: Python :: 3.8",
|
||||
"Programming Language :: Python :: 3.9",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"setproctitle",
|
||||
"aiohttp",
|
||||
"orjson",
|
||||
"uvicorn",
|
||||
"fastapi",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"requests>=2.25.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
smg = "sglang_router.cli:main"
|
||||
amg = "sglang_router.cli:main"
|
||||
sglang-router = "sglang_router.cli:main"
|
||||
|
||||
|
||||
[tool.maturin]
|
||||
python-source = "."
|
||||
module-name = "sglang_router.sglang_router_rs"
|
||||
# Exclude bindings/python/README.md to use root README only
|
||||
exclude = ["README.md"]
|
||||
@@ -0,0 +1,28 @@
|
||||
import os
|
||||
import warnings
|
||||
|
||||
from setuptools import setup
|
||||
|
||||
with_rust = os.environ.get("SGLANG_ROUTER_BUILD_WITH_RUST", None)
|
||||
with_rust = with_rust is None or (not with_rust.lower() in ["0", "false", "no"])
|
||||
|
||||
rust_extensions = []
|
||||
if with_rust:
|
||||
from setuptools_rust import Binding, RustExtension
|
||||
|
||||
rust_extensions.append(
|
||||
RustExtension(
|
||||
target="sglang_router_rs",
|
||||
path="Cargo.toml",
|
||||
binding=Binding.PyO3,
|
||||
)
|
||||
)
|
||||
else:
|
||||
warnings.warn(
|
||||
"Building 'sglang-router' without Rust support. Performance may be degraded."
|
||||
)
|
||||
|
||||
setup(
|
||||
rust_extensions=rust_extensions,
|
||||
zip_safe=False,
|
||||
)
|
||||
@@ -0,0 +1,3 @@
|
||||
from sglang_router.version import __version__
|
||||
|
||||
__all__ = ["__version__"]
|
||||
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
Allow running the CLI via: python -m sglang_router
|
||||
"""
|
||||
|
||||
from sglang_router.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SGLang Model Gateway CLI
|
||||
|
||||
Provides convenient command-line interface for launching the router and server.
|
||||
|
||||
Usage:
|
||||
smg launch [args] # Launch router only
|
||||
smg server [args] # Launch router + server
|
||||
smg --help # Show help
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from typing import List, Optional
|
||||
|
||||
from sglang_router.sglang_router_rs import (
|
||||
get_verbose_version_string,
|
||||
get_version_string,
|
||||
)
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create the main CLI parser with subcommands."""
|
||||
prog_name = os.path.basename(sys.argv[0]) if sys.argv else "smg"
|
||||
parser = argparse.ArgumentParser(
|
||||
prog=prog_name,
|
||||
description="SGLang Model Gateway - High-performance inference router",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", help="Available commands")
|
||||
|
||||
# Launch router subcommand
|
||||
launch_parser = subparsers.add_parser(
|
||||
"launch",
|
||||
help="Launch router only (requires existing worker URLs)",
|
||||
description="Launch the SGLang router with existing worker instances",
|
||||
add_help=False, # Let router handle --help
|
||||
)
|
||||
|
||||
# Launch server + router subcommand
|
||||
server_parser = subparsers.add_parser(
|
||||
"server",
|
||||
help="Launch router and server processes together",
|
||||
description="Launch both SGLang router and server processes",
|
||||
add_help=False, # Let server handle --help
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Optional[List[str]] = None) -> None:
|
||||
"""Main CLI entry point."""
|
||||
if argv is None:
|
||||
argv = sys.argv[1:]
|
||||
|
||||
# Handle version flags before parsing
|
||||
if argv and argv[0] in ["--version", "-V", "--version-verbose"]:
|
||||
if argv[0] == "--version-verbose":
|
||||
print(get_verbose_version_string())
|
||||
else:
|
||||
print(get_version_string())
|
||||
sys.exit(0)
|
||||
|
||||
# Handle empty command - show help
|
||||
if not argv or argv[0] not in ["launch", "server", "-h", "--help"]:
|
||||
parser = create_parser()
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
parser = create_parser()
|
||||
args, unknown = parser.parse_known_args(argv)
|
||||
|
||||
if args.command == "launch":
|
||||
# Import and call launch_router functions directly
|
||||
from sglang_router.launch_router import launch_router, parse_router_args
|
||||
|
||||
# All router args are in unknown
|
||||
router_args = parse_router_args(unknown)
|
||||
launch_router(router_args)
|
||||
|
||||
elif args.command == "server":
|
||||
# Import and call launch_server main with proper argv
|
||||
# Note: launch_server.main() uses argparse internally which reads sys.argv
|
||||
# We need to temporarily set sys.argv for compatibility
|
||||
import sglang_router.launch_server as launch_server_module
|
||||
|
||||
# Preserve original sys.argv
|
||||
original_argv = sys.argv
|
||||
try:
|
||||
# All server args are in unknown
|
||||
prog_name = os.path.basename(sys.argv[0]) if sys.argv else "smg"
|
||||
sys.argv = [f"{prog_name} server"] + unknown
|
||||
launch_server_module.main()
|
||||
finally:
|
||||
# Restore original sys.argv
|
||||
sys.argv = original_argv
|
||||
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,113 @@
|
||||
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:
|
||||
# TODO: support tracing for router(Rust).
|
||||
del router_args.enable_trace
|
||||
del router_args.otlp_traces_endpoint
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,213 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
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.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 in appropriate mode (HTTP or gRPC)
|
||||
if server_args.grpc_mode:
|
||||
from sglang.srt.entrypoints.grpc_server import serve_grpc
|
||||
|
||||
asyncio.run(serve_grpc(server_args))
|
||||
else:
|
||||
from sglang.srt.entrypoints.http_server import launch_server
|
||||
|
||||
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
|
||||
# Use grpc:// protocol if server is in gRPC mode, otherwise http://
|
||||
protocol = "grpc" if server_args.grpc_mode else "http"
|
||||
router_args.worker_urls = [
|
||||
f"{protocol}://{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()
|
||||
@@ -0,0 +1,488 @@
|
||||
"""
|
||||
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
|
||||
|
||||
try:
|
||||
from sglang.srt.tracing.trace import (
|
||||
process_tracing_init,
|
||||
trace_get_remote_propagate_context,
|
||||
trace_req_finish,
|
||||
trace_req_start,
|
||||
trace_set_thread_info,
|
||||
trace_slice_end,
|
||||
trace_slice_start,
|
||||
)
|
||||
|
||||
trace_package_imported = True
|
||||
except ImportError:
|
||||
trace_package_imported = False
|
||||
|
||||
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
|
||||
self.otlp_traces_endpoint = router_args.otlp_traces_endpoint
|
||||
self.enable_trace = router_args.enable_trace
|
||||
if self.enable_trace and not trace_package_imported:
|
||||
logger.warning(
|
||||
"Tracing is not supported in this environment. Please install sglang."
|
||||
)
|
||||
self.enable_trace = False
|
||||
|
||||
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
|
||||
if self.enable_trace:
|
||||
process_tracing_init(self.otlp_traces_endpoint, "sglang")
|
||||
trace_set_thread_info("Mini lb")
|
||||
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:
|
||||
headers = {}
|
||||
bootstrap_room_list = []
|
||||
if self.enable_trace:
|
||||
bootstrap_room_list = (
|
||||
modified_request["bootstrap_room"]
|
||||
if isinstance(modified_request["bootstrap_room"], list)
|
||||
else [modified_request["bootstrap_room"]]
|
||||
)
|
||||
trace_context = trace_get_remote_propagate_context(bootstrap_room_list)
|
||||
headers = {"trace_context": trace_context}
|
||||
|
||||
tasks = [
|
||||
session.post(
|
||||
f"{prefill_server}/{endpoint}",
|
||||
json=modified_request,
|
||||
headers=headers,
|
||||
),
|
||||
session.post(
|
||||
f"{decode_server}/{endpoint}",
|
||||
json=modified_request,
|
||||
headers=headers,
|
||||
),
|
||||
]
|
||||
|
||||
for bootstrap_room in bootstrap_room_list:
|
||||
trace_slice_end("mini_lb_launch", bootstrap_room, auto_next_anon=True)
|
||||
|
||||
# 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()
|
||||
|
||||
for bootstrap_room in bootstrap_room_list:
|
||||
trace_slice_end(
|
||||
"wait_PD_finish",
|
||||
bootstrap_room,
|
||||
thread_finish_flag=True,
|
||||
)
|
||||
trace_req_finish(bootstrap_room)
|
||||
|
||||
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
|
||||
headers = {}
|
||||
bootstrap_room_list = []
|
||||
if self.enable_trace:
|
||||
bootstrap_room_list = (
|
||||
modified_request["bootstrap_room"]
|
||||
if isinstance(modified_request["bootstrap_room"], list)
|
||||
else [modified_request["bootstrap_room"]]
|
||||
)
|
||||
trace_context = trace_get_remote_propagate_context(
|
||||
bootstrap_room_list
|
||||
)
|
||||
headers = {"trace_context": trace_context}
|
||||
|
||||
tasks = [
|
||||
session.post(
|
||||
f"{prefill_server}/{endpoint}",
|
||||
json=modified_request,
|
||||
headers=headers,
|
||||
),
|
||||
session.post(
|
||||
f"{decode_server}/{endpoint}",
|
||||
json=modified_request,
|
||||
headers=headers,
|
||||
),
|
||||
]
|
||||
|
||||
for bootstrap_room in bootstrap_room_list:
|
||||
trace_slice_end(
|
||||
"mini_lb_launch", bootstrap_room, auto_next_anon=True
|
||||
)
|
||||
# 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
|
||||
|
||||
for bootstrap_room in bootstrap_room_list:
|
||||
trace_slice_end(
|
||||
"wait_PD_finish",
|
||||
bootstrap_room,
|
||||
thread_finish_flag=True,
|
||||
)
|
||||
trace_req_finish(bootstrap_room)
|
||||
|
||||
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():
|
||||
bootstrap_room = random.randint(0, 2**63 - 1)
|
||||
if lb.enable_trace:
|
||||
trace_req_start(bootstrap_room, bootstrap_room, role="router")
|
||||
trace_slice_start("mini_lb_launch", bootstrap_room)
|
||||
return bootstrap_room
|
||||
|
||||
|
||||
# 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))
|
||||
@@ -0,0 +1,226 @@
|
||||
from typing import Optional
|
||||
|
||||
from sglang_router.router_args import RouterArgs
|
||||
from sglang_router.sglang_router_rs import (
|
||||
BackendType,
|
||||
HistoryBackendType,
|
||||
PolicyType,
|
||||
PyOracleConfig,
|
||||
PyPostgresConfig,
|
||||
)
|
||||
from sglang_router.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,
|
||||
"bucket": PolicyType.Bucket,
|
||||
}
|
||||
return policy_map[policy_str]
|
||||
|
||||
|
||||
def backend_from_str(backend_str: Optional[str]) -> BackendType:
|
||||
"""Convert backend string to BackendType enum."""
|
||||
if isinstance(backend_str, BackendType):
|
||||
return backend_str
|
||||
if backend_str is None:
|
||||
return BackendType.Sglang
|
||||
backend_map = {"sglang": BackendType.Sglang, "openai": BackendType.Openai}
|
||||
backend_lower = backend_str.lower()
|
||||
if backend_lower not in backend_map:
|
||||
raise ValueError(
|
||||
f"Unknown backend: {backend_str}. Valid options: {', '.join(backend_map.keys())}"
|
||||
)
|
||||
return backend_map[backend_lower]
|
||||
|
||||
|
||||
def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType:
|
||||
"""Convert history backend string to HistoryBackendType enum."""
|
||||
if isinstance(backend_str, HistoryBackendType):
|
||||
return backend_str
|
||||
if backend_str is None:
|
||||
return HistoryBackendType.Memory
|
||||
backend_lower = backend_str.lower()
|
||||
if backend_lower == "memory":
|
||||
return HistoryBackendType.Memory
|
||||
elif backend_lower == "none":
|
||||
# Use getattr to access 'None' which is a Python keyword
|
||||
return getattr(HistoryBackendType, "None")
|
||||
elif backend_lower == "oracle":
|
||||
return HistoryBackendType.Oracle
|
||||
elif backend_lower == "postgres":
|
||||
return HistoryBackendType.Postgres
|
||||
else:
|
||||
raise ValueError(f"Unknown history backend: {backend_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. Supports IPv4, IPv6 (e.g., ::, ::1), or 0.0.0.0 for all interfaces. Default: '0.0.0.0'
|
||||
port: Port number to bind the router server. Default: 3001
|
||||
worker_startup_timeout_secs: Timeout in seconds for worker startup and registration. Large models can take significant time to load into GPU memory. Default: 1800 (30 minutes)
|
||||
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
|
||||
enable_igw: Enable IGW (Inference-Gateway) mode for multi-model support. When enabled,
|
||||
the router can manage multiple models simultaneously with per-model load balancing
|
||||
policies. 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', 'warn', 'error'.
|
||||
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"])
|
||||
|
||||
# Convert backend
|
||||
args_dict["backend"] = backend_from_str(args_dict.get("backend"))
|
||||
|
||||
# Convert history_backend to enum first
|
||||
history_backend_raw = args_dict.get("history_backend", "memory")
|
||||
history_backend = history_backend_from_str(history_backend_raw)
|
||||
|
||||
# Convert Oracle config if needed
|
||||
oracle_config = None
|
||||
if history_backend == HistoryBackendType.Oracle:
|
||||
# Prioritize TNS alias over connect descriptor
|
||||
tns_alias = args_dict.get("oracle_tns_alias")
|
||||
connect_descriptor = args_dict.get("oracle_connect_descriptor")
|
||||
|
||||
# Use TNS alias if provided, otherwise use connect descriptor
|
||||
final_descriptor = tns_alias if tns_alias else connect_descriptor
|
||||
|
||||
oracle_config = PyOracleConfig(
|
||||
password=args_dict.get("oracle_password"),
|
||||
username=args_dict.get("oracle_username"),
|
||||
connect_descriptor=final_descriptor,
|
||||
wallet_path=args_dict.get("oracle_wallet_path"),
|
||||
pool_min=args_dict.get("oracle_pool_min", 1),
|
||||
pool_max=args_dict.get("oracle_pool_max", 16),
|
||||
pool_timeout_secs=args_dict.get("oracle_pool_timeout_secs", 30),
|
||||
)
|
||||
args_dict["oracle_config"] = oracle_config
|
||||
args_dict["history_backend"] = history_backend
|
||||
|
||||
# Convert Postgres config if needed
|
||||
postgres_config = None
|
||||
if history_backend == HistoryBackendType.Postgres:
|
||||
postgres_config = PyPostgresConfig(
|
||||
db_url=args_dict.get("postgres_db_url"),
|
||||
pool_max=args_dict.get("postgres_pool_max", 16),
|
||||
)
|
||||
args_dict["postgres_config"] = postgres_config
|
||||
|
||||
# Remove fields that shouldn't be passed to Rust Router constructor
|
||||
fields_to_remove = [
|
||||
"mini_lb",
|
||||
"oracle_wallet_path",
|
||||
"oracle_tns_alias",
|
||||
"oracle_connect_descriptor",
|
||||
"oracle_username",
|
||||
"oracle_password",
|
||||
"oracle_pool_min",
|
||||
"oracle_pool_max",
|
||||
"oracle_pool_timeout_secs",
|
||||
"postgres_db_url",
|
||||
"postgres_pool_max",
|
||||
]
|
||||
for field in fields_to_remove:
|
||||
args_dict.pop(field, None)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,782 @@
|
||||
import argparse
|
||||
import dataclasses
|
||||
import logging
|
||||
import os
|
||||
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 = "0.0.0.0"
|
||||
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 = 1800
|
||||
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
|
||||
bucket_adjust_interval_secs: int = 5
|
||||
dp_aware: bool = False
|
||||
enable_igw: bool = False # Enable IGW (Inter-Gateway) mode for multi-model support
|
||||
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 (-1 to disable)
|
||||
max_concurrent_requests: int = -1
|
||||
# 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
|
||||
model_path: Optional[str] = None
|
||||
tokenizer_path: Optional[str] = None
|
||||
chat_template: Optional[str] = None
|
||||
# Tokenizer cache configuration
|
||||
tokenizer_cache_enable_l0: bool = False
|
||||
tokenizer_cache_l0_max_entries: int = 10000
|
||||
tokenizer_cache_enable_l1: bool = False
|
||||
tokenizer_cache_l1_max_memory: int = 50 * 1024 * 1024 # 50MB
|
||||
reasoning_parser: Optional[str] = None
|
||||
tool_call_parser: Optional[str] = None
|
||||
# MCP server configuration
|
||||
mcp_config_path: Optional[str] = None
|
||||
# Backend selection
|
||||
backend: str = "sglang"
|
||||
# History backend configuration
|
||||
history_backend: str = "memory"
|
||||
oracle_wallet_path: Optional[str] = None
|
||||
oracle_tns_alias: Optional[str] = None
|
||||
oracle_connect_descriptor: Optional[str] = None
|
||||
oracle_username: Optional[str] = None
|
||||
oracle_password: Optional[str] = None
|
||||
oracle_pool_min: int = 1
|
||||
oracle_pool_max: int = 16
|
||||
oracle_pool_timeout_secs: int = 30
|
||||
postgres_db_url: Optional[str] = None
|
||||
postgres_pool_max: int = 16
|
||||
# mTLS configuration for worker communication
|
||||
client_cert_path: Optional[str] = None
|
||||
client_key_path: Optional[str] = None
|
||||
ca_cert_paths: List[str] = dataclasses.field(default_factory=list)
|
||||
# Trace
|
||||
enable_trace: bool = False
|
||||
otlp_traces_endpoint: str = "localhost:4317"
|
||||
|
||||
@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. Supports IPv4, IPv6 (e.g., ::, ::1), or 0.0.0.0 for all interfaces",
|
||||
)
|
||||
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. Supports IPv4 and IPv6 addresses (use brackets for IPv6, e.g., http://[::1]:8000 http://192.168.1.1: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", "bucket"],
|
||||
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 and registration (default: 1800 / 30 minutes). Large models can take significant time to load into GPU memory.",
|
||||
)
|
||||
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}bucket-adjust-interval-secs",
|
||||
type=int,
|
||||
default=RouterArgs.bucket_adjust_interval_secs,
|
||||
help="Interval in seconds between bucket boundary adjustment operations",
|
||||
)
|
||||
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}enable-igw",
|
||||
action="store_true",
|
||||
help="Enable IGW (Inference-Gateway) mode for multi-model support",
|
||||
)
|
||||
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", "warn", "error"],
|
||||
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="0.0.0.0",
|
||||
help="Host address to bind the Prometheus metrics server. Supports IPv4, IPv6 (e.g., ::, ::1), or 0.0.0.0 for all interfaces",
|
||||
)
|
||||
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). Set to -1 to disable 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)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}chat-template",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Chat template path (optional)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}tokenizer-cache-enable-l0",
|
||||
action="store_true",
|
||||
default=RouterArgs.tokenizer_cache_enable_l0,
|
||||
help="Enable L0 (whole-string exact match) tokenizer cache (default: False)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}tokenizer-cache-l0-max-entries",
|
||||
type=int,
|
||||
default=RouterArgs.tokenizer_cache_l0_max_entries,
|
||||
help="Maximum number of entries in L0 tokenizer cache (default: 10000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}tokenizer-cache-enable-l1",
|
||||
action="store_true",
|
||||
default=RouterArgs.tokenizer_cache_enable_l1,
|
||||
help="Enable L1 (prefix matching) tokenizer cache (default: False)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}tokenizer-cache-l1-max-memory",
|
||||
type=int,
|
||||
default=RouterArgs.tokenizer_cache_l1_max_memory,
|
||||
help="Maximum memory for L1 tokenizer cache in bytes (default: 50MB)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}reasoning-parser",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Specify the parser for reasoning models (e.g., deepseek-r1, qwen3)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}tool-call-parser",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Specify the parser for handling tool-call interactions",
|
||||
)
|
||||
# MCP server configuration
|
||||
parser.add_argument(
|
||||
f"--{prefix}mcp-config-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to MCP (Model Context Protocol) server configuration file",
|
||||
)
|
||||
# Backend selection
|
||||
parser.add_argument(
|
||||
f"--{prefix}backend",
|
||||
type=str,
|
||||
default=RouterArgs.backend,
|
||||
choices=["sglang", "openai"],
|
||||
help="Backend runtime to use (default: sglang)",
|
||||
)
|
||||
# History backend configuration
|
||||
parser.add_argument(
|
||||
f"--{prefix}history-backend",
|
||||
type=str,
|
||||
default=RouterArgs.history_backend,
|
||||
choices=["memory", "none", "oracle", "postgres"],
|
||||
help="History storage backend for conversations and responses (default: memory)",
|
||||
)
|
||||
# Oracle configuration
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-wallet-path",
|
||||
type=str,
|
||||
default=os.getenv("ATP_WALLET_PATH"),
|
||||
help="Path to Oracle ATP wallet directory (env: ATP_WALLET_PATH)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-tns-alias",
|
||||
type=str,
|
||||
default=os.getenv("ATP_TNS_ALIAS"),
|
||||
help="Oracle TNS alias from tnsnames.ora (env: ATP_TNS_ALIAS).",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-connect-descriptor",
|
||||
type=str,
|
||||
default=os.getenv("ATP_DSN"),
|
||||
help="Oracle connection descriptor/DSN (full connection string) (env: ATP_DSN)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-username",
|
||||
type=str,
|
||||
default=os.getenv("ATP_USER"),
|
||||
help="Oracle database username (env: ATP_USER)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-password",
|
||||
type=str,
|
||||
default=os.getenv("ATP_PASSWORD"),
|
||||
help="Oracle database password (env: ATP_PASSWORD)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-pool-min",
|
||||
type=int,
|
||||
default=int(os.getenv("ATP_POOL_MIN", RouterArgs.oracle_pool_min)),
|
||||
help="Minimum Oracle connection pool size (default: 1, env: ATP_POOL_MIN)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-pool-max",
|
||||
type=int,
|
||||
default=int(os.getenv("ATP_POOL_MAX", RouterArgs.oracle_pool_max)),
|
||||
help="Maximum Oracle connection pool size (default: 16, env: ATP_POOL_MAX)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}oracle-pool-timeout-secs",
|
||||
type=int,
|
||||
default=int(
|
||||
os.getenv("ATP_POOL_TIMEOUT_SECS", RouterArgs.oracle_pool_timeout_secs)
|
||||
),
|
||||
help="Oracle connection pool timeout in seconds (default: 30, env: ATP_POOL_TIMEOUT_SECS)",
|
||||
)
|
||||
# Postgres configuration
|
||||
parser.add_argument(
|
||||
f"--{prefix}postgres-db-url",
|
||||
type=str,
|
||||
default=os.getenv("POSTGRES_DB_URL"),
|
||||
help="PostgreSQL database connection URL (env: POSTGRES_DB_URL)",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}postgres-pool-max",
|
||||
type=int,
|
||||
default=int(os.getenv("POSTGRES_POOL_MAX", RouterArgs.postgres_pool_max)),
|
||||
help="Maximum PostgreSQL connection pool size (default: 16, env: POSTGRES_POOL_MAX)",
|
||||
)
|
||||
# mTLS configuration
|
||||
parser.add_argument(
|
||||
f"--{prefix}client-cert-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to client certificate for mTLS authentication with workers",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}client-key-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to client private key for mTLS authentication with workers",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}ca-cert-paths",
|
||||
type=str,
|
||||
nargs="*",
|
||||
default=[],
|
||||
help="Path(s) to CA certificate(s) for verifying worker TLS certificates. Can specify multiple CAs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}enable-trace",
|
||||
action="store_true",
|
||||
help="Enable opentelemetry trace",
|
||||
)
|
||||
parser.add_argument(
|
||||
f"--{prefix}otlp-traces-endpoint",
|
||||
type=str,
|
||||
default="localhost:4317",
|
||||
help="Config opentelemetry collector endpoint if --enable-trace is set. format: <ip>:<port>",
|
||||
)
|
||||
|
||||
@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:
|
||||
# Allow empty URLs even without service discovery to support dynamic worker addition
|
||||
# URLs will be validated separately if provided
|
||||
pass
|
||||
|
||||
# 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 {}
|
||||
|
||||
# Support `- --selector\n- a=b c=d` case
|
||||
if len(selector_list) == 1 and (" " in selector_list[0]):
|
||||
selector_list = selector_list[0].split(" ")
|
||||
|
||||
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]
|
||||
@@ -0,0 +1 @@
|
||||
__version__ = "0.2.3"
|
||||
@@ -0,0 +1,730 @@
|
||||
use pyo3::prelude::*;
|
||||
use sgl_model_gateway::*;
|
||||
use std::collections::HashMap;
|
||||
|
||||
// Define the enums with PyO3 bindings
|
||||
#[pyclass(eq)]
|
||||
#[derive(Clone, PartialEq, Debug)]
|
||||
pub enum PolicyType {
|
||||
Random,
|
||||
RoundRobin,
|
||||
CacheAware,
|
||||
PowerOfTwo,
|
||||
Bucket,
|
||||
}
|
||||
|
||||
#[pyclass(eq)]
|
||||
#[derive(Clone, PartialEq, Debug)]
|
||||
pub enum BackendType {
|
||||
Sglang,
|
||||
Openai,
|
||||
}
|
||||
|
||||
#[pyclass(eq)]
|
||||
#[derive(Clone, PartialEq, Debug)]
|
||||
pub enum HistoryBackendType {
|
||||
Memory,
|
||||
None,
|
||||
Oracle,
|
||||
Postgres,
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct PyOracleConfig {
|
||||
#[pyo3(get, set)]
|
||||
pub wallet_path: Option<String>,
|
||||
#[pyo3(get, set)]
|
||||
pub connect_descriptor: Option<String>,
|
||||
#[pyo3(get, set)]
|
||||
pub username: Option<String>,
|
||||
#[pyo3(get, set)]
|
||||
pub password: Option<String>,
|
||||
#[pyo3(get, set)]
|
||||
pub pool_min: usize,
|
||||
#[pyo3(get, set)]
|
||||
pub pool_max: usize,
|
||||
#[pyo3(get, set)]
|
||||
pub pool_timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for PyOracleConfig {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("PyOracleConfig")
|
||||
.field("wallet_path", &self.wallet_path)
|
||||
.field("connect_descriptor", &"<redacted>")
|
||||
.field("username", &self.username)
|
||||
.field("password", &"<redacted>")
|
||||
.field("pool_min", &self.pool_min)
|
||||
.field("pool_max", &self.pool_max)
|
||||
.field("pool_timeout_secs", &self.pool_timeout_secs)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyOracleConfig {
|
||||
#[new]
|
||||
#[pyo3(signature = (
|
||||
password = None,
|
||||
username = None,
|
||||
connect_descriptor = None,
|
||||
wallet_path = None,
|
||||
pool_min = 1,
|
||||
pool_max = 16,
|
||||
pool_timeout_secs = 30,
|
||||
))]
|
||||
fn new(
|
||||
password: Option<String>,
|
||||
username: Option<String>,
|
||||
connect_descriptor: Option<String>,
|
||||
wallet_path: Option<String>,
|
||||
pool_min: usize,
|
||||
pool_max: usize,
|
||||
pool_timeout_secs: u64,
|
||||
) -> PyResult<Self> {
|
||||
if pool_min == 0 {
|
||||
return Err(pyo3::exceptions::PyValueError::new_err(
|
||||
"pool_min must be at least 1",
|
||||
));
|
||||
}
|
||||
if pool_max < pool_min {
|
||||
return Err(pyo3::exceptions::PyValueError::new_err(
|
||||
"pool_max must be >= pool_min",
|
||||
));
|
||||
}
|
||||
|
||||
Ok(PyOracleConfig {
|
||||
wallet_path,
|
||||
connect_descriptor,
|
||||
username,
|
||||
password,
|
||||
pool_min,
|
||||
pool_max,
|
||||
pool_timeout_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl PyOracleConfig {
|
||||
pub fn to_config_oracle(&self) -> config::OracleConfig {
|
||||
config::OracleConfig {
|
||||
wallet_path: self.wallet_path.clone(),
|
||||
connect_descriptor: self.connect_descriptor.clone().unwrap_or_default(),
|
||||
username: self.username.clone().unwrap_or_default(),
|
||||
password: self.password.clone().unwrap_or_default(),
|
||||
pool_min: self.pool_min,
|
||||
pool_max: self.pool_max,
|
||||
pool_timeout_secs: self.pool_timeout_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PyPostgresConfig {
|
||||
#[pyo3(get, set)]
|
||||
pub db_url: Option<String>,
|
||||
|
||||
#[pyo3(get, set)]
|
||||
pub pool_max: usize,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyPostgresConfig {
|
||||
#[new]
|
||||
#[pyo3(signature = (db_url = None,pool_max = 16,))]
|
||||
fn new(db_url: Option<String>, pool_max: usize) -> PyResult<Self> {
|
||||
Ok(PyPostgresConfig { db_url, pool_max })
|
||||
}
|
||||
}
|
||||
|
||||
impl PyPostgresConfig {
|
||||
pub fn to_config_postgres(&self) -> config::PostgresConfig {
|
||||
config::PostgresConfig {
|
||||
db_url: self.db_url.clone().unwrap_or_default(),
|
||||
pool_max: self.pool_max,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
struct Router {
|
||||
host: String,
|
||||
port: u16,
|
||||
worker_urls: Vec<String>,
|
||||
policy: PolicyType,
|
||||
worker_startup_timeout_secs: u64,
|
||||
worker_startup_check_interval: u64,
|
||||
cache_threshold: f32,
|
||||
balance_abs_threshold: usize,
|
||||
balance_rel_threshold: f32,
|
||||
eviction_interval_secs: u64,
|
||||
max_tree_size: usize,
|
||||
max_payload_size: usize,
|
||||
dp_aware: bool,
|
||||
api_key: Option<String>,
|
||||
log_dir: Option<String>,
|
||||
log_level: Option<String>,
|
||||
service_discovery: bool,
|
||||
selector: HashMap<String, String>,
|
||||
service_discovery_port: u16,
|
||||
service_discovery_namespace: Option<String>,
|
||||
prefill_selector: HashMap<String, String>,
|
||||
decode_selector: HashMap<String, String>,
|
||||
bootstrap_port_annotation: String,
|
||||
prometheus_port: Option<u16>,
|
||||
prometheus_host: Option<String>,
|
||||
request_timeout_secs: u64,
|
||||
request_id_headers: Option<Vec<String>>,
|
||||
pd_disaggregation: bool,
|
||||
bucket_adjust_interval_secs: usize,
|
||||
prefill_urls: Option<Vec<(String, Option<u16>)>>,
|
||||
decode_urls: Option<Vec<String>>,
|
||||
prefill_policy: Option<PolicyType>,
|
||||
decode_policy: Option<PolicyType>,
|
||||
max_concurrent_requests: i32,
|
||||
cors_allowed_origins: Vec<String>,
|
||||
retry_max_retries: u32,
|
||||
retry_initial_backoff_ms: u64,
|
||||
retry_max_backoff_ms: u64,
|
||||
retry_backoff_multiplier: f32,
|
||||
retry_jitter_factor: f32,
|
||||
disable_retries: bool,
|
||||
cb_failure_threshold: u32,
|
||||
cb_success_threshold: u32,
|
||||
cb_timeout_duration_secs: u64,
|
||||
cb_window_duration_secs: u64,
|
||||
disable_circuit_breaker: bool,
|
||||
health_failure_threshold: u32,
|
||||
health_success_threshold: u32,
|
||||
health_check_timeout_secs: u64,
|
||||
health_check_interval_secs: u64,
|
||||
health_check_endpoint: String,
|
||||
enable_igw: bool,
|
||||
queue_size: usize,
|
||||
queue_timeout_secs: u64,
|
||||
rate_limit_tokens_per_second: Option<i32>,
|
||||
connection_mode: core::ConnectionMode,
|
||||
model_path: Option<String>,
|
||||
tokenizer_path: Option<String>,
|
||||
chat_template: Option<String>,
|
||||
tokenizer_cache_enable_l0: bool,
|
||||
tokenizer_cache_l0_max_entries: usize,
|
||||
tokenizer_cache_enable_l1: bool,
|
||||
tokenizer_cache_l1_max_memory: usize,
|
||||
reasoning_parser: Option<String>,
|
||||
tool_call_parser: Option<String>,
|
||||
mcp_config_path: Option<String>,
|
||||
backend: BackendType,
|
||||
history_backend: HistoryBackendType,
|
||||
oracle_config: Option<PyOracleConfig>,
|
||||
postgres_config: Option<PyPostgresConfig>,
|
||||
client_cert_path: Option<String>,
|
||||
client_key_path: Option<String>,
|
||||
ca_cert_paths: Vec<String>,
|
||||
}
|
||||
|
||||
impl Router {
|
||||
fn determine_connection_mode(worker_urls: &[String]) -> core::ConnectionMode {
|
||||
for url in worker_urls {
|
||||
if url.starts_with("grpc://") || url.starts_with("grpcs://") {
|
||||
return core::ConnectionMode::Grpc { port: None };
|
||||
}
|
||||
}
|
||||
core::ConnectionMode::Http
|
||||
}
|
||||
|
||||
pub fn to_router_config(&self) -> config::ConfigResult<config::RouterConfig> {
|
||||
use config::{
|
||||
DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode,
|
||||
};
|
||||
|
||||
let convert_policy = |policy: &PolicyType| -> ConfigPolicyConfig {
|
||||
match policy {
|
||||
PolicyType::Random => ConfigPolicyConfig::Random,
|
||||
PolicyType::RoundRobin => ConfigPolicyConfig::RoundRobin,
|
||||
PolicyType::CacheAware => ConfigPolicyConfig::CacheAware {
|
||||
cache_threshold: self.cache_threshold,
|
||||
balance_abs_threshold: self.balance_abs_threshold,
|
||||
balance_rel_threshold: self.balance_rel_threshold,
|
||||
eviction_interval_secs: self.eviction_interval_secs,
|
||||
max_tree_size: self.max_tree_size,
|
||||
},
|
||||
PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
|
||||
load_check_interval_secs: 5,
|
||||
},
|
||||
PolicyType::Bucket => ConfigPolicyConfig::Bucket {
|
||||
balance_abs_threshold: self.balance_abs_threshold,
|
||||
balance_rel_threshold: self.balance_rel_threshold,
|
||||
bucket_adjust_interval_secs: self.bucket_adjust_interval_secs,
|
||||
},
|
||||
}
|
||||
};
|
||||
|
||||
let mode = if self.enable_igw {
|
||||
RoutingMode::Regular {
|
||||
worker_urls: vec![],
|
||||
}
|
||||
} else if matches!(self.backend, BackendType::Openai) {
|
||||
RoutingMode::OpenAI {
|
||||
worker_urls: self.worker_urls.clone(),
|
||||
}
|
||||
} else if self.pd_disaggregation {
|
||||
RoutingMode::PrefillDecode {
|
||||
prefill_urls: self.prefill_urls.clone().unwrap_or_default(),
|
||||
decode_urls: self.decode_urls.clone().unwrap_or_default(),
|
||||
prefill_policy: self.prefill_policy.as_ref().map(convert_policy),
|
||||
decode_policy: self.decode_policy.as_ref().map(convert_policy),
|
||||
}
|
||||
} else {
|
||||
RoutingMode::Regular {
|
||||
worker_urls: self.worker_urls.clone(),
|
||||
}
|
||||
};
|
||||
|
||||
let policy = convert_policy(&self.policy);
|
||||
|
||||
let discovery = if self.service_discovery {
|
||||
Some(DiscoveryConfig {
|
||||
enabled: true,
|
||||
namespace: self.service_discovery_namespace.clone(),
|
||||
port: self.service_discovery_port,
|
||||
check_interval_secs: 60,
|
||||
selector: self.selector.clone(),
|
||||
prefill_selector: self.prefill_selector.clone(),
|
||||
decode_selector: self.decode_selector.clone(),
|
||||
bootstrap_port_annotation: self.bootstrap_port_annotation.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let metrics = match (self.prometheus_port, self.prometheus_host.as_ref()) {
|
||||
(Some(port), Some(host)) => Some(MetricsConfig {
|
||||
port,
|
||||
host: host.clone(),
|
||||
}),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let history_backend = match self.history_backend {
|
||||
HistoryBackendType::Memory => config::HistoryBackend::Memory,
|
||||
HistoryBackendType::None => config::HistoryBackend::None,
|
||||
HistoryBackendType::Oracle => config::HistoryBackend::Oracle,
|
||||
HistoryBackendType::Postgres => config::HistoryBackend::Postgres,
|
||||
};
|
||||
|
||||
let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) {
|
||||
self.oracle_config
|
||||
.as_ref()
|
||||
.map(|cfg| cfg.to_config_oracle())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let postgres_config = if matches!(self.history_backend, HistoryBackendType::Postgres) {
|
||||
self.postgres_config
|
||||
.as_ref()
|
||||
.map(|cfg| cfg.to_config_postgres())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
config::RouterConfig::builder()
|
||||
.mode(mode)
|
||||
.policy(policy)
|
||||
.host(&self.host)
|
||||
.port(self.port)
|
||||
.connection_mode(self.connection_mode.clone())
|
||||
.max_payload_size(self.max_payload_size)
|
||||
.request_timeout_secs(self.request_timeout_secs)
|
||||
.worker_startup_timeout_secs(self.worker_startup_timeout_secs)
|
||||
.worker_startup_check_interval_secs(self.worker_startup_check_interval)
|
||||
.max_concurrent_requests(self.max_concurrent_requests)
|
||||
.queue_size(self.queue_size)
|
||||
.queue_timeout_secs(self.queue_timeout_secs)
|
||||
.cors_allowed_origins(self.cors_allowed_origins.clone())
|
||||
.retry_config(config::RetryConfig {
|
||||
max_retries: self.retry_max_retries,
|
||||
initial_backoff_ms: self.retry_initial_backoff_ms,
|
||||
max_backoff_ms: self.retry_max_backoff_ms,
|
||||
backoff_multiplier: self.retry_backoff_multiplier,
|
||||
jitter_factor: self.retry_jitter_factor,
|
||||
})
|
||||
.circuit_breaker_config(config::CircuitBreakerConfig {
|
||||
failure_threshold: self.cb_failure_threshold,
|
||||
success_threshold: self.cb_success_threshold,
|
||||
timeout_duration_secs: self.cb_timeout_duration_secs,
|
||||
window_duration_secs: self.cb_window_duration_secs,
|
||||
})
|
||||
.health_check_config(config::HealthCheckConfig {
|
||||
failure_threshold: self.health_failure_threshold,
|
||||
success_threshold: self.health_success_threshold,
|
||||
timeout_secs: self.health_check_timeout_secs,
|
||||
check_interval_secs: self.health_check_interval_secs,
|
||||
endpoint: self.health_check_endpoint.clone(),
|
||||
})
|
||||
.tokenizer_cache(config::TokenizerCacheConfig {
|
||||
enable_l0: self.tokenizer_cache_enable_l0,
|
||||
l0_max_entries: self.tokenizer_cache_l0_max_entries,
|
||||
enable_l1: self.tokenizer_cache_enable_l1,
|
||||
l1_max_memory: self.tokenizer_cache_l1_max_memory,
|
||||
})
|
||||
.history_backend(history_backend)
|
||||
.maybe_api_key(self.api_key.as_ref())
|
||||
.maybe_discovery(discovery)
|
||||
.maybe_metrics(metrics)
|
||||
.maybe_log_dir(self.log_dir.as_ref())
|
||||
.maybe_log_level(self.log_level.as_ref())
|
||||
.maybe_request_id_headers(self.request_id_headers.clone())
|
||||
.maybe_rate_limit_tokens_per_second(self.rate_limit_tokens_per_second)
|
||||
.maybe_model_path(self.model_path.as_ref())
|
||||
.maybe_tokenizer_path(self.tokenizer_path.as_ref())
|
||||
.maybe_chat_template(self.chat_template.as_ref())
|
||||
.maybe_oracle(oracle)
|
||||
.maybe_postgres(postgres_config)
|
||||
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
|
||||
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
|
||||
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
|
||||
.dp_aware(self.dp_aware)
|
||||
.retries(!self.disable_retries)
|
||||
.circuit_breaker(!self.disable_circuit_breaker)
|
||||
.igw(self.enable_igw)
|
||||
.maybe_client_cert_and_key(
|
||||
self.client_cert_path.as_ref(),
|
||||
self.client_key_path.as_ref(),
|
||||
)
|
||||
.add_ca_certificates(self.ca_cert_paths.clone())
|
||||
.build()
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Router {
|
||||
#[new]
|
||||
#[pyo3(signature = (
|
||||
worker_urls,
|
||||
policy = PolicyType::RoundRobin,
|
||||
host = String::from("0.0.0.0"),
|
||||
port = 3001,
|
||||
worker_startup_timeout_secs = 600,
|
||||
worker_startup_check_interval = 30,
|
||||
cache_threshold = 0.3,
|
||||
balance_abs_threshold = 64,
|
||||
balance_rel_threshold = 1.5,
|
||||
eviction_interval_secs = 120,
|
||||
max_tree_size = 2usize.pow(26),
|
||||
max_payload_size = 512 * 1024 * 1024,
|
||||
dp_aware = false,
|
||||
api_key = None,
|
||||
log_dir = None,
|
||||
log_level = None,
|
||||
service_discovery = false,
|
||||
selector = HashMap::new(),
|
||||
service_discovery_port = 80,
|
||||
service_discovery_namespace = None,
|
||||
prefill_selector = HashMap::new(),
|
||||
decode_selector = HashMap::new(),
|
||||
bootstrap_port_annotation = String::from("sglang.ai/bootstrap-port"),
|
||||
prometheus_port = None,
|
||||
prometheus_host = None,
|
||||
request_timeout_secs = 1800,
|
||||
request_id_headers = None,
|
||||
pd_disaggregation = false,
|
||||
bucket_adjust_interval_secs = 5,
|
||||
prefill_urls = None,
|
||||
decode_urls = None,
|
||||
prefill_policy = None,
|
||||
decode_policy = None,
|
||||
max_concurrent_requests = -1,
|
||||
cors_allowed_origins = vec![],
|
||||
retry_max_retries = 5,
|
||||
retry_initial_backoff_ms = 50,
|
||||
retry_max_backoff_ms = 30_000,
|
||||
retry_backoff_multiplier = 1.5,
|
||||
retry_jitter_factor = 0.2,
|
||||
disable_retries = false,
|
||||
cb_failure_threshold = 10,
|
||||
cb_success_threshold = 3,
|
||||
cb_timeout_duration_secs = 60,
|
||||
cb_window_duration_secs = 120,
|
||||
disable_circuit_breaker = false,
|
||||
health_failure_threshold = 3,
|
||||
health_success_threshold = 2,
|
||||
health_check_timeout_secs = 5,
|
||||
health_check_interval_secs = 60,
|
||||
health_check_endpoint = String::from("/health"),
|
||||
enable_igw = false,
|
||||
queue_size = 100,
|
||||
queue_timeout_secs = 60,
|
||||
rate_limit_tokens_per_second = None,
|
||||
model_path = None,
|
||||
tokenizer_path = None,
|
||||
chat_template = None,
|
||||
tokenizer_cache_enable_l0 = false,
|
||||
tokenizer_cache_l0_max_entries = 10000,
|
||||
tokenizer_cache_enable_l1 = false,
|
||||
tokenizer_cache_l1_max_memory = 52428800,
|
||||
reasoning_parser = None,
|
||||
tool_call_parser = None,
|
||||
mcp_config_path = None,
|
||||
backend = BackendType::Sglang,
|
||||
history_backend = HistoryBackendType::Memory,
|
||||
oracle_config = None,
|
||||
postgres_config = None,
|
||||
client_cert_path = None,
|
||||
client_key_path = None,
|
||||
ca_cert_paths = vec![],
|
||||
))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn new(
|
||||
worker_urls: Vec<String>,
|
||||
policy: PolicyType,
|
||||
host: String,
|
||||
port: u16,
|
||||
worker_startup_timeout_secs: u64,
|
||||
worker_startup_check_interval: u64,
|
||||
cache_threshold: f32,
|
||||
balance_abs_threshold: usize,
|
||||
balance_rel_threshold: f32,
|
||||
eviction_interval_secs: u64,
|
||||
max_tree_size: usize,
|
||||
max_payload_size: usize,
|
||||
dp_aware: bool,
|
||||
api_key: Option<String>,
|
||||
log_dir: Option<String>,
|
||||
log_level: Option<String>,
|
||||
service_discovery: bool,
|
||||
selector: HashMap<String, String>,
|
||||
service_discovery_port: u16,
|
||||
service_discovery_namespace: Option<String>,
|
||||
prefill_selector: HashMap<String, String>,
|
||||
decode_selector: HashMap<String, String>,
|
||||
bootstrap_port_annotation: String,
|
||||
prometheus_port: Option<u16>,
|
||||
prometheus_host: Option<String>,
|
||||
request_timeout_secs: u64,
|
||||
request_id_headers: Option<Vec<String>>,
|
||||
pd_disaggregation: bool,
|
||||
bucket_adjust_interval_secs: usize,
|
||||
prefill_urls: Option<Vec<(String, Option<u16>)>>,
|
||||
decode_urls: Option<Vec<String>>,
|
||||
prefill_policy: Option<PolicyType>,
|
||||
decode_policy: Option<PolicyType>,
|
||||
max_concurrent_requests: i32,
|
||||
cors_allowed_origins: Vec<String>,
|
||||
retry_max_retries: u32,
|
||||
retry_initial_backoff_ms: u64,
|
||||
retry_max_backoff_ms: u64,
|
||||
retry_backoff_multiplier: f32,
|
||||
retry_jitter_factor: f32,
|
||||
disable_retries: bool,
|
||||
cb_failure_threshold: u32,
|
||||
cb_success_threshold: u32,
|
||||
cb_timeout_duration_secs: u64,
|
||||
cb_window_duration_secs: u64,
|
||||
disable_circuit_breaker: bool,
|
||||
health_failure_threshold: u32,
|
||||
health_success_threshold: u32,
|
||||
health_check_timeout_secs: u64,
|
||||
health_check_interval_secs: u64,
|
||||
health_check_endpoint: String,
|
||||
enable_igw: bool,
|
||||
queue_size: usize,
|
||||
queue_timeout_secs: u64,
|
||||
rate_limit_tokens_per_second: Option<i32>,
|
||||
model_path: Option<String>,
|
||||
tokenizer_path: Option<String>,
|
||||
chat_template: Option<String>,
|
||||
tokenizer_cache_enable_l0: bool,
|
||||
tokenizer_cache_l0_max_entries: usize,
|
||||
tokenizer_cache_enable_l1: bool,
|
||||
tokenizer_cache_l1_max_memory: usize,
|
||||
reasoning_parser: Option<String>,
|
||||
tool_call_parser: Option<String>,
|
||||
mcp_config_path: Option<String>,
|
||||
backend: BackendType,
|
||||
history_backend: HistoryBackendType,
|
||||
oracle_config: Option<PyOracleConfig>,
|
||||
postgres_config: Option<PyPostgresConfig>,
|
||||
client_cert_path: Option<String>,
|
||||
client_key_path: Option<String>,
|
||||
ca_cert_paths: Vec<String>,
|
||||
) -> PyResult<Self> {
|
||||
let mut all_urls = worker_urls.clone();
|
||||
|
||||
if let Some(ref prefill_urls) = prefill_urls {
|
||||
for (url, _) in prefill_urls {
|
||||
all_urls.push(url.clone());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref decode_urls) = decode_urls {
|
||||
all_urls.extend(decode_urls.clone());
|
||||
}
|
||||
|
||||
let connection_mode = Self::determine_connection_mode(&all_urls);
|
||||
|
||||
Ok(Router {
|
||||
host,
|
||||
port,
|
||||
worker_urls,
|
||||
policy,
|
||||
worker_startup_timeout_secs,
|
||||
worker_startup_check_interval,
|
||||
cache_threshold,
|
||||
balance_abs_threshold,
|
||||
balance_rel_threshold,
|
||||
eviction_interval_secs,
|
||||
max_tree_size,
|
||||
max_payload_size,
|
||||
dp_aware,
|
||||
api_key,
|
||||
log_dir,
|
||||
log_level,
|
||||
service_discovery,
|
||||
selector,
|
||||
service_discovery_port,
|
||||
service_discovery_namespace,
|
||||
prefill_selector,
|
||||
decode_selector,
|
||||
bootstrap_port_annotation,
|
||||
prometheus_port,
|
||||
prometheus_host,
|
||||
request_timeout_secs,
|
||||
request_id_headers,
|
||||
pd_disaggregation,
|
||||
bucket_adjust_interval_secs,
|
||||
prefill_urls,
|
||||
decode_urls,
|
||||
prefill_policy,
|
||||
decode_policy,
|
||||
max_concurrent_requests,
|
||||
cors_allowed_origins,
|
||||
retry_max_retries,
|
||||
retry_initial_backoff_ms,
|
||||
retry_max_backoff_ms,
|
||||
retry_backoff_multiplier,
|
||||
retry_jitter_factor,
|
||||
disable_retries,
|
||||
cb_failure_threshold,
|
||||
cb_success_threshold,
|
||||
cb_timeout_duration_secs,
|
||||
cb_window_duration_secs,
|
||||
disable_circuit_breaker,
|
||||
health_failure_threshold,
|
||||
health_success_threshold,
|
||||
health_check_timeout_secs,
|
||||
health_check_interval_secs,
|
||||
health_check_endpoint,
|
||||
enable_igw,
|
||||
queue_size,
|
||||
queue_timeout_secs,
|
||||
rate_limit_tokens_per_second,
|
||||
connection_mode,
|
||||
model_path,
|
||||
tokenizer_path,
|
||||
chat_template,
|
||||
tokenizer_cache_enable_l0,
|
||||
tokenizer_cache_l0_max_entries,
|
||||
tokenizer_cache_enable_l1,
|
||||
tokenizer_cache_l1_max_memory,
|
||||
reasoning_parser,
|
||||
tool_call_parser,
|
||||
mcp_config_path,
|
||||
backend,
|
||||
history_backend,
|
||||
oracle_config,
|
||||
postgres_config,
|
||||
client_cert_path,
|
||||
client_key_path,
|
||||
ca_cert_paths,
|
||||
})
|
||||
}
|
||||
|
||||
fn start(&self) -> PyResult<()> {
|
||||
use metrics::PrometheusConfig;
|
||||
|
||||
let router_config = self.to_router_config().map_err(|e| {
|
||||
pyo3::exceptions::PyValueError::new_err(format!("Configuration error: {}", e))
|
||||
})?;
|
||||
|
||||
router_config.validate().map_err(|e| {
|
||||
pyo3::exceptions::PyValueError::new_err(format!(
|
||||
"Configuration validation failed: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
|
||||
let service_discovery_config = if self.service_discovery {
|
||||
Some(service_discovery::ServiceDiscoveryConfig {
|
||||
enabled: true,
|
||||
selector: self.selector.clone(),
|
||||
check_interval: std::time::Duration::from_secs(60),
|
||||
port: self.service_discovery_port,
|
||||
namespace: self.service_discovery_namespace.clone(),
|
||||
pd_mode: self.pd_disaggregation,
|
||||
prefill_selector: self.prefill_selector.clone(),
|
||||
decode_selector: self.decode_selector.clone(),
|
||||
bootstrap_port_annotation: self.bootstrap_port_annotation.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let prometheus_config = Some(PrometheusConfig {
|
||||
port: self.prometheus_port.unwrap_or(29000),
|
||||
host: self
|
||||
.prometheus_host
|
||||
.clone()
|
||||
.unwrap_or_else(|| "127.0.0.1".to_string()),
|
||||
});
|
||||
|
||||
let runtime = tokio::runtime::Runtime::new()
|
||||
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
|
||||
|
||||
runtime.block_on(async move {
|
||||
server::startup(server::ServerConfig {
|
||||
host: self.host.clone(),
|
||||
port: self.port,
|
||||
router_config,
|
||||
max_payload_size: self.max_payload_size,
|
||||
log_dir: self.log_dir.clone(),
|
||||
log_level: self.log_level.clone(),
|
||||
service_discovery_config,
|
||||
prometheus_config,
|
||||
request_timeout_secs: self.request_timeout_secs,
|
||||
request_id_headers: self.request_id_headers.clone(),
|
||||
})
|
||||
.await
|
||||
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Get simple version string (default for --version)
|
||||
#[pyfunction]
|
||||
fn get_version_string() -> String {
|
||||
version::get_version_string()
|
||||
}
|
||||
|
||||
/// Get verbose version information string with full build details (for --version-verbose)
|
||||
#[pyfunction]
|
||||
fn get_verbose_version_string() -> String {
|
||||
version::get_verbose_version_string()
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PolicyType>()?;
|
||||
m.add_class::<BackendType>()?;
|
||||
m.add_class::<HistoryBackendType>()?;
|
||||
m.add_class::<PyOracleConfig>()?;
|
||||
m.add_class::<PyPostgresConfig>()?;
|
||||
m.add_class::<Router>()?;
|
||||
m.add_function(wrap_pyfunction!(get_version_string, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(get_verbose_version_string, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
use std::process::Command;
|
||||
|
||||
const DEFAULT_VERSION: &str = "0.0.0";
|
||||
const DEFAULT_PROJECT_NAME: &str = "sgl-model-gateway";
|
||||
|
||||
/// Set a compile-time environment variable with the SGL_MODEL_GATEWAY_ prefix
|
||||
macro_rules! set_env {
|
||||
($name:expr, $value:expr) => {
|
||||
println!("cargo:rustc-env=SGL_MODEL_GATEWAY_{}={}", $name, $value);
|
||||
};
|
||||
}
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Rebuild triggers
|
||||
println!("cargo:rerun-if-changed=src/proto/sglang_scheduler.proto");
|
||||
println!("cargo:rerun-if-changed=src/proto/vllm_engine.proto");
|
||||
println!("cargo:rerun-if-changed=Cargo.toml");
|
||||
|
||||
// Compile protobuf files
|
||||
tonic_prost_build::configure()
|
||||
.build_server(true)
|
||||
.build_client(true)
|
||||
.type_attribute("GetModelInfoResponse", "#[derive(serde::Serialize)]")
|
||||
.protoc_arg("--experimental_allow_proto3_optional")
|
||||
.compile_protos(
|
||||
&[
|
||||
"src/proto/sglang_scheduler.proto",
|
||||
"src/proto/vllm_engine.proto",
|
||||
],
|
||||
&["src/proto"],
|
||||
)?;
|
||||
|
||||
// Set version info environment variables
|
||||
let version = read_cargo_version().unwrap_or_else(|_| DEFAULT_VERSION.to_string());
|
||||
let target = std::env::var("TARGET").unwrap_or_else(|_| get_rustc_host().unwrap_or_default());
|
||||
let profile = std::env::var("PROFILE").unwrap_or_default();
|
||||
|
||||
set_env!("PROJECT_NAME", DEFAULT_PROJECT_NAME);
|
||||
set_env!("VERSION", version);
|
||||
set_env!(
|
||||
"BUILD_TIME",
|
||||
chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC")
|
||||
);
|
||||
set_env!(
|
||||
"BUILD_MODE",
|
||||
if profile == "release" {
|
||||
"release"
|
||||
} else {
|
||||
"debug"
|
||||
}
|
||||
);
|
||||
set_env!("TARGET_TRIPLE", target);
|
||||
set_env!(
|
||||
"GIT_BRANCH",
|
||||
git_branch().unwrap_or_else(|| "unknown".into())
|
||||
);
|
||||
set_env!(
|
||||
"GIT_COMMIT",
|
||||
git_commit().unwrap_or_else(|| "unknown".into())
|
||||
);
|
||||
set_env!(
|
||||
"GIT_STATUS",
|
||||
git_status().unwrap_or_else(|| "unknown".into())
|
||||
);
|
||||
set_env!(
|
||||
"RUSTC_VERSION",
|
||||
rustc_version().unwrap_or_else(|| "unknown".into())
|
||||
);
|
||||
set_env!(
|
||||
"CARGO_VERSION",
|
||||
cargo_version().unwrap_or_else(|| "unknown".into())
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn read_cargo_version() -> Result<String, Box<dyn std::error::Error>> {
|
||||
let content = std::fs::read_to_string("Cargo.toml")?;
|
||||
let toml: toml::Value = toml::from_str(&content)?;
|
||||
toml.get("package")
|
||||
.and_then(|p| p.get("version"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
.ok_or_else(|| "Missing version in Cargo.toml".into())
|
||||
}
|
||||
|
||||
fn run_cmd(cmd: &str, args: &[&str]) -> Option<String> {
|
||||
Command::new(cmd)
|
||||
.args(args)
|
||||
.output()
|
||||
.ok()
|
||||
.filter(|o| o.status.success())
|
||||
.and_then(|o| String::from_utf8(o.stdout).ok())
|
||||
.map(|s| s.trim().to_string())
|
||||
}
|
||||
|
||||
fn git_branch() -> Option<String> {
|
||||
run_cmd("git", &["rev-parse", "--abbrev-ref", "HEAD"])
|
||||
}
|
||||
|
||||
fn git_commit() -> Option<String> {
|
||||
run_cmd("git", &["rev-parse", "--short", "HEAD"])
|
||||
}
|
||||
|
||||
fn git_status() -> Option<String> {
|
||||
run_cmd("git", &["status", "--porcelain"])
|
||||
.map(|s| if s.is_empty() { "clean" } else { "dirty" }.into())
|
||||
}
|
||||
|
||||
fn rustc_version() -> Option<String> {
|
||||
run_cmd("rustc", &["--version"])
|
||||
}
|
||||
|
||||
fn cargo_version() -> Option<String> {
|
||||
run_cmd("cargo", &["--version"])
|
||||
}
|
||||
|
||||
fn get_rustc_host() -> Option<String> {
|
||||
run_cmd("rustc", &["-vV"])?
|
||||
.lines()
|
||||
.find(|l| l.starts_with("host: "))
|
||||
.and_then(|l| l.strip_prefix("host: "))
|
||||
.map(|s| s.trim().to_string())
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
# Rust build artifacts
|
||||
target/
|
||||
**/target/
|
||||
|
||||
# Cargo lock files (examples don't need locked dependencies)
|
||||
Cargo.lock
|
||||
**/Cargo.lock
|
||||
|
||||
# Generated WASM files
|
||||
*.wasm
|
||||
*.component.wasm
|
||||
**/*.wasm
|
||||
**/*.component.wasm
|
||||
|
||||
# Build scripts output
|
||||
build/
|
||||
|
||||
# IDE files
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS files
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
@@ -0,0 +1,102 @@
|
||||
# WASM Guest Examples for sgl-router
|
||||
|
||||
This directory contains example WASM middleware components demonstrating how to implement custom middleware for sgl-router using the WebAssembly Component Model.
|
||||
|
||||
## Examples Overview
|
||||
|
||||
### [wasm-guest-auth](./wasm-guest-auth/)
|
||||
|
||||
API key authentication middleware that validates API keys for requests to `/api` and `/v1` paths.
|
||||
|
||||
**Features:**
|
||||
- Validates API keys from `Authorization` header or `x-api-key` header
|
||||
- Returns `401 Unauthorized` for missing or invalid keys
|
||||
- Attach point: `OnRequest` only
|
||||
|
||||
**Use case:** Protect API endpoints with API key authentication.
|
||||
|
||||
### [wasm-guest-logging](./wasm-guest-logging/)
|
||||
|
||||
Request tracking and status code conversion middleware.
|
||||
|
||||
**Features:**
|
||||
- Adds tracking headers (`x-request-id`, `x-wasm-processed`, `x-processed-at`, `x-api-route`)
|
||||
- Converts `500` errors to `503` for better client handling
|
||||
- Attach points: `OnRequest` and `OnResponse`
|
||||
|
||||
**Use case:** Request tracing and error status code conversion.
|
||||
|
||||
### [wasm-guest-ratelimit](./wasm-guest-ratelimit/)
|
||||
|
||||
Rate limiting middleware with configurable limits.
|
||||
|
||||
**Features:**
|
||||
- Rate limiting per identifier (API Key, IP, or Request ID)
|
||||
- Default: 60 requests per minute
|
||||
- Returns `429 Too Many Requests` when limit exceeded
|
||||
- Attach point: `OnRequest` only
|
||||
|
||||
**Note:** This is a simplified demonstration with per-instance state. For production, use router-level rate limiting with shared state.
|
||||
|
||||
**Use case:** Protect against request flooding and abuse.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Each example includes its own README with detailed build and deployment instructions. See individual example directories for:
|
||||
|
||||
- Build instructions
|
||||
- Deployment configuration
|
||||
- Customization options
|
||||
- Testing examples
|
||||
|
||||
## Common Prerequisites
|
||||
|
||||
All examples require:
|
||||
|
||||
- Rust toolchain (latest stable)
|
||||
- `wasm32-wasip2` target: `rustup target add wasm32-wasip2`
|
||||
- `wasm-tools`: `cargo install wasm-tools`
|
||||
- sgl-router running with WASM enabled (`--enable-wasm`)
|
||||
|
||||
## Building All Examples
|
||||
|
||||
```bash
|
||||
cd examples/wasm
|
||||
for example in wasm-guest-auth wasm-guest-logging wasm-guest-ratelimit; do
|
||||
echo "Building $example..."
|
||||
cd $example && ./build.sh && cd ..
|
||||
done
|
||||
```
|
||||
|
||||
## Deploying Multiple Modules
|
||||
|
||||
You can deploy all three modules together:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:3000/wasm \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"modules": [
|
||||
{
|
||||
"name": "auth-middleware",
|
||||
"file_path": "/path/to/wasm_guest_auth.component.wasm",
|
||||
"module_type": "Middleware",
|
||||
"attach_points": [{"Middleware": "OnRequest"}]
|
||||
},
|
||||
{
|
||||
"name": "logging-middleware",
|
||||
"file_path": "/path/to/wasm_guest_logging.component.wasm",
|
||||
"module_type": "Middleware",
|
||||
"attach_points": [{"Middleware": "OnRequest"}, {"Middleware": "OnResponse"}]
|
||||
},
|
||||
{
|
||||
"name": "ratelimit-middleware",
|
||||
"file_path": "/path/to/wasm_guest_ratelimit.component.wasm",
|
||||
"module_type": "Middleware",
|
||||
"attach_points": [{"Middleware": "OnRequest"}]
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Modules execute in the order they are deployed. If a module returns `Reject`, subsequent modules won't execute.
|
||||
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "wasm-guest-auth"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
wit-bindgen = { version = "0.21", features = ["macros"] }
|
||||
@@ -0,0 +1,62 @@
|
||||
# WASM Auth Example for sgl-router
|
||||
|
||||
This example demonstrates API key authentication middleware for sgl-router using the WebAssembly Component Model.
|
||||
|
||||
## Overview
|
||||
|
||||
This middleware validates API keys for requests to `/api` and `/v1` paths:
|
||||
|
||||
- Supports `Authorization: Bearer <key>` header
|
||||
- Supports `Authorization: ApiKey <key>` header
|
||||
- Supports `x-api-key` header
|
||||
- Returns `401 Unauthorized` for missing or invalid keys
|
||||
|
||||
**Default API Key**: `secret-api-key-12345`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Build and Deploy
|
||||
|
||||
```bash
|
||||
# Build
|
||||
cd examples/wasm-guest-auth
|
||||
./build.sh
|
||||
|
||||
# Deploy (replace file_path with actual path)
|
||||
curl -X POST http://localhost:3000/wasm \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"modules": [{
|
||||
"name": "auth-middleware",
|
||||
"file_path": "/absolute/path/to/wasm_guest_auth.component.wasm",
|
||||
"module_type": "Middleware",
|
||||
"attach_points": [{"Middleware": "OnRequest"}]
|
||||
}]
|
||||
}'
|
||||
```
|
||||
|
||||
### Customization
|
||||
|
||||
Modify `EXPECTED_API_KEY` in `src/lib.rs`:
|
||||
|
||||
```rust
|
||||
const EXPECTED_API_KEY: &str = "your-secret-key";
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Test unauthorized (returns 401)
|
||||
curl -v http://localhost:3000/api/test
|
||||
|
||||
# Test authorized (passes)
|
||||
curl -v http://localhost:3000/api/test \
|
||||
-H "Authorization: Bearer secret-api-key-12345"
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- Verify API key matches `EXPECTED_API_KEY` in code
|
||||
- Check request header format and path (`/api` or `/v1`)
|
||||
- Verify module is attached to `OnRequest` phase
|
||||
- Check router logs for errors
|
||||
@@ -0,0 +1,77 @@
|
||||
#!/bin/bash
|
||||
# Build script for WASM guest auth example
|
||||
# This script simplifies the build process for the WASM middleware component
|
||||
|
||||
set -e
|
||||
|
||||
echo "Building WASM guest auth example..."
|
||||
|
||||
# Check if we're in the right directory
|
||||
if [ ! -f "Cargo.toml" ]; then
|
||||
echo "Error: Cargo.toml not found. Please run this script from the wasm-guest-auth directory."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check for required tools
|
||||
command -v cargo >/dev/null 2>&1 || { echo "Error: cargo is required but not installed. Aborting." >&2; exit 1; }
|
||||
|
||||
# Check and install wasm32-wasip2 target
|
||||
echo "Checking for wasm32-wasip2 target..."
|
||||
if ! rustup target list --installed | grep -q "wasm32-wasip2"; then
|
||||
echo "wasm32-wasip2 target not found. Installing..."
|
||||
rustup target add wasm32-wasip2
|
||||
echo "✓ wasm32-wasip2 target installed"
|
||||
else
|
||||
echo "✓ wasm32-wasip2 target already installed"
|
||||
fi
|
||||
|
||||
# Check for wasm-tools
|
||||
if ! command -v wasm-tools >/dev/null 2>&1; then
|
||||
echo "Error: wasm-tools is required but not installed."
|
||||
echo "Install it with: cargo install wasm-tools"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Build with cargo (wit-bindgen uses cargo, not wasm-pack)
|
||||
echo "Running cargo build..."
|
||||
cargo build --target wasm32-wasip2 --release
|
||||
|
||||
# Output locations
|
||||
WASM_MODULE="target/wasm32-wasip2/release/wasm_guest_auth.wasm"
|
||||
WASM_COMPONENT="target/wasm32-wasip2/release/wasm_guest_auth.component.wasm"
|
||||
|
||||
if [ ! -f "$WASM_MODULE" ]; then
|
||||
echo "Error: Build failed - WASM module not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if the file is already a component
|
||||
echo "Checking WASM file format..."
|
||||
if wasm-tools print "$WASM_MODULE" 2>/dev/null | grep -q "^(\s*component"; then
|
||||
echo "✓ WASM file is already in component format"
|
||||
# Copy to component path for consistency
|
||||
cp "$WASM_MODULE" "$WASM_COMPONENT"
|
||||
else
|
||||
# Wrap the WASM module into a component format
|
||||
echo "Wrapping WASM module into component format..."
|
||||
wasm-tools component new "$WASM_MODULE" -o "$WASM_COMPONENT"
|
||||
if [ ! -f "$WASM_COMPONENT" ]; then
|
||||
echo "Error: Failed to create component file"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -f "$WASM_COMPONENT" ]; then
|
||||
echo ""
|
||||
echo "✓ Build successful!"
|
||||
echo " WASM module: $WASM_MODULE"
|
||||
echo " WASM component: $WASM_COMPONENT"
|
||||
echo ""
|
||||
echo "Next steps:"
|
||||
echo "1. Use the component file ($WASM_COMPONENT) when adding the module"
|
||||
echo "2. Prepare the module configuration (see README.md for JSON format)"
|
||||
echo "3. Use the API endpoint to add the module (see README.md for details)"
|
||||
else
|
||||
echo "Error: Component file not found"
|
||||
exit 1
|
||||
fi
|
||||
@@ -0,0 +1,70 @@
|
||||
//! WASM Guest Auth Example for sgl-router
|
||||
//!
|
||||
//! This example demonstrates API key authentication middleware
|
||||
//! for sgl-router using the WebAssembly Component Model.
|
||||
//!
|
||||
//! Features:
|
||||
//! - API Key authentication
|
||||
|
||||
wit_bindgen::generate!({
|
||||
path: "../../../src/wasm/interface",
|
||||
world: "sgl-router",
|
||||
});
|
||||
|
||||
use exports::sgl::router::{
|
||||
middleware_on_request::Guest as OnRequestGuest,
|
||||
middleware_on_response::Guest as OnResponseGuest,
|
||||
};
|
||||
use sgl::router::middleware_types::{Action, Request, Response};
|
||||
|
||||
/// Expected API Key (in production, this should be passed as configuration)
|
||||
const EXPECTED_API_KEY: &str = "secret-api-key-12345";
|
||||
|
||||
/// Main middleware implementation
|
||||
struct Middleware;
|
||||
|
||||
// Helper function to find header value
|
||||
fn find_header_value(
|
||||
headers: &[sgl::router::middleware_types::Header],
|
||||
name: &str,
|
||||
) -> Option<String> {
|
||||
headers
|
||||
.iter()
|
||||
.find(|h| h.name.eq_ignore_ascii_case(name))
|
||||
.map(|h| h.value.clone())
|
||||
}
|
||||
|
||||
// Implement on-request interface
|
||||
impl OnRequestGuest for Middleware {
|
||||
fn on_request(req: Request) -> Action {
|
||||
// API Key Authentication
|
||||
// Check for API key in Authorization header for /api routes
|
||||
if req.path.starts_with("/api") || req.path.starts_with("/v1") {
|
||||
let api_key = find_header_value(&req.headers, "authorization")
|
||||
.and_then(|h| {
|
||||
h.strip_prefix("Bearer ")
|
||||
.or_else(|| h.strip_prefix("ApiKey "))
|
||||
.map(|s| s.to_string())
|
||||
})
|
||||
.or_else(|| find_header_value(&req.headers, "x-api-key"));
|
||||
|
||||
// Reject if API key is missing or invalid
|
||||
if api_key.as_deref() != Some(EXPECTED_API_KEY) {
|
||||
return Action::Reject(401);
|
||||
}
|
||||
}
|
||||
|
||||
// Authentication passed, continue processing
|
||||
Action::Continue
|
||||
}
|
||||
}
|
||||
|
||||
// Implement on-response interface (empty - not used for auth)
|
||||
impl OnResponseGuest for Middleware {
|
||||
fn on_response(_resp: Response) -> Action {
|
||||
Action::Continue
|
||||
}
|
||||
}
|
||||
|
||||
// Export the component
|
||||
export!(Middleware);
|
||||
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "wasm-guest-logging"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
wit-bindgen = { version = "0.21", features = ["macros"] }
|
||||
@@ -0,0 +1,53 @@
|
||||
# WASM Logging Example for sgl-router
|
||||
|
||||
This example demonstrates logging and tracing middleware for sgl-router using the WebAssembly Component Model.
|
||||
|
||||
## Overview
|
||||
|
||||
This middleware provides:
|
||||
|
||||
- **Request Tracking** - Adds tracking headers (`x-request-id`, `x-wasm-processed`, `x-processed-at`, `x-api-route`)
|
||||
- **Status Code Conversion** - Converts `500` errors to `503`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Build and Deploy
|
||||
|
||||
```bash
|
||||
# Build
|
||||
cd examples/wasm-guest-logging
|
||||
./build.sh
|
||||
|
||||
# Deploy (replace file_path with actual path)
|
||||
curl -X POST http://localhost:3000/wasm \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"modules": [{
|
||||
"name": "logging-middleware",
|
||||
"file_path": "/absolute/path/to/wasm_guest_logging.component.wasm",
|
||||
"module_type": "Middleware",
|
||||
"attach_points": [{"Middleware": "OnRequest"}, {"Middleware": "OnResponse"}]
|
||||
}]
|
||||
}'
|
||||
```
|
||||
|
||||
### Customization
|
||||
|
||||
Modify `on_request` or `on_response` functions in `src/lib.rs` to add custom tracking headers or status code conversions.
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Check tracking headers
|
||||
curl -v http://localhost:3000/v1/models 2>&1 | \
|
||||
grep -E "(x-request-id|x-wasm-processed|x-processed-at)"
|
||||
|
||||
# Test status code conversion (requires endpoint returning 500)
|
||||
curl -v http://localhost:3000/some-endpoint 2>&1 | grep -E "(< HTTP|500|503)"
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- Verify module attached to both `OnRequest` and `OnResponse` phases
|
||||
- Check router logs for execution errors
|
||||
- Ensure module built successfully
|
||||
@@ -0,0 +1,77 @@
|
||||
#!/bin/bash
|
||||
# Build script for WASM guest logging example
|
||||
# This script simplifies the build process for the WASM middleware component
|
||||
|
||||
set -e
|
||||
|
||||
echo "Building WASM guest logging example..."
|
||||
|
||||
# Check if we're in the right directory
|
||||
if [ ! -f "Cargo.toml" ]; then
|
||||
echo "Error: Cargo.toml not found. Please run this script from the wasm-guest-logging directory."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check for required tools
|
||||
command -v cargo >/dev/null 2>&1 || { echo "Error: cargo is required but not installed. Aborting." >&2; exit 1; }
|
||||
|
||||
# Check and install wasm32-wasip2 target
|
||||
echo "Checking for wasm32-wasip2 target..."
|
||||
if ! rustup target list --installed | grep -q "wasm32-wasip2"; then
|
||||
echo "wasm32-wasip2 target not found. Installing..."
|
||||
rustup target add wasm32-wasip2
|
||||
echo "✓ wasm32-wasip2 target installed"
|
||||
else
|
||||
echo "✓ wasm32-wasip2 target already installed"
|
||||
fi
|
||||
|
||||
# Check for wasm-tools
|
||||
if ! command -v wasm-tools >/dev/null 2>&1; then
|
||||
echo "Error: wasm-tools is required but not installed."
|
||||
echo "Install it with: cargo install wasm-tools"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Build with cargo (wit-bindgen uses cargo, not wasm-pack)
|
||||
echo "Running cargo build..."
|
||||
cargo build --target wasm32-wasip2 --release
|
||||
|
||||
# Output locations
|
||||
WASM_MODULE="target/wasm32-wasip2/release/wasm_guest_logging.wasm"
|
||||
WASM_COMPONENT="target/wasm32-wasip2/release/wasm_guest_logging.component.wasm"
|
||||
|
||||
if [ ! -f "$WASM_MODULE" ]; then
|
||||
echo "Error: Build failed - WASM module not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if the file is already a component
|
||||
echo "Checking WASM file format..."
|
||||
if wasm-tools print "$WASM_MODULE" 2>/dev/null | grep -q "^(\s*component"; then
|
||||
echo "✓ WASM file is already in component format"
|
||||
# Copy to component path for consistency
|
||||
cp "$WASM_MODULE" "$WASM_COMPONENT"
|
||||
else
|
||||
# Wrap the WASM module into a component format
|
||||
echo "Wrapping WASM module into component format..."
|
||||
wasm-tools component new "$WASM_MODULE" -o "$WASM_COMPONENT"
|
||||
if [ ! -f "$WASM_COMPONENT" ]; then
|
||||
echo "Error: Failed to create component file"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -f "$WASM_COMPONENT" ]; then
|
||||
echo ""
|
||||
echo "✓ Build successful!"
|
||||
echo " WASM module: $WASM_MODULE"
|
||||
echo " WASM component: $WASM_COMPONENT"
|
||||
echo ""
|
||||
echo "Next steps:"
|
||||
echo "1. Use the component file ($WASM_COMPONENT) when adding the module"
|
||||
echo "2. Prepare the module configuration (see README.md for JSON format)"
|
||||
echo "3. Use the API endpoint to add the module (see README.md for details)"
|
||||
else
|
||||
echo "Error: Component file not found"
|
||||
exit 1
|
||||
fi
|
||||
@@ -0,0 +1,88 @@
|
||||
//! WASM Guest Logging Example for sgl-router
|
||||
//!
|
||||
//! This example demonstrates logging and tracing middleware
|
||||
//! for sgl-router using the WebAssembly Component Model.
|
||||
//!
|
||||
//! Features:
|
||||
//! - Request tracking and tracing headers
|
||||
//! - Response status code conversion
|
||||
|
||||
wit_bindgen::generate!({
|
||||
path: "../../../src/wasm/interface",
|
||||
world: "sgl-router",
|
||||
});
|
||||
|
||||
use exports::sgl::router::{
|
||||
middleware_on_request::Guest as OnRequestGuest,
|
||||
middleware_on_response::Guest as OnResponseGuest,
|
||||
};
|
||||
use sgl::router::middleware_types::{Action, Header, ModifyAction, Request, Response};
|
||||
|
||||
/// Main middleware implementation
|
||||
struct Middleware;
|
||||
|
||||
// Helper function to create header
|
||||
fn create_header(name: &str, value: &str) -> Header {
|
||||
Header {
|
||||
name: name.to_string(),
|
||||
value: value.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
// Implement on-request interface
|
||||
impl OnRequestGuest for Middleware {
|
||||
fn on_request(req: Request) -> Action {
|
||||
let mut modify_action = ModifyAction {
|
||||
status: None,
|
||||
headers_set: vec![],
|
||||
headers_add: vec![],
|
||||
headers_remove: vec![],
|
||||
body_replace: None,
|
||||
};
|
||||
|
||||
// Request Logging and Tracing
|
||||
// Add tracing headers with request ID
|
||||
modify_action
|
||||
.headers_add
|
||||
.push(create_header("x-request-id", &req.request_id));
|
||||
modify_action
|
||||
.headers_add
|
||||
.push(create_header("x-wasm-processed", "true"));
|
||||
modify_action.headers_add.push(create_header(
|
||||
"x-processed-at",
|
||||
&req.now_epoch_ms.to_string(),
|
||||
));
|
||||
|
||||
// Add custom header for API requests
|
||||
if req.path.starts_with("/api") || req.path.starts_with("/v1") {
|
||||
modify_action
|
||||
.headers_add
|
||||
.push(create_header("x-api-route", "true"));
|
||||
}
|
||||
|
||||
Action::Modify(modify_action)
|
||||
}
|
||||
}
|
||||
|
||||
// Implement on-response interface
|
||||
impl OnResponseGuest for Middleware {
|
||||
fn on_response(resp: Response) -> Action {
|
||||
// Status code conversion: Convert 500 to 503 for better client handling
|
||||
if resp.status == 500 {
|
||||
let modify_action = ModifyAction {
|
||||
status: Some(503),
|
||||
headers_set: vec![],
|
||||
headers_add: vec![],
|
||||
headers_remove: vec![],
|
||||
body_replace: None,
|
||||
};
|
||||
Action::Modify(modify_action)
|
||||
} else {
|
||||
// No modification needed
|
||||
Action::Continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Export the component
|
||||
export!(Middleware);
|
||||
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "wasm-guest-ratelimit"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
wit-bindgen = { version = "0.21", features = ["macros"] }
|
||||
@@ -0,0 +1,68 @@
|
||||
# WASM Rate Limit Example for sgl-router
|
||||
|
||||
This example demonstrates rate limiting middleware for sgl-router using the WebAssembly Component Model.
|
||||
|
||||
## Overview
|
||||
|
||||
This middleware provides rate limiting:
|
||||
|
||||
- **Default**: 60 requests per minute per identifier
|
||||
- **Identifier Priority**: API Key > IP Address > Request ID
|
||||
- **Response**: Returns `429 Too Many Requests` when limit exceeded
|
||||
|
||||
**Important**: This is a simplified demonstration. Since WASM components are stateless, each worker thread maintains its own counter. For production, implement rate limiting at the router/host level with shared state.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Build and Deploy
|
||||
|
||||
```bash
|
||||
# Build
|
||||
cd examples/wasm-guest-ratelimit
|
||||
./build.sh
|
||||
|
||||
# Deploy (replace file_path with actual path)
|
||||
curl -X POST http://localhost:3000/wasm \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"modules": [{
|
||||
"name": "ratelimit-middleware",
|
||||
"file_path": "/absolute/path/to/wasm_guest_ratelimit.component.wasm",
|
||||
"module_type": "Middleware",
|
||||
"attach_points": [{"Middleware": "OnRequest"}]
|
||||
}]
|
||||
}'
|
||||
```
|
||||
|
||||
### Customization
|
||||
|
||||
Modify constants in `src/lib.rs`:
|
||||
|
||||
```rust
|
||||
const RATE_LIMIT_REQUESTS: u64 = 100; // requests per window
|
||||
const RATE_LIMIT_WINDOW_MS: u64 = 60_000; // time window in ms
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Send multiple requests (first 60 succeed, then 429)
|
||||
for i in {1..65}; do
|
||||
curl -s -o /dev/null -w "%{http_code}\n" \
|
||||
http://localhost:3000/v1/models \
|
||||
-H "Authorization: Bearer secret-api-key-12345"
|
||||
done
|
||||
```
|
||||
|
||||
## Limitations
|
||||
|
||||
- Per-instance state (not shared across workers)
|
||||
- No cross-process state sharing
|
||||
- Memory growth with unique identifiers
|
||||
- State lost on instance restart
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- Verify module attached to `OnRequest` phase
|
||||
- Check identifier extraction logic matches request format
|
||||
- Note: Each WASM worker has separate counter
|
||||
@@ -0,0 +1,77 @@
|
||||
#!/bin/bash
|
||||
# Build script for WASM guest rate limit example
|
||||
# This script simplifies the build process for the WASM middleware component
|
||||
|
||||
set -e
|
||||
|
||||
echo "Building WASM guest rate limit example..."
|
||||
|
||||
# Check if we're in the right directory
|
||||
if [ ! -f "Cargo.toml" ]; then
|
||||
echo "Error: Cargo.toml not found. Please run this script from the wasm-guest-ratelimit directory."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check for required tools
|
||||
command -v cargo >/dev/null 2>&1 || { echo "Error: cargo is required but not installed. Aborting." >&2; exit 1; }
|
||||
|
||||
# Check and install wasm32-wasip2 target
|
||||
echo "Checking for wasm32-wasip2 target..."
|
||||
if ! rustup target list --installed | grep -q "wasm32-wasip2"; then
|
||||
echo "wasm32-wasip2 target not found. Installing..."
|
||||
rustup target add wasm32-wasip2
|
||||
echo "✓ wasm32-wasip2 target installed"
|
||||
else
|
||||
echo "✓ wasm32-wasip2 target already installed"
|
||||
fi
|
||||
|
||||
# Check for wasm-tools
|
||||
if ! command -v wasm-tools >/dev/null 2>&1; then
|
||||
echo "Error: wasm-tools is required but not installed."
|
||||
echo "Install it with: cargo install wasm-tools"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Build with cargo (wit-bindgen uses cargo, not wasm-pack)
|
||||
echo "Running cargo build..."
|
||||
cargo build --target wasm32-wasip2 --release
|
||||
|
||||
# Output locations
|
||||
WASM_MODULE="target/wasm32-wasip2/release/wasm_guest_ratelimit.wasm"
|
||||
WASM_COMPONENT="target/wasm32-wasip2/release/wasm_guest_ratelimit.component.wasm"
|
||||
|
||||
if [ ! -f "$WASM_MODULE" ]; then
|
||||
echo "Error: Build failed - WASM module not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if the file is already a component
|
||||
echo "Checking WASM file format..."
|
||||
if wasm-tools print "$WASM_MODULE" 2>/dev/null | grep -q "^(\s*component"; then
|
||||
echo "✓ WASM file is already in component format"
|
||||
# Copy to component path for consistency
|
||||
cp "$WASM_MODULE" "$WASM_COMPONENT"
|
||||
else
|
||||
# Wrap the WASM module into a component format
|
||||
echo "Wrapping WASM module into component format..."
|
||||
wasm-tools component new "$WASM_MODULE" -o "$WASM_COMPONENT"
|
||||
if [ ! -f "$WASM_COMPONENT" ]; then
|
||||
echo "Error: Failed to create component file"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -f "$WASM_COMPONENT" ]; then
|
||||
echo ""
|
||||
echo "✓ Build successful!"
|
||||
echo " WASM module: $WASM_MODULE"
|
||||
echo " WASM component: $WASM_COMPONENT"
|
||||
echo ""
|
||||
echo "Next steps:"
|
||||
echo "1. Use the component file ($WASM_COMPONENT) when adding the module"
|
||||
echo "2. Prepare the module configuration (see README.md for JSON format)"
|
||||
echo "3. Use the API endpoint to add the module (see README.md for details)"
|
||||
else
|
||||
echo "Error: Component file not found"
|
||||
exit 1
|
||||
fi
|
||||
@@ -0,0 +1,155 @@
|
||||
//! WASM Guest Rate Limit Example for sgl-router
|
||||
//!
|
||||
//! This example demonstrates rate limiting middleware
|
||||
//! for sgl-router using the WebAssembly Component Model.
|
||||
//!
|
||||
//! Features:
|
||||
//! - Rate limiting based on API Key or IP address
|
||||
//! - Fixed time window (e.g., 60 requests per minute)
|
||||
//! - Returns 429 Too Many Requests when limit exceeded
|
||||
//!
|
||||
//! Note: This is a simplified implementation. Since WASM components are stateless,
|
||||
//! each instance maintains its own counters. For production use, consider
|
||||
//! implementing rate limiting at the host/router level with shared state.
|
||||
|
||||
wit_bindgen::generate!({
|
||||
path: "../../../src/wasm/interface",
|
||||
world: "sgl-router",
|
||||
});
|
||||
|
||||
use std::cell::RefCell;
|
||||
|
||||
use exports::sgl::router::{
|
||||
middleware_on_request::Guest as OnRequestGuest,
|
||||
middleware_on_response::Guest as OnResponseGuest,
|
||||
};
|
||||
use sgl::router::middleware_types::{Action, Request, Response};
|
||||
|
||||
/// Main middleware implementation
|
||||
struct Middleware;
|
||||
|
||||
// Rate limit configuration
|
||||
const RATE_LIMIT_REQUESTS: u64 = 60; // Maximum requests per window
|
||||
const RATE_LIMIT_WINDOW_MS: u64 = 60_000; // Time window in milliseconds (1 minute)
|
||||
|
||||
// Simple in-memory counter (per WASM instance)
|
||||
// In a real implementation, this would be shared across all instances
|
||||
// This is a simplified example for demonstration purposes
|
||||
struct RateLimitState {
|
||||
requests: Vec<(String, u64)>, // (identifier, timestamp_ms)
|
||||
}
|
||||
|
||||
impl RateLimitState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
requests: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up old entries outside the time window
|
||||
fn cleanup(&mut self, current_time_ms: u64) {
|
||||
let cutoff = current_time_ms.saturating_sub(RATE_LIMIT_WINDOW_MS);
|
||||
self.requests.retain(|(_, timestamp)| *timestamp > cutoff);
|
||||
}
|
||||
|
||||
// Check if identifier has exceeded rate limit
|
||||
fn check_limit(&mut self, identifier: &str, current_time_ms: u64) -> bool {
|
||||
self.cleanup(current_time_ms);
|
||||
|
||||
// Count requests in current window for this identifier
|
||||
let count = self
|
||||
.requests
|
||||
.iter()
|
||||
.filter(|(id, timestamp)| {
|
||||
id == identifier
|
||||
&& *timestamp > current_time_ms.saturating_sub(RATE_LIMIT_WINDOW_MS)
|
||||
})
|
||||
.count() as u64;
|
||||
|
||||
if count >= RATE_LIMIT_REQUESTS {
|
||||
return false; // Limit exceeded
|
||||
}
|
||||
|
||||
// Add new request
|
||||
self.requests
|
||||
.push((identifier.to_string(), current_time_ms));
|
||||
true // Within limit
|
||||
}
|
||||
}
|
||||
|
||||
// Thread-local state (per WASM instance thread)
|
||||
// Using thread_local! is safer than static mut as it avoids unsafe blocks
|
||||
// and provides separate state for each thread automatically
|
||||
thread_local! {
|
||||
static RATE_LIMIT_STATE: RefCell<RateLimitState> = RefCell::new(RateLimitState::new());
|
||||
}
|
||||
|
||||
fn get_identifier(req: &Request) -> String {
|
||||
// Helper function to find header value
|
||||
let find_header_value =
|
||||
|headers: &[sgl::router::middleware_types::Header], name: &str| -> Option<String> {
|
||||
headers
|
||||
.iter()
|
||||
.find(|h| h.name.eq_ignore_ascii_case(name))
|
||||
.map(|h| h.value.clone())
|
||||
};
|
||||
|
||||
// Prefer API Key as identifier (more stable than IP)
|
||||
if let Some(auth_header) = find_header_value(&req.headers, "authorization") {
|
||||
if auth_header.starts_with("Bearer ") {
|
||||
return format!("api_key:{}", &auth_header[7..]);
|
||||
} else if auth_header.starts_with("ApiKey ") {
|
||||
return format!("api_key:{}", &auth_header[7..]);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(api_key) = find_header_value(&req.headers, "x-api-key") {
|
||||
return format!("api_key:{}", api_key);
|
||||
}
|
||||
|
||||
// Fall back to IP address from forwarded headers
|
||||
if let Some(forwarded_for) = find_header_value(&req.headers, "x-forwarded-for") {
|
||||
// Take first IP from comma-separated list
|
||||
let ip = forwarded_for.split(',').next().unwrap_or("").trim();
|
||||
if !ip.is_empty() {
|
||||
return format!("ip:{}", ip);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(real_ip) = find_header_value(&req.headers, "x-real-ip") {
|
||||
return format!("ip:{}", real_ip);
|
||||
}
|
||||
|
||||
// Last resort: use request ID (not ideal, but better than nothing)
|
||||
format!("req_id:{}", req.request_id)
|
||||
}
|
||||
|
||||
// Implement on-request interface
|
||||
impl OnRequestGuest for Middleware {
|
||||
fn on_request(req: Request) -> Action {
|
||||
let identifier = get_identifier(&req);
|
||||
let current_time_ms = req.now_epoch_ms;
|
||||
|
||||
// Access thread-local state safely without unsafe blocks
|
||||
// Each thread gets its own RateLimitState instance
|
||||
RATE_LIMIT_STATE.with(|state| {
|
||||
let mut state = state.borrow_mut();
|
||||
if !state.check_limit(&identifier, current_time_ms) {
|
||||
// Rate limit exceeded
|
||||
return Action::Reject(429);
|
||||
}
|
||||
// Within rate limit, continue processing
|
||||
Action::Continue
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Implement on-response interface (empty - not used for rate limiting)
|
||||
impl OnResponseGuest for Middleware {
|
||||
fn on_response(_resp: Response) -> Action {
|
||||
Action::Continue
|
||||
}
|
||||
}
|
||||
|
||||
// Export the component
|
||||
export!(Middleware);
|
||||
@@ -0,0 +1 @@
|
||||
"""Test package root for router Python tests."""
|
||||
@@ -0,0 +1,15 @@
|
||||
import sys
|
||||
from importlib.util import find_spec
|
||||
from pathlib import Path
|
||||
|
||||
# Only add bindings/python to path if the wheel is not installed (for local development)
|
||||
# This ensures CI tests use the installed wheel which contains the Rust extension
|
||||
_ROOT = Path(__file__).resolve().parents[1]
|
||||
_SRC = _ROOT / "bindings" / "python"
|
||||
|
||||
# Check if sglang_router is already installed with the Rust extension
|
||||
_wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
|
||||
|
||||
# Only add bindings/python if wheel is not installed (development mode)
|
||||
if not _wheel_installed and str(_SRC) not in sys.path:
|
||||
sys.path.insert(0, str(_SRC))
|
||||
@@ -0,0 +1,343 @@
|
||||
"""
|
||||
gRPC Router E2E Test - OpenAI Server API Compatibility
|
||||
|
||||
This test file is REUSED from test/srt/openai_server/basic/test_openai_server.py
|
||||
with minimal changes:
|
||||
- Swap popen_launch_server() → popen_launch_workers_and_router()
|
||||
- Update teardown to cleanup router + workers
|
||||
- All test logic and assertions remain identical
|
||||
|
||||
Run with:
|
||||
python3 -m pytest e2e_grpc/basic/test_openai_server.py -v
|
||||
python3 -m unittest e2e_grpc.basic.test_openai_server.TestOpenAIServer.test_completion
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_GPT_OSS_MODEL_PATH,
|
||||
DEFAULT_MODEL_PATH,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
get_tokenizer,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestOpenAIServer(CustomTestCase):
|
||||
"""
|
||||
Test OpenAI API through gRPC router.
|
||||
|
||||
REUSED from test/srt/openai_server/basic/test_openai_server.py
|
||||
ONLY CHANGE: Server launch mechanism
|
||||
- Launches SGLang workers with --enable-grpc
|
||||
- Launches gRPC router pointing to those workers
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
policy="round_robin",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
# ALL TEST METHODS BELOW ARE UNCHANGED FROM ORIGINAL
|
||||
# They validate that the router maintains OpenAI API compatibility
|
||||
def run_chat_completion(self, logprobs, parallel_sample_num):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the capital of France? Answer in a few words.",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
logprobs=logprobs is not None and logprobs > 0,
|
||||
top_logprobs=logprobs,
|
||||
n=parallel_sample_num,
|
||||
)
|
||||
|
||||
if logprobs:
|
||||
assert isinstance(
|
||||
response.choices[0].logprobs.content[0].top_logprobs[0].token, str
|
||||
)
|
||||
|
||||
ret_num_top_logprobs = len(
|
||||
response.choices[0].logprobs.content[0].top_logprobs
|
||||
)
|
||||
assert (
|
||||
ret_num_top_logprobs == logprobs
|
||||
), f"{ret_num_top_logprobs} vs {logprobs}"
|
||||
|
||||
assert len(response.choices) == parallel_sample_num
|
||||
assert response.choices[0].message.role == "assistant"
|
||||
assert isinstance(response.choices[0].message.content, str)
|
||||
assert response.id
|
||||
assert response.created
|
||||
assert response.usage.prompt_tokens > 0
|
||||
assert response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
def run_chat_completion_stream(self, logprobs, parallel_sample_num=1):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
generator = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
],
|
||||
temperature=0,
|
||||
logprobs=logprobs is not None and logprobs > 0,
|
||||
top_logprobs=logprobs,
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
n=parallel_sample_num,
|
||||
)
|
||||
|
||||
is_firsts = {}
|
||||
is_finished = {}
|
||||
finish_reason_counts = {}
|
||||
for response in generator:
|
||||
usage = response.usage
|
||||
if usage is not None:
|
||||
assert usage.prompt_tokens > 0, f"usage.prompt_tokens was zero"
|
||||
assert usage.completion_tokens > 0, f"usage.completion_tokens was zero"
|
||||
assert usage.total_tokens > 0, f"usage.total_tokens was zero"
|
||||
continue
|
||||
|
||||
index = response.choices[0].index
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
if finish_reason is not None:
|
||||
is_finished[index] = True
|
||||
finish_reason_counts[index] = finish_reason_counts.get(index, 0) + 1
|
||||
|
||||
data = response.choices[0].delta
|
||||
|
||||
if is_firsts.get(index, True):
|
||||
assert (
|
||||
data.role == "assistant"
|
||||
), f"data.role was not 'assistant' for first chunk"
|
||||
is_firsts[index] = False
|
||||
continue
|
||||
|
||||
if logprobs and not is_finished.get(index, False):
|
||||
assert response.choices[0].logprobs, f"logprobs was not returned"
|
||||
assert isinstance(
|
||||
response.choices[0].logprobs.content[0].top_logprobs[0].token, str
|
||||
), f"top_logprobs token was not a string"
|
||||
assert isinstance(
|
||||
response.choices[0].logprobs.content[0].top_logprobs, list
|
||||
), f"top_logprobs was not a list"
|
||||
ret_num_top_logprobs = len(
|
||||
response.choices[0].logprobs.content[0].top_logprobs
|
||||
)
|
||||
assert (
|
||||
ret_num_top_logprobs == logprobs
|
||||
), f"{ret_num_top_logprobs} vs {logprobs}"
|
||||
|
||||
assert (
|
||||
isinstance(data.content, str)
|
||||
or isinstance(data.reasoning_content, str)
|
||||
or (isinstance(data.tool_calls, list) and len(data.tool_calls) > 0)
|
||||
or response.choices[0].finish_reason
|
||||
)
|
||||
assert response.id
|
||||
assert response.created
|
||||
|
||||
for index in [i for i in range(parallel_sample_num)]:
|
||||
assert not is_firsts.get(
|
||||
index, True
|
||||
), f"index {index} is not found in the response"
|
||||
|
||||
for index in range(parallel_sample_num):
|
||||
assert (
|
||||
index in finish_reason_counts
|
||||
), f"No finish_reason found for index {index}"
|
||||
assert (
|
||||
finish_reason_counts[index] == 1
|
||||
), f"Expected 1 finish_reason chunk for index {index}, got {finish_reason_counts[index]}"
|
||||
|
||||
def test_chat_completion(self):
|
||||
for logprobs in [None, 5]:
|
||||
for parallel_sample_num in [1, 2]:
|
||||
self.run_chat_completion(logprobs, parallel_sample_num)
|
||||
|
||||
def test_chat_completion_stream(self):
|
||||
for logprobs in [None, 5]:
|
||||
for parallel_sample_num in [1, 2]:
|
||||
self.run_chat_completion_stream(logprobs, parallel_sample_num)
|
||||
|
||||
def test_regex(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
regex = (
|
||||
r"""\{\n"""
|
||||
+ r""" "name": "[\w]+",\n"""
|
||||
+ r""" "population": [\d]+\n"""
|
||||
+ r"""\}"""
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{"role": "user", "content": "Introduce the capital of France."},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=128,
|
||||
extra_body={"regex": regex},
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
|
||||
try:
|
||||
js_obj = json.loads(text)
|
||||
except (TypeError, json.decoder.JSONDecodeError):
|
||||
raise
|
||||
assert isinstance(js_obj["name"], str)
|
||||
assert isinstance(js_obj["population"], int)
|
||||
|
||||
def test_penalty(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{"role": "user", "content": "Introduce the capital of France."},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=32,
|
||||
frequency_penalty=1.0,
|
||||
)
|
||||
text = response.choices[0].message.content
|
||||
assert isinstance(text, str)
|
||||
|
||||
def test_response_prefill(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": """
|
||||
Extract the name, size, price, and color from this product description as a JSON object:
|
||||
|
||||
<description>
|
||||
The SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. At just 5 inches wide, it lets you control lights, thermostats, and other connected devices via voice or app—no matter where you place it in your home. This affordable little hub brings convenient hands-free control to your smart devices.
|
||||
</description>
|
||||
""",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "{\n",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
extra_body={"continue_final_message": True},
|
||||
)
|
||||
|
||||
assert (
|
||||
response.choices[0]
|
||||
.message.content.strip()
|
||||
.startswith('"name": "SmartHome Mini",')
|
||||
)
|
||||
|
||||
def test_model_list(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
models = list(client.models.list().data)
|
||||
assert len(models) == 1
|
||||
# assert isinstance(getattr(models[0], "max_model_len", None), int)
|
||||
|
||||
@unittest.skip("Skipping retrieve model test as it is not supported by the router")
|
||||
def test_retrieve_model(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
retrieved_model = client.models.retrieve(self.model)
|
||||
self.assertEqual(retrieved_model.id, self.model)
|
||||
self.assertEqual(retrieved_model.root, self.model)
|
||||
|
||||
with self.assertRaises(openai.NotFoundError):
|
||||
client.models.retrieve("non-existent-model")
|
||||
|
||||
|
||||
class TestOpenAIServerGptOss(TestOpenAIServer):
|
||||
"""
|
||||
Test OpenAI API through gRPC router with openai/gpt-oss-20b model.
|
||||
Extends TestOpenAIServer and only changes the model.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_GPT_OSS_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
policy="round_robin",
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
def test_chat_completion(self):
|
||||
for parallel_sample_num in [1, 2]:
|
||||
self.run_chat_completion(None, parallel_sample_num)
|
||||
|
||||
def test_chat_completion_stream(self):
|
||||
for parallel_sample_num in [1, 2]:
|
||||
self.run_chat_completion_stream(None, parallel_sample_num)
|
||||
|
||||
@unittest.skip("Skipping for OSS models")
|
||||
def test_regex(self):
|
||||
super().test_regex()
|
||||
|
||||
@unittest.skip("Skipping for OSS models")
|
||||
def test_response_prefill(self):
|
||||
super().test_response_prefill()
|
||||
|
||||
@unittest.skip("Skipping for OSS models")
|
||||
def test_penalty(self):
|
||||
super().test_penalty()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
Pytest configuration for gRPC router e2e tests.
|
||||
|
||||
This module provides shared fixtures that can be used across all gRPC router tests.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest # noqa: F401
|
||||
|
||||
# Ensure router bindings/python is importable
|
||||
_ROUTER_ROOT = Path(__file__).resolve().parents[2]
|
||||
_ROUTER_SRC = _ROUTER_ROOT / "bindings" / "python"
|
||||
if str(_ROUTER_SRC) not in sys.path:
|
||||
sys.path.insert(0, str(_ROUTER_SRC))
|
||||
|
||||
# Ensure e2e_grpc test utilities are importable
|
||||
_E2E_GRPC_DIR = Path(__file__).parent
|
||||
if str(_E2E_GRPC_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_E2E_GRPC_DIR))
|
||||
|
||||
|
||||
# Pytest markers for test organization
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line("markers", "e2e: end-to-end tests with real workers")
|
||||
config.addinivalue_line("markers", "grpc: gRPC-specific tests")
|
||||
config.addinivalue_line("markers", "slow: slow-running tests")
|
||||
config.addinivalue_line("markers", "pd: prefill-decode disaggregation tests")
|
||||
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_chat_completion_with_reasoning
|
||||
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_chat_completion_without_reasoning
|
||||
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_stream_chat_completion_with_reasoning
|
||||
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_stream_chat_completion_without_reasoning
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_ENABLE_THINKING_MODEL_PATH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestEnableThinking(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
cls.model = DEFAULT_ENABLE_THINKING_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-1234"
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=120,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
"--reasoning-parser",
|
||||
"qwen3",
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=4,
|
||||
)
|
||||
cls.additional_chat_kwargs = {}
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
def test_chat_completion_with_reasoning(self):
|
||||
# Test non-streaming with "enable_thinking": True, reasoning_content should not be empty
|
||||
client = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"temperature": 0,
|
||||
"separate_reasoning": True,
|
||||
"chat_template_kwargs": {"enable_thinking": True},
|
||||
**self.additional_chat_kwargs,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(client.status_code, 200, f"Failed with: {client.text}")
|
||||
data = client.json()
|
||||
|
||||
self.assertIn("choices", data)
|
||||
self.assertTrue(len(data["choices"]) > 0)
|
||||
self.assertIn("message", data["choices"][0])
|
||||
self.assertIn("reasoning_content", data["choices"][0]["message"])
|
||||
self.assertIsNotNone(data["choices"][0]["message"]["reasoning_content"])
|
||||
|
||||
def test_chat_completion_without_reasoning(self):
|
||||
# Test non-streaming with "enable_thinking": False, reasoning_content should be empty
|
||||
client = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"temperature": 0,
|
||||
"separate_reasoning": True,
|
||||
"chat_template_kwargs": {"enable_thinking": False},
|
||||
**self.additional_chat_kwargs,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(client.status_code, 200, f"Failed with: {client.text}")
|
||||
data = client.json()
|
||||
|
||||
self.assertIn("choices", data)
|
||||
self.assertTrue(len(data["choices"]) > 0)
|
||||
self.assertIn("message", data["choices"][0])
|
||||
|
||||
if "reasoning_content" in data["choices"][0]["message"]:
|
||||
self.assertIsNone(data["choices"][0]["message"]["reasoning_content"])
|
||||
|
||||
def test_stream_chat_completion_with_reasoning(self):
|
||||
# Test streaming with "enable_thinking": True, reasoning_content should not be empty
|
||||
response = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"temperature": 0,
|
||||
"separate_reasoning": True,
|
||||
"stream": True,
|
||||
"chat_template_kwargs": {"enable_thinking": True},
|
||||
**self.additional_chat_kwargs,
|
||||
},
|
||||
stream=True,
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200, f"Failed with: {response.text}")
|
||||
|
||||
has_reasoning = False
|
||||
has_content = False
|
||||
|
||||
for line in response.iter_lines():
|
||||
if line:
|
||||
line = line.decode("utf-8")
|
||||
if line.startswith("data:") and not line.startswith("data: [DONE]"):
|
||||
data = json.loads(line[6:])
|
||||
if "choices" in data and len(data["choices"]) > 0:
|
||||
delta = data["choices"][0].get("delta", {})
|
||||
|
||||
if "reasoning_content" in delta and delta["reasoning_content"]:
|
||||
has_reasoning = True
|
||||
|
||||
if "content" in delta and delta["content"]:
|
||||
has_content = True
|
||||
|
||||
self.assertTrue(
|
||||
has_reasoning,
|
||||
"The reasoning content is not included in the stream response",
|
||||
)
|
||||
self.assertTrue(
|
||||
has_content, "The stream response does not contain normal content"
|
||||
)
|
||||
|
||||
def test_stream_chat_completion_without_reasoning(self):
|
||||
# Test streaming with "enable_thinking": False, reasoning_content should be empty
|
||||
response = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"temperature": 0,
|
||||
"separate_reasoning": True,
|
||||
"stream": True,
|
||||
"chat_template_kwargs": {"enable_thinking": False},
|
||||
**self.additional_chat_kwargs,
|
||||
},
|
||||
stream=True,
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200, f"Failed with: {response.text}")
|
||||
|
||||
has_reasoning = False
|
||||
has_content = False
|
||||
|
||||
for line in response.iter_lines():
|
||||
if line:
|
||||
line = line.decode("utf-8")
|
||||
if line.startswith("data:") and not line.startswith("data: [DONE]"):
|
||||
data = json.loads(line[6:])
|
||||
if "choices" in data and len(data["choices"]) > 0:
|
||||
delta = data["choices"][0].get("delta", {})
|
||||
|
||||
if "reasoning_content" in delta and delta["reasoning_content"]:
|
||||
has_reasoning = True
|
||||
|
||||
if "content" in delta and delta["content"]:
|
||||
has_content = True
|
||||
|
||||
self.assertFalse(
|
||||
has_reasoning,
|
||||
"The reasoning content should not be included in the stream response",
|
||||
)
|
||||
self.assertTrue(
|
||||
has_content, "The stream response does not contain normal content"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_streaming_separate_reasoning_false
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_streaming_separate_reasoning_true
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_streaming_separate_reasoning_true_stream_reasoning_false
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_nonstreaming_separate_reasoning_false
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_nonstreaming_separate_reasoning_true
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentStartup.test_nonstreaming
|
||||
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentStartup.test_streaming
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_REASONING_MODEL_PATH,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestReasoningContentAPI(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
cls.model = DEFAULT_REASONING_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-1234"
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
"--reasoning-parser",
|
||||
"deepseek_r1",
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
def test_streaming_separate_reasoning_false(self):
|
||||
# Test streaming with separate_reasoning=False, reasoning_content should be empty
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 1+3?",
|
||||
}
|
||||
],
|
||||
"max_tokens": 100,
|
||||
"stream": True,
|
||||
"extra_body": {"separate_reasoning": False},
|
||||
}
|
||||
response = client.chat.completions.create(**payload)
|
||||
|
||||
reasoning_content = ""
|
||||
content = ""
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.content:
|
||||
content += chunk.choices[0].delta.content
|
||||
elif chunk.choices[0].delta.reasoning_content:
|
||||
reasoning_content += chunk.choices[0].delta.reasoning_content
|
||||
|
||||
assert len(reasoning_content) == 0
|
||||
assert len(content) > 0
|
||||
|
||||
def test_streaming_separate_reasoning_true(self):
|
||||
# Test streaming with separate_reasoning=True, reasoning_content should not be empty
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 1+3?",
|
||||
}
|
||||
],
|
||||
"max_tokens": 100,
|
||||
"stream": True,
|
||||
"extra_body": {"separate_reasoning": True},
|
||||
}
|
||||
response = client.chat.completions.create(**payload)
|
||||
|
||||
reasoning_content = ""
|
||||
content = ""
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.content:
|
||||
content += chunk.choices[0].delta.content
|
||||
elif chunk.choices[0].delta.reasoning_content:
|
||||
reasoning_content += chunk.choices[0].delta.reasoning_content
|
||||
|
||||
assert len(reasoning_content) > 0
|
||||
assert len(content) > 0
|
||||
|
||||
def test_streaming_separate_reasoning_true_stream_reasoning_false(self):
|
||||
# Test streaming with separate_reasoning=True, reasoning_content should not be empty
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 1+3?",
|
||||
}
|
||||
],
|
||||
"max_tokens": 100,
|
||||
"stream": True,
|
||||
"extra_body": {"separate_reasoning": True, "stream_reasoning": False},
|
||||
}
|
||||
response = client.chat.completions.create(**payload)
|
||||
|
||||
reasoning_content = ""
|
||||
content = ""
|
||||
first_chunk = False
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.reasoning_content:
|
||||
reasoning_content = chunk.choices[0].delta.reasoning_content
|
||||
first_chunk = True
|
||||
if chunk.choices[0].delta.content:
|
||||
content += chunk.choices[0].delta.content
|
||||
if not first_chunk:
|
||||
reasoning_content = chunk.choices[0].delta.reasoning_content
|
||||
first_chunk = True
|
||||
if not first_chunk:
|
||||
assert (
|
||||
not chunk.choices[0].delta.reasoning_content
|
||||
or len(chunk.choices[0].delta.reasoning_content) == 0
|
||||
)
|
||||
assert len(reasoning_content) > 0
|
||||
assert len(content) > 0
|
||||
|
||||
def test_nonstreaming_separate_reasoning_false(self):
|
||||
# Test non-streaming with separate_reasoning=False, reasoning_content should be empty
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 1+3?",
|
||||
}
|
||||
],
|
||||
"max_tokens": 100,
|
||||
"extra_body": {"separate_reasoning": False},
|
||||
}
|
||||
response = client.chat.completions.create(**payload)
|
||||
|
||||
assert (
|
||||
not response.choices[0].message.reasoning_content
|
||||
or len(response.choices[0].message.reasoning_content) == 0
|
||||
)
|
||||
assert len(response.choices[0].message.content) > 0
|
||||
|
||||
def test_nonstreaming_separate_reasoning_true(self):
|
||||
# Test non-streaming with separate_reasoning=True, reasoning_content should not be empty
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 1+3?",
|
||||
}
|
||||
],
|
||||
"max_tokens": 100,
|
||||
"extra_body": {"separate_reasoning": True},
|
||||
}
|
||||
response = client.chat.completions.create(**payload)
|
||||
|
||||
assert len(response.choices[0].message.reasoning_content) > 0
|
||||
assert len(response.choices[0].message.content) > 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,342 @@
|
||||
"""
|
||||
Fixtures for launching gRPC router + workers for e2e testing.
|
||||
|
||||
This module provides fixtures for launching SGLang workers and gRPC router separately:
|
||||
1. Launch N SGLang workers with gRPC enabled
|
||||
2. Launch router pointing to those workers
|
||||
|
||||
This approach gives more control and matches production deployment patterns.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def find_free_port() -> int:
|
||||
"""Find an available 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]
|
||||
|
||||
|
||||
def wait_for_workers_ready(
|
||||
router_url: str,
|
||||
expected_workers: int,
|
||||
timeout: int = 300,
|
||||
api_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Wait for router to have all workers connected.
|
||||
|
||||
Polls the /workers endpoint until the 'total' field matches expected_workers.
|
||||
|
||||
Example response from /workers endpoint:
|
||||
{"workers":[],"total":0,"stats":{"prefill_count":0,"decode_count":0,"regular_count":0}}
|
||||
|
||||
Args:
|
||||
router_url: Base URL of router (e.g., "http://127.0.0.1:30000")
|
||||
expected_workers: Number of workers expected to be connected
|
||||
timeout: Max seconds to wait
|
||||
api_key: Optional API key for authentication
|
||||
"""
|
||||
start_time = time.time()
|
||||
last_error = None
|
||||
attempt = 0
|
||||
|
||||
headers = {}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
|
||||
with requests.Session() as session:
|
||||
while time.time() - start_time < timeout:
|
||||
attempt += 1
|
||||
elapsed = int(time.time() - start_time)
|
||||
|
||||
# Log progress every 10 seconds
|
||||
if elapsed > 0 and elapsed % 10 == 0 and attempt % 10 == 0:
|
||||
logger.info(
|
||||
f" Still waiting for workers... ({elapsed}/{timeout}s elapsed)"
|
||||
)
|
||||
|
||||
try:
|
||||
response = session.get(
|
||||
f"{router_url}/workers", headers=headers, timeout=5
|
||||
)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
total_workers = data.get("total", 0)
|
||||
|
||||
if total_workers == expected_workers:
|
||||
logger.info(
|
||||
f" All {expected_workers} workers connected after {elapsed}s"
|
||||
)
|
||||
return
|
||||
else:
|
||||
last_error = f"Workers: {total_workers}/{expected_workers}"
|
||||
else:
|
||||
last_error = f"HTTP {response.status_code}"
|
||||
except requests.ConnectionError:
|
||||
last_error = "Connection refused (router not ready yet)"
|
||||
except requests.Timeout:
|
||||
last_error = "Timeout"
|
||||
except requests.RequestException as e:
|
||||
last_error = str(e)
|
||||
except (ValueError, KeyError) as e:
|
||||
last_error = f"Invalid response: {e}"
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
raise TimeoutError(
|
||||
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s.\n"
|
||||
f"Last status: {last_error}\n"
|
||||
f"Hint: Run with SHOW_ROUTER_LOGS=1 to see startup logs"
|
||||
)
|
||||
|
||||
|
||||
def popen_launch_workers_and_router(
|
||||
model: str,
|
||||
base_url: str,
|
||||
timeout: int = 300,
|
||||
num_workers: int = 2,
|
||||
policy: str = "round_robin",
|
||||
api_key: Optional[str] = None,
|
||||
worker_args: Optional[list] = None,
|
||||
router_args: Optional[list] = None,
|
||||
tp_size: int = 1,
|
||||
env: Optional[dict] = None,
|
||||
stdout=None,
|
||||
stderr=None,
|
||||
) -> dict:
|
||||
"""
|
||||
Launch SGLang workers and gRPC router separately.
|
||||
|
||||
This approach:
|
||||
1. Starts N SGLang workers with --grpc-mode flag
|
||||
2. Waits for workers to initialize (process startup)
|
||||
3. Starts a gRPC router pointing to those workers
|
||||
4. Waits for router health check to pass (router validates worker connectivity)
|
||||
|
||||
This matches production deployment patterns better than the integrated approach.
|
||||
|
||||
Args:
|
||||
model: Model path (e.g., /home/ubuntu/models/llama-3.1-8b-instruct)
|
||||
base_url: Base URL for router (e.g., "http://127.0.0.1:8080")
|
||||
timeout: Timeout for server startup (default: 300s)
|
||||
num_workers: Number of workers to launch
|
||||
policy: Routing policy (round_robin, random, power_of_two, cache_aware)
|
||||
api_key: Optional API key for router
|
||||
worker_args: Additional arguments for workers (e.g., ["--context-len", "8192"])
|
||||
router_args: Additional arguments for router (e.g., ["--max-total-token", "1536"])
|
||||
tp_size: Tensor parallelism size for workers (default: 1)
|
||||
env: Optional environment variables for workers (e.g., {"SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION": "256"})
|
||||
stdout: Optional file handle for worker stdout (default: subprocess.PIPE)
|
||||
stderr: Optional file handle for worker stderr (default: subprocess.PIPE)
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
- workers: list of worker process objects
|
||||
- worker_urls: list of gRPC worker URLs
|
||||
- router: router process object
|
||||
- base_url: router URL (HTTP endpoint)
|
||||
|
||||
Example:
|
||||
>>> cluster = popen_launch_workers_and_router(model, base_url, num_workers=2)
|
||||
>>> # Use cluster['base_url'] for HTTP requests
|
||||
>>> # Cleanup:
|
||||
>>> for worker in cluster['workers']:
|
||||
>>> kill_process_tree(worker.pid)
|
||||
>>> kill_process_tree(cluster['router'].pid)
|
||||
"""
|
||||
import os
|
||||
|
||||
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
|
||||
|
||||
# Note: timeout parameter is used for router health check below
|
||||
|
||||
# Parse router port from base_url
|
||||
if ":" in base_url.split("//")[-1]:
|
||||
router_port = int(base_url.split(":")[-1])
|
||||
else:
|
||||
router_port = find_free_port()
|
||||
|
||||
logger.info(f"\n{'='*70}")
|
||||
logger.info(f"Launching gRPC cluster (separate workers + router)")
|
||||
logger.info(f"{'='*70}")
|
||||
logger.info(f" Model: {model}")
|
||||
logger.info(f" Router port: {router_port}")
|
||||
logger.info(f" Workers: {num_workers}")
|
||||
logger.info(f" TP size: {tp_size}")
|
||||
logger.info(f" Policy: {policy}")
|
||||
|
||||
# Step 1: Launch workers with gRPC enabled
|
||||
workers = []
|
||||
worker_urls = []
|
||||
|
||||
for i in range(num_workers):
|
||||
worker_port = find_free_port()
|
||||
worker_url = f"grpc://127.0.0.1:{worker_port}"
|
||||
worker_urls.append(worker_url)
|
||||
|
||||
logger.info(f"\n[Worker {i+1}/{num_workers}]")
|
||||
logger.info(f" Port: {worker_port}")
|
||||
logger.info(f" URL: {worker_url}")
|
||||
|
||||
# Build worker command
|
||||
worker_cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang.launch_server",
|
||||
"--model-path",
|
||||
model,
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(worker_port),
|
||||
"--grpc-mode", # Enable gRPC for this worker
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
]
|
||||
|
||||
# Add TP size
|
||||
if tp_size > 1:
|
||||
worker_cmd.extend(["--tp-size", str(tp_size)])
|
||||
|
||||
# Add worker-specific args
|
||||
if worker_args:
|
||||
worker_cmd.extend(worker_args)
|
||||
|
||||
# Launch worker with optional environment variables
|
||||
if show_output:
|
||||
worker_proc = subprocess.Popen(
|
||||
worker_cmd,
|
||||
env=env,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
else:
|
||||
worker_proc = subprocess.Popen(
|
||||
worker_cmd,
|
||||
stdout=stdout if stdout is not None else subprocess.PIPE,
|
||||
stderr=stderr if stderr is not None else subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
|
||||
workers.append(worker_proc)
|
||||
logger.info(f" PID: {worker_proc.pid}")
|
||||
|
||||
# Give workers a moment to start binding to ports
|
||||
# The router will check worker health when it starts
|
||||
logger.info(f"\nWaiting for {num_workers} workers to initialize (20s)...")
|
||||
time.sleep(20)
|
||||
|
||||
# Quick check: make sure worker processes are still alive
|
||||
for i, worker in enumerate(workers):
|
||||
if worker.poll() is not None:
|
||||
logger.error(
|
||||
f" ✗ Worker {i+1} died during startup (exit code: {worker.poll()})"
|
||||
)
|
||||
# Cleanup: kill all workers
|
||||
for w in workers:
|
||||
try:
|
||||
w.kill()
|
||||
except:
|
||||
pass
|
||||
raise RuntimeError(f"Worker {i+1} failed to start")
|
||||
|
||||
logger.info(
|
||||
f"✓ All {num_workers} workers started (router will verify connectivity)"
|
||||
)
|
||||
|
||||
# Step 2: Launch router pointing to workers
|
||||
logger.info(f"\n[Router]")
|
||||
logger.info(f" Port: {router_port}")
|
||||
logger.info(f" Worker URLs: {', '.join(worker_urls)}")
|
||||
|
||||
# Build router command
|
||||
router_cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang_router.launch_router",
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(router_port),
|
||||
"--prometheus-port",
|
||||
"9321",
|
||||
"--policy",
|
||||
policy,
|
||||
"--model-path",
|
||||
model,
|
||||
"--log-level",
|
||||
"warn",
|
||||
]
|
||||
|
||||
# Add worker URLs
|
||||
router_cmd.append("--worker-urls")
|
||||
router_cmd.extend(worker_urls)
|
||||
|
||||
# Add API key
|
||||
if api_key:
|
||||
router_cmd.extend(["--api-key", api_key])
|
||||
|
||||
# Add router-specific args
|
||||
if router_args:
|
||||
router_cmd.extend(router_args)
|
||||
|
||||
if show_output:
|
||||
logger.info(f" Command: {' '.join(router_cmd)}")
|
||||
|
||||
# Launch router
|
||||
if show_output:
|
||||
router_proc = subprocess.Popen(router_cmd)
|
||||
else:
|
||||
router_proc = subprocess.Popen(
|
||||
router_cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
|
||||
logger.info(f" PID: {router_proc.pid}")
|
||||
|
||||
# Wait for router to be ready
|
||||
router_url = f"http://127.0.0.1:{router_port}"
|
||||
logger.info(f"\nWaiting for router to start at {router_url}...")
|
||||
|
||||
try:
|
||||
wait_for_workers_ready(
|
||||
router_url, expected_workers=num_workers, timeout=180, api_key=api_key
|
||||
)
|
||||
logger.info(f"✓ Router ready at {router_url}")
|
||||
except TimeoutError:
|
||||
logger.error(f"✗ Router failed to start")
|
||||
# Cleanup: kill router and all workers
|
||||
try:
|
||||
router_proc.kill()
|
||||
except:
|
||||
pass
|
||||
for worker in workers:
|
||||
try:
|
||||
worker.kill()
|
||||
except:
|
||||
pass
|
||||
raise
|
||||
|
||||
logger.info(f"\n{'='*70}")
|
||||
logger.info(f"✓ gRPC cluster ready!")
|
||||
logger.info(f" Router: {router_url}")
|
||||
logger.info(f" Workers: {len(workers)}")
|
||||
logger.info(f"{'='*70}\n")
|
||||
|
||||
return {
|
||||
"workers": workers,
|
||||
"worker_urls": worker_urls,
|
||||
"router": router_proc,
|
||||
"base_url": router_url,
|
||||
}
|
||||
@@ -0,0 +1,950 @@
|
||||
"""
|
||||
gRPC Router E2E Test - Test Openai Function Calling
|
||||
|
||||
This test file is REUSED from test/srt/openai_server/function_call/test_openai_function_calling.py
|
||||
with minimal changes:
|
||||
num_workers=2,
|
||||
- Swap popen_launch_server() → popen_launch_workers_and_router()
|
||||
- Update teardown to cleanup router + workers
|
||||
- All test logic and assertions remain identical
|
||||
|
||||
Run with:
|
||||
pytest py_test/e2e_grpc/e2e_grpc/function_call/test_openai_function_calling.py -v
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_MODEL_PATH,
|
||||
DEFAULT_SMALL_MODEL_PATH,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
get_tokenizer,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
# NOTE: this system_message is for Llama3.2 system prompt. Without this,
|
||||
# sometimes Llama3.2 gives a different tool call format such as:
|
||||
# '<|python_tag|>{"type": "function", "function": "add", "parameters": {"a": "3", "b": "5"}}'
|
||||
SYSTEM_MESSAGE = (
|
||||
"You are a helpful assistant with tool calling capabilities. "
|
||||
"Only reply with a tool call if the function exists in the library provided by the user. "
|
||||
"If it doesn't exist, just reply directly in natural language. "
|
||||
"When you receive a tool call response, use the output to format an answer to the original user question. "
|
||||
"You have access to the following functions. "
|
||||
"To call a function, please respond with JSON for a function call. "
|
||||
'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. '
|
||||
"Do not use variables.\n\n"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
# Using small model for function calling tests
|
||||
cls.model = DEFAULT_SMALL_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
# Start the local OpenAI Server. If necessary, you can add other parameters such as --enable-tools.
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
# If your server needs extra parameters to test function calling, please add them here.
|
||||
"--tool-call-parser",
|
||||
"llama",
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
def test_function_calling_format(self):
|
||||
"""
|
||||
Test: Whether the function call format returned by the AI is correct.
|
||||
When returning a tool call, message.content should be None, and tool_calls should be a list.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Compute the sum of two numbers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
"b": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{"role": "user", "content": "Compute (3+5)"},
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
|
||||
assert (
|
||||
isinstance(tool_calls, list) and len(tool_calls) > 0
|
||||
), "tool_calls should be a non-empty list"
|
||||
|
||||
function_name = tool_calls[0].function.name
|
||||
assert function_name == "add", "Function name should be 'add'"
|
||||
|
||||
# This unit test is too difficult for default model. Mark it as optional unit tests so it won't trigger unless specified.
|
||||
def _test_function_calling_multiturn(self):
|
||||
"""
|
||||
Test: Whether the function call format returned by the AI is correct.
|
||||
When returning a tool call, message.content should be None, and tool_calls should be a list.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Compute the sum of two numbers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
"b": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "Compute (3+5)"}]
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
tool_call = response.choices[0].message.tool_calls[0]
|
||||
function_name = tool_call.function.name
|
||||
assert function_name == "add", "Function name should be 'add'"
|
||||
function_arguments = tool_call.function.arguments
|
||||
function_arguments = json.loads(tool_call.function.arguments)
|
||||
assert function_arguments in [
|
||||
{"a": 3, "b": 5},
|
||||
{"a": "3", "b": "5"},
|
||||
], f"Unexpected function arguments: {function_arguments}"
|
||||
|
||||
messages.append(response.choices[0].message)
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": "8",
|
||||
"name": function_name,
|
||||
}
|
||||
)
|
||||
|
||||
final_response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
assert (
|
||||
"8" in final_response.choices[0].message.content
|
||||
), "tool_call response should have the sum 8 in the content"
|
||||
|
||||
def test_function_calling_streaming_simple(self):
|
||||
"""
|
||||
Test: Whether the function name can be correctly recognized in streaming mode.
|
||||
- Expect a function call to be found, and the function name to be correct.
|
||||
- Verify that streaming mode returns at least multiple chunks.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to find the weather for",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "Weather unit (celsius or fahrenheit)",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["city", "unit"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the temperature in Paris in celsius??",
|
||||
},
|
||||
]
|
||||
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
self.assertTrue(len(chunks) > 0, "Streaming should return at least one chunk")
|
||||
|
||||
found_function_name = False
|
||||
for chunk in chunks:
|
||||
choice = chunk.choices[0]
|
||||
# Check whether the current chunk contains tool_calls
|
||||
if choice.delta.tool_calls:
|
||||
tool_call = choice.delta.tool_calls[0]
|
||||
if tool_call.function.name:
|
||||
self.assertEqual(
|
||||
tool_call.function.name,
|
||||
"get_current_weather",
|
||||
"Function name should be 'get_current_weather'",
|
||||
)
|
||||
found_function_name = True
|
||||
break
|
||||
|
||||
self.assertTrue(
|
||||
found_function_name,
|
||||
"Target function name 'get_current_weather' was not found in the streaming chunks",
|
||||
)
|
||||
|
||||
finish_reason = chunks[-1].choices[0].finish_reason
|
||||
self.assertEqual(
|
||||
finish_reason,
|
||||
"tool_calls",
|
||||
"Final response of function calling should have finish_reason 'tool_calls'",
|
||||
)
|
||||
|
||||
def test_function_calling_streaming_args_parsing(self):
|
||||
"""
|
||||
Test: Whether the function call arguments returned in streaming mode can be correctly concatenated into valid JSON.
|
||||
- The user request requires multiple parameters.
|
||||
- AI may return the arguments in chunks that need to be concatenated.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Compute the sum of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
"strict": True, # Llama-3.2-1B is flaky in tool call. It won't always respond with parameters unless we set strict.
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{"role": "user", "content": "Please sum 5 and 7, just call the function."},
|
||||
]
|
||||
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.9,
|
||||
top_p=0.9,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
argument_fragments = []
|
||||
chunks = list(response_stream)
|
||||
function_name = None
|
||||
for chunk in chunks:
|
||||
choice = chunk.choices[0]
|
||||
if choice.delta.tool_calls:
|
||||
tool_call = choice.delta.tool_calls[0]
|
||||
# Record the function name on first occurrence
|
||||
function_name = tool_call.function.name or function_name
|
||||
# In case of multiple chunks, JSON fragments may need to be concatenated
|
||||
if tool_call.function.arguments is not None:
|
||||
argument_fragments.append(tool_call.function.arguments)
|
||||
|
||||
self.assertEqual(function_name, "add", "Function name should be 'add'")
|
||||
joined_args = "".join(argument_fragments)
|
||||
self.assertTrue(
|
||||
len(joined_args) > 0,
|
||||
"No parameter fragments were returned in the function call",
|
||||
)
|
||||
|
||||
finish_reason = chunks[-1].choices[0].finish_reason
|
||||
self.assertEqual(
|
||||
finish_reason,
|
||||
"tool_calls",
|
||||
"Final response of function calling should have finish_reason 'tool_calls'",
|
||||
)
|
||||
|
||||
# Check whether the concatenated JSON is valid
|
||||
try:
|
||||
args_obj = json.loads(joined_args)
|
||||
except json.JSONDecodeError:
|
||||
self.fail(
|
||||
"The concatenated tool call arguments are not valid JSON, parsing failed"
|
||||
)
|
||||
|
||||
self.assertIn("a", args_obj, "Missing parameter 'a'")
|
||||
self.assertIn("b", args_obj, "Missing parameter 'b'")
|
||||
self.assertEqual(str(args_obj["a"]), "5", "Parameter a should be 5")
|
||||
self.assertEqual(str(args_obj["b"]), "7", "Parameter b should be 7")
|
||||
|
||||
@unittest.skip(
|
||||
"Skipping function call strict test as it is not supported by the router"
|
||||
)
|
||||
def test_function_call_strict(self):
|
||||
"""
|
||||
Test: Whether the strict mode of function calling works as expected.
|
||||
- When strict mode is enabled, the AI should not return a function call if the function name is not recognized.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "sub",
|
||||
"description": "Compute the difference of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"int_a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"int_b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["int_a", "int_b"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "Please compute 5 - 7, using your tool."}
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
function_name = tool_calls[0].function.name
|
||||
arguments = tool_calls[0].function.arguments
|
||||
args_obj = json.loads(arguments)
|
||||
|
||||
self.assertEqual(function_name, "sub", "Function name should be 'sub'")
|
||||
self.assertEqual(str(args_obj["int_a"]), "5", "Parameter int_a should be 5")
|
||||
self.assertEqual(str(args_obj["int_b"]), "7", "Parameter int_b should be 7")
|
||||
|
||||
def test_function_call_required(self):
|
||||
"""
|
||||
Test: Whether tool_choice: "required" works as expected
|
||||
- When tool_choice == "required", the model should return one or more tool_calls.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "sub",
|
||||
"description": "Compute the difference of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"int_a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"int_b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["int_a", "int_b"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "use this to get latest weather information for a city given its name",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "name of the city to get weather for",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "What is the capital of France?"}]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls, "No tool_calls in the response")
|
||||
function_name = tool_calls[0].function.name
|
||||
arguments = tool_calls[0].function.arguments
|
||||
args_obj = json.loads(arguments)
|
||||
|
||||
self.assertEqual(
|
||||
function_name,
|
||||
"get_weather",
|
||||
f"Function name should be 'get_weather', got: {function_name}",
|
||||
)
|
||||
self.assertIn(
|
||||
"city", args_obj, f"Function arguments should have 'city', got: {args_obj}"
|
||||
)
|
||||
|
||||
# Make the test more robust by checking type and accepting valid responses
|
||||
city_value = args_obj["city"]
|
||||
self.assertIsInstance(
|
||||
city_value,
|
||||
str,
|
||||
f"Parameter city should be a string, got: {type(city_value)}",
|
||||
)
|
||||
self.assertTrue(
|
||||
"Paris" in city_value or "France" in city_value,
|
||||
f"Parameter city should contain either 'Paris' or 'France', got: {city_value}",
|
||||
)
|
||||
|
||||
def test_function_call_specific(self):
|
||||
"""
|
||||
Test: Whether tool_choice: ToolChoice works as expected
|
||||
- When tool_choice is a specific ToolChoice, the model should return one or more tool_calls.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "sub",
|
||||
"description": "Compute the difference of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"int_a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"int_b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["int_a", "int_b"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "use this to get latest weather information for a city given its name",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "name of the city to get weather for",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "What is the capital of France?"}]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "function": {"name": "get_weather"}},
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls, "No tool_calls in the response")
|
||||
function_name = tool_calls[0].function.name
|
||||
arguments = tool_calls[0].function.arguments
|
||||
args_obj = json.loads(arguments)
|
||||
|
||||
self.assertEqual(
|
||||
function_name, "get_weather", "Function name should be 'get_weather'"
|
||||
)
|
||||
self.assertIn("city", args_obj, "Function arguments should have 'city'")
|
||||
|
||||
def test_streaming_multiple_choices_finish_reason(self):
|
||||
"""
|
||||
Test: Verify that each choice gets its own finish_reason chunk in streaming mode with n > 1.
|
||||
This tests the fix for the bug where only the last index got a finish_reason chunk.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "What is the weather like in Los Angeles?"}
|
||||
]
|
||||
|
||||
# Request with n=2 to get multiple choices
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
temperature=0.8,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
tool_choice="required", # Force tool calls
|
||||
n=2, # Multiple choices
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
|
||||
# Track finish_reason chunks for each index
|
||||
finish_reason_chunks = {}
|
||||
for chunk in chunks:
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
if choice.finish_reason is not None:
|
||||
index = choice.index
|
||||
if index not in finish_reason_chunks:
|
||||
finish_reason_chunks[index] = []
|
||||
finish_reason_chunks[index].append(choice.finish_reason)
|
||||
|
||||
# Verify we got finish_reason chunks for both indices
|
||||
self.assertEqual(
|
||||
len(finish_reason_chunks),
|
||||
2,
|
||||
f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}",
|
||||
)
|
||||
|
||||
# Verify both index 0 and 1 have finish_reason
|
||||
self.assertIn(
|
||||
0, finish_reason_chunks, "Missing finish_reason chunk for index 0"
|
||||
)
|
||||
self.assertIn(
|
||||
1, finish_reason_chunks, "Missing finish_reason chunk for index 1"
|
||||
)
|
||||
|
||||
# Verify the finish_reason is "tool_calls" since we forced tool calls
|
||||
for index, reasons in finish_reason_chunks.items():
|
||||
self.assertEqual(
|
||||
reasons[-1], # Last finish_reason for this index
|
||||
"tool_calls",
|
||||
f"Expected finish_reason 'tool_calls' for index {index}, got {reasons[-1]}",
|
||||
)
|
||||
|
||||
def test_function_calling_streaming_no_tool_call(self):
|
||||
"""
|
||||
Test: Whether the finish_reason is stop in streaming mode when no tool call is given.
|
||||
- Expect no function call to be found.
|
||||
- Verify that finish_reason is stop
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to find the weather for",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "Weather unit (celsius or fahrenheit)",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["city", "unit"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "Who are you?"}]
|
||||
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
tool_choice="none",
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
self.assertTrue(len(chunks) > 0, "Streaming should return at least one chunk")
|
||||
|
||||
found_tool_call = False
|
||||
for chunk in chunks:
|
||||
choice = chunk.choices[0]
|
||||
# Check whether the current chunk contains tool_calls
|
||||
found_tool_call = choice.delta.tool_calls is not None
|
||||
|
||||
self.assertFalse(
|
||||
found_tool_call,
|
||||
"Shouldn't have any tool_call in the streaming chunks",
|
||||
)
|
||||
|
||||
finish_reason = chunks[-1].choices[0].finish_reason
|
||||
self.assertEqual(
|
||||
finish_reason,
|
||||
"stop",
|
||||
"Final response of no function calling should have finish_reason 'stop'",
|
||||
)
|
||||
|
||||
def test_streaming_multiple_choices_without_tools(self):
|
||||
"""
|
||||
Test: Verify that each choice gets its own finish_reason chunk without tool calls.
|
||||
This tests the fix for regular content streaming with multiple choices.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
messages = [{"role": "user", "content": "Say hello in one word."}]
|
||||
|
||||
# Request with n=2 to get multiple choices, no tools
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
stream=True,
|
||||
max_tokens=10, # Keep it short
|
||||
n=2, # Multiple choices
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
|
||||
# Track finish_reason chunks for each index
|
||||
finish_reason_chunks = {}
|
||||
for chunk in chunks:
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
if choice.finish_reason is not None:
|
||||
index = choice.index
|
||||
if index not in finish_reason_chunks:
|
||||
finish_reason_chunks[index] = []
|
||||
finish_reason_chunks[index].append(choice.finish_reason)
|
||||
|
||||
# Verify we got finish_reason chunks for both indices
|
||||
self.assertEqual(
|
||||
len(finish_reason_chunks),
|
||||
2,
|
||||
f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}",
|
||||
)
|
||||
|
||||
# Verify both index 0 and 1 have finish_reason
|
||||
self.assertIn(
|
||||
0, finish_reason_chunks, "Missing finish_reason chunk for index 0"
|
||||
)
|
||||
self.assertIn(
|
||||
1, finish_reason_chunks, "Missing finish_reason chunk for index 1"
|
||||
)
|
||||
|
||||
# Verify the finish_reason is "stop" (regular completion)
|
||||
for index, reasons in finish_reason_chunks.items():
|
||||
self.assertIn(
|
||||
reasons[-1],
|
||||
["stop", "length"], # Could be either depending on how model responds
|
||||
f"Expected finish_reason 'stop' or 'length' for index {index}, got {reasons[-1]}",
|
||||
)
|
||||
|
||||
|
||||
class TestOpenAIPythonicFunctionCalling(CustomTestCase):
|
||||
PYTHONIC_TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather for a given location.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The name of the city or location.",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_tourist_attractions",
|
||||
"description": "Get a list of top tourist attractions for a given city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The name of the city to find attractions for.",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
PYTHONIC_MESSAGES = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a travel assistant. "
|
||||
"When asked to call functions, ALWAYS respond ONLY with a python list of function calls, "
|
||||
"using this format: [func_name1(param1=value1, param2=value2), func_name2(param=value)]. "
|
||||
"Do NOT use JSON, do NOT use variables, do NOT use any other format. "
|
||||
"Here is an example:\n"
|
||||
'[get_weather(location="Paris"), get_tourist_attractions(city="Paris")]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"I'm planning a trip to Tokyo next week. What's the weather like and what are some top tourist attractions? "
|
||||
"Propose parallel tool calls at once, using the python list of function calls format as shown above."
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
cls.model = DEFAULT_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
"--tool-call-parser",
|
||||
"pythonic",
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
def test_pythonic_tool_call_prompt(self):
|
||||
"""
|
||||
Test: Explicit prompt for pythonic tool call format without chat template.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=self.PYTHONIC_MESSAGES,
|
||||
tools=self.PYTHONIC_TOOLS,
|
||||
temperature=0.1,
|
||||
stream=False,
|
||||
)
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsInstance(tool_calls, list, "No tool_calls found")
|
||||
self.assertGreaterEqual(len(tool_calls), 1)
|
||||
names = [tc.function.name for tc in tool_calls]
|
||||
self.assertTrue(
|
||||
"get_weather" in names or "get_tourist_attractions" in names,
|
||||
f"Function name '{names}' should container either 'get_weather' or 'get_tourist_attractions'",
|
||||
)
|
||||
|
||||
def test_pythonic_tool_call_streaming(self):
|
||||
"""
|
||||
Test: Streaming pythonic tool call format; assert tool_call index is present.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=self.PYTHONIC_MESSAGES,
|
||||
tools=self.PYTHONIC_TOOLS,
|
||||
temperature=0.1,
|
||||
stream=True,
|
||||
)
|
||||
found_tool_calls = False
|
||||
found_index = False
|
||||
found_names = set()
|
||||
for chunk in response_stream:
|
||||
choice = chunk.choices[0]
|
||||
if getattr(choice.delta, "tool_calls", None):
|
||||
found_tool_calls = True
|
||||
tool_call = choice.delta.tool_calls[0]
|
||||
if hasattr(tool_call, "index") or (
|
||||
isinstance(tool_call, dict) and "index" in tool_call
|
||||
):
|
||||
found_index = True
|
||||
found_names.add(str(tool_call.function.name))
|
||||
|
||||
self.assertTrue(found_tool_calls, "No tool_calls found in streaming response")
|
||||
self.assertTrue(found_index, "No index field found in any streamed tool_call")
|
||||
self.assertTrue(
|
||||
"get_weather" in found_names or "get_tourist_attractions" in found_names,
|
||||
f"Function name '{found_names}' should container either 'get_weather' or 'get_tourist_attractions'",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,795 @@
|
||||
"""
|
||||
Test script for tool_choice functionality in SGLang
|
||||
Tests: required, auto, and specific function choices in both streaming and non-streaming modes
|
||||
|
||||
# To run the tests, use the following command:
|
||||
#
|
||||
# python3 -m unittest openai_server.function_call.test_tool_choice
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH,
|
||||
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH,
|
||||
DEFAULT_SMALL_MODEL_PATH,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
get_tokenizer,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestToolChoiceLlama32(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
# Mark flaky tests for this model
|
||||
cls.flaky_tests = {
|
||||
"test_multi_tool_scenario_auto",
|
||||
"test_multi_tool_scenario_required",
|
||||
}
|
||||
|
||||
# Use a model that supports function calling
|
||||
cls.model = DEFAULT_SMALL_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
# Start the local OpenAI Server with tool calling support
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
"--tool-call-parser",
|
||||
"llama", # Default parser for the test model
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
def setUp(self):
|
||||
self.client = openai.Client(base_url=self.base_url, api_key=self.api_key)
|
||||
self.model_name = self.client.models.list().data[0].id
|
||||
|
||||
def _is_flaky_test(self):
|
||||
"""Check if the current test is marked as flaky for this class"""
|
||||
return (
|
||||
hasattr(self.__class__, "flaky_tests")
|
||||
and self._testMethodName in self.__class__.flaky_tests
|
||||
)
|
||||
|
||||
def get_test_tools(self):
|
||||
"""Get the test tools for function calling"""
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "use this to get latest weather information for a city given its name",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "name of the city to get weather for",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_pokemon_info",
|
||||
"description": "get detailed information about a pokemon given its name",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "name of the pokemon to get info for",
|
||||
}
|
||||
},
|
||||
"required": ["name"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "make_next_step_decision",
|
||||
"description": "You will be given a trace of thinking process in the following format.\n\nQuestion: the input question you must answer\nTOOL: think about what to do, and choose a tool to use ONLY IF there are defined tools. \n You should never call the same tool with the same input twice in a row.\n If the previous conversation history already contains the information that can be retrieved from the tool, you should not call the tool again.\nOBSERVATION: the result of the tool call, NEVER include this in your response, this information will be provided\n... (this TOOL/OBSERVATION can repeat N times)\nANSWER: If you know the answer to the original question, require for more information,\n or you don't know the answer and there are no defined tools or all available tools are not helpful, respond with the answer without mentioning anything else.\n If the previous conversation history already contains the answer, respond with the answer right away.\n\n If no tools are configured, naturally mention this limitation while still being helpful. Briefly note that adding tools in the agent configuration would expand capabilities.\n\nYour task is to respond with the next step to take, based on the traces, \nor answer the question if you have enough information.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"decision": {
|
||||
"type": "string",
|
||||
"description": 'The next step to take, it must be either "TOOL" or "ANSWER". If the previous conversation history already contains the information that can be retrieved from the tool, you should not call the tool again. If there are no defined tools, you should not return "TOOL" in your response.',
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": 'The content of the next step. If the decision is "TOOL", this should be a short and concise reasoning of why you chose the tool, MUST include the tool name. If the decision is "ANSWER", this should be the answer to the question. If no tools are available, integrate this limitation conversationally without sounding scripted.',
|
||||
},
|
||||
},
|
||||
"required": ["decision", "content"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
def get_test_messages(self):
|
||||
"""Get test messages that should trigger tool usage"""
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Answer the following questions as best you can:\n\nYou will be given a trace of thinking process in the following format.\n\nQuestion: the input question you must answer\nTOOL: think about what to do, and choose a tool to use ONLY IF there are defined tools\nOBSERVATION: the result of the tool call or the observation of the current task, NEVER include this in your response, this information will be provided\n... (this TOOL/OBSERVATION can repeat N times)\nANSWER: If you know the answer to the original question, require for more information, \nif the previous conversation history already contains the answer, \nor you don't know the answer and there are no defined tools or all available tools are not helpful, respond with the answer without mentioning anything else.\nYou may use light Markdown formatting to improve clarity (e.g. lists, **bold**, *italics*), but keep it minimal and unobtrusive.\n\nYour task is to respond with the next step to take, based on the traces, \nor answer the question if you have enough information.\n\nQuestion: what is the weather in top 5 populated cities in the US in celsius?\n\nTraces:\n\n\nThese are some additional instructions that you should follow:",
|
||||
}
|
||||
]
|
||||
|
||||
def get_travel_tools(self):
|
||||
"""Get tools for travel assistant scenario that should trigger multiple tool calls"""
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather for a given location.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The name of the city or location.",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_tourist_attractions",
|
||||
"description": "Get a list of top tourist attractions for a given city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The name of the city to find attractions for.",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
def get_travel_messages(self):
|
||||
"""Get travel assistant messages that should trigger multiple tool calls"""
|
||||
return [
|
||||
{
|
||||
"content": "You are a travel assistant providing real-time weather updates and top tourist attractions.",
|
||||
"role": "system",
|
||||
},
|
||||
{
|
||||
"content": "I'm planning a trip to Tokyo next week. What's the weather like? What are the most amazing sights?",
|
||||
"role": "user",
|
||||
},
|
||||
]
|
||||
|
||||
def test_tool_choice_auto_non_streaming(self):
|
||||
"""Test tool_choice='auto' in non-streaming mode"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(response.choices[0].message)
|
||||
# With auto, tool calls are optional
|
||||
|
||||
def test_tool_choice_auto_streaming(self):
|
||||
"""Test tool_choice='auto' in streaming mode"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect streaming response
|
||||
content_chunks = []
|
||||
tool_call_chunks = []
|
||||
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.content:
|
||||
content_chunks.append(chunk.choices[0].delta.content)
|
||||
elif chunk.choices[0].delta.tool_calls:
|
||||
tool_call_chunks.extend(chunk.choices[0].delta.tool_calls)
|
||||
|
||||
# Should complete without errors
|
||||
self.assertIsInstance(content_chunks, list)
|
||||
self.assertIsInstance(tool_call_chunks, list)
|
||||
|
||||
def test_tool_choice_required_non_streaming(self):
|
||||
"""Test tool_choice='required' in non-streaming mode"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
temperature=0.2,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# With required, we should get tool calls
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls)
|
||||
self.assertGreater(len(tool_calls), 0)
|
||||
|
||||
def test_tool_choice_required_streaming(self):
|
||||
"""Test tool_choice='required' in streaming mode"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect streaming response
|
||||
tool_call_chunks = []
|
||||
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.tool_calls:
|
||||
tool_call_chunks.extend(chunk.choices[0].delta.tool_calls)
|
||||
|
||||
# With required, we should get tool call chunks
|
||||
self.assertGreater(len(tool_call_chunks), 0)
|
||||
|
||||
def test_tool_choice_specific_function_non_streaming(self):
|
||||
"""Test tool_choice with specific function in non-streaming mode"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
tool_choice = {"type": "function", "function": {"name": "get_weather"}}
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Should call the specific function
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls)
|
||||
# Our messages ask the top 5 populated cities in the US, so the model could get 5 tool calls
|
||||
self.assertGreaterEqual(len(tool_calls), 1)
|
||||
for tool_call in tool_calls:
|
||||
self.assertEqual(tool_call.function.name, "get_weather")
|
||||
|
||||
def test_tool_choice_specific_function_streaming(self):
|
||||
"""Test tool_choice with specific function in streaming mode"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
tool_choice = {"type": "function", "function": {"name": "get_weather"}}
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect streaming response
|
||||
tool_call_chunks = []
|
||||
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.tool_calls:
|
||||
tool_call_chunks.extend(chunk.choices[0].delta.tool_calls)
|
||||
|
||||
# Should get tool call chunks for the specific function
|
||||
self.assertGreater(len(tool_call_chunks), 0)
|
||||
|
||||
# Find function name in chunks
|
||||
found_name = None
|
||||
for chunk in tool_call_chunks:
|
||||
if chunk.function and chunk.function.name:
|
||||
found_name = chunk.function.name
|
||||
break
|
||||
|
||||
self.assertEqual(found_name, "get_weather")
|
||||
|
||||
def test_required_streaming_arguments_chunks_json(self):
|
||||
"""In streaming required mode, complete tool call arguments should be valid JSON when all chunks are combined"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=1024,
|
||||
temperature=0.1,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Collect all tool call chunks and reconstruct complete tool calls
|
||||
tool_calls_by_index = {}
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.tool_calls:
|
||||
for tool_call_delta in chunk.choices[0].delta.tool_calls:
|
||||
tool_index = tool_call_delta.index
|
||||
|
||||
# Initialize tool call if not seen before
|
||||
if tool_index not in tool_calls_by_index:
|
||||
tool_calls_by_index[tool_index] = {
|
||||
"id": tool_call_delta.id,
|
||||
"type": "function",
|
||||
"function": {"name": "", "arguments": ""},
|
||||
}
|
||||
|
||||
# Update function name if present (first chunk)
|
||||
if tool_call_delta.function and tool_call_delta.function.name:
|
||||
tool_calls_by_index[tool_index]["function"][
|
||||
"name"
|
||||
] = tool_call_delta.function.name
|
||||
|
||||
# Accumulate arguments (all chunks)
|
||||
if tool_call_delta.function and tool_call_delta.function.arguments:
|
||||
tool_calls_by_index[tool_index]["function"][
|
||||
"arguments"
|
||||
] += tool_call_delta.function.arguments
|
||||
|
||||
self.assertGreater(len(tool_calls_by_index), 0)
|
||||
|
||||
# Validate that complete tool calls have valid JSON arguments
|
||||
for tool_call in tool_calls_by_index.values():
|
||||
self.assertIsNotNone(tool_call["function"]["name"])
|
||||
self.assertIsNotNone(tool_call["function"]["arguments"])
|
||||
|
||||
# The complete arguments should be valid JSON
|
||||
try:
|
||||
args = json.loads(tool_call["function"]["arguments"])
|
||||
self.assertIsInstance(args, dict)
|
||||
except json.JSONDecodeError:
|
||||
self.fail(
|
||||
f"Invalid JSON in complete tool call arguments: {tool_call['function']['arguments']}"
|
||||
)
|
||||
|
||||
def test_complex_parameters_required_non_streaming(self):
|
||||
"""Validate complex nested parameter schemas in non-streaming required mode"""
|
||||
complex_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "analyze_data",
|
||||
"description": "Analyze complex data structures",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"data": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"metrics": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
"config": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"threshold": {"type": "number"},
|
||||
"enabled": {"type": "boolean"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["metrics"],
|
||||
},
|
||||
"options": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"value": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["data"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Analyze some data with metrics and configuration",
|
||||
}
|
||||
]
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=1024,
|
||||
temperature=0.1,
|
||||
tools=complex_tools,
|
||||
tool_choice="required",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls)
|
||||
self.assertGreater(len(tool_calls), 0)
|
||||
|
||||
for tool_call in tool_calls:
|
||||
self.assertEqual(tool_call.function.name, "analyze_data")
|
||||
try:
|
||||
args = json.loads(tool_call.function.arguments)
|
||||
self.assertIsInstance(args, dict)
|
||||
self.assertIn("data", args)
|
||||
self.assertIsInstance(args["data"], dict)
|
||||
except json.JSONDecodeError:
|
||||
self.fail(
|
||||
f"Invalid JSON in complex tool call arguments: {tool_call.function.arguments}"
|
||||
)
|
||||
|
||||
def test_multi_tool_scenario_auto(self):
|
||||
"""Test multi-tool scenario with tool_choice='auto'"""
|
||||
tools = self.get_travel_tools()
|
||||
messages = self.get_travel_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
temperature=0.2,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Should complete without errors
|
||||
self.assertIsNotNone(response.choices[0].message)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
expected_functions = {"get_weather", "get_tourist_attractions"}
|
||||
|
||||
if self._is_flaky_test():
|
||||
# For flaky tests, just verify all called functions are available tools
|
||||
if tool_calls:
|
||||
available_names = [tool["function"]["name"] for tool in tools]
|
||||
for call in tool_calls:
|
||||
self.assertIn(call.function.name, available_names)
|
||||
else:
|
||||
# For non-flaky tests, enforce strict requirements
|
||||
self.assertIsNotNone(tool_calls, "Expected tool calls but got none")
|
||||
self.assertEqual(
|
||||
len(tool_calls), 2, f"Expected 2 tool calls, got {len(tool_calls)}"
|
||||
)
|
||||
|
||||
called_functions = {call.function.name for call in tool_calls}
|
||||
self.assertEqual(
|
||||
called_functions,
|
||||
expected_functions,
|
||||
f"Expected functions {expected_functions}, got {called_functions}",
|
||||
)
|
||||
|
||||
def test_multi_tool_scenario_required(self):
|
||||
"""Test multi-tool scenario with tool_choice='required'"""
|
||||
tools = self.get_travel_tools()
|
||||
messages = self.get_travel_messages()
|
||||
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
temperature=0.2,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# With required, we should get at least one tool call
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls)
|
||||
self.assertGreater(len(tool_calls), 0)
|
||||
|
||||
# Verify all called functions are available tools
|
||||
available_names = [tool["function"]["name"] for tool in tools]
|
||||
expected_functions = {"get_weather", "get_tourist_attractions"}
|
||||
|
||||
for tool_call in tool_calls:
|
||||
self.assertIsNotNone(tool_call.function.name)
|
||||
self.assertIsNotNone(tool_call.function.arguments)
|
||||
|
||||
if self._is_flaky_test():
|
||||
# For flaky tests, just ensure basic functionality works
|
||||
self.assertGreater(
|
||||
len(tool_calls),
|
||||
0,
|
||||
f"Expected at least 1 tool call, got {len(tool_calls)}",
|
||||
)
|
||||
for call in tool_calls:
|
||||
self.assertIn(call.function.name, available_names)
|
||||
else:
|
||||
# For non-flaky tests, enforce strict requirements
|
||||
self.assertEqual(
|
||||
len(tool_calls), 2, f"Expected 2 tool calls, got {len(tool_calls)}"
|
||||
)
|
||||
|
||||
called_functions = {call.function.name for call in tool_calls}
|
||||
self.assertEqual(
|
||||
called_functions,
|
||||
expected_functions,
|
||||
f"Expected functions {expected_functions}, got {called_functions}",
|
||||
)
|
||||
|
||||
def test_error_handling_invalid_tool_choice(self):
|
||||
"""Test error handling for invalid tool_choice"""
|
||||
tools = self.get_test_tools()
|
||||
messages = self.get_test_messages()
|
||||
|
||||
# Test with invalid function name
|
||||
tool_choice = {"type": "function", "function": {"name": "nonexistent_function"}}
|
||||
|
||||
# Expect a 400 BadRequestError to be raised for invalid tool_choice
|
||||
with self.assertRaises(openai.BadRequestError) as context:
|
||||
self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Verify the error message contains the expected text
|
||||
self.assertIn(
|
||||
"function 'nonexistent_function' not found in",
|
||||
str(context.exception),
|
||||
)
|
||||
|
||||
def test_invalid_tool_missing_name(self):
|
||||
"""Test what happens when user doesn't provide a tool name in request"""
|
||||
# Test with malformed JSON in tool parameters - missing required "name" field
|
||||
invalid_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
# Missing required "name" field
|
||||
"description": "Test function with invalid schema",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"test_field": {
|
||||
"type": "string",
|
||||
"description": "Test field",
|
||||
}
|
||||
},
|
||||
"required": ["test_field"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Test the function",
|
||||
}
|
||||
]
|
||||
|
||||
# Should raise BadRequestError due to missing required 'name' field
|
||||
with self.assertRaises(openai.BadRequestError) as context:
|
||||
self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=100,
|
||||
temperature=0.1,
|
||||
tools=invalid_tools,
|
||||
tool_choice="required",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Verify the error message indicates missing name field
|
||||
error_msg = str(context.exception).lower()
|
||||
self.assertIn("name", error_msg)
|
||||
|
||||
def test_conflicting_defs_required_tool_choice(self):
|
||||
"""Test that conflicting $defs with required tool_choice returns 400 error"""
|
||||
conflicting_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "tool1",
|
||||
"description": "Tool 1 with conflicting $defs",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"data": {"$ref": "#/$defs/DataType"},
|
||||
},
|
||||
"required": ["data"],
|
||||
"$defs": {
|
||||
"DataType": {
|
||||
"type": "object",
|
||||
"properties": {"value": {"type": "string"}},
|
||||
"required": ["value"],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "tool2",
|
||||
"description": "Tool 2 with conflicting $defs",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"data": {"$ref": "#/$defs/DataType"},
|
||||
},
|
||||
"required": ["data"],
|
||||
"$defs": {
|
||||
"DataType": { # Different definition for DataType
|
||||
"type": "object",
|
||||
"properties": {"value": {"type": "number"}},
|
||||
"required": ["value"],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Test the conflicting tools",
|
||||
}
|
||||
]
|
||||
|
||||
# Should raise BadRequestError due to conflicting $defs
|
||||
with self.assertRaises(openai.BadRequestError) as context:
|
||||
self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
max_tokens=100,
|
||||
temperature=0.1,
|
||||
tools=conflicting_tools,
|
||||
tool_choice="required",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Verify the error message indicates conflicting tool definitions
|
||||
error_msg = str(context.exception).lower()
|
||||
self.assertIn("invalid tool configuration", error_msg)
|
||||
self.assertIn("not supported", error_msg)
|
||||
|
||||
|
||||
class TestToolChoiceQwen25(TestToolChoiceLlama32):
|
||||
"""Test tool_choice functionality with Qwen2.5 model"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
cls.flaky_tests = {}
|
||||
|
||||
cls.model = DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
"--tool-call-parser",
|
||||
"qwen",
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
|
||||
class TestToolChoiceMistral(TestToolChoiceLlama32):
|
||||
"""Test tool_choice functionality with Mistral model"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
# Mark flaky tests for this model
|
||||
cls.flaky_tests = {
|
||||
"test_multi_tool_scenario_auto",
|
||||
"test_multi_tool_scenario_required",
|
||||
}
|
||||
|
||||
cls.model = DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
router_args=[
|
||||
"--tool-call-parser",
|
||||
"mistral",
|
||||
],
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@unittest.skip("Fails due to whitespace issue with Mistral - skipping")
|
||||
def test_complex_parameters_required_non_streaming(self):
|
||||
"""Validate complex nested parameter schemas in non-streaming required mode"""
|
||||
super().test_complex_parameters_required_non_streaming()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,13 @@
|
||||
[pytest]
|
||||
# Show print statements and logs
|
||||
log_cli = true
|
||||
log_cli_level = INFO
|
||||
log_cli_format = %(asctime)s [%(levelname)8s] %(message)s
|
||||
log_cli_date_format = %Y-%m-%d %H:%M:%S
|
||||
|
||||
# Show stdout/stderr
|
||||
addopts = -v -s --tb=short
|
||||
|
||||
# Capture settings
|
||||
# -s means don't capture stdout (show print statements)
|
||||
# --tb=short means short traceback format
|
||||
@@ -0,0 +1,260 @@
|
||||
"""
|
||||
Standalone utilities for e2e_grpc tests.
|
||||
|
||||
This module provides all necessary utilities without depending on sglang Python package.
|
||||
Extracted and adapted from:
|
||||
- sglang.srt.utils.kill_process_tree
|
||||
- sglang.srt.utils.hf_transformers_utils.get_tokenizer
|
||||
- sglang.test.test_utils (constants and CustomTestCase)
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
|
||||
import psutil
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
PreTrainedTokenizer,
|
||||
PreTrainedTokenizerBase,
|
||||
PreTrainedTokenizerFast,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"transformers is required for tokenizer utilities. "
|
||||
"Install with: pip install transformers"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Constants
|
||||
# ============================================================================
|
||||
|
||||
# Server and timeout constants
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH = 600
|
||||
DEFAULT_PORT_FOR_SRT_TEST_RUNNER = 20000
|
||||
DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
|
||||
|
||||
# File name constants for test output
|
||||
STDOUT_FILENAME = "/tmp/sglang_test_stdout.txt"
|
||||
STDERR_FILENAME = "/tmp/sglang_test_stderr.txt"
|
||||
|
||||
# Model base path - can be overridden via environment variable
|
||||
# By default, use HuggingFace model identifiers (no local path prefix)
|
||||
# Set ROUTER_LOCAL_MODEL_PATH to use local models (e.g., "/home/ubuntu/models")
|
||||
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
|
||||
|
||||
|
||||
# Helper function to build model paths
|
||||
def _get_model_path(model_identifier: str) -> str:
|
||||
"""
|
||||
Build model path from base path and model identifier.
|
||||
|
||||
If ROUTER_LOCAL_MODEL_PATH is set, prepend it to the identifier.
|
||||
Otherwise, return the identifier as-is (for HuggingFace download).
|
||||
"""
|
||||
if ROUTER_LOCAL_MODEL_PATH:
|
||||
return os.path.join(ROUTER_LOCAL_MODEL_PATH, model_identifier)
|
||||
return model_identifier
|
||||
|
||||
|
||||
# Model paths used in e2e_grpc tests
|
||||
# These can be either HuggingFace identifiers or local paths (depending on ROUTER_LOCAL_MODEL_PATH)
|
||||
|
||||
# Main test model - Llama 3.1 8B Instruct
|
||||
DEFAULT_MODEL_PATH = _get_model_path("meta-llama/Llama-3.1-8B-Instruct")
|
||||
|
||||
# Small models for function calling tests
|
||||
DEFAULT_SMALL_MODEL_PATH = _get_model_path("meta-llama/Llama-3.2-1B-Instruct")
|
||||
|
||||
# Reasoning models
|
||||
DEFAULT_REASONING_MODEL_PATH = _get_model_path(
|
||||
"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
|
||||
)
|
||||
|
||||
# Thinking-enabled models
|
||||
DEFAULT_ENABLE_THINKING_MODEL_PATH = _get_model_path("Qwen/Qwen3-30B-A3B")
|
||||
|
||||
# Function calling models
|
||||
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH = _get_model_path("Qwen/Qwen2.5-7B-Instruct")
|
||||
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH = _get_model_path(
|
||||
"mistralai/Mistral-7B-Instruct-v0.3"
|
||||
)
|
||||
|
||||
# GPT-OSS models
|
||||
DEFAULT_GPT_OSS_MODEL_PATH = _get_model_path("openai/gpt-oss-20b")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Process Management
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
|
||||
"""
|
||||
Kill the process and all its child processes.
|
||||
|
||||
Args:
|
||||
parent_pid: PID of the parent process
|
||||
include_parent: Whether to kill the parent process itself
|
||||
skip_pid: Optional PID to skip during cleanup
|
||||
"""
|
||||
# Remove sigchld handler to avoid spammy logs
|
||||
if threading.current_thread() is threading.main_thread():
|
||||
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
|
||||
|
||||
if parent_pid is None:
|
||||
parent_pid = os.getpid()
|
||||
include_parent = False
|
||||
|
||||
try:
|
||||
itself = psutil.Process(parent_pid)
|
||||
except psutil.NoSuchProcess:
|
||||
return
|
||||
|
||||
children = itself.children(recursive=True)
|
||||
for child in children:
|
||||
if child.pid == skip_pid:
|
||||
continue
|
||||
try:
|
||||
child.kill()
|
||||
except psutil.NoSuchProcess:
|
||||
pass
|
||||
|
||||
if include_parent:
|
||||
try:
|
||||
itself.kill()
|
||||
except psutil.NoSuchProcess:
|
||||
pass
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tokenizer Utilities
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def check_gguf_file(model_path: str) -> bool:
|
||||
"""Check if the model path points to a GGUF file."""
|
||||
if not isinstance(model_path, str):
|
||||
return False
|
||||
return model_path.endswith(".gguf")
|
||||
|
||||
|
||||
def is_remote_url(path: str) -> bool:
|
||||
"""Check if the path is a remote URL."""
|
||||
if not isinstance(path, str):
|
||||
return False
|
||||
return path.startswith("http://") or path.startswith("https://")
|
||||
|
||||
|
||||
def get_tokenizer(
|
||||
tokenizer_name: str,
|
||||
*args,
|
||||
tokenizer_mode: str = "auto",
|
||||
trust_remote_code: bool = False,
|
||||
tokenizer_revision: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
|
||||
"""
|
||||
Gets a tokenizer for the given model name via Huggingface.
|
||||
|
||||
Args:
|
||||
tokenizer_name: Name or path of the tokenizer
|
||||
tokenizer_mode: Mode for tokenizer loading ("auto", "slow")
|
||||
trust_remote_code: Whether to trust remote code
|
||||
tokenizer_revision: Specific revision to use
|
||||
**kwargs: Additional arguments passed to AutoTokenizer.from_pretrained
|
||||
|
||||
Returns:
|
||||
Loaded tokenizer instance
|
||||
"""
|
||||
if tokenizer_mode == "slow":
|
||||
if kwargs.get("use_fast", False):
|
||||
raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
|
||||
kwargs["use_fast"] = False
|
||||
|
||||
# Handle special model name mapping
|
||||
if tokenizer_name == "mistralai/Devstral-Small-2505":
|
||||
tokenizer_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
|
||||
|
||||
is_gguf = check_gguf_file(tokenizer_name)
|
||||
if is_gguf:
|
||||
kwargs["gguf_file"] = tokenizer_name
|
||||
tokenizer_name = Path(tokenizer_name).parent
|
||||
|
||||
# Note: Removed remote URL handling and local directory download
|
||||
# as they depend on sglang-specific utilities
|
||||
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_name,
|
||||
*args,
|
||||
trust_remote_code=trust_remote_code,
|
||||
tokenizer_revision=tokenizer_revision,
|
||||
**kwargs,
|
||||
)
|
||||
except TypeError as e:
|
||||
# Handle specific errors
|
||||
err_msg = (
|
||||
"Failed to load the tokenizer. If you are running a model with "
|
||||
"a custom tokenizer, please set the --trust-remote-code flag."
|
||||
)
|
||||
raise RuntimeError(err_msg) from e
|
||||
|
||||
if not isinstance(tokenizer, PreTrainedTokenizerFast):
|
||||
logger.warning(
|
||||
f"Using a slow tokenizer. This might cause a performance "
|
||||
f"degradation. Consider using a fast tokenizer instead."
|
||||
)
|
||||
|
||||
return tokenizer
|
||||
|
||||
|
||||
def get_tokenizer_from_processor(processor):
|
||||
"""Extract tokenizer from a processor object."""
|
||||
if isinstance(processor, PreTrainedTokenizerBase):
|
||||
return processor
|
||||
return processor.tokenizer
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test Utilities
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class CustomTestCase(unittest.TestCase):
|
||||
"""
|
||||
Custom test case base class with retry support.
|
||||
|
||||
This provides automatic test retry functionality based on environment variables.
|
||||
"""
|
||||
|
||||
def _callTestMethod(self, method):
|
||||
"""Override to add retry logic."""
|
||||
max_retry = int(os.environ.get("SGLANG_TEST_MAX_RETRY", "0"))
|
||||
|
||||
if max_retry == 0:
|
||||
# No retry, just run once
|
||||
return super(CustomTestCase, self)._callTestMethod(method)
|
||||
|
||||
# Retry logic
|
||||
for attempt in range(max_retry + 1):
|
||||
try:
|
||||
return super(CustomTestCase, self)._callTestMethod(method)
|
||||
except Exception as e:
|
||||
if attempt < max_retry:
|
||||
logger.info(
|
||||
f"Test failed on attempt {attempt + 1}/{max_retry + 1}, retrying..."
|
||||
)
|
||||
continue
|
||||
else:
|
||||
# Last attempt, re-raise the exception
|
||||
raise
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
python3 -m unittest openai_server.validation.test_large_max_new_tokens.TestLargeMaxNewTokens.test_chat_completion
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_MODEL_PATH,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
STDERR_FILENAME,
|
||||
STDOUT_FILENAME,
|
||||
CustomTestCase,
|
||||
get_tokenizer,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestLargeMaxNewTokens(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
cls.stdout = open(STDOUT_FILENAME, "w")
|
||||
cls.stderr = open(STDERR_FILENAME, "w")
|
||||
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
worker_args=(
|
||||
"--max-total-token",
|
||||
"1536",
|
||||
"--context-len",
|
||||
"8192",
|
||||
"--decode-log-interval",
|
||||
"2",
|
||||
),
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
env={"SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION": "256", **os.environ},
|
||||
stdout=cls.stdout,
|
||||
stderr=cls.stderr,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
cls.stdout.close()
|
||||
cls.stderr.close()
|
||||
os.remove(STDOUT_FILENAME)
|
||||
os.remove(STDERR_FILENAME)
|
||||
|
||||
def run_chat_completion(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Please repeat the world 'hello' for 10000 times.",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
)
|
||||
return response
|
||||
|
||||
def test_chat_completion(self):
|
||||
num_requests = 4
|
||||
all_requests_running = False
|
||||
|
||||
futures = []
|
||||
with ThreadPoolExecutor(num_requests) as executor:
|
||||
# Send multiple requests
|
||||
for i in range(num_requests):
|
||||
futures.append(executor.submit(self.run_chat_completion))
|
||||
|
||||
# Ensure that they are running concurrently
|
||||
pt = 0
|
||||
while pt >= 0:
|
||||
time.sleep(5)
|
||||
# Flush stderr to ensure logs are written
|
||||
self.stderr.flush()
|
||||
lines = open(STDERR_FILENAME).readlines()
|
||||
for line in lines[pt:]:
|
||||
if f"#running-req: {num_requests}" in line:
|
||||
all_requests_running = True
|
||||
pt = -1
|
||||
break
|
||||
pt += 1
|
||||
|
||||
assert all_requests_running
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
gRPC Router E2E Test - Test Openai Server Ignore Eos
|
||||
|
||||
This test file is REUSED from test/srt/openai_server/validation/test_openai_server_ignore_eos.py
|
||||
with minimal changes:
|
||||
num_workers=2,
|
||||
- Swap popen_launch_server() → popen_launch_workers_and_router()
|
||||
- Update teardown to cleanup router + workers
|
||||
- All test logic and assertions remain identical
|
||||
|
||||
Run with:
|
||||
pytest py_test/e2e_grpc/e2e_grpc/validation/test_openai_server_ignore_eos.py -v
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR.parent))
|
||||
from fixtures import popen_launch_workers_and_router
|
||||
from util import (
|
||||
DEFAULT_MODEL_PATH,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
get_tokenizer,
|
||||
kill_process_tree,
|
||||
)
|
||||
|
||||
|
||||
class TestOpenAIServerIgnoreEOS(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# CHANGE: Launch gRPC router with integrated workers (single command)
|
||||
cls.model = DEFAULT_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.cluster = popen_launch_workers_and_router(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
num_workers=1,
|
||||
tp_size=2,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
# Cleanup router and workers
|
||||
kill_process_tree(cls.cluster["router"].pid)
|
||||
for worker in cls.cluster.get("workers", []):
|
||||
kill_process_tree(worker.pid)
|
||||
|
||||
def test_ignore_eos(self):
|
||||
"""
|
||||
Test that ignore_eos=True allows generation to continue beyond EOS token
|
||||
and reach the max_tokens limit.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
max_tokens = 200
|
||||
|
||||
response_default = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Count from 1 to 20."},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=max_tokens,
|
||||
extra_body={"ignore_eos": False},
|
||||
)
|
||||
|
||||
response_ignore_eos = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Count from 1 to 20."},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=max_tokens,
|
||||
extra_body={"ignore_eos": True},
|
||||
)
|
||||
|
||||
default_tokens = len(
|
||||
self.tokenizer.encode(response_default.choices[0].message.content)
|
||||
)
|
||||
ignore_eos_tokens = len(
|
||||
self.tokenizer.encode(response_ignore_eos.choices[0].message.content)
|
||||
)
|
||||
|
||||
# Check if ignore_eos resulted in more tokens or exactly max_tokens
|
||||
# The ignore_eos response should either:
|
||||
# 1. Have more tokens than the default response (if default stopped at EOS before max_tokens)
|
||||
# 2. Have exactly max_tokens (if it reached the max_tokens limit)
|
||||
self.assertTrue(
|
||||
ignore_eos_tokens > default_tokens or ignore_eos_tokens >= max_tokens,
|
||||
f"ignore_eos did not generate more tokens: {ignore_eos_tokens} vs {default_tokens}",
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
response_ignore_eos.choices[0].finish_reason,
|
||||
"length",
|
||||
f"Expected finish_reason='length' for ignore_eos=True, got {response_ignore_eos.choices[0].finish_reason}",
|
||||
)
|
||||
@@ -0,0 +1,807 @@
|
||||
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",
|
||||
"--router-log-level",
|
||||
"warn",
|
||||
]
|
||||
|
||||
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),
|
||||
"--log-level",
|
||||
"warning",
|
||||
]
|
||||
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,
|
||||
enable_igw: 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 enable_igw:
|
||||
cmd += ["--enable-igw"]
|
||||
if api_key is not None:
|
||||
cmd += ["--api-key", api_key]
|
||||
cmd += [
|
||||
"--prometheus-port",
|
||||
str(prom_port),
|
||||
"--prometheus-host",
|
||||
"127.0.0.1",
|
||||
"--log-level",
|
||||
"warn",
|
||||
]
|
||||
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
|
||||
|
||||
|
||||
def _gpu_monitor_should_run(thresholds: Optional[dict]) -> bool:
|
||||
"""Decide whether to enable the GPU monitor.
|
||||
|
||||
Runs if thresholds request GPU checks or if GPU_UTIL_LOG is truthy.
|
||||
"""
|
||||
want = False
|
||||
try:
|
||||
mean_th = None if thresholds is None else thresholds.get("gpu_util_mean_min")
|
||||
p50_th = None if thresholds is None else thresholds.get("gpu_util_p50_min")
|
||||
want = bool(mean_th is not None or p50_th is not None)
|
||||
except Exception:
|
||||
want = False
|
||||
if not want:
|
||||
env_flag = os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")
|
||||
want = want or env_flag
|
||||
return want
|
||||
|
||||
|
||||
def _gpu_monitor_path(experiment_folder: str) -> str:
|
||||
"""Return the JSON path for storing GPU monitor results."""
|
||||
base = Path.cwd() / experiment_folder
|
||||
return str(base / "gpu_utilization.json")
|
||||
|
||||
|
||||
def _launch_gpu_monitor(bench_pid: int, experiment_folder: str, interval: float):
|
||||
"""Start the GPU monitor process. Returns (proc, path) or (None, None)."""
|
||||
try:
|
||||
from multiprocessing import Process
|
||||
|
||||
out_path = _gpu_monitor_path(experiment_folder)
|
||||
proc = Process(
|
||||
target=_gpu_monitor_proc_entry,
|
||||
args=(bench_pid, out_path, interval),
|
||||
daemon=True,
|
||||
)
|
||||
proc.start()
|
||||
return proc, out_path
|
||||
except Exception as e:
|
||||
logger.warning("Failed to launch GPU monitor: %s", e)
|
||||
return None, None
|
||||
|
||||
|
||||
def _read_gpu_monitor_result(path: Optional[str]) -> Optional[dict]:
|
||||
try:
|
||||
if path and os.path.exists(path):
|
||||
with open(path, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read GPU monitor result from %r: %s", path, e)
|
||||
return None
|
||||
|
||||
|
||||
def _log_and_assert_gpu_thresholds(
|
||||
result: Optional[dict], thresholds: Optional[dict]
|
||||
) -> None:
|
||||
if not result or not isinstance(result, dict) or result.get("count", 0) <= 0:
|
||||
logger.warning("GPU utilization monitor produced no samples.")
|
||||
return
|
||||
|
||||
overall = result.get("overall", {}) if isinstance(result, dict) else {}
|
||||
count = int(result.get("count", 0))
|
||||
mean_th = None if thresholds is None else thresholds.get("gpu_util_mean_min")
|
||||
p50_th = None if thresholds is None else thresholds.get("gpu_util_p50_min")
|
||||
|
||||
mean_v = float(overall.get("mean", 0.0))
|
||||
p50_v = overall.get("p50")
|
||||
|
||||
logger.info(
|
||||
"GPU utilization overall: mean=%.2f%% p50=%s (samples=%d)",
|
||||
mean_v,
|
||||
(f"{float(p50_v):.2f}%" if p50_v is not None else "n/a"),
|
||||
count,
|
||||
)
|
||||
|
||||
if mean_th is not None:
|
||||
assert mean_v >= float(
|
||||
mean_th
|
||||
), f"GPU utilization mean below threshold: {mean_v:.2f}% < {mean_th}%"
|
||||
if p50_th is not None and p50_v is not None:
|
||||
p50_f = float(p50_v)
|
||||
assert p50_f >= float(
|
||||
p50_th
|
||||
), f"GPU utilization p50 below threshold: {p50_f:.2f}% < {p50_th}%"
|
||||
|
||||
|
||||
def _gpu_monitor_proc_entry(bench_pid: int, out_file: str, interval: float) -> None:
|
||||
"""Low-impact GPU utilization monitor using NVML in a separate process.
|
||||
|
||||
Writes JSON to out_file that includes overall and per-GPU raw samples and summary stats.
|
||||
"""
|
||||
try:
|
||||
try:
|
||||
os.nice(10)
|
||||
except Exception:
|
||||
pass
|
||||
total = 0.0
|
||||
n = 0
|
||||
try:
|
||||
import pynvml # type: ignore
|
||||
|
||||
pynvml.nvmlInit()
|
||||
except Exception:
|
||||
with open(out_file, "w") as f:
|
||||
os.makedirs(os.path.dirname(out_file), exist_ok=True)
|
||||
json.dump(
|
||||
{
|
||||
"count": 0,
|
||||
"overall": {"mean": 0.0},
|
||||
"per_gpu": {},
|
||||
"raw": {},
|
||||
},
|
||||
f,
|
||||
)
|
||||
return
|
||||
try:
|
||||
import pynvml # type: ignore
|
||||
|
||||
count = pynvml.nvmlDeviceGetCount()
|
||||
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
|
||||
except Exception:
|
||||
with open(out_file, "w") as f:
|
||||
os.makedirs(os.path.dirname(out_file), exist_ok=True)
|
||||
json.dump(
|
||||
{
|
||||
"count": 0,
|
||||
"overall": {"mean": 0.0},
|
||||
"per_gpu": {},
|
||||
"raw": {},
|
||||
},
|
||||
f,
|
||||
)
|
||||
return
|
||||
|
||||
# Prepare per-GPU and overall raw collectors
|
||||
per_gpu_samples: dict[str, list[float]] = {}
|
||||
overall_samples: list[float] = []
|
||||
|
||||
while True:
|
||||
if not os.path.exists(f"/proc/{bench_pid}"):
|
||||
break
|
||||
try:
|
||||
vals = []
|
||||
import pynvml # type: ignore
|
||||
|
||||
for idx, h in enumerate(handles):
|
||||
try:
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(h).gpu
|
||||
vals.append(float(util))
|
||||
key = str(idx)
|
||||
per_gpu_samples.setdefault(key, []).append(float(util))
|
||||
except Exception:
|
||||
continue
|
||||
if vals:
|
||||
avg = sum(vals) / len(vals)
|
||||
overall_samples.append(avg)
|
||||
total += avg
|
||||
n += 1
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(interval)
|
||||
finally:
|
||||
try:
|
||||
os.makedirs(os.path.dirname(out_file), exist_ok=True)
|
||||
with open(out_file, "w") as f:
|
||||
|
||||
def pct_from(samples: list[float], p: float) -> float:
|
||||
if not samples:
|
||||
return 0.0
|
||||
srt = sorted(samples)
|
||||
i = max(
|
||||
0, min(len(srt) - 1, int(round((p / 100.0) * (len(srt) - 1))))
|
||||
)
|
||||
return float(srt[i])
|
||||
|
||||
overall_mean = (total / n) if n > 0 else 0.0
|
||||
|
||||
per_gpu_summary: dict[str, dict] = {}
|
||||
for key, arr in per_gpu_samples.items():
|
||||
per_gpu_summary[key] = {
|
||||
"mean": float(sum(arr) / len(arr)) if arr else 0.0,
|
||||
"p5": pct_from(arr, 5),
|
||||
"p10": pct_from(arr, 10),
|
||||
"p25": pct_from(arr, 25),
|
||||
"p50": pct_from(arr, 50),
|
||||
"p75": pct_from(arr, 75),
|
||||
"p90": pct_from(arr, 90),
|
||||
"p95": pct_from(arr, 95),
|
||||
"min": float(min(arr)) if arr else 0.0,
|
||||
"max": float(max(arr)) if arr else 0.0,
|
||||
"count": len(arr),
|
||||
}
|
||||
|
||||
out_payload = {
|
||||
"bench_pid": bench_pid,
|
||||
"interval_sec": interval,
|
||||
"count": n,
|
||||
"overall": {
|
||||
"mean": float(overall_mean),
|
||||
"p5": pct_from(overall_samples, 5),
|
||||
"p10": pct_from(overall_samples, 10),
|
||||
"p25": pct_from(overall_samples, 25),
|
||||
"p50": pct_from(overall_samples, 50),
|
||||
"p75": pct_from(overall_samples, 75),
|
||||
"p90": pct_from(overall_samples, 90),
|
||||
"p95": pct_from(overall_samples, 95),
|
||||
"min": float(min(overall_samples)) if overall_samples else 0.0,
|
||||
"max": float(max(overall_samples)) if overall_samples else 0.0,
|
||||
},
|
||||
"per_gpu": per_gpu_summary,
|
||||
"raw": {
|
||||
"overall": overall_samples,
|
||||
"per_gpu": per_gpu_samples,
|
||||
},
|
||||
}
|
||||
json.dump(out_payload, f)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import pynvml # type: ignore
|
||||
|
||||
pynvml.nvmlShutdown()
|
||||
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 * 5
|
||||
)
|
||||
|
||||
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",
|
||||
"3",
|
||||
"--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
|
||||
)
|
||||
# Optional GPU utilization monitor in a low-priority child process (pynvml only)
|
||||
# Enabled only when gpu_util_mean_min is provided in thresholds.
|
||||
monitor_path = None
|
||||
monitor_proc = None
|
||||
gpu_util_result: dict | None = None
|
||||
want_gpu_monitor = _gpu_monitor_should_run(thresholds)
|
||||
if want_gpu_monitor:
|
||||
interval = float(os.environ.get("GPU_UTIL_SAMPLE_INTERVAL", "2.0"))
|
||||
monitor_proc, monitor_path = _launch_gpu_monitor(
|
||||
bench_pid=proc.pid,
|
||||
experiment_folder=experiment_folder,
|
||||
interval=interval,
|
||||
)
|
||||
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 = []
|
||||
for _ in range(10):
|
||||
json_files = [
|
||||
p
|
||||
for p in actual_folder.rglob("*.json")
|
||||
if "experiment_metadata" not in p.name
|
||||
]
|
||||
if json_files:
|
||||
break
|
||||
time.sleep(1)
|
||||
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})"
|
||||
|
||||
# Validate optional GPU utilization threshold if provided
|
||||
if want_gpu_monitor:
|
||||
try:
|
||||
if monitor_proc is not None:
|
||||
monitor_proc.join(timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
gpu_util_result = _read_gpu_monitor_result(monitor_path)
|
||||
_log_and_assert_gpu_thresholds(gpu_util_result, thresholds)
|
||||
|
||||
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
|
||||
# Ensure GPU monitor process is cleaned up
|
||||
if monitor_proc is not None and monitor_proc.is_alive():
|
||||
try:
|
||||
monitor_proc.terminate()
|
||||
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 os.getenv("E2E_PRIMARY_MODEL", 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_embedding_model() -> str:
|
||||
"""Embedding model to use for E2E tests.
|
||||
|
||||
Defaults to an E5 Mistral model, can be overridden via E2E_EMBEDDING_MODEL env var.
|
||||
"""
|
||||
import os
|
||||
|
||||
return os.getenv("E2E_EMBEDDING_MODEL", "intfloat/e5-mistral-7b-instruct")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def e2e_primary_embedding_worker(e2e_embedding_model: str):
|
||||
"""Launch a single embedding worker using the specified model."""
|
||||
port = _find_available_port()
|
||||
base_url = f"http://127.0.0.1:{port}"
|
||||
proc = _popen_launch_worker(e2e_embedding_model, base_url)
|
||||
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)
|
||||
@@ -0,0 +1,62 @@
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
|
||||
def _wait_for_workers(
|
||||
base_url: str, expected_count: int, timeout: float = 60.0, headers: dict = None
|
||||
) -> None:
|
||||
"""Poll /workers endpoint until expected number of workers are registered."""
|
||||
start = time.perf_counter()
|
||||
with requests.Session() as session:
|
||||
while time.perf_counter() - start < timeout:
|
||||
try:
|
||||
r = session.get(f"{base_url}/workers", headers=headers, timeout=5)
|
||||
if r.status_code == 200:
|
||||
workers = r.json().get("workers", [])
|
||||
if len(workers) >= expected_count:
|
||||
return
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(
|
||||
f"Expected {expected_count} workers at {base_url}, timed out after {timeout}s"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
def test_embeddings_basic(
|
||||
e2e_router_only_rr, e2e_primary_embedding_worker, e2e_embedding_model
|
||||
):
|
||||
base = e2e_router_only_rr.url
|
||||
worker_url = e2e_primary_embedding_worker.url
|
||||
|
||||
# Attach embedding worker to router-only instance
|
||||
r = requests.post(f"{base}/workers", json={"url": worker_url}, timeout=180)
|
||||
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
# Wait for worker to be registered
|
||||
_wait_for_workers(base, expected_count=1, timeout=60.0)
|
||||
|
||||
# Simple embedding request with two inputs
|
||||
payload = {
|
||||
"model": e2e_embedding_model,
|
||||
"input": [
|
||||
"the quick brown fox",
|
||||
"jumps over the lazy dog",
|
||||
],
|
||||
}
|
||||
r = requests.post(f"{base}/v1/embeddings", json=payload, timeout=120)
|
||||
|
||||
assert r.status_code == 200, f"unexpected status: {r.status_code} {r.text}"
|
||||
|
||||
data = r.json()
|
||||
assert "data" in data and isinstance(data["data"], list)
|
||||
assert len(data["data"]) == 2
|
||||
|
||||
# Validate shape of embedding objects
|
||||
for item in data["data"]:
|
||||
assert "embedding" in item and isinstance(item["embedding"], list)
|
||||
# Ensure non-empty vectors
|
||||
assert len(item["embedding"]) > 0
|
||||
@@ -0,0 +1,264 @@
|
||||
import logging
|
||||
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",
|
||||
"--log-level",
|
||||
"warn",
|
||||
]
|
||||
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": 13,
|
||||
"e2e_latency_mean_max": 16,
|
||||
"input_throughput_mean_min": 350,
|
||||
"output_throughput_mean_min": 18,
|
||||
"gpu_util_p50_min": 99,
|
||||
},
|
||||
kill_procs=pd_cluster.workers,
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from sglang.test.run_eval import run_eval
|
||||
|
||||
|
||||
def _wait_for_workers(
|
||||
base_url: str, expected_count: int, timeout: float = 60.0, headers: dict = None
|
||||
) -> None:
|
||||
"""Poll /workers endpoint until expected number of workers are registered."""
|
||||
start = time.perf_counter()
|
||||
with requests.Session() as session:
|
||||
while time.perf_counter() - start < timeout:
|
||||
try:
|
||||
r = session.get(f"{base_url}/workers", headers=headers, timeout=5)
|
||||
if r.status_code == 200:
|
||||
workers = r.json().get("workers", [])
|
||||
if len(workers) >= expected_count:
|
||||
return
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(
|
||||
f"Expected {expected_count} workers at {base_url}, timed out after {timeout}s"
|
||||
)
|
||||
|
||||
|
||||
@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}/workers", json={"url": w.url}, timeout=180)
|
||||
assert (
|
||||
r.status_code == 202
|
||||
), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
# Wait for workers to be registered
|
||||
_wait_for_workers(base, expected_count=2, timeout=60.0)
|
||||
|
||||
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}/workers", json={"url": w.url}, timeout=180)
|
||||
assert (
|
||||
r.status_code == 202
|
||||
), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
# Wait for workers to be registered
|
||||
_wait_for_workers(base, expected_count=2, timeout=60.0)
|
||||
|
||||
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": 800, # temp relax from 1000 to 800 for now
|
||||
"output_throughput_mean_min": 12,
|
||||
# Enforce GPU utilization p50 >= 99% during the run.
|
||||
"gpu_util_p50_min": 99,
|
||||
},
|
||||
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}/workers", json={"url": worker_url}, timeout=180)
|
||||
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
# Wait for worker to be registered
|
||||
_wait_for_workers(base, expected_count=1, timeout=60.0)
|
||||
|
||||
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
|
||||
from urllib.parse import quote
|
||||
|
||||
encoded_url = quote(worker_url, safe="")
|
||||
r = requests.delete(f"{base}/workers/{encoded_url}", timeout=60)
|
||||
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
|
||||
@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}/workers", json={"url": worker.url}, timeout=180)
|
||||
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
# Wait for worker to be registered
|
||||
_wait_for_workers(base, expected_count=1, timeout=60.0)
|
||||
|
||||
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}/workers",
|
||||
json={"url": worker_url, "api_key": api_key},
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=180,
|
||||
)
|
||||
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
|
||||
|
||||
# Wait for workers to be registered and expanded
|
||||
_wait_for_workers(
|
||||
router_url,
|
||||
expected_count=2,
|
||||
timeout=60.0,
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
)
|
||||
|
||||
# Verify the expanded workers have correct URLs
|
||||
r = requests.get(
|
||||
f"{router_url}/workers",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=30,
|
||||
)
|
||||
r.raise_for_status()
|
||||
workers = r.json().get("workers", [])
|
||||
urls = [w["url"] for w in workers]
|
||||
assert len(urls) == 2
|
||||
assert set(urls) == {f"{worker_url}@0", f"{worker_url}@1"}
|
||||
|
||||
# Verify API key enforcement
|
||||
# 1) Without Authorization -> Should get 401 Unauthorized
|
||||
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
|
||||
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
pytest configuration for e2e_response_api tests.
|
||||
|
||||
This configures pytest to not collect base test classes that are meant to be inherited.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import openai
|
||||
import pytest # noqa: F401
|
||||
from router_fixtures import (
|
||||
popen_launch_openai_xai_router,
|
||||
popen_launch_workers_and_router,
|
||||
)
|
||||
from util import kill_process_tree
|
||||
|
||||
# ------------------------------
|
||||
# Backend Configuration Map
|
||||
# ------------------------------
|
||||
BACKENDS = {
|
||||
"openai": {
|
||||
"model": "gpt-5-nano",
|
||||
"base_url_port": "http://127.0.0.1:30010",
|
||||
"launcher": popen_launch_openai_xai_router,
|
||||
"launcher_kwargs": {
|
||||
"backend": "openai",
|
||||
"history_backend": "memory",
|
||||
},
|
||||
"api_key_env": "OPENAI_API_KEY",
|
||||
"needs_workers": False,
|
||||
},
|
||||
"xai": {
|
||||
"model": "grok-4-fast",
|
||||
"base_url_port": "http://127.0.0.1:30023",
|
||||
"launcher": popen_launch_openai_xai_router,
|
||||
"launcher_kwargs": {
|
||||
"backend": "xai",
|
||||
"history_backend": "memory",
|
||||
},
|
||||
"api_key_env": "XAI_API_KEY",
|
||||
"needs_workers": False,
|
||||
},
|
||||
"grpc": {
|
||||
"model": "/home/ubuntu/models/Qwen/Qwen2.5-14B-Instruct",
|
||||
"base_url_port": "http://127.0.0.1:30030",
|
||||
"launcher": popen_launch_workers_and_router,
|
||||
"launcher_kwargs": {
|
||||
"timeout": 90,
|
||||
"num_workers": 1,
|
||||
"tp_size": 2,
|
||||
"policy": "round_robin",
|
||||
"worker_args": ["--context-length=1000"],
|
||||
"router_args": [
|
||||
"--history-backend",
|
||||
"memory",
|
||||
"--tool-call-parser",
|
||||
"qwen",
|
||||
],
|
||||
},
|
||||
"api_key_env": None, # grpc does not use API keys
|
||||
"needs_workers": True,
|
||||
},
|
||||
"grpc_harmony": {
|
||||
"model": "/home/ubuntu/models/openai/gpt-oss-20b",
|
||||
"base_url_port": "http://127.0.0.1:30030",
|
||||
"launcher": popen_launch_workers_and_router,
|
||||
"launcher_kwargs": {
|
||||
"timeout": 90,
|
||||
"num_workers": 1,
|
||||
"tp_size": 2,
|
||||
"policy": "round_robin",
|
||||
"worker_args": ["--reasoning-parser=gpt-oss"],
|
||||
"router_args": ["--history-backend", "memory"],
|
||||
},
|
||||
"api_key_env": None,
|
||||
"needs_workers": True,
|
||||
},
|
||||
"oracle_store": {
|
||||
"model": "gpt-5-nano",
|
||||
"base_url_port": "http://127.0.0.1:30040",
|
||||
"launcher": popen_launch_openai_xai_router,
|
||||
"launcher_kwargs": {
|
||||
"backend": "openai",
|
||||
"history_backend": "oracle",
|
||||
},
|
||||
"api_key_env": "OPENAI_API_KEY",
|
||||
"needs_workers": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def setup_backend(request):
|
||||
backend = request.param
|
||||
if backend not in BACKENDS:
|
||||
raise RuntimeError(f"Unknown backend {backend}")
|
||||
|
||||
cfg = BACKENDS[backend]
|
||||
|
||||
# Launch cluster
|
||||
cluster = (
|
||||
cfg["launcher"](
|
||||
cfg["model"],
|
||||
cfg["base_url_port"],
|
||||
**cfg["launcher_kwargs"],
|
||||
)
|
||||
if cfg["launcher"] is popen_launch_workers_and_router
|
||||
else cfg["launcher"](
|
||||
backend=cfg["launcher_kwargs"]["backend"],
|
||||
base_url=cfg["base_url_port"],
|
||||
history_backend=cfg["launcher_kwargs"]["history_backend"],
|
||||
)
|
||||
)
|
||||
|
||||
# Build client
|
||||
api_key = os.environ.get(cfg["api_key_env"]) if cfg["api_key_env"] else None
|
||||
client = openai.Client(
|
||||
api_key=api_key,
|
||||
base_url=cluster["base_url"] + "/v1",
|
||||
)
|
||||
|
||||
# Yield data to test
|
||||
try:
|
||||
yield backend, cfg["model"], client
|
||||
finally:
|
||||
# Always kill router
|
||||
kill_process_tree(cluster["router"].pid)
|
||||
|
||||
# If workers exist, kill them as well
|
||||
if cfg["needs_workers"]:
|
||||
for w in cluster.get("workers", []):
|
||||
kill_process_tree(w.pid)
|
||||
@@ -0,0 +1,241 @@
|
||||
"""
|
||||
Base test class for Response API e2e tests.
|
||||
|
||||
This module provides base test classes that can be reused across different backends
|
||||
(OpenAI, XAI, gRPC) with common test logic.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from openai import OpenAI
|
||||
from openai.types import responses
|
||||
|
||||
# Add current directory for local imports
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setup_backend", ["openai", "oracle_store"], indirect=True)
|
||||
class TestResponseCRUD:
|
||||
"""Base class for Response API CRUD tests."""
|
||||
|
||||
def test_create_and_get_response(self, setup_backend):
|
||||
"""Test creating response and retrieving it."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create response
|
||||
create_resp = client.responses.create(model=model, input="Hello, world!")
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.error is None
|
||||
assert create_resp.status == "completed"
|
||||
assert len(create_resp.output_text) > 0
|
||||
response_id = create_resp.id
|
||||
|
||||
# Get response
|
||||
get_resp = client.responses.retrieve(response_id=response_id)
|
||||
assert get_resp.error is None
|
||||
assert get_resp.id == response_id
|
||||
assert get_resp.status == "completed"
|
||||
|
||||
input_resp = client.responses.input_items.list(response_id=get_resp.id)
|
||||
assert input_resp.data is not None
|
||||
assert len(input_resp.data) > 0
|
||||
|
||||
@pytest.mark.skip(reason="TODO: Add delete response feature")
|
||||
def test_delete_response(self, setup_backend):
|
||||
"""Test deleting response."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create response
|
||||
create_resp = client.responses.create(model=model, input="Test deletion")
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.error is None
|
||||
assert create_resp.status == "completed"
|
||||
assert len(create_resp.output_text) > 0
|
||||
|
||||
response_id = create_resp.id
|
||||
|
||||
# Delete response
|
||||
client.responses.delete(response_id=response_id)
|
||||
|
||||
# Verify it's deleted (should return 404)
|
||||
with pytest.raises(openai.NotFoundError):
|
||||
client.responses.retrieve(response_id=response_id)
|
||||
|
||||
@pytest.mark.skip(reason="TODO: Add background response feature")
|
||||
def test_background_response(self, setup_backend):
|
||||
"""Test background response execution."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create background response
|
||||
create_resp = client.responses.create(
|
||||
model=model,
|
||||
input="Write a short story",
|
||||
background=True,
|
||||
max_output_tokens=100,
|
||||
)
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.error is None
|
||||
assert create_resp.status in ["in_progress", "queued"]
|
||||
|
||||
response_id = create_resp.id
|
||||
|
||||
# Wait for completion
|
||||
final_data = wait_for_background_task(client, response_id, timeout=60)
|
||||
assert final_data.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setup_backend", ["openai", "oracle_store"], indirect=True)
|
||||
class TestConversationCRUD:
|
||||
"""Base class for Conversation API CRUD tests."""
|
||||
|
||||
def test_create_and_get_conversation(self, setup_backend):
|
||||
"""Test creating and retrieving conversation."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create conversation
|
||||
create_resp = client.conversations.create(metadata={"user": "test_user"})
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.created_at is not None
|
||||
|
||||
create_data = create_resp.metadata
|
||||
assert create_data["user"] == "test_user"
|
||||
conversation_id = create_resp.id
|
||||
|
||||
# Get conversation
|
||||
get_resp = client.conversations.retrieve(conversation_id=conversation_id)
|
||||
assert get_resp.id is not None
|
||||
assert get_resp.created_at is not None
|
||||
|
||||
get_data = get_resp.metadata
|
||||
assert get_resp.id == conversation_id
|
||||
assert get_data["user"] == "test_user"
|
||||
|
||||
def test_update_conversation(self, setup_backend):
|
||||
"""Test updating conversation metadata."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create conversation
|
||||
create_resp = client.conversations.create(metadata={"key1": "value1"})
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.created_at is not None
|
||||
|
||||
create_data = create_resp.metadata
|
||||
assert create_data["key1"] == "value1"
|
||||
assert "key2" not in create_data
|
||||
conversation_id = create_resp.id
|
||||
|
||||
# Update conversation
|
||||
update_resp = client.conversations.update(
|
||||
conversation_id=conversation_id,
|
||||
metadata={"key1": "value1", "key2": "value2"},
|
||||
)
|
||||
assert update_resp.id == conversation_id
|
||||
update_data = update_resp.metadata
|
||||
assert update_data["key1"] == "value1"
|
||||
assert update_data["key2"] == "value2"
|
||||
|
||||
# Verify update
|
||||
get_resp = client.conversations.retrieve(conversation_id=conversation_id)
|
||||
get_data = get_resp.metadata
|
||||
assert update_data["key1"] == "value1"
|
||||
assert update_data["key2"] == "value2"
|
||||
|
||||
def test_delete_conversation(self, setup_backend):
|
||||
"""Test deleting conversation."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create conversation
|
||||
create_resp = client.conversations.create()
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.created_at is not None
|
||||
conversation_id = create_resp.id
|
||||
|
||||
# Delete conversation
|
||||
delete_resp = client.conversations.delete(conversation_id=conversation_id)
|
||||
assert delete_resp.id is not None
|
||||
assert delete_resp.deleted
|
||||
|
||||
# Verify deletion
|
||||
with pytest.raises(openai.NotFoundError):
|
||||
client.conversations.retrieve(conversation_id=conversation_id)
|
||||
|
||||
def test_list_conversation_items(self, setup_backend):
|
||||
"""Test listing conversation items."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create conversation
|
||||
conv_resp = client.conversations.create()
|
||||
assert conv_resp.id is not None
|
||||
conversation_id = conv_resp.id
|
||||
|
||||
# Create response with conversation
|
||||
resp1 = client.responses.create(
|
||||
model=model,
|
||||
input="First message",
|
||||
conversation=conversation_id,
|
||||
max_output_tokens=50,
|
||||
)
|
||||
assert resp1.error is None
|
||||
resp2 = client.responses.create(
|
||||
model=model,
|
||||
input="Second message",
|
||||
conversation=conversation_id,
|
||||
max_output_tokens=50,
|
||||
)
|
||||
assert resp2.error is None
|
||||
|
||||
# List items
|
||||
list_resp = client.conversations.items.list(conversation_id=conversation_id)
|
||||
assert list_resp is not None
|
||||
assert list_resp.data is not None
|
||||
|
||||
list_data = list_resp.data
|
||||
# Should have at least 4 items (2 inputs + 2 outputs)
|
||||
assert len(list_data) >= 4
|
||||
|
||||
|
||||
def wait_for_background_task(
|
||||
client: OpenAI, response_id: str, timeout: int = 30, poll_interval: float = 0.5
|
||||
) -> responses.Response:
|
||||
"""
|
||||
Wait for background task to complete.
|
||||
|
||||
Args:
|
||||
client: openai client
|
||||
response_id: Response ID to poll
|
||||
timeout: Max seconds to wait
|
||||
poll_interval: Seconds between polls
|
||||
|
||||
Returns:
|
||||
Final response data
|
||||
|
||||
Raises:
|
||||
TimeoutError: If task doesn't complete in time
|
||||
AssertionError: If task fails
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
while time.time() - start_time < timeout:
|
||||
resp = client.responses.retrieve(response_id=response_id)
|
||||
assert resp.error is None
|
||||
assert resp.id == response_id
|
||||
|
||||
status = resp.status
|
||||
|
||||
if status == "completed":
|
||||
return resp
|
||||
elif status == "failed":
|
||||
raise AssertionError(f"Background task failed: {resp.error}")
|
||||
elif status == "cancelled":
|
||||
raise AssertionError("Background task was cancelled")
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
raise TimeoutError(
|
||||
f"Background task {response_id} did not complete within {timeout}s"
|
||||
)
|
||||
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
State management tests for Response API.
|
||||
|
||||
Tests both previous_response_id and conversation-based state management.
|
||||
These tests should work across all backends (OpenAI, XAI, gRPC).
|
||||
"""
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setup_backend", ["openai", "xai", "grpc", "grpc_harmony"], indirect=True
|
||||
)
|
||||
class TestStateManagement:
|
||||
"""Tests for state management using previous_response_id and conversation."""
|
||||
|
||||
def test_basic_response_creation(self, setup_backend):
|
||||
"""Test basic response creation without state."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
resp = client.responses.create(model=model, input="What is 2+2?")
|
||||
|
||||
assert resp.id is not None
|
||||
assert resp.error is None
|
||||
assert resp.status == "completed"
|
||||
assert len(resp.output_text) > 0
|
||||
assert resp.usage is not None
|
||||
|
||||
def test_streaming_response(self, setup_backend):
|
||||
"""Test streaming response."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model, input="Count to 5", stream=True, max_output_tokens=50
|
||||
)
|
||||
|
||||
# Check for response.created event
|
||||
events = [event for event in resp]
|
||||
created_events = [event for event in events if event.type == "response.created"]
|
||||
assert len(created_events) > 0
|
||||
|
||||
# Check for final completed event or in_progress events
|
||||
assert any(
|
||||
event.type in ["response.completed", "response.in_progress"]
|
||||
for event in events
|
||||
)
|
||||
|
||||
def test_previous_response_id_chaining(self, setup_backend):
|
||||
"""Test chaining responses using previous_response_id."""
|
||||
_, model, client = setup_backend
|
||||
# First response
|
||||
resp1 = client.responses.create(
|
||||
model=model, input="My name is Alice and my friend is Bob. Remember it."
|
||||
)
|
||||
assert resp1.error is None
|
||||
assert resp1.status == "completed"
|
||||
response1_id = resp1.id
|
||||
|
||||
# Second response referencing first
|
||||
resp2 = client.responses.create(
|
||||
model=model, input="What is my name", previous_response_id=response1_id
|
||||
)
|
||||
assert resp2.error is None
|
||||
assert resp2.status == "completed"
|
||||
|
||||
# The model should remember the name from previous response
|
||||
assert "Alice" in resp2.output_text
|
||||
|
||||
# Third response referencing second
|
||||
resp3 = client.responses.create(
|
||||
model=model,
|
||||
input="What is my friend name?",
|
||||
previous_response_id=resp2.id,
|
||||
)
|
||||
assert resp3.error is None
|
||||
assert resp3.status == "completed"
|
||||
assert "Bob" in resp3.output_text
|
||||
|
||||
@pytest.mark.skip(reason="TODO: Add the invalid previous_response_id check")
|
||||
def test_previous_response_id_invalid(self, setup_backend):
|
||||
"""Test using invalid previous_response_id."""
|
||||
_, model, client = setup_backend
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
client.responses.create(
|
||||
model=model,
|
||||
input="Test",
|
||||
previous_response_id="resp_invalid123",
|
||||
max_output_tokens=50,
|
||||
)
|
||||
|
||||
def test_conversation_with_multiple_turns(self, setup_backend):
|
||||
"""Test state management using conversation ID."""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["grpc", "grpc_harmony"]:
|
||||
pytest.skip("TODO: 501 Not Implemented")
|
||||
|
||||
# Create conversation
|
||||
conv_resp = client.conversations.create(metadata={"topic": "math"})
|
||||
assert conv_resp.id is not None
|
||||
assert conv_resp.created_at is not None
|
||||
|
||||
conversation_id = conv_resp.id
|
||||
|
||||
# First response in conversation
|
||||
resp1 = client.responses.create(
|
||||
model=model, input="I have 5 apples.", conversation=conversation_id
|
||||
)
|
||||
assert resp1.error is None
|
||||
assert resp1.status == "completed"
|
||||
|
||||
# Second response in same conversation
|
||||
resp2 = client.responses.create(
|
||||
model=model,
|
||||
input="How many apples do I have?",
|
||||
conversation=conversation_id,
|
||||
)
|
||||
assert resp2.error is None
|
||||
assert resp2.status == "completed"
|
||||
output_text = resp2.output_text
|
||||
|
||||
# Should remember "5 apples"
|
||||
assert "5" in output_text or "five" in output_text.lower()
|
||||
|
||||
# Third response in same conversation
|
||||
resp3 = client.responses.create(
|
||||
model=model,
|
||||
input="If I get 3 more, how many total?",
|
||||
conversation=conversation_id,
|
||||
)
|
||||
assert resp3.error is None
|
||||
assert resp3.status == "completed"
|
||||
output_text = resp3.output_text
|
||||
|
||||
# Should calculate 5 + 3 = 8
|
||||
assert "8" in output_text or "eight" in output_text.lower()
|
||||
list_resp = client.conversations.items.list(conversation_id)
|
||||
assert list_resp.data is not None
|
||||
items = list_resp.data
|
||||
# Should have at least 6 items (3 inputs + 3 outputs)
|
||||
assert len(items) >= 6
|
||||
|
||||
def test_mutually_exclusive_parameters(self, setup_backend):
|
||||
"""Test that previous_response_id and conversation are mutually exclusive."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# TODO: Remove this once the conversation API is implemented for GRPC backend
|
||||
conversation_id = "conv_123"
|
||||
|
||||
resp1 = client.responses.create(model=model, input="Test")
|
||||
response1_id = resp1.id
|
||||
|
||||
# Try to use both parameters
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
client.responses.create(
|
||||
model=model,
|
||||
input="This should fail",
|
||||
previous_response_id=response1_id,
|
||||
conversation=conversation_id,
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
"""
|
||||
Streaming events tests for Response API.
|
||||
Tests for streaming event validation including:
|
||||
- Zero-based output_index for reasoning content
|
||||
- OutputItemDone event emission and output array construction
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setup_backend", ["grpc", "grpc_harmony"], indirect=True)
|
||||
class TestStreamingEvents:
|
||||
"""Tests for streaming event validation."""
|
||||
|
||||
def test_output_item_event_emitted(self, setup_backend):
|
||||
"""
|
||||
Test that output_index is zero-based in streaming responses.
|
||||
Verifies that the first output item has output_index: 0.
|
||||
"""
|
||||
_, model, client = setup_backend
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="Count from 1 to 3",
|
||||
stream=True,
|
||||
max_output_tokens=50,
|
||||
)
|
||||
|
||||
events = [event for event in resp]
|
||||
assert len(events) > 0
|
||||
|
||||
# Find output_item.added events
|
||||
output_item_added_events = [
|
||||
event for event in events if event.type == "response.output_item.added"
|
||||
]
|
||||
assert len(output_item_added_events) > 0, "Should have output_item.added events"
|
||||
|
||||
# Verify first output item has output_index: 0
|
||||
first_item_event = output_item_added_events[0]
|
||||
assert first_item_event.item is not None
|
||||
assert first_item_event.output_index is not None
|
||||
assert (
|
||||
first_item_event.output_index == 0
|
||||
), "First output item must have output_index: 0 (zero-based indexing)"
|
||||
|
||||
# Verify subsequent items increment correctly
|
||||
for i, event in enumerate(output_item_added_events):
|
||||
assert (
|
||||
event.output_index == i
|
||||
), f"Output item {i} should have output_index: {i}"
|
||||
|
||||
# Verify output_item.done event exists
|
||||
output_item_done_events = [
|
||||
event for event in events if event.type == "response.output_item.done"
|
||||
]
|
||||
assert len(output_item_done_events) > 0
|
||||
|
||||
# Verify output_item.done event structure
|
||||
for event in output_item_done_events:
|
||||
assert event.item is not None
|
||||
assert event.output_index is not None
|
||||
assert event.item.type is not None
|
||||
|
||||
# Find response.completed event
|
||||
completed_events = [
|
||||
event for event in events if event.type == "response.completed"
|
||||
]
|
||||
assert len(completed_events) == 1, "Should have exactly one completed event"
|
||||
|
||||
# Verify output array exists and contains items
|
||||
completed_event = completed_events[0]
|
||||
|
||||
assert completed_event.response.output is not None
|
||||
output_array = completed_event.response.output
|
||||
assert isinstance(output_array, list)
|
||||
assert len(output_array) > 0, "Output array should contain at least one item"
|
||||
|
||||
# Verify each item in output array has proper structure
|
||||
for i, item in enumerate(output_array):
|
||||
assert item.type is not None
|
||||
|
||||
# Verify output_item.added events match items in final output array
|
||||
output_item_added_events = [
|
||||
event for event in events if event.type == "response.output_item.added"
|
||||
]
|
||||
|
||||
assert len(output_item_added_events) == len(
|
||||
output_array
|
||||
), "Number of output_item.added events should match output array length"
|
||||
|
||||
def test_reasoning_content(self, setup_backend):
|
||||
"""
|
||||
Test that reasoning content has correct zero-based output_index.
|
||||
Specifically tests that reasoning item has output_index: 0
|
||||
and message item has output_index: 1.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
if backend in ["grpc"]:
|
||||
pytest.skip("skip test_reasoning_content for grpc")
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="What is the capital of France? Think step by step.",
|
||||
stream=True,
|
||||
max_output_tokens=200,
|
||||
)
|
||||
|
||||
events = [event for event in resp]
|
||||
assert len(events) > 0
|
||||
|
||||
# Find output_item.added events
|
||||
output_item_added_events = [
|
||||
event for event in events if event.type == "response.output_item.added"
|
||||
]
|
||||
assert len(output_item_added_events) > 0
|
||||
|
||||
reasoning_items = [
|
||||
item for item in output_item_added_events if item.item.type == "reasoning"
|
||||
]
|
||||
message_items = [
|
||||
item for item in output_item_added_events if item.item.type == "message"
|
||||
]
|
||||
|
||||
# If reasoning is present, verify it has output_index: 0
|
||||
if reasoning_items:
|
||||
reasoning_item = reasoning_items[0]
|
||||
assert (
|
||||
reasoning_item.output_index == 0
|
||||
), "Reasoning item should have output_index: 0"
|
||||
|
||||
# If message is present after reasoning, verify it has output_index: 1
|
||||
if reasoning_items and message_items:
|
||||
message_item = message_items[0]
|
||||
assert (
|
||||
message_item.output_index == 1
|
||||
), "Message item after reasoning should have output_index: 1"
|
||||
|
||||
# Find response.completed event
|
||||
completed_events = [
|
||||
event for event in events if event.type == "response.completed"
|
||||
]
|
||||
assert len(completed_events) == 1
|
||||
|
||||
# Get output array from completed event
|
||||
output_array = completed_events[0].response.output
|
||||
assert len(output_array) > 0
|
||||
|
||||
# Check if reasoning items are in output array
|
||||
reasoning_items_in_output = [
|
||||
item for item in output_array if item.type == "reasoning"
|
||||
]
|
||||
assert len(reasoning_items_in_output) > 0
|
||||
@@ -0,0 +1,172 @@
|
||||
"""
|
||||
Structured output tests for Response API.
|
||||
|
||||
Tests for text.format field with json_object and json_schema formats.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Add current directory for local imports
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setup_backend", ["openai", "grpc_harmony"], indirect=True)
|
||||
class TestStructuredOutput:
|
||||
|
||||
def test_structured_output_json_schema(self, setup_backend):
|
||||
"""Test structured output with json_schema format."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
# Create response with structured output
|
||||
params = {
|
||||
"model": model,
|
||||
"input": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful math tutor. Guide the user through the solution step by step.",
|
||||
},
|
||||
{"role": "user", "content": "how can I solve 8x + 7 = -23"},
|
||||
],
|
||||
"text": {
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"name": "math_reasoning",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"steps": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"explanation": {"type": "string"},
|
||||
"output": {"type": "string"},
|
||||
},
|
||||
"required": ["explanation", "output"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"final_answer": {"type": "string"},
|
||||
},
|
||||
"required": ["steps", "final_answer"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"strict": True,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
create_resp = client.responses.create(**params)
|
||||
assert create_resp.error is None
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.output is not None
|
||||
assert create_resp.text is not None
|
||||
|
||||
# Verify text format was echoed back correctly
|
||||
assert create_resp.text.format is not None
|
||||
assert create_resp.text.format.type == "json_schema"
|
||||
assert create_resp.text.format.name == "math_reasoning"
|
||||
assert create_resp.text.format.schema_ is not None
|
||||
assert create_resp.text.format.strict
|
||||
|
||||
# Find the message output (output[0] may be reasoning, output[1] is message)
|
||||
output_text = next(
|
||||
(
|
||||
content.text
|
||||
for item in create_resp.output
|
||||
if item.type == "message"
|
||||
for content in item.content
|
||||
if content.type == "output_text"
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert output_text is not None, "No output_text found in response"
|
||||
assert output_text.strip(), "output_text is empty"
|
||||
|
||||
# Parse JSON output
|
||||
output_json = json.loads(output_text)
|
||||
|
||||
# Verify schema structure
|
||||
assert "steps" in output_json
|
||||
assert "final_answer" in output_json
|
||||
assert isinstance(output_json["steps"], list)
|
||||
assert len(output_json["steps"]) > 0
|
||||
|
||||
# Verify each step has required fields
|
||||
for step in output_json["steps"]:
|
||||
assert "explanation" in step
|
||||
assert "output" in step
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
|
||||
class TestSimpleSchemaStructuredOutput:
|
||||
|
||||
def test_structured_output_json_schema(self, setup_backend):
|
||||
"""Override with simpler schema for Llama model (complex schemas not well supported)."""
|
||||
_, model, client = setup_backend
|
||||
|
||||
params = {
|
||||
"model": model,
|
||||
"input": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a math solver. Return ONLY a JSON object that matches the schema—no extra text.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 1 + 1?",
|
||||
},
|
||||
],
|
||||
"text": {
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"name": "math_answer",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"answer": {"type": "string"}},
|
||||
"required": ["answer"],
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
create_resp = client.responses.create(**params)
|
||||
assert create_resp.error is None
|
||||
assert create_resp.id is not None
|
||||
assert create_resp.output is not None
|
||||
assert create_resp.text is not None
|
||||
|
||||
# Verify text format was echoed back correctly
|
||||
assert create_resp.text.format is not None
|
||||
assert create_resp.text.format.type == "json_schema"
|
||||
assert create_resp.text.format.name == "math_answer"
|
||||
assert create_resp.text.format.schema_ is not None
|
||||
|
||||
# Find the message output
|
||||
output_text = next(
|
||||
(
|
||||
content.text
|
||||
for item in create_resp.output
|
||||
if item.type == "message"
|
||||
for content in item.content
|
||||
if content.type == "output_text"
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert output_text is not None, "No output_text found in response"
|
||||
assert output_text.strip(), "output_text is empty"
|
||||
|
||||
# Parse JSON output
|
||||
output_json = json.loads(output_text)
|
||||
|
||||
# Verify simple schema structure (just answer field)
|
||||
assert "answer" in output_json
|
||||
assert isinstance(output_json["answer"], str)
|
||||
assert output_json["answer"], "Answer is empty"
|
||||
@@ -0,0 +1,765 @@
|
||||
"""
|
||||
Test class for tool calling tests.
|
||||
|
||||
This module provides test cases for function calling functionality, tool choices
|
||||
and mcp calling functionality across different backends.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Add current directory for local imports
|
||||
_TEST_DIR = Path(__file__).parent
|
||||
sys.path.insert(0, str(_TEST_DIR))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setup_backend", ["openai", "grpc", "grpc_harmony"], indirect=True
|
||||
)
|
||||
class TestToolCalling:
|
||||
|
||||
# Shared function tool definitions
|
||||
SYSTEM_DIAGNOSTICS_FUNCTION = {
|
||||
"type": "function",
|
||||
"name": "get_system_diagnostics",
|
||||
"description": "Retrieve real-time diagnostics for a spacecraft system.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"system_name": {
|
||||
"type": "string",
|
||||
"description": "Name of the spacecraft system to query. "
|
||||
"Example: 'Astra-7 Core Reactor'.",
|
||||
}
|
||||
},
|
||||
"required": ["system_name"],
|
||||
},
|
||||
}
|
||||
|
||||
GET_WEATHER_FUNCTION = {
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city name, e.g., San Francisco",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
}
|
||||
|
||||
CALCULATE_FUNCTION = {
|
||||
"type": "function",
|
||||
"name": "calculate",
|
||||
"description": "Perform a mathematical calculation",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "The mathematical expression to evaluate",
|
||||
}
|
||||
},
|
||||
"required": ["expression"],
|
||||
},
|
||||
}
|
||||
|
||||
SEARCH_WEB_FUNCTION = {
|
||||
"type": "function",
|
||||
"name": "search_web",
|
||||
"description": "Search the web for information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
}
|
||||
|
||||
LOCAL_SEARCH_FUNCTION = {
|
||||
"type": "function",
|
||||
"name": "local_search",
|
||||
"description": "Search local database",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
}
|
||||
|
||||
# Shared constants for MCP tests
|
||||
BRAVE_MCP_TOOL = {
|
||||
"type": "mcp",
|
||||
"server_label": "brave",
|
||||
"server_description": "A Tool to do web search",
|
||||
"server_url": "http://localhost:8001/sse",
|
||||
"require_approval": "never",
|
||||
}
|
||||
|
||||
DEEPWIKI_MCP_TOOL = {
|
||||
"type": "mcp",
|
||||
"server_label": "deepwiki",
|
||||
"server_url": "https://mcp.deepwiki.com/mcp",
|
||||
"require_approval": "never",
|
||||
}
|
||||
|
||||
MCP_TEST_PROMPT = (
|
||||
"show me some news about sglang router, use the tool to just search "
|
||||
"one result and return one sentence response"
|
||||
)
|
||||
|
||||
# Test cases for basic function calling functionality
|
||||
|
||||
def test_basic_function_call(self, setup_backend):
|
||||
"""
|
||||
Test basic function calling workflow.
|
||||
|
||||
This test follows the pattern from function_call_test.py:
|
||||
1. Define a function tool (get_horoscope)
|
||||
2. Send user message asking for horoscope
|
||||
3. Model should return function_call
|
||||
4. Execute function locally and provide output
|
||||
5. Model should generate final response using the function output
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["grpc"]:
|
||||
pytest.skip("skip for grpc")
|
||||
|
||||
# 1. Define a list of callable tools for the model
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_horoscope",
|
||||
"description": "Get today's horoscope for an astrological sign.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"sign": {
|
||||
"type": "string",
|
||||
"description": "An astrological sign like Taurus or Aquarius",
|
||||
},
|
||||
},
|
||||
"required": ["sign"],
|
||||
},
|
||||
},
|
||||
]
|
||||
system_prompt = (
|
||||
"You are a helpful assistant that can call functions. "
|
||||
"When a user asks for horoscope information, call the function. "
|
||||
"IMPORTANT: Don't reply directly to the user, only call the function. "
|
||||
)
|
||||
|
||||
# Create a running input list we will add to over time
|
||||
input_list = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "What is my horoscope? I am an Aquarius."},
|
||||
]
|
||||
|
||||
# 2. Prompt the model with tools defined
|
||||
resp = client.responses.create(model=model, input=input_list, tools=tools)
|
||||
|
||||
# Should successfully make the request
|
||||
assert resp.error is None
|
||||
|
||||
# Basic response structure
|
||||
assert resp.id is not None
|
||||
assert resp.status == "completed"
|
||||
assert resp.output is not None
|
||||
|
||||
# Verify output array is not empty
|
||||
output = resp.output
|
||||
assert isinstance(output, list)
|
||||
assert len(output) > 0
|
||||
|
||||
# Check for function_call in output
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert (
|
||||
len(function_calls) > 0
|
||||
), "Response should contain at least one function_call"
|
||||
|
||||
# Verify function_call structure
|
||||
function_call = function_calls[0]
|
||||
assert function_call.call_id is not None
|
||||
assert function_call.name is not None
|
||||
assert function_call.name == "get_horoscope"
|
||||
assert function_call.arguments is not None
|
||||
|
||||
# Parse arguments
|
||||
args = json.loads(function_call.arguments)
|
||||
assert "sign" in args
|
||||
assert args["sign"].lower() == "aquarius"
|
||||
|
||||
# 3. Save function call outputs for subsequent requests
|
||||
input_list.append(function_call)
|
||||
|
||||
# 4. Execute the function logic for get_horoscope
|
||||
horoscope = f"{args['sign']}: Next Tuesday you will befriend a baby otter."
|
||||
|
||||
# 5. Provide function call results to the model
|
||||
input_list.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": function_call.call_id,
|
||||
"output": json.dumps({"horoscope": horoscope}),
|
||||
}
|
||||
)
|
||||
|
||||
# 6. Make second request with function output
|
||||
resp2 = client.responses.create(
|
||||
model=model,
|
||||
input=input_list,
|
||||
instructions="Respond only with a horoscope generated by a tool.",
|
||||
tools=tools,
|
||||
)
|
||||
assert resp2.error is None
|
||||
assert resp2.status == "completed"
|
||||
|
||||
# The model should be able to give a response using the function output
|
||||
output2 = resp2.output
|
||||
assert len(output2) > 0
|
||||
|
||||
# Find message output
|
||||
messages = [item for item in output2 if item.type == "message"]
|
||||
assert len(messages) > 0, "Response should contain at least one message"
|
||||
|
||||
# Verify message contains the horoscope
|
||||
message = messages[0]
|
||||
assert message.content is not None
|
||||
content_parts = message.content
|
||||
assert len(content_parts) > 0
|
||||
|
||||
# Get text from content
|
||||
text_parts = [part.text for part in content_parts if part.type == "output_text"]
|
||||
full_text = " ".join(text_parts).lower()
|
||||
|
||||
# Should mention the horoscope or baby otter
|
||||
assert (
|
||||
"baby otter" in full_text or "aquarius" in full_text
|
||||
), "Response should reference the horoscope content"
|
||||
|
||||
# Test cases for tool_choice parameter support, these tests require --reasoning-parser
|
||||
|
||||
def test_tool_choice_auto(self, setup_backend):
|
||||
"""
|
||||
Test tool_choice="auto" allows model to decide whether to use tools.
|
||||
|
||||
The model should be able to choose to call a tool or not.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip("skip for openai")
|
||||
|
||||
tools = [self.GET_WEATHER_FUNCTION]
|
||||
|
||||
# Query that should trigger tool use
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="What is the weather in Seattle?",
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert resp.id is not None
|
||||
assert resp.error is None
|
||||
|
||||
output = resp.output
|
||||
assert len(output) > 0
|
||||
|
||||
# With auto, model should choose to call get_weather for this query
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert (
|
||||
len(function_calls) > 0
|
||||
), "Model should choose to call function with tool_choice='auto'"
|
||||
|
||||
def test_tool_choice_required(self, setup_backend):
|
||||
"""
|
||||
Test tool_choice="required" forces the model to call at least one tool.
|
||||
|
||||
The model must make at least one function call.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip("skip for openai")
|
||||
|
||||
tools = [self.CALCULATE_FUNCTION]
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="What is 15 * 23?",
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert resp.id is not None
|
||||
assert resp.error is None
|
||||
|
||||
output = resp.output
|
||||
|
||||
# Must have at least one function call
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert (
|
||||
len(function_calls) > 0
|
||||
), "tool_choice='required' must force at least one function call"
|
||||
|
||||
def test_tool_choice_specific_function(self, setup_backend):
|
||||
"""
|
||||
Test tool_choice with specific function name forces that function to be called.
|
||||
|
||||
The model must call the specified function.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip("skip for openai")
|
||||
|
||||
tools = [self.SEARCH_WEB_FUNCTION, self.GET_WEATHER_FUNCTION]
|
||||
|
||||
# Force specific function call
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="What's happening in the news today?",
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "function": {"name": "search_web"}},
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert resp.id is not None
|
||||
assert resp.error is None
|
||||
|
||||
output = resp.output
|
||||
|
||||
# Must have function call
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert len(function_calls) > 0, "Must call the specified function"
|
||||
|
||||
# Must be the specified function
|
||||
called_function = function_calls[0]
|
||||
assert (
|
||||
called_function.name == "search_web"
|
||||
), "Must call the function specified in tool_choice"
|
||||
|
||||
def test_tool_choice_streaming(self, setup_backend):
|
||||
"""
|
||||
Test tool_choice parameter works correctly with streaming.
|
||||
|
||||
Verifies that tool_choice constraints are applied in streaming mode.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai", "grpc"]:
|
||||
pytest.skip("skip for openai")
|
||||
|
||||
tools = [self.CALCULATE_FUNCTION]
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="Calculate 42 * 17",
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
stream=True,
|
||||
)
|
||||
|
||||
events = [event for event in resp]
|
||||
assert len(events) > 0
|
||||
|
||||
event_types = [e.type for e in events]
|
||||
|
||||
# Should have function call events
|
||||
assert (
|
||||
"response.function_call_arguments.delta" in event_types
|
||||
), "Should have function_call_arguments.delta events"
|
||||
|
||||
# Verify completed event has function call
|
||||
completed_events = [e for e in events if e.type == "response.completed"]
|
||||
assert len(completed_events) == 1
|
||||
|
||||
output = completed_events[0].response.output
|
||||
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert (
|
||||
len(function_calls) > 0
|
||||
), "Streaming with tool_choice='required' must produce function call"
|
||||
|
||||
def test_tool_choice_with_mcp_tools(self, setup_backend):
|
||||
"""
|
||||
Test tool_choice parameter works with MCP tools.
|
||||
|
||||
Verifies that tool_choice can control MCP tool usage.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip("skip for openai")
|
||||
|
||||
tools = [self.DEEPWIKI_MCP_TOOL]
|
||||
|
||||
# With tool_choice="auto", should allow MCP tool calls
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="What transport protocols does the 2025-03-26 version of the MCP spec (modelcontextprotocol/modelcontextprotocol) support?",
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert resp.id is not None
|
||||
assert resp.error is None
|
||||
|
||||
output = resp.output
|
||||
|
||||
# Should have mcp_call with auto
|
||||
mcp_calls = [item for item in output if item.type == "mcp_call"]
|
||||
assert len(mcp_calls) > 0, "tool_choice='auto' should allow MCP tool calls"
|
||||
|
||||
def test_tool_choice_mixed_function_and_mcp(self, setup_backend):
|
||||
"""
|
||||
Test tool_choice with mixed function and MCP tools.
|
||||
|
||||
Verifies tool_choice can select specific tools when both function and MCP tools are available.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip("skip for openai")
|
||||
|
||||
tools = [self.DEEPWIKI_MCP_TOOL, self.LOCAL_SEARCH_FUNCTION]
|
||||
|
||||
# Force specific function call
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="Search for information about Python",
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "function": {"name": "local_search"}},
|
||||
stream=False,
|
||||
)
|
||||
|
||||
assert resp.id is not None
|
||||
assert resp.error is None
|
||||
|
||||
output = resp.output
|
||||
|
||||
# Must call local_search, not MCP
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert len(function_calls) > 0
|
||||
assert function_calls[0].name == "local_search"
|
||||
|
||||
# Should not have mcp_call
|
||||
mcp_calls = [item for item in output if item.type == "mcp_call"]
|
||||
assert len(mcp_calls) == 0, "Should only call specified function, not MCP tools"
|
||||
|
||||
# Tests for MCP tool calling in both streaming and non-streaming modes.
|
||||
|
||||
def test_mcp_basic_tool_call(self, setup_backend):
|
||||
"""
|
||||
Test basic MCP tool call (non-streaming).
|
||||
Validation strictness is controlled by parameter `backend` from setup_backend fixture.
|
||||
Set to "strict" if backend is http.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
# To avoid being rate-limited by brave search server
|
||||
time.sleep(2)
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input=self.MCP_TEST_PROMPT,
|
||||
tools=[self.BRAVE_MCP_TOOL],
|
||||
stream=False,
|
||||
reasoning={"effort": "low"},
|
||||
)
|
||||
|
||||
# Should successfully make the request
|
||||
assert resp.error is None
|
||||
|
||||
# Basic response structure
|
||||
assert resp.id is not None
|
||||
assert resp.status == "completed"
|
||||
assert resp.model is not None
|
||||
assert resp.output is not None
|
||||
|
||||
# Verify output array is not empty
|
||||
assert len(resp.output_text) > 0
|
||||
|
||||
# Check for MCP-specific output types
|
||||
output_types = [item.type for item in resp.output]
|
||||
|
||||
# Should have mcp_list_tools - tools are listed before calling
|
||||
assert (
|
||||
"mcp_list_tools" in output_types
|
||||
), "Response should contain mcp_list_tools"
|
||||
|
||||
# Should have at least one mcp_call
|
||||
mcp_calls = [item for item in resp.output if item.type == "mcp_call"]
|
||||
assert len(mcp_calls) > 0, "Response should contain at least one mcp_call"
|
||||
|
||||
# Verify mcp_call structure
|
||||
for mcp_call in mcp_calls:
|
||||
assert mcp_call.id is not None
|
||||
assert mcp_call.error is None
|
||||
assert mcp_call.status == "completed"
|
||||
assert mcp_call.server_label == "brave"
|
||||
assert mcp_call.name is not None
|
||||
assert mcp_call.arguments is not None
|
||||
assert mcp_call.output is not None
|
||||
|
||||
# Strict mode: additional validation for HTTP backends
|
||||
if backend == "openai":
|
||||
# Should have final message output
|
||||
messages = [item for item in resp.output if item.type == "message"]
|
||||
assert len(messages) > 0, "Response should contain at least one message"
|
||||
# Verify message structure
|
||||
for msg in messages:
|
||||
assert msg.content is not None
|
||||
assert isinstance(msg.content, list)
|
||||
|
||||
# Check content has text
|
||||
for content_item in msg.content:
|
||||
if content_item.type == "output_text":
|
||||
assert content_item.text is not None
|
||||
assert isinstance(content_item.text, str)
|
||||
assert len(content_item.text) > 0
|
||||
|
||||
def test_mcp_basic_tool_call_streaming(self, setup_backend):
|
||||
"""Test basic MCP tool call (streaming).
|
||||
|
||||
Validation strictness is controlled by the class attribute `mcp_validation_mode`.
|
||||
Set to "strict" in subclasses for additional HTTP-specific validation.
|
||||
"""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
# To avoid being rate-limited by brave search server
|
||||
time.sleep(2)
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input=self.MCP_TEST_PROMPT,
|
||||
tools=[self.BRAVE_MCP_TOOL],
|
||||
stream=True,
|
||||
reasoning={"effort": "low"},
|
||||
)
|
||||
|
||||
# Should successfully make the request
|
||||
events = [event for event in resp]
|
||||
assert len(events) > 0
|
||||
|
||||
event_types = [event.type for event in events]
|
||||
# Check for lifecycle events
|
||||
assert "response.created" in event_types, "Should have response.created event"
|
||||
assert (
|
||||
"response.completed" in event_types
|
||||
), "Should have response.completed event"
|
||||
|
||||
# Check for MCP list tools events
|
||||
assert (
|
||||
"response.output_item.added" in event_types
|
||||
), "Should have output_item.added events"
|
||||
assert (
|
||||
"response.mcp_list_tools.in_progress" in event_types
|
||||
), "Should have mcp_list_tools.in_progress event"
|
||||
assert (
|
||||
"response.mcp_list_tools.completed" in event_types
|
||||
), "Should have mcp_list_tools.completed event"
|
||||
|
||||
# Check for MCP call events
|
||||
assert (
|
||||
"response.mcp_call.in_progress" in event_types
|
||||
), "Should have mcp_call.in_progress event"
|
||||
assert (
|
||||
"response.mcp_call_arguments.delta" in event_types
|
||||
), "Should have mcp_call_arguments.delta event"
|
||||
assert (
|
||||
"response.mcp_call_arguments.done" in event_types
|
||||
), "Should have mcp_call_arguments.done event"
|
||||
assert (
|
||||
"response.mcp_call.completed" in event_types
|
||||
), "Should have mcp_call.completed event"
|
||||
|
||||
# Verify final completed event has full response
|
||||
completed_events = [e for e in events if e.type == "response.completed"]
|
||||
assert len(completed_events) == 1
|
||||
|
||||
final_response = completed_events[0].response
|
||||
assert final_response.id is not None
|
||||
assert final_response.status == "completed"
|
||||
assert final_response.output is not None
|
||||
|
||||
# Verify final output contains expected items
|
||||
final_output = final_response.output
|
||||
final_output_types = [item.type for item in final_output]
|
||||
|
||||
assert "mcp_list_tools" in final_output_types
|
||||
assert "mcp_call" in final_output_types
|
||||
|
||||
# Verify mcp_call items in final output
|
||||
mcp_calls = [item for item in final_output if item.type == "mcp_call"]
|
||||
assert len(mcp_calls) > 0
|
||||
|
||||
for mcp_call in mcp_calls:
|
||||
assert mcp_call.error is None
|
||||
assert mcp_call.status == "completed"
|
||||
assert mcp_call.server_label == "brave"
|
||||
assert mcp_call.name is not None
|
||||
assert mcp_call.arguments is not None
|
||||
assert mcp_call.output is not None
|
||||
|
||||
# Strict mode: additional validation for HTTP backends
|
||||
if backend == "openai":
|
||||
# Check for text output events
|
||||
assert (
|
||||
"response.content_part.added" in event_types
|
||||
), "Should have content_part.added event"
|
||||
assert (
|
||||
"response.output_text.delta" in event_types
|
||||
), "Should have output_text.delta events"
|
||||
assert (
|
||||
"response.output_text.done" in event_types
|
||||
), "Should have output_text.done event"
|
||||
assert (
|
||||
"response.content_part.done" in event_types
|
||||
), "Should have content_part.done event"
|
||||
|
||||
assert "message" in final_output_types
|
||||
|
||||
# Verify text deltas combine to final message
|
||||
text_deltas = [
|
||||
e.delta for e in events if e.type == "response.output_text.delta"
|
||||
]
|
||||
assert len(text_deltas) > 0, "Should have text deltas"
|
||||
|
||||
# Get final text from output_text.done event
|
||||
text_done_events = [
|
||||
e for e in events if e.type == "response.output_text.done"
|
||||
]
|
||||
assert len(text_done_events) > 0
|
||||
|
||||
final_text = text_done_events[0].text
|
||||
assert len(final_text) > 0, "Final text should not be empty"
|
||||
|
||||
def test_mixed_mcp_and_function_tools(self, setup_backend):
|
||||
"""Test mixed MCP and function tools (non-streaming)."""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip(
|
||||
"Requires external MCP server (deepwiki) - may not be accessible in CI"
|
||||
)
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="Give me diagnostics for the Astra-7 Core Reactor.",
|
||||
tools=[self.BRAVE_MCP_TOOL, self.SYSTEM_DIAGNOSTICS_FUNCTION],
|
||||
stream=False,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
# Should successfully make the request
|
||||
assert resp.error is None
|
||||
|
||||
# Basic response structure
|
||||
assert resp.id is not None
|
||||
assert resp.status is not None
|
||||
assert resp.output is not None
|
||||
|
||||
# Verify output array is not empty
|
||||
output = resp.output
|
||||
assert isinstance(output, list)
|
||||
assert len(output) > 0
|
||||
|
||||
# Check for function_call (not mcp_call for get_system_diagnostics)
|
||||
function_calls = [item for item in output if item.type == "function_call"]
|
||||
assert (
|
||||
len(function_calls) > 0
|
||||
), "Response should contain at least one function_call"
|
||||
|
||||
# Verify function_call structure for get_system_diagnostics
|
||||
system_diagnostics_call = function_calls[0]
|
||||
assert system_diagnostics_call.name == "get_system_diagnostics"
|
||||
assert system_diagnostics_call.call_id is not None
|
||||
assert system_diagnostics_call.arguments is not None
|
||||
assert system_diagnostics_call.status is not None
|
||||
|
||||
# Parse and verify arguments
|
||||
args = json.loads(system_diagnostics_call.arguments)
|
||||
assert "system_name" in args
|
||||
assert "astra-7" in args["system_name"].lower()
|
||||
|
||||
def test_mixed_mcp_and_function_tools_streaming(self, setup_backend):
|
||||
"""Test mixed MCP and function tools (streaming)."""
|
||||
backend, model, client = setup_backend
|
||||
|
||||
if backend in ["openai"]:
|
||||
pytest.skip(
|
||||
"Requires external MCP server (deepwiki) - may not be accessible in CI"
|
||||
)
|
||||
|
||||
resp = client.responses.create(
|
||||
model=model,
|
||||
input="Give me diagnostics for the Astra-7 Core Reactor.",
|
||||
tools=[self.BRAVE_MCP_TOOL, self.SYSTEM_DIAGNOSTICS_FUNCTION],
|
||||
stream=True,
|
||||
tool_choice="auto", # Encourage tool usage
|
||||
)
|
||||
|
||||
# Should successfully make the request
|
||||
events = [event for event in resp]
|
||||
assert len(events) > 0
|
||||
|
||||
event_types = [e.type for e in events]
|
||||
|
||||
# Check for lifecycle events
|
||||
assert "response.created" in event_types, "Should have response.created event"
|
||||
|
||||
# Should have mcp_list_tools events
|
||||
assert (
|
||||
"response.mcp_list_tools.completed" in event_types
|
||||
), "Should have mcp_list_tools.completed event"
|
||||
|
||||
# Should have function_call_arguments events (not mcp_call_arguments)
|
||||
assert (
|
||||
"response.function_call_arguments.delta" in event_types
|
||||
), "Should have function_call_arguments.delta event for function tools"
|
||||
assert (
|
||||
"response.function_call_arguments.done" in event_types
|
||||
), "Should have function_call_arguments.done event for function tools"
|
||||
|
||||
# Should NOT have mcp_call_arguments events for function tools
|
||||
# (get_system_diagnostics should use function_call_arguments, not mcp_call_arguments)
|
||||
mcp_call_arg_events = [
|
||||
e
|
||||
for e in events
|
||||
if e.type == "response.mcp_call_arguments.delta"
|
||||
and "get_system_diagnostics" in str(e.delta)
|
||||
]
|
||||
assert (
|
||||
len(mcp_call_arg_events) == 0
|
||||
), "Should NOT emit mcp_call_arguments.delta for function tools (get_system_diagnostics)"
|
||||
|
||||
# Verify function_call_arguments.delta event structure
|
||||
func_arg_deltas = [
|
||||
e for e in events if e.type == "response.function_call_arguments.delta"
|
||||
]
|
||||
assert (
|
||||
len(func_arg_deltas) > 0
|
||||
), "Should have function_call_arguments.delta events"
|
||||
|
||||
# Check that delta event contains system_name arguments
|
||||
full_delta_event = ""
|
||||
for event in func_arg_deltas:
|
||||
full_delta_event += event.delta
|
||||
|
||||
assert (
|
||||
"system_name" in full_delta_event.lower()
|
||||
and "astra-7" in full_delta_event.lower()
|
||||
), "function_call_arguments.delta should contain system_name and astra-7"
|
||||
@@ -0,0 +1,565 @@
|
||||
"""
|
||||
Fixtures for launching OpenAI/XAI router for response API e2e testing.
|
||||
|
||||
This module provides fixtures for launching SGLang router with OpenAI or XAI backends:
|
||||
1. Launch router with --backend openai pointing to OpenAI or XAI API
|
||||
2. Configure history backend (memory or oracle)
|
||||
|
||||
This supports testing the Response API against real cloud providers.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def wait_for_workers_ready(
|
||||
router_url: str,
|
||||
expected_workers: int,
|
||||
timeout: int = 300,
|
||||
api_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Wait for router to have all workers connected.
|
||||
|
||||
Polls the /workers endpoint until the 'total' field matches expected_workers.
|
||||
|
||||
Example response from /workers endpoint:
|
||||
{"workers":[],"total":0,"stats":{"prefill_count":0,"decode_count":0,"regular_count":0}}
|
||||
|
||||
Args:
|
||||
router_url: Base URL of router (e.g., "http://127.0.0.1:30000")
|
||||
expected_workers: Number of workers expected to be connected
|
||||
timeout: Max seconds to wait
|
||||
api_key: Optional API key for authentication
|
||||
"""
|
||||
start_time = time.time()
|
||||
last_error = None
|
||||
attempt = 0
|
||||
|
||||
headers = {}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
|
||||
with requests.Session() as session:
|
||||
while time.time() - start_time < timeout:
|
||||
attempt += 1
|
||||
elapsed = int(time.time() - start_time)
|
||||
|
||||
# Log progress every 10 seconds
|
||||
if elapsed > 0 and elapsed % 10 == 0 and attempt % 10 == 0:
|
||||
logger.info(
|
||||
f" Still waiting for workers... ({elapsed}/{timeout}s elapsed)"
|
||||
)
|
||||
|
||||
try:
|
||||
response = session.get(
|
||||
f"{router_url}/workers", headers=headers, timeout=5
|
||||
)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
total_workers = data.get("total", 0)
|
||||
|
||||
if total_workers == expected_workers:
|
||||
logger.info(
|
||||
f" All {expected_workers} workers connected after {elapsed}s"
|
||||
)
|
||||
return
|
||||
else:
|
||||
last_error = f"Workers: {total_workers}/{expected_workers}"
|
||||
else:
|
||||
last_error = f"HTTP {response.status_code}"
|
||||
except requests.ConnectionError:
|
||||
last_error = "Connection refused (router not ready yet)"
|
||||
except requests.Timeout:
|
||||
last_error = "Timeout"
|
||||
except requests.RequestException as e:
|
||||
last_error = str(e)
|
||||
except (ValueError, KeyError) as e:
|
||||
last_error = f"Invalid response: {e}"
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
raise TimeoutError(
|
||||
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s.\n"
|
||||
f"Last status: {last_error}\n"
|
||||
f"Hint: Run with SHOW_ROUTER_LOGS=1 to see startup logs"
|
||||
)
|
||||
|
||||
|
||||
def find_free_port() -> int:
|
||||
"""Find an available 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]
|
||||
|
||||
|
||||
def wait_for_router_ready(
|
||||
router_url: str,
|
||||
timeout: int = 60,
|
||||
api_key: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Wait for router to be ready.
|
||||
|
||||
Polls the /health endpoint until it returns 200.
|
||||
|
||||
Args:
|
||||
router_url: Base URL of router (e.g., "http://127.0.0.1:30000")
|
||||
timeout: Max seconds to wait
|
||||
api_key: Optional API key for authentication
|
||||
"""
|
||||
start_time = time.time()
|
||||
last_error = None
|
||||
attempt = 0
|
||||
|
||||
headers = {}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
|
||||
with requests.Session() as session:
|
||||
while time.time() - start_time < timeout:
|
||||
attempt += 1
|
||||
elapsed = int(time.time() - start_time)
|
||||
|
||||
# Log progress every 10 seconds
|
||||
if elapsed > 0 and elapsed % 10 == 0 and attempt % 10 == 0:
|
||||
logger.info(
|
||||
f" Still waiting for router... ({elapsed}/{timeout}s elapsed)"
|
||||
)
|
||||
|
||||
try:
|
||||
response = session.get(
|
||||
f"{router_url}/health", headers=headers, timeout=5
|
||||
)
|
||||
if response.status_code == 200:
|
||||
logger.info(f" Router ready after {elapsed}s")
|
||||
return
|
||||
else:
|
||||
last_error = f"HTTP {response.status_code}"
|
||||
except requests.ConnectionError:
|
||||
last_error = "Connection refused (router not ready yet)"
|
||||
except requests.Timeout:
|
||||
last_error = "Timeout"
|
||||
except requests.RequestException as e:
|
||||
last_error = str(e)
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
raise TimeoutError(
|
||||
f"Router at {router_url} did not become ready within {timeout}s.\n"
|
||||
f"Last status: {last_error}\n"
|
||||
f"Hint: Run with SHOW_ROUTER_LOGS=1 to see startup logs"
|
||||
)
|
||||
|
||||
|
||||
def popen_launch_openai_xai_router(
|
||||
backend: str, # "openai" or "xai"
|
||||
base_url: str,
|
||||
timeout: int = 60,
|
||||
history_backend: str = "memory",
|
||||
api_key: Optional[str] = None,
|
||||
router_args: Optional[list] = None,
|
||||
stdout=None,
|
||||
stderr=None,
|
||||
prometheus_port: Optional[int] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Launch SGLang router with OpenAI or XAI backend.
|
||||
|
||||
This approach:
|
||||
1. Starts router with --backend openai
|
||||
2. Points to OpenAI or XAI API via --worker-urls
|
||||
3. Configures history backend (memory or oracle)
|
||||
4. Waits for router health check to pass
|
||||
|
||||
Args:
|
||||
backend: "openai" or "xai"
|
||||
base_url: Base URL for router (e.g., "http://127.0.0.1:30000")
|
||||
timeout: Timeout for router startup (default: 60s)
|
||||
history_backend: "memory" or "oracle" (default: memory)
|
||||
api_key: Optional API key for router authentication
|
||||
router_args: Additional arguments for router
|
||||
stdout: Optional file handle for router stdout
|
||||
stderr: Optional file handle for router stderr
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
- router: router process object
|
||||
- base_url: router URL (HTTP endpoint)
|
||||
|
||||
Example:
|
||||
>>> cluster = popen_launch_openai_xai_router(
|
||||
... "openai", "http://127.0.0.1:30000"
|
||||
... )
|
||||
>>> # Use cluster['base_url'] for HTTP requests
|
||||
>>> # Cleanup:
|
||||
>>> kill_process_tree(cluster['router'].pid)
|
||||
"""
|
||||
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
|
||||
|
||||
# Parse router port from base_url
|
||||
if ":" in base_url.split("//")[-1]:
|
||||
router_port = int(base_url.split(":")[-1])
|
||||
else:
|
||||
router_port = find_free_port()
|
||||
|
||||
logger.info(f"\n{'='*70}")
|
||||
logger.info(f"Launching {backend.upper()} router")
|
||||
logger.info(f"{'='*70}")
|
||||
logger.info(f" Backend: {backend}")
|
||||
logger.info(f" Router port: {router_port}")
|
||||
logger.info(f" History backend: {history_backend}")
|
||||
|
||||
# Determine worker URL based on backend
|
||||
if backend == "openai":
|
||||
worker_url = "https://api.openai.com"
|
||||
# Get API key from environment
|
||||
backend_api_key = os.environ.get("OPENAI_API_KEY")
|
||||
if not backend_api_key:
|
||||
raise ValueError(
|
||||
"OPENAI_API_KEY environment variable must be set for OpenAI backend"
|
||||
)
|
||||
elif backend == "xai":
|
||||
worker_url = "https://api.x.ai"
|
||||
# Get API key from environment
|
||||
backend_api_key = os.environ.get("XAI_API_KEY")
|
||||
if not backend_api_key:
|
||||
raise ValueError(
|
||||
"XAI_API_KEY environment variable must be set for XAI backend"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported backend: {backend}")
|
||||
|
||||
logger.info(f" Worker URL: {worker_url}")
|
||||
|
||||
# Build router command
|
||||
router_cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang_router.launch_router",
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(router_port),
|
||||
"--backend",
|
||||
"openai",
|
||||
"--worker-urls",
|
||||
worker_url,
|
||||
"--history-backend",
|
||||
history_backend,
|
||||
"--log-level",
|
||||
"warn",
|
||||
]
|
||||
|
||||
# Note: Not adding --api-key to router command for local testing
|
||||
# The router will not require authentication
|
||||
|
||||
# Add Prometheus port to avoid conflicts (use unique port or disable)
|
||||
if prometheus_port is None:
|
||||
# Auto-assign a unique prometheus port based on router port
|
||||
prometheus_port = router_port + 1000
|
||||
router_cmd.extend(["--prometheus-port", str(prometheus_port)])
|
||||
|
||||
# Add router-specific args
|
||||
if router_args:
|
||||
router_cmd.extend(router_args)
|
||||
|
||||
if show_output:
|
||||
logger.info(f" Command: {' '.join(router_cmd)}")
|
||||
|
||||
# Set up environment with backend API key
|
||||
env = os.environ.copy()
|
||||
if backend == "openai":
|
||||
env["OPENAI_API_KEY"] = backend_api_key
|
||||
else:
|
||||
env["XAI_API_KEY"] = backend_api_key
|
||||
|
||||
# Launch router
|
||||
if show_output:
|
||||
router_proc = subprocess.Popen(
|
||||
router_cmd,
|
||||
env=env,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
else:
|
||||
router_proc = subprocess.Popen(
|
||||
router_cmd,
|
||||
stdout=stdout if stdout is not None else subprocess.PIPE,
|
||||
stderr=stderr if stderr is not None else subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
|
||||
print(f" PID: {router_proc.pid}")
|
||||
|
||||
# Wait for router to be ready
|
||||
router_url = f"http://127.0.0.1:{router_port}"
|
||||
print(f"\nWaiting for router to start at {router_url}...")
|
||||
|
||||
try:
|
||||
wait_for_router_ready(router_url, timeout=timeout, api_key=None)
|
||||
logger.info(f"✓ Router ready at {router_url}")
|
||||
except TimeoutError:
|
||||
logger.error(f"✗ Router failed to start")
|
||||
# Cleanup: kill router
|
||||
try:
|
||||
router_proc.kill()
|
||||
except:
|
||||
pass
|
||||
raise
|
||||
|
||||
logger.info(f"\n{'='*70}")
|
||||
logger.info(f"✓ {backend.upper()} router ready!")
|
||||
logger.info(f" Router: {router_url}")
|
||||
logger.info(f"{'='*70}\n")
|
||||
|
||||
return {
|
||||
"router": router_proc,
|
||||
"base_url": router_url,
|
||||
}
|
||||
|
||||
|
||||
def popen_launch_workers_and_router(
|
||||
model: str,
|
||||
base_url: str,
|
||||
timeout: int = 300,
|
||||
num_workers: int = 2,
|
||||
policy: str = "round_robin",
|
||||
api_key: Optional[str] = None,
|
||||
worker_args: Optional[list] = None,
|
||||
router_args: Optional[list] = None,
|
||||
tp_size: int = 1,
|
||||
env: Optional[dict] = None,
|
||||
stdout=None,
|
||||
stderr=None,
|
||||
) -> dict:
|
||||
"""
|
||||
Launch SGLang workers and gRPC router separately.
|
||||
|
||||
This approach:
|
||||
1. Starts N SGLang workers with --grpc-mode flag
|
||||
2. Waits for workers to initialize (process startup)
|
||||
3. Starts a gRPC router pointing to those workers
|
||||
4. Waits for router health check to pass (router validates worker connectivity)
|
||||
|
||||
This matches production deployment patterns better than the integrated approach.
|
||||
|
||||
Args:
|
||||
model: Model path (e.g., /home/ubuntu/models/llama-3.1-8b-instruct)
|
||||
base_url: Base URL for router (e.g., "http://127.0.0.1:8080")
|
||||
timeout: Timeout for server startup (default: 300s)
|
||||
num_workers: Number of workers to launch
|
||||
policy: Routing policy (round_robin, random, power_of_two, cache_aware)
|
||||
api_key: Optional API key for router
|
||||
worker_args: Additional arguments for workers (e.g., ["--context-len", "8192"])
|
||||
router_args: Additional arguments for router (e.g., ["--max-total-token", "1536"])
|
||||
tp_size: Tensor parallelism size for workers (default: 1)
|
||||
env: Optional environment variables for workers (e.g., {"SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION": "256"})
|
||||
stdout: Optional file handle for worker stdout (default: subprocess.PIPE)
|
||||
stderr: Optional file handle for worker stderr (default: subprocess.PIPE)
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
- workers: list of worker process objects
|
||||
- worker_urls: list of gRPC worker URLs
|
||||
- router: router process object
|
||||
- base_url: router URL (HTTP endpoint)
|
||||
|
||||
Example:
|
||||
>>> cluster = popen_launch_workers_and_router(model, base_url, num_workers=2)
|
||||
>>> # Use cluster['base_url'] for HTTP requests
|
||||
>>> # Cleanup:
|
||||
>>> for worker in cluster['workers']:
|
||||
>>> kill_process_tree(worker.pid)
|
||||
>>> kill_process_tree(cluster['router'].pid)
|
||||
"""
|
||||
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
|
||||
|
||||
# Parse router port from base_url
|
||||
if ":" in base_url.split("//")[-1]:
|
||||
router_port = int(base_url.split(":")[-1])
|
||||
else:
|
||||
router_port = find_free_port()
|
||||
|
||||
logger.info(f"\n{'='*70}")
|
||||
logger.info(f"Launching gRPC cluster (separate workers + router)")
|
||||
logger.info(f"{'='*70}")
|
||||
logger.info(f" Model: {model}")
|
||||
logger.info(f" Router port: {router_port}")
|
||||
logger.info(f" Workers: {num_workers}")
|
||||
logger.info(f" TP size: {tp_size}")
|
||||
logger.info(f" Policy: {policy}")
|
||||
|
||||
# Step 1: Launch workers with gRPC enabled
|
||||
workers = []
|
||||
worker_urls = []
|
||||
|
||||
for i in range(num_workers):
|
||||
worker_port = find_free_port()
|
||||
worker_url = f"grpc://127.0.0.1:{worker_port}"
|
||||
worker_urls.append(worker_url)
|
||||
|
||||
logger.info(f"\n[Worker {i+1}/{num_workers}]")
|
||||
logger.info(f" Port: {worker_port}")
|
||||
logger.info(f" URL: {worker_url}")
|
||||
|
||||
# Build worker command
|
||||
worker_cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang.launch_server",
|
||||
"--model-path",
|
||||
model,
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(worker_port),
|
||||
"--grpc-mode", # Enable gRPC for this worker
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
]
|
||||
|
||||
# Add TP size
|
||||
if tp_size > 1:
|
||||
worker_cmd.extend(["--tp-size", str(tp_size)])
|
||||
|
||||
# Add worker-specific args
|
||||
if worker_args:
|
||||
worker_cmd.extend(worker_args)
|
||||
|
||||
# Launch worker with optional environment variables
|
||||
if show_output:
|
||||
worker_proc = subprocess.Popen(
|
||||
worker_cmd,
|
||||
env=env,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
else:
|
||||
worker_proc = subprocess.Popen(
|
||||
worker_cmd,
|
||||
stdout=stdout if stdout is not None else subprocess.PIPE,
|
||||
stderr=stderr if stderr is not None else subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
|
||||
workers.append(worker_proc)
|
||||
logger.info(f" PID: {worker_proc.pid}")
|
||||
|
||||
# Give workers a moment to start binding to ports
|
||||
# The router will check worker health when it starts
|
||||
logger.info(f"\nWaiting for {num_workers} workers to initialize (20s)...")
|
||||
time.sleep(20)
|
||||
|
||||
# Quick check: make sure worker processes are still alive
|
||||
for i, worker in enumerate(workers):
|
||||
if worker.poll() is not None:
|
||||
logger.error(
|
||||
f" ✗ Worker {i+1} died during startup (exit code: {worker.poll()})"
|
||||
)
|
||||
# Cleanup: kill all workers
|
||||
for w in workers:
|
||||
try:
|
||||
w.kill()
|
||||
except:
|
||||
pass
|
||||
raise RuntimeError(f"Worker {i+1} failed to start")
|
||||
|
||||
logger.info(
|
||||
f"✓ All {num_workers} workers started (router will verify connectivity)"
|
||||
)
|
||||
|
||||
# Step 2: Launch router pointing to workers
|
||||
logger.info(f"\n[Router]")
|
||||
logger.info(f" Port: {router_port}")
|
||||
logger.info(f" Worker URLs: {', '.join(worker_urls)}")
|
||||
|
||||
# Build router command
|
||||
router_cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang_router.launch_router",
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(router_port),
|
||||
"--prometheus-port",
|
||||
"9321",
|
||||
"--policy",
|
||||
policy,
|
||||
"--model-path",
|
||||
model,
|
||||
"--log-level",
|
||||
"warn",
|
||||
]
|
||||
|
||||
# Add worker URLs
|
||||
router_cmd.append("--worker-urls")
|
||||
router_cmd.extend(worker_urls)
|
||||
|
||||
# Add API key
|
||||
if api_key:
|
||||
router_cmd.extend(["--api-key", api_key])
|
||||
|
||||
# Add router-specific args
|
||||
if router_args:
|
||||
router_cmd.extend(router_args)
|
||||
|
||||
if show_output:
|
||||
logger.info(f" Command: {' '.join(router_cmd)}")
|
||||
|
||||
# Launch router
|
||||
if show_output:
|
||||
router_proc = subprocess.Popen(router_cmd)
|
||||
else:
|
||||
router_proc = subprocess.Popen(
|
||||
router_cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
|
||||
logger.info(f" PID: {router_proc.pid}")
|
||||
|
||||
# Wait for router to be ready
|
||||
router_url = f"http://127.0.0.1:{router_port}"
|
||||
logger.info(f"\nWaiting for router to start at {router_url}...")
|
||||
|
||||
try:
|
||||
wait_for_workers_ready(
|
||||
router_url, expected_workers=num_workers, timeout=180, api_key=api_key
|
||||
)
|
||||
logger.info(f"✓ Router ready at {router_url}")
|
||||
except TimeoutError:
|
||||
logger.error(f"✗ Router failed to start")
|
||||
# Cleanup: kill router and all workers
|
||||
try:
|
||||
router_proc.kill()
|
||||
except:
|
||||
pass
|
||||
for worker in workers:
|
||||
try:
|
||||
worker.kill()
|
||||
except:
|
||||
pass
|
||||
raise
|
||||
|
||||
logger.info(f"\n{'='*70}")
|
||||
logger.info(f"✓ gRPC cluster ready!")
|
||||
logger.info(f" Router: {router_url}")
|
||||
logger.info(f" Workers: {len(workers)}")
|
||||
logger.info(f"{'='*70}\n")
|
||||
|
||||
return {
|
||||
"workers": workers,
|
||||
"worker_urls": worker_urls,
|
||||
"router": router_proc,
|
||||
"base_url": router_url,
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
Utility functions for Response API e2e tests.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
import psutil
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
|
||||
"""
|
||||
Kill the process and all its child processes.
|
||||
|
||||
Args:
|
||||
parent_pid: PID of the parent process
|
||||
include_parent: Whether to kill the parent process itself
|
||||
skip_pid: Optional PID to skip during cleanup
|
||||
"""
|
||||
# Remove sigchld handler to avoid spammy logs
|
||||
if threading.current_thread() is threading.main_thread():
|
||||
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
|
||||
|
||||
if parent_pid is None:
|
||||
parent_pid = os.getpid()
|
||||
include_parent = False
|
||||
|
||||
try:
|
||||
itself = psutil.Process(parent_pid)
|
||||
except psutil.NoSuchProcess:
|
||||
return
|
||||
|
||||
children = itself.children(recursive=True)
|
||||
for child in children:
|
||||
if child.pid == skip_pid:
|
||||
continue
|
||||
try:
|
||||
child.kill()
|
||||
except psutil.NoSuchProcess:
|
||||
pass
|
||||
|
||||
if include_parent:
|
||||
try:
|
||||
itself.kill()
|
||||
except psutil.NoSuchProcess:
|
||||
pass
|
||||
|
||||
|
||||
class CustomTestCase(unittest.TestCase):
|
||||
"""
|
||||
Custom test case base class with retry support.
|
||||
|
||||
This provides automatic test retry functionality based on environment variables.
|
||||
"""
|
||||
|
||||
def _callTestMethod(self, method):
|
||||
"""Override to add retry logic."""
|
||||
max_retry = int(os.environ.get("SGLANG_TEST_MAX_RETRY", "0"))
|
||||
|
||||
if max_retry == 0:
|
||||
# No retry, just run once
|
||||
return super(CustomTestCase, self)._callTestMethod(method)
|
||||
|
||||
# Retry logic
|
||||
for attempt in range(max_retry + 1):
|
||||
try:
|
||||
return super(CustomTestCase, self)._callTestMethod(method)
|
||||
except Exception as e:
|
||||
if attempt < max_retry:
|
||||
logger.info(
|
||||
f"Test failed on attempt {attempt + 1}/{max_retry + 1}, retrying..."
|
||||
)
|
||||
continue
|
||||
else:
|
||||
# Last attempt, re-raise the exception
|
||||
raise
|
||||
@@ -0,0 +1 @@
|
||||
"""Shared fixtures for router integration tests."""
|
||||
@@ -0,0 +1,236 @@
|
||||
"""
|
||||
Generate self-signed certificates for mTLS integration testing.
|
||||
Creates a Certificate Authority (CA), server certificates, and client certificates.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import ipaddress
|
||||
from pathlib import Path
|
||||
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.x509.oid import NameOID
|
||||
|
||||
|
||||
def generate_private_key():
|
||||
"""Generate an RSA private key."""
|
||||
return rsa.generate_private_key(
|
||||
public_exponent=65537,
|
||||
key_size=2048,
|
||||
)
|
||||
|
||||
|
||||
def generate_ca_certificate():
|
||||
"""Generate a self-signed CA certificate."""
|
||||
private_key = generate_private_key()
|
||||
|
||||
subject = issuer = x509.Name(
|
||||
[
|
||||
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
|
||||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, "Test CA"),
|
||||
]
|
||||
)
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(issuer)
|
||||
.public_key(private_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.utcnow())
|
||||
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
|
||||
.add_extension(
|
||||
x509.BasicConstraints(ca=True, path_length=None),
|
||||
critical=True,
|
||||
)
|
||||
.add_extension(
|
||||
x509.KeyUsage(
|
||||
digital_signature=True,
|
||||
key_cert_sign=True,
|
||||
crl_sign=True,
|
||||
key_encipherment=False,
|
||||
content_commitment=False,
|
||||
data_encipherment=False,
|
||||
key_agreement=False,
|
||||
encipher_only=False,
|
||||
decipher_only=False,
|
||||
),
|
||||
critical=True,
|
||||
)
|
||||
.sign(private_key, hashes.SHA256())
|
||||
)
|
||||
|
||||
return private_key, cert
|
||||
|
||||
|
||||
def generate_server_certificate(ca_key, ca_cert):
|
||||
"""Generate a server certificate signed by the CA."""
|
||||
private_key = generate_private_key()
|
||||
|
||||
subject = x509.Name(
|
||||
[
|
||||
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
|
||||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
|
||||
]
|
||||
)
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(ca_cert.subject)
|
||||
.public_key(private_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.utcnow())
|
||||
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
|
||||
.add_extension(
|
||||
x509.SubjectAlternativeName(
|
||||
[
|
||||
x509.DNSName("localhost"),
|
||||
x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
|
||||
]
|
||||
),
|
||||
critical=False,
|
||||
)
|
||||
.add_extension(
|
||||
x509.KeyUsage(
|
||||
digital_signature=True,
|
||||
key_encipherment=True,
|
||||
key_cert_sign=False,
|
||||
crl_sign=False,
|
||||
content_commitment=False,
|
||||
data_encipherment=False,
|
||||
key_agreement=False,
|
||||
encipher_only=False,
|
||||
decipher_only=False,
|
||||
),
|
||||
critical=True,
|
||||
)
|
||||
.add_extension(
|
||||
x509.ExtendedKeyUsage(
|
||||
[
|
||||
x509.oid.ExtendedKeyUsageOID.SERVER_AUTH,
|
||||
]
|
||||
),
|
||||
critical=False,
|
||||
)
|
||||
.sign(ca_key, hashes.SHA256())
|
||||
)
|
||||
|
||||
return private_key, cert
|
||||
|
||||
|
||||
def generate_client_certificate(ca_key, ca_cert):
|
||||
"""Generate a client certificate signed by the CA."""
|
||||
private_key = generate_private_key()
|
||||
|
||||
subject = x509.Name(
|
||||
[
|
||||
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
|
||||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, "test-client"),
|
||||
]
|
||||
)
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(ca_cert.subject)
|
||||
.public_key(private_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.utcnow())
|
||||
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
|
||||
.add_extension(
|
||||
x509.KeyUsage(
|
||||
digital_signature=True,
|
||||
key_encipherment=True,
|
||||
key_cert_sign=False,
|
||||
crl_sign=False,
|
||||
content_commitment=False,
|
||||
data_encipherment=False,
|
||||
key_agreement=False,
|
||||
encipher_only=False,
|
||||
decipher_only=False,
|
||||
),
|
||||
critical=True,
|
||||
)
|
||||
.add_extension(
|
||||
x509.ExtendedKeyUsage(
|
||||
[
|
||||
x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH,
|
||||
]
|
||||
),
|
||||
critical=False,
|
||||
)
|
||||
.sign(ca_key, hashes.SHA256())
|
||||
)
|
||||
|
||||
return private_key, cert
|
||||
|
||||
|
||||
def save_key(key, path: Path):
|
||||
"""Save private key to PEM file."""
|
||||
with open(path, "wb") as f:
|
||||
f.write(
|
||||
key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.TraditionalOpenSSL,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def save_cert(cert, path: Path):
|
||||
"""Save certificate to PEM file."""
|
||||
with open(path, "wb") as f:
|
||||
f.write(cert.public_bytes(serialization.Encoding.PEM))
|
||||
|
||||
|
||||
def generate_all_certificates(output_dir: Path):
|
||||
"""Generate all certificates and keys for mTLS testing."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print("==> Generating CA certificate...")
|
||||
ca_key, ca_cert = generate_ca_certificate()
|
||||
save_key(ca_key, output_dir / "ca-key.pem")
|
||||
save_cert(ca_cert, output_dir / "ca-cert.pem")
|
||||
|
||||
print("==> Generating server certificate...")
|
||||
server_key, server_cert = generate_server_certificate(ca_key, ca_cert)
|
||||
save_key(server_key, output_dir / "server-key.pem")
|
||||
save_cert(server_cert, output_dir / "server-cert.pem")
|
||||
|
||||
print("==> Generating client certificate...")
|
||||
client_key, client_cert = generate_client_certificate(ca_key, ca_cert)
|
||||
save_key(client_key, output_dir / "client-key.pem")
|
||||
save_cert(client_cert, output_dir / "client-cert.pem")
|
||||
|
||||
print(f"==> Certificates generated successfully in {output_dir}")
|
||||
print()
|
||||
print("Files created:")
|
||||
print(" - ca-cert.pem : CA certificate (for verifying server/client certs)")
|
||||
print(" - ca-key.pem : CA private key")
|
||||
print(" - server-cert.pem : Server certificate")
|
||||
print(" - server-key.pem : Server private key")
|
||||
print(" - client-cert.pem : Client certificate")
|
||||
print(" - client-key.pem : Client private key")
|
||||
print()
|
||||
print("Test server can use: server-cert.pem + server-key.pem")
|
||||
print("Test router can use: client-cert.pem + client-key.pem + ca-cert.pem")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
script_dir = Path(__file__).parent
|
||||
certs_dir = script_dir / "test_certs"
|
||||
generate_all_certificates(certs_dir)
|
||||
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
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)
|
||||
# TLS/mTLS configuration
|
||||
p.add_argument(
|
||||
"--ssl-certfile", type=str, default=None, help="Path to SSL certificate file"
|
||||
)
|
||||
p.add_argument("--ssl-keyfile", type=str, default=None, help="Path to SSL key file")
|
||||
p.add_argument(
|
||||
"--ssl-ca-certs",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to CA certificates for client verification",
|
||||
)
|
||||
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(request: Request):
|
||||
check_api_key(request)
|
||||
# Return format matching real workers: array of load info per DP rank
|
||||
return JSONResponse(
|
||||
[
|
||||
{
|
||||
"dp_rank": 0,
|
||||
"num_reqs": _inflight,
|
||||
"num_waiting_reqs": 0,
|
||||
"num_tokens": _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))
|
||||
|
||||
# Configure SSL if certificates are provided
|
||||
ssl_config = {}
|
||||
if args.ssl_certfile and args.ssl_keyfile:
|
||||
ssl_config["ssl_certfile"] = args.ssl_certfile
|
||||
ssl_config["ssl_keyfile"] = args.ssl_keyfile
|
||||
# If CA certs provided, require client certificates (mTLS)
|
||||
if args.ssl_ca_certs:
|
||||
ssl_config["ssl_ca_certs"] = args.ssl_ca_certs
|
||||
ssl_config["ssl_cert_reqs"] = 2 # ssl.CERT_REQUIRED
|
||||
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="warning", **ssl_config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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]
|
||||
@@ -0,0 +1,238 @@
|
||||
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",
|
||||
# mTLS configuration
|
||||
"client_cert_path": "--client-cert-path",
|
||||
"client_key_path": "--client-key-path",
|
||||
"ca_cert_paths": "--ca-cert-paths",
|
||||
}
|
||||
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)
|
||||
elif isinstance(v, list):
|
||||
# Handle list arguments (e.g., ca_cert_paths)
|
||||
if v: # Only add if list is not empty
|
||||
cmd.append(flag)
|
||||
cmd.extend([str(item) for item in v])
|
||||
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, timeout: float = 30.0) -> None:
|
||||
r = requests.post(f"{base_url}/workers", json={"url": worker_url})
|
||||
assert (
|
||||
r.status_code == 202
|
||||
), f"add_worker failed: {r.status_code} {r.text}" # ACCEPTED status
|
||||
|
||||
# Poll until worker is actually added and healthy
|
||||
from urllib.parse import quote
|
||||
|
||||
encoded_url = quote(worker_url, safe="")
|
||||
start = time.time()
|
||||
with requests.Session() as s:
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
r = s.get(f"{base_url}/workers/{encoded_url}", timeout=2)
|
||||
if r.status_code == 200:
|
||||
data = r.json()
|
||||
# Check if registration job failed
|
||||
job_status = data.get("job_status")
|
||||
if job_status and job_status.get("state") == "failed":
|
||||
raise RuntimeError(
|
||||
f"Worker registration failed: {job_status.get('message', 'Unknown error')}"
|
||||
)
|
||||
# Check if worker is healthy and registered (not just in job queue)
|
||||
if data.get("is_healthy", False):
|
||||
return
|
||||
# Worker not ready yet, continue polling
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.1)
|
||||
raise TimeoutError(
|
||||
f"Worker {worker_url} was not added and healthy after {timeout}s"
|
||||
)
|
||||
|
||||
def remove_worker(
|
||||
self, base_url: str, worker_url: str, timeout: float = 30.0
|
||||
) -> None:
|
||||
# URL encode the worker_url for path parameter
|
||||
from urllib.parse import quote
|
||||
|
||||
encoded_url = quote(worker_url, safe="")
|
||||
r = requests.delete(f"{base_url}/workers/{encoded_url}")
|
||||
assert (
|
||||
r.status_code == 202
|
||||
), f"remove_worker failed: {r.status_code} {r.text}" # ACCEPTED status
|
||||
|
||||
# Poll until worker is actually removed (GET returns 404) or timeout
|
||||
start = time.time()
|
||||
last_status = None
|
||||
with requests.Session() as s:
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
r = s.get(f"{base_url}/workers/{encoded_url}", timeout=2)
|
||||
if r.status_code == 404:
|
||||
# Worker successfully removed
|
||||
return
|
||||
elif r.status_code == 200:
|
||||
# Check if removal job failed
|
||||
data = r.json()
|
||||
job_status = data.get("job_status")
|
||||
if job_status:
|
||||
last_status = job_status
|
||||
if job_status.get("state") == "failed":
|
||||
raise RuntimeError(
|
||||
f"Worker removal failed: {job_status.get('message', 'Unknown error')}"
|
||||
)
|
||||
# Worker still being processed, continue polling
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.1)
|
||||
|
||||
# Provide detailed timeout error with last known status
|
||||
error_msg = f"Worker {worker_url} was not removed after {timeout}s"
|
||||
if last_status:
|
||||
error_msg += f". Last job status: {last_status}"
|
||||
raise TimeoutError(error_msg)
|
||||
|
||||
def list_workers(self, base_url: str) -> list[str]:
|
||||
r = requests.get(f"{base_url}/workers")
|
||||
assert r.status_code == 200, f"list_workers failed: {r.status_code} {r.text}"
|
||||
data = r.json()
|
||||
# Extract URLs from WorkerInfo objects
|
||||
workers = data.get("workers", [])
|
||||
return [w["url"] for w in workers]
|
||||
|
||||
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()
|
||||
@@ -0,0 +1 @@
|
||||
"""Integration test package for the router."""
|
||||
@@ -0,0 +1,128 @@
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Iterable, List, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from ..fixtures.generate_test_certs import generate_all_certificates
|
||||
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: Optional[List[str]] = 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()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def test_certificates():
|
||||
"""Generate test certificates for mTLS tests, clean up after session."""
|
||||
# Get the test_certs directory path
|
||||
fixtures_dir = Path(__file__).parent.parent / "fixtures"
|
||||
certs_dir = fixtures_dir / "test_certs"
|
||||
|
||||
# Generate certificates
|
||||
generate_all_certificates(certs_dir)
|
||||
|
||||
# Yield the path to the certificates directory
|
||||
yield certs_dir
|
||||
|
||||
# Cleanup: remove the generated certificates
|
||||
if certs_dir.exists():
|
||||
shutil.rmtree(certs_dir)
|
||||
@@ -0,0 +1 @@
|
||||
"""Load balancing integration tests."""
|
||||
@@ -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
|
||||
@@ -0,0 +1,99 @@
|
||||
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]
|
||||
slow_url = urls_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
|
||||
|
||||
# Start background load in a non-blocking way to keep slow worker busy
|
||||
background_executor = concurrent.futures.ThreadPoolExecutor(max_workers=8)
|
||||
background_futures = []
|
||||
for i in range(32):
|
||||
future = background_executor.submit(_direct_load, i)
|
||||
background_futures.append(future)
|
||||
|
||||
# Wait longer for the load monitor to update (at least 2 monitor intervals)
|
||||
time.sleep(3)
|
||||
|
||||
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
|
||||
|
||||
# Clean up background executor
|
||||
background_executor.shutdown(wait=False)
|
||||
|
||||
# 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
|
||||
@@ -0,0 +1,32 @@
|
||||
import collections
|
||||
|
||||
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
|
||||
@@ -0,0 +1,33 @@
|
||||
import collections
|
||||
|
||||
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
|
||||
@@ -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
|
||||
@@ -0,0 +1,228 @@
|
||||
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,
|
||||
)
|
||||
|
||||
# should see 500 when worker actually starts, before that should see 503
|
||||
saw_500 = False
|
||||
for _ in range(8):
|
||||
r = post_once()
|
||||
if r.status_code == 500:
|
||||
# Worker starts, continue to circuit breaker test
|
||||
saw_500 = True
|
||||
break
|
||||
assert (
|
||||
r.status_code == 503
|
||||
), "Should only see 503 when waiting for worker to start"
|
||||
assert saw_500, "Worker didn't start after 8 requests"
|
||||
|
||||
saw_503 = False
|
||||
for _ in range(4):
|
||||
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,
|
||||
)
|
||||
|
||||
# should see 500 when worker actually starts, before that should see 503
|
||||
saw_500 = False
|
||||
for _ in range(8):
|
||||
r = post_once()
|
||||
if r.status_code == 500:
|
||||
# Worker starts, continue to circuit breaker test
|
||||
saw_500 = True
|
||||
break
|
||||
assert (
|
||||
r.status_code == 503
|
||||
), "Should only see 503 when waiting for worker to start"
|
||||
assert saw_500, "Worker didn't start after 8 requests"
|
||||
|
||||
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,
|
||||
},
|
||||
)
|
||||
|
||||
saw_500 = False
|
||||
for _ in range(8):
|
||||
r = requests.post(
|
||||
f"{rh.url}/v1/completions",
|
||||
json={
|
||||
"model": "test-model",
|
||||
"prompt": "x",
|
||||
"max_tokens": 1,
|
||||
"stream": False,
|
||||
},
|
||||
timeout=3,
|
||||
)
|
||||
if r.status_code == 500:
|
||||
# Worker starts, continue to check
|
||||
saw_500 = True
|
||||
break
|
||||
assert (
|
||||
r.status_code == 503
|
||||
), "Should only see 503 when waiting for worker to start"
|
||||
|
||||
assert saw_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
|
||||
@@ -0,0 +1,32 @@
|
||||
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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user