Skip to content

GitLab

"test/vscode:/vscode.git/clone" did not exist on "6b242c29af317498ea51b465529cf8e68c2c88fd"
Unverified Commit c8f31042 authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[router] Refactor router and policy traits with dependency injection (#7987) 


Co-authored-by: default avatarJin Pan <jpan236@wisc.edu>
Co-authored-by: default avatarKeru Yang <rukeyang@gmail.com>
Co-authored-by: default avatarYingyi Huang <yingyihuang2000@outlook.com>
Co-authored-by: default avatarPhilip Zhu <phlipzhux@gmail.com>
parent 1f76fc87
Show whitespace changes
Inline Side-by-side
Showing with 3076 additions and 499 deletions
.github/workflows/pr-test-pd-router.yml
View file @ c8f31042
......@@ -131,36 +131,83 @@ jobs:
SERVER_PID=$!
echo "server_pid=$SERVER_PID" >> $GITHUB_OUTPUT
# Wait for all 8 servers to be healthy (script already does this)
wait_count=0
while [ $wait_count -lt 30 ]; do
if ps -p $SERVER_PID > /dev/null; then
# Check if the startup script printed success message
sleep 2
wait_count=$((wait_count + 1))
else
# Script exited - check if it was successful
wait $SERVER_PID
exit_code=$?
if [ $exit_code -eq 0 ]; then
echo "✓ All disaggregation servers are healthy"
break
else
echo "Error: Server startup failed with code $exit_code"
exit 1
fi
fi
done
echo "✓ Servers started (PID: $SERVER_PID)"
- name: Test all policies sequentially
timeout-minutes: 30
run: |
POLICIES=("random" "round_robin" "cache_aware" "power_of_two")
BASE_URL="http://127.0.0.9:8000"
for policy in "${POLICIES[@]}"; do
echo ""
echo "=================================================="
echo "Testing policy: $policy"
echo "=================================================="
# Start router with the current policy
echo "Starting router with policy: $policy..."
python3 -m sglang_router.launch_router \
--pd-disaggregation \
--policy "$policy" \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 9002 \
--prefill http://127.0.0.3:30003 9003 \
--prefill http://127.0.0.4:30004 9004 \
--decode http://127.0.0.5:30005 \
--decode http://127.0.0.6:30006 \
--decode http://127.0.0.7:30007 \
--decode http://127.0.0.8:30008 \
--host 127.0.0.9 \
--port 8000 &
ROUTER_PID=$!
# Wait for router to become healthy
echo "Waiting for router to become healthy..."
TIMEOUT=300
TIMEOUT=60
ELAPSED=0
while [ $ELAPSED -lt $TIMEOUT ]; do
if curl --connect-timeout 5 --silent http://127.0.0.9:8000 > /dev/null 2>&1; then
echo "✓ Router is reachable"
break
fi
if ! ps -p $SERVER_PID > /dev/null; then
echo "Error: Server processes failed to start"
if ! ps -p $ROUTER_PID > /dev/null; then
echo "Error: Router process died"
exit 1
fi
echo "Waiting for router... (${ELAPSED}s/${TIMEOUT}s)"
sleep 10
ELAPSED=$((ELAPSED + 10))
sleep 5
ELAPSED=$((ELAPSED + 5))
done
if [ $ELAPSED -ge $TIMEOUT ]; then
echo "Error: Router health check timeout after ${TIMEOUT}s"
echo "Error: Router health check timeout"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
echo "✓ Servers started and healthy (PID: $SERVER_PID)"
- name: Test API functionality
timeout-minutes: 5
run: |
BASE_URL="http://127.0.0.9:8000"
echo "Testing API completions..."
# Test API functionality
echo "Testing API completions for $policy..."
response=$(curl -s -X POST "$BASE_URL/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer test-token" \
......@@ -174,13 +221,15 @@ jobs:
}')
if echo "$response" | jq -e '.choices[0].message.content' > /dev/null 2>&1; then
echo "✓ API test passed"
echo "✓ API test passed for $policy"
else
echo "✗ API test failed: $response"
echo "✗ API test failed for $policy: $response"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
echo "Testing streaming API..."
# Test streaming
echo "Testing streaming API for $policy..."
stream_response=$(timeout 30 curl -s -X POST "$BASE_URL/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer test-token" \
......@@ -194,16 +243,15 @@ jobs:
}')
if echo "$stream_response" | grep -q "data:"; then
echo "✓ Streaming API test passed"
echo "✓ Streaming API test passed for $policy"
else
echo "✗ Streaming API test failed"
echo "✗ Streaming API test failed for $policy"
kill $ROUTER_PID 2>/dev/null || true
exit 1
fi
- name: Run benchmark test
timeout-minutes: 5
run: |
echo "Running benchmark test..."
# Run benchmark
echo "Running benchmark for $policy..."
benchmark_output=$(python3 -m sglang.bench_one_batch_server \
--model-path "/raid/models/meta-llama/Llama-3.1-8B-Instruct" \
--base-url "http://127.0.0.9:8000" \
......@@ -214,28 +262,69 @@ jobs:
echo "$benchmark_output"
# Extract metrics from output
# Save benchmark output
echo "$benchmark_output" > "benchmark_${policy}.txt"
# Extract and validate metrics
latency=$(echo "$benchmark_output" | grep "latency:" | awk '{print $2}' | sed 's/s//')
input_throughput=$(echo "$benchmark_output" | grep "input throughput:" | awk '{print $3}')
output_throughput=$(echo "$benchmark_output" | grep "output throughput:" | awk '{print $3}')
# Validate performance (latency<1.5s, input>20k, output>1k)
command -v bc >/dev/null || (apt-get update && apt-get install -y bc)
echo "Performance: ${latency}s | ${input_throughput} | ${output_throughput} tok/s"
echo "Performance for $policy: ${latency}s | ${input_throughput} | ${output_throughput} tok/s"
# Validate performance
fail=""
(( $(echo "$latency > 1.5" | bc -l) )) && fail="Latency too high (${latency}s>1.5s) "
(( $(echo "$input_throughput < 20000" | bc -l) )) && fail="${fail}Input too low (${input_throughput}<20k) "
(( $(echo "$output_throughput < 1000" | bc -l) )) && fail="${fail}Output too low (${output_throughput}<1k) "
if [ -n "$fail" ]; then
echo "✗ Benchmark failed: $fail"
echo "✗ Benchmark failed for $policy: $fail"
kill $ROUTER_PID 2>/dev/null || true
exit 1
else
echo "✓ Performance validation passed"
echo "✓ Performance validation passed for $policy"
fi
# Stop router before testing next policy
echo "Stopping router for $policy..."
# First try graceful shutdown
kill $ROUTER_PID 2>/dev/null || true
# Wait up to 5 seconds for graceful shutdown
for i in {1..5}; do
if ! ps -p $ROUTER_PID > /dev/null 2>&1; then
echo "Router stopped gracefully"
break
fi
sleep 1
done
# Force kill if still running
if ps -p $ROUTER_PID > /dev/null 2>&1; then
echo "Force killing router..."
kill -9 $ROUTER_PID 2>/dev/null || true
fi
# Short delay to ensure port is released
sleep 2
echo "✓ Completed testing for $policy"
done
echo ""
echo "✅ All policies tested successfully!"
- name: Upload benchmark results
if: success()
uses: actions/upload-artifact@v4
with:
name: benchmark-results-all-policies
path: benchmark_*.txt
- name: Cleanup servers
if: always()
run: |
......@@ -247,3 +336,34 @@ jobs:
sleep 5
remaining=$(ps aux | grep -c "sglang.launch_server" || echo "0")
echo "Cleanup completed. Remaining processes: $remaining"
summarize-benchmarks:
needs: test-disaggregation
runs-on: ubuntu-latest
if: success()
steps:
- name: Download benchmark results
uses: actions/download-artifact@v4
with:
name: benchmark-results-all-policies
- name: Create benchmark summary
run: |
echo "## PD Router Benchmark Results Summary" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "| Policy | Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY
echo "|--------|-------------|-------------------------|--------------------------|" >> $GITHUB_STEP_SUMMARY
for policy in random round_robin cache_aware power_of_two; do
if [ -f "benchmark_${policy}.txt" ]; then
latency=$(grep "latency:" "benchmark_${policy}.txt" | awk '{print $2}')
input_throughput=$(grep "input throughput:" "benchmark_${policy}.txt" | awk '{print $3}')
output_throughput=$(grep "output throughput:" "benchmark_${policy}.txt" | awk '{print $3}')
echo "| ${policy} | ${latency} | ${input_throughput} | ${output_throughput} |" >> $GITHUB_STEP_SUMMARY
fi
done
echo "" >> $GITHUB_STEP_SUMMARY
echo "✅ All policies tested successfully!" >> $GITHUB_STEP_SUMMARY
scripts/ci_start_disaggregation_servers.sh
View file @ c8f31042
......@@ -87,20 +87,8 @@ while true; do
fi
done
# Launch the router
echo "Launching router at 127.0.0.9:8000..."
python3 -m sglang_router.launch_router \
--pd-disaggregation \
--policy power_of_two \
--prefill http://127.0.0.1:30001 9001 \
--prefill http://127.0.0.2:30002 9002 \
--prefill http://127.0.0.3:30003 9003 \
--prefill http://127.0.0.4:30004 9004 \
--decode http://127.0.0.5:30005 \
--decode http://127.0.0.6:30006 \
--decode http://127.0.0.7:30007 \
--decode http://127.0.0.8:30008 \
--host 127.0.0.9 \
--port 8000 &
# Don't launch router here - just keep servers running
echo "✅ All disaggregation servers are ready and waiting for router connections"
wait # Wait for all background jobs to finish
# Keep the script running
wait # Wait for all background server jobs
sgl-router/benches/request_processing.rs
View file @ c8f31042
......@@ -6,7 +6,7 @@ use sglang_router_rs::openai_api_types::{
ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateParameters, GenerateRequest,
SamplingParams, StringOrArray, UserMessageContent,
};
use sglang_router_rs::request_adapter::{RouteableRequest, ToPdRequest};
use sglang_router_rs::routers::request_adapter::{RouteableRequest, ToPdRequest};
// Sample request data for benchmarks
fn create_sample_generate_request() -> GenerateRequest {
......
sgl-router/py_test/test_launch_router.py
View file @ c8f31042
......@@ -164,56 +164,47 @@ class TestLaunchRouter(unittest.TestCase):
"""Test that policy validation works correctly for PD and regular modes."""
from sglang_router.launch_router import RouterArgs, launch_router
# Test 1: PowerOfTwo is only valid in PD mode
# Test 1: PowerOfTwo requires at least 2 workers
args = self.create_router_args(
pd_disaggregation=False,
policy="power_of_two",
worker_urls=["http://localhost:8000"],
worker_urls=["http://localhost:8000"], # Only 1 worker
)
# Should raise error
with self.assertRaises(ValueError) as cm:
launch_router(args)
self.assertIn(
"PowerOfTwo policy is only supported in PD disaggregated mode",
"Power-of-two policy requires at least 2 workers",
str(cm.exception),
)
# Test 2: RoundRobin is not valid in PD mode
# Test 2: PowerOfTwo with sufficient workers should succeed
args = self.create_router_args(
pd_disaggregation=True,
policy="round_robin",
prefill=[["http://prefill1:8080", "9000"]],
decode=[["http://decode1:8081"]],
worker_urls=[],
)
# Should raise error
with self.assertRaises(ValueError) as cm:
launch_router(args)
self.assertIn(
"RoundRobin policy is not supported in PD disaggregated mode",
str(cm.exception),
pd_disaggregation=False,
policy="power_of_two",
worker_urls=["http://localhost:8000", "http://localhost:8001"], # 2 workers
)
# This should not raise an error (validation passes)
# Test 3: Valid combinations should not raise errors
# Test 3: All policies now work in both modes
# Regular mode with RoundRobin
args = self.create_router_args(
pd_disaggregation=False,
policy="round_robin",
worker_urls=["http://localhost:8000"],
)
# This should not raise (though it may fail to connect)
# This should not raise validation error
# PD mode with PowerOfTwo
# PD mode with RoundRobin (now supported!)
args = self.create_router_args(
pd_disaggregation=True,
policy="power_of_two",
policy="round_robin",
prefill=[["http://prefill1:8080", "9000"]],
decode=[["http://decode1:8081"]],
worker_urls=[],
)
# This should not raise (though it may fail to connect)
# This should not raise validation error
def test_pd_service_discovery_args_parsing(self):
"""Test PD service discovery CLI argument parsing."""
......
sgl-router/src/config/types.rs
View file @ c8f31042
use super::{ConfigError, ConfigResult};
use super::ConfigResult;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
......@@ -215,6 +215,7 @@ impl RouterConfig {
self.metrics.is_some()
}
/* Commented out - no longer needed without compatibility layer
/// Convert to routing PolicyConfig for internal use
pub fn to_routing_policy_config(&self) -> ConfigResult<crate::router::PolicyConfig> {
match (&self.mode, &self.policy) {
......@@ -291,4 +292,5 @@ impl RouterConfig {
}
}
}
*/
}
sgl-router/src/config/validation.rs
View file @ c8f31042
......@@ -255,29 +255,8 @@ impl ConfigValidator {
/// Validate compatibility between different configuration sections
fn validate_compatibility(config: &RouterConfig) -> ConfigResult<()> {
// Check mode and policy compatibility
match (&config.mode, &config.policy) {
(RoutingMode::Regular { .. }, PolicyConfig::PowerOfTwo { .. }) => {
// PowerOfTwo is only supported in PD mode
return Err(ConfigError::IncompatibleConfig {
reason: "PowerOfTwo policy is only supported in PD disaggregated mode"
.to_string(),
});
}
(RoutingMode::PrefillDecode { .. }, PolicyConfig::RoundRobin) => {
return Err(ConfigError::IncompatibleConfig {
reason: "RoundRobin policy is not supported in PD disaggregated mode"
.to_string(),
});
}
(RoutingMode::PrefillDecode { .. }, PolicyConfig::CacheAware { .. }) => {
return Err(ConfigError::IncompatibleConfig {
reason: "CacheAware policy is not supported in PD disaggregated mode"
.to_string(),
});
}
_ => {}
}
// All policies are now supported for both router types thanks to the unified trait design
// No mode/policy restrictions needed anymore
// Check if service discovery is enabled for worker count validation
let has_service_discovery = config.discovery.as_ref().map_or(false, |d| d.enabled);
......@@ -459,8 +438,8 @@ mod tests {
}
#[test]
fn test_validate_incompatible_policy() {
// RoundRobin with PD mode
fn test_validate_roundrobin_with_pd_mode() {
// RoundRobin with PD mode is now supported
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![("http://prefill:8000".to_string(), None)],
......@@ -470,16 +449,12 @@ mod tests {
);
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("RoundRobin policy is not supported in PD disaggregated mode"));
assert!(result.is_ok());
}
#[test]
fn test_validate_cache_aware_with_pd_mode() {
// CacheAware with PD mode should fail
// CacheAware with PD mode is now supported
let config = RouterConfig::new(
RoutingMode::PrefillDecode {
prefill_urls: vec![("http://prefill:8000".to_string(), None)],
......@@ -495,16 +470,12 @@ mod tests {
);
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("CacheAware policy is not supported in PD disaggregated mode"));
assert!(result.is_ok());
}
#[test]
fn test_validate_power_of_two_with_regular_mode() {
// PowerOfTwo with Regular mode should fail
// PowerOfTwo with Regular mode is now supported
let config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec![
......@@ -518,10 +489,6 @@ mod tests {
);
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("PowerOfTwo policy is only supported in PD disaggregated mode"));
assert!(result.is_ok());
}
}
sgl-router/src/lib.rs
View file @ c8f31042
......@@ -4,11 +4,9 @@ pub mod logging;
use std::collections::HashMap;
pub mod core;
pub mod openai_api_types;
pub mod pd_router;
pub mod pd_types;
pub mod policies;
pub mod prometheus;
pub mod request_adapter;
pub mod router;
pub mod routers;
pub mod server;
pub mod service_discovery;
pub mod tree;
......@@ -241,11 +239,6 @@ impl Router {
))
})?;
// Convert to internal policy config
let policy_config = router_config
.to_routing_policy_config()
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
// Create service discovery config if enabled
let service_discovery_config = if self.service_discovery {
Some(service_discovery::ServiceDiscoveryConfig {
......@@ -282,8 +275,7 @@ impl Router {
server::startup(server::ServerConfig {
host: self.host.clone(),
port: self.port,
worker_urls: self.worker_urls.clone(),
policy_config,
router_config,
max_payload_size: self.max_payload_size,
log_dir: self.log_dir.clone(),
log_level: self.log_level.clone(),
......
sgl-router/src/policies/cache_aware.rs 0 → 100644
View file @ c8f31042
/*
Cache-Aware Load Balancing Router
This router combines two strategies to optimize both cache utilization and request distribution:
1. Cache-Aware Routing (Approximate Tree)
2. Load Balancing (Shortest Queue with Balance Thresholds)
The router dynamically switches between these strategies based on load conditions:
- Uses load balancing when the system is imbalanced
- Uses cache-aware routing when the system is balanced
A system is considered imbalanced if both conditions are met:
1. (max - min) > abs_threshold
2. max > rel_threshold * min
Strategy Details:
1. Cache-Aware Routing (Approximate Tree)
-------------------------------------------
This strategy maintains an approximate radix tree for each worker based on request history,
eliminating the need for direct cache state queries. The tree stores raw text characters
instead of token IDs to avoid tokenization overhead.
Process:
a. For each request, find the worker with the highest prefix match
b. If match rate > cache_threshold:
Route to the worker with highest match (likely has relevant data cached)
c. If match rate ≤ cache_threshold:
Route to the worker with smallest tree size (most available cache capacity)
d. Background maintenance:
Periodically evict least recently used leaf nodes to prevent memory overflow
2. Load Balancing (Shortest Queue)
-------------------------------------------
This strategy tracks pending request counts per worker and routes new requests
to the least busy worker when the system is detected to be imbalanced.
Configuration Parameters:
------------------------
1. cache_threshold: (float, 0.0 to 1.0)
Minimum prefix match ratio to use highest-match routing.
Below this threshold, routes to worker with most available cache space.
2. balance_abs_threshold: (integer)
Absolute difference threshold for load imbalance detection.
System is potentially imbalanced if (max_load - min_load) > abs_threshold
3. balance_rel_threshold: (float)
Relative ratio threshold for load imbalance detection.
System is potentially imbalanced if max_load > min_load * rel_threshold
Used in conjunction with abs_threshold to determine final imbalance state.
4. eviction_interval_secs: (integer)
Interval between LRU eviction cycles for the approximate trees.
5. max_tree_size: (integer)
Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted
during the next eviction cycle.
*/
use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy};
use crate::core::Worker;
use crate::tree::Tree;
use metrics::{counter, gauge};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use tracing::{debug, info};
/// Cache-aware routing policy
///
/// Routes requests based on cache affinity when load is balanced,
/// switches to shortest-queue routing when load is imbalanced.
#[derive(Debug)]
pub struct CacheAwarePolicy {
config: CacheAwareConfig,
tree: Arc<Mutex<Tree>>,
eviction_handle: Option<thread::JoinHandle<()>>,
}
impl CacheAwarePolicy {
pub fn new() -> Self {
Self::with_config(CacheAwareConfig::default())
}
pub fn with_config(config: CacheAwareConfig) -> Self {
let tree = Arc::new(Mutex::new(Tree::new()));
// Start background eviction thread if configured
let eviction_handle = if config.eviction_interval_secs > 0 {
let tree_clone = Arc::clone(&tree);
let max_tree_size = config.max_tree_size;
let interval = config.eviction_interval_secs;
Some(thread::spawn(move || loop {
thread::sleep(Duration::from_secs(interval));
if let Ok(tree_guard) = tree_clone.lock() {
tree_guard.evict_tenant_by_size(max_tree_size);
debug!("Cache eviction completed, max_size: {}", max_tree_size);
}
}))
} else {
None
};
Self {
config,
tree,
eviction_handle,
}
}
/// Initialize the tree with worker URLs
pub fn init_workers(&self, workers: &[Box<dyn Worker>]) {
if let Ok(tree) = self.tree.lock() {
for worker in workers {
tree.insert("", worker.url());
}
}
}
/// Remove a worker from the tree
pub fn remove_worker(&self, url: &str) {
if let Ok(tree) = self.tree.lock() {
tree.remove_tenant(url);
}
}
/// Run cache eviction to prevent unbounded growth
pub fn evict_cache(&self, max_size: usize) {
if let Ok(tree) = self.tree.lock() {
tree.evict_tenant_by_size(max_size);
}
}
}
impl LoadBalancingPolicy for CacheAwarePolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
// Get current load statistics
let loads: Vec<usize> = workers.iter().map(|w| w.load()).collect();
let max_load = *loads.iter().max().unwrap_or(&0);
let min_load = *loads.iter().min().unwrap_or(&0);
// Check if load is imbalanced
let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold
&& (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold);
if is_imbalanced {
// Log load balancing trigger
let worker_loads: Vec<(String, usize)> = workers
.iter()
.map(|w| (w.url().to_string(), w.load()))
.collect();
info!(
"Load balancing triggered due to workload imbalance:\n\
Max load: {}, Min load: {}\n\
Current worker loads: {:?}",
max_load, min_load, worker_loads
);
counter!("sgl_router_load_balancing_events_total").increment(1);
gauge!("sgl_router_max_load").set(max_load as f64);
gauge!("sgl_router_min_load").set(min_load as f64);
// Use shortest queue when imbalanced
let min_load_idx = healthy_indices
.iter()
.min_by_key(|&&idx| workers[idx].load())
.copied()?;
// Increment processed counter
workers[min_load_idx].increment_processed();
counter!("sgl_router_processed_requests_total", "worker" => workers[min_load_idx].url().to_string())
.increment(1);
return Some(min_load_idx);
}
// Use cache-aware routing when balanced
let text = request_text.unwrap_or("");
if let Ok(tree) = self.tree.lock() {
let (matched_text, matched_worker) = tree.prefix_match(text);
let match_rate = if text.is_empty() {
0.0
} else {
matched_text.chars().count() as f32 / text.chars().count() as f32
};
let selected_url = if match_rate > self.config.cache_threshold {
counter!("sgl_router_cache_hits_total").increment(1);
matched_worker.to_string()
} else {
counter!("sgl_router_cache_misses_total").increment(1);
tree.get_smallest_tenant()
};
// Find the index of the selected worker
let selected_idx = workers.iter().position(|w| w.url() == selected_url)?;
// Only proceed if the worker is healthy
if !workers[selected_idx].is_healthy() {
return healthy_indices.first().copied();
}
// Update the tree with this request
tree.insert(text, &selected_url);
// Increment processed counter
workers[selected_idx].increment_processed();
counter!("sgl_router_processed_requests_total", "worker" => selected_url).increment(1);
return Some(selected_idx);
}
// Fallback to first healthy worker if tree operations fail
healthy_indices.first().copied()
}
fn name(&self) -> &'static str {
"cache_aware"
}
fn on_request_complete(&self, worker_url: &str, success: bool) {
// Could track success rates per worker for more intelligent routing
if !success {
// Optionally reduce affinity for failed requests
tracing::debug!(
"Request to {} completed with success={}",
worker_url,
success
);
}
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn select_worker_pair(
&self,
prefill_workers: &[Box<dyn Worker>],
decode_workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<(usize, usize)> {
// In PD mode:
// - Prefill: Use cache-aware routing for better cache utilization
// - Decode: Use least-load routing for better load distribution
// Select prefill worker using cache-aware logic
let prefill_idx = self.select_worker(prefill_workers, request_text)?;
// Select decode worker using least-load logic
let healthy_decode = get_healthy_worker_indices(decode_workers);
if healthy_decode.is_empty() {
return None;
}
let decode_idx = healthy_decode
.iter()
.min_by_key(|&&idx| decode_workers[idx].load())
.copied()?;
Some((prefill_idx, decode_idx))
}
}
impl Default for CacheAwarePolicy {
fn default() -> Self {
Self::new()
}
}
impl Drop for CacheAwarePolicy {
fn drop(&mut self) {
// Note: We can't properly stop the eviction thread since it's in an infinite loop
// In a production system, we'd use a channel or atomic flag to signal shutdown
if let Some(handle) = self.eviction_handle.take() {
// The thread will continue running until the program exits
// This is acceptable for now since the router typically runs for the lifetime of the program
drop(handle);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_cache_aware_with_balanced_load() {
// Create policy without eviction thread for testing
let config = CacheAwareConfig {
eviction_interval_secs: 0, // Disable eviction thread
..Default::default()
};
let policy = CacheAwarePolicy::with_config(config);
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Initialize the policy with workers
policy.init_workers(&workers);
// First request should be distributed
let idx1 = policy.select_worker(&workers, Some("hello world")).unwrap();
// Same request should go to same worker (cache hit)
let idx2 = policy.select_worker(&workers, Some("hello world")).unwrap();
assert_eq!(idx1, idx2);
// Similar request should also go to same worker
let idx3 = policy.select_worker(&workers, Some("hello")).unwrap();
assert_eq!(idx1, idx3);
}
#[test]
fn test_cache_aware_with_imbalanced_load() {
let policy = CacheAwarePolicy::with_config(CacheAwareConfig {
cache_threshold: 0.5,
balance_abs_threshold: 5,
balance_rel_threshold: 2.0,
eviction_interval_secs: 0, // Disable eviction thread
max_tree_size: 10000,
});
let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular);
let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular);
// Create significant load imbalance
for _ in 0..20 {
worker1.increment_load();
}
// worker2 has load 0
let workers: Vec<Box<dyn Worker>> = vec![Box::new(worker1), Box::new(worker2)];
policy.init_workers(&workers);
// Should select worker2 (lower load) despite cache affinity
for _ in 0..5 {
let idx = policy.select_worker(&workers, Some("test")).unwrap();
assert_eq!(idx, 1); // Should always pick worker2
}
}
#[test]
fn test_cache_aware_worker_removal() {
let config = CacheAwareConfig {
eviction_interval_secs: 0, // Disable eviction thread
..Default::default()
};
let policy = CacheAwarePolicy::with_config(config);
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
policy.init_workers(&workers);
// Route some requests
policy.select_worker(&workers, Some("test1"));
policy.select_worker(&workers, Some("test2"));
// Remove a worker
policy.remove_worker("http://w1:8000");
workers[0].set_healthy(false);
// All requests should now go to worker2
let idx = policy.select_worker(&workers, Some("test1")).unwrap();
assert_eq!(idx, 1);
}
}
sgl-router/src/policies/factory.rs 0 → 100644
View file @ c8f31042
//! Factory for creating load balancing policies
use super::{
CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy,
RoundRobinPolicy,
};
use crate::config::PolicyConfig;
use std::sync::Arc;
/// Factory for creating policy instances
pub struct PolicyFactory;
impl PolicyFactory {
/// Create a policy from configuration
pub fn create_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
match config {
PolicyConfig::Random => Arc::new(RandomPolicy::new()),
PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
PolicyConfig::CacheAware {
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
eviction_interval_secs,
max_tree_size,
} => {
let config = CacheAwareConfig {
cache_threshold: *cache_threshold,
balance_abs_threshold: *balance_abs_threshold,
balance_rel_threshold: *balance_rel_threshold,
eviction_interval_secs: *eviction_interval_secs,
max_tree_size: *max_tree_size,
};
Arc::new(CacheAwarePolicy::with_config(config))
}
}
}
/// Create a policy by name (for dynamic loading)
pub fn create_by_name(name: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
match name.to_lowercase().as_str() {
"random" => Some(Arc::new(RandomPolicy::new())),
"round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())),
"power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())),
"cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_from_config() {
// Test Random
let policy = PolicyFactory::create_from_config(&PolicyConfig::Random);
assert_eq!(policy.name(), "random");
// Test RoundRobin
let policy = PolicyFactory::create_from_config(&PolicyConfig::RoundRobin);
assert_eq!(policy.name(), "round_robin");
// Test PowerOfTwo
let policy = PolicyFactory::create_from_config(&PolicyConfig::PowerOfTwo {
load_check_interval_secs: 60,
});
assert_eq!(policy.name(), "power_of_two");
// Test CacheAware
let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAware {
cache_threshold: 0.7,
balance_abs_threshold: 10,
balance_rel_threshold: 1.5,
eviction_interval_secs: 30,
max_tree_size: 1000,
});
assert_eq!(policy.name(), "cache_aware");
}
#[test]
fn test_create_by_name() {
assert!(PolicyFactory::create_by_name("random").is_some());
assert!(PolicyFactory::create_by_name("RANDOM").is_some());
assert!(PolicyFactory::create_by_name("round_robin").is_some());
assert!(PolicyFactory::create_by_name("RoundRobin").is_some());
assert!(PolicyFactory::create_by_name("power_of_two").is_some());
assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some());
assert!(PolicyFactory::create_by_name("cache_aware").is_some());
assert!(PolicyFactory::create_by_name("CacheAware").is_some());
assert!(PolicyFactory::create_by_name("unknown").is_none());
}
}
sgl-router/src/policies/mod.rs 0 → 100644
View file @ c8f31042
//! Load balancing policies for SGLang router
//!
//! This module provides a unified abstraction for routing policies that work
//! across both regular and prefill-decode (PD) routing modes.
use crate::core::Worker;
use std::fmt::Debug;
mod cache_aware;
mod factory;
mod power_of_two;
mod random;
mod round_robin;
pub use cache_aware::CacheAwarePolicy;
pub use factory::PolicyFactory;
pub use power_of_two::PowerOfTwoPolicy;
pub use random::RandomPolicy;
pub use round_robin::RoundRobinPolicy;
/// Core trait for load balancing policies
///
/// This trait provides a unified interface for implementing routing algorithms
/// that can work with both regular single-worker selection and PD dual-worker selection.
pub trait LoadBalancingPolicy: Send + Sync + Debug {
/// Select a single worker from the available workers
///
/// This is used for regular routing mode where requests go to a single worker.
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize>;
/// Select a pair of workers (prefill and decode) for PD routing
///
/// Returns indices of (prefill_worker, decode_worker) from their respective arrays.
/// Default implementation uses select_worker for each array independently.
fn select_worker_pair(
&self,
prefill_workers: &[Box<dyn Worker>],
decode_workers: &[Box<dyn Worker>],
request_text: Option<&str>,
) -> Option<(usize, usize)> {
// Default implementation: independently select from each pool
let prefill_idx = self.select_worker(prefill_workers, request_text)?;
let decode_idx = self.select_worker(decode_workers, request_text)?;
Some((prefill_idx, decode_idx))
}
/// Update policy state after request completion
///
/// This is called when a request completes (successfully or not) to allow
/// policies to update their internal state.
fn on_request_complete(&self, _worker_url: &str, _success: bool) {
// Default: no-op for stateless policies
}
/// Get policy name for metrics and debugging
fn name(&self) -> &'static str;
/// Update worker load information
///
/// This is called periodically with current load information for load-aware policies.
fn update_loads(&self, _loads: &std::collections::HashMap<String, isize>) {
// Default: no-op for policies that don't use load information
}
/// Reset any internal state
///
/// This is useful for policies that maintain state (e.g., round-robin counters).
fn reset(&self) {
// Default: no-op for stateless policies
}
/// Get as Any for downcasting
fn as_any(&self) -> &dyn std::any::Any;
}
/// Configuration for cache-aware policy
#[derive(Debug, Clone)]
pub struct CacheAwareConfig {
pub cache_threshold: f32,
pub balance_abs_threshold: usize,
pub balance_rel_threshold: f32,
pub eviction_interval_secs: u64,
pub max_tree_size: usize,
}
impl Default for CacheAwareConfig {
fn default() -> Self {
Self {
cache_threshold: 0.5,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
eviction_interval_secs: 30,
max_tree_size: 10000,
}
}
}
/// Helper function to filter healthy workers and return their indices
pub(crate) fn get_healthy_worker_indices(workers: &[Box<dyn Worker>]) -> Vec<usize> {
workers
.iter()
.enumerate()
.filter(|(_, w)| w.is_healthy())
.map(|(idx, _)| idx)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_get_healthy_worker_indices() {
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// All healthy initially
let indices = get_healthy_worker_indices(&workers);
assert_eq!(indices, vec![0, 1, 2]);
// Mark one unhealthy
workers[1].set_healthy(false);
let indices = get_healthy_worker_indices(&workers);
assert_eq!(indices, vec![0, 2]);
}
}
sgl-router/src/policies/power_of_two.rs 0 → 100644
View file @ c8f31042
//! Power-of-two choices load balancing policy
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
use metrics::counter;
use rand::Rng;
use std::collections::HashMap;
use std::sync::RwLock;
use tracing::info;
/// Power-of-two choices policy
///
/// Randomly selects two workers and routes to the one with lower load.
/// This provides good load distribution with minimal coordination overhead.
#[derive(Debug)]
pub struct PowerOfTwoPolicy {
/// Cached load information from external monitoring
cached_loads: RwLock<HashMap<String, isize>>,
}
impl PowerOfTwoPolicy {
pub fn new() -> Self {
Self {
cached_loads: RwLock::new(HashMap::new()),
}
}
fn get_worker_load(&self, worker: &dyn Worker) -> isize {
// First check cached loads (from external monitoring)
if let Ok(loads) = self.cached_loads.read() {
if let Some(&load) = loads.get(worker.url()) {
return load;
}
}
// Fall back to local load counter
worker.load() as isize
}
}
impl LoadBalancingPolicy for PowerOfTwoPolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
_request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
if healthy_indices.len() == 1 {
return Some(healthy_indices[0]);
}
// Select two random workers
let mut rng = rand::thread_rng();
let idx1 = rng.gen_range(0..healthy_indices.len());
let mut idx2 = rng.gen_range(0..healthy_indices.len());
// Ensure we pick two different workers
while idx2 == idx1 {
idx2 = rng.gen_range(0..healthy_indices.len());
}
let worker_idx1 = healthy_indices[idx1];
let worker_idx2 = healthy_indices[idx2];
// Compare loads and select the less loaded one
let load1 = self.get_worker_load(workers[worker_idx1].as_ref());
let load2 = self.get_worker_load(workers[worker_idx2].as_ref());
// Log selection for debugging
let selected_idx = if load1 <= load2 {
worker_idx1
} else {
worker_idx2
};
info!(
"Power-of-two selection: {}={} vs {}={} -> selected {}",
workers[worker_idx1].url(),
load1,
workers[worker_idx2].url(),
load2,
workers[selected_idx].url()
);
// Increment processed counter
workers[selected_idx].increment_processed();
counter!("sgl_router_processed_requests_total", "worker" => workers[selected_idx].url().to_string())
.increment(1);
Some(selected_idx)
}
fn name(&self) -> &'static str {
"power_of_two"
}
fn update_loads(&self, loads: &HashMap<String, isize>) {
if let Ok(mut cached) = self.cached_loads.write() {
*cached = loads.clone();
}
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
impl Default for PowerOfTwoPolicy {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_power_of_two_selection() {
let policy = PowerOfTwoPolicy::new();
let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular);
let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular);
let worker3 = BasicWorker::new("http://w3:8000".to_string(), WorkerType::Regular);
// Set different loads
for _ in 0..10 {
worker1.increment_load();
}
for _ in 0..5 {
worker2.increment_load();
}
// worker3 has load 0
let workers: Vec<Box<dyn Worker>> =
vec![Box::new(worker1), Box::new(worker2), Box::new(worker3)];
// Run multiple selections
let mut selected_counts = vec![0; 3];
for _ in 0..100 {
if let Some(idx) = policy.select_worker(&workers, None) {
selected_counts[idx] += 1;
}
}
// Worker with lowest load (worker3) should be selected most often
assert!(selected_counts[2] > selected_counts[1]);
assert!(selected_counts[1] > selected_counts[0]);
}
#[test]
fn test_power_of_two_with_cached_loads() {
let policy = PowerOfTwoPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Update cached loads
let mut loads = HashMap::new();
loads.insert("http://w1:8000".to_string(), 100);
loads.insert("http://w2:8000".to_string(), 10);
policy.update_loads(&loads);
// Should prefer worker2 with lower cached load
let mut w2_selected = 0;
for _ in 0..50 {
if let Some(idx) = policy.select_worker(&workers, None) {
if idx == 1 {
w2_selected += 1;
}
}
}
// Worker2 should be selected significantly more often
assert!(w2_selected > 35); // Should win most of the time
}
#[test]
fn test_power_of_two_single_worker() {
let policy = PowerOfTwoPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
))];
// With single worker, should always select it
assert_eq!(policy.select_worker(&workers, None), Some(0));
}
}
sgl-router/src/policies/random.rs 0 → 100644
View file @ c8f31042
//! Random load balancing policy
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
use rand::Rng;
/// Random selection policy
///
/// Selects workers randomly with uniform distribution among healthy workers.
#[derive(Debug, Default)]
pub struct RandomPolicy;
impl RandomPolicy {
pub fn new() -> Self {
Self
}
}
impl LoadBalancingPolicy for RandomPolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
_request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
let mut rng = rand::thread_rng();
let random_idx = rng.gen_range(0..healthy_indices.len());
Some(healthy_indices[random_idx])
}
fn name(&self) -> &'static str {
"random"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use std::collections::HashMap;
#[test]
fn test_random_selection() {
let policy = RandomPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// Test multiple selections to ensure randomness
let mut counts = HashMap::new();
for _ in 0..100 {
if let Some(idx) = policy.select_worker(&workers, None) {
*counts.entry(idx).or_insert(0) += 1;
}
}
// All workers should be selected at least once
assert_eq!(counts.len(), 3);
assert!(counts.values().all(|&count| count > 0));
}
#[test]
fn test_random_with_unhealthy_workers() {
let policy = RandomPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Mark first worker as unhealthy
workers[0].set_healthy(false);
// Should always select the healthy worker (index 1)
for _ in 0..10 {
assert_eq!(policy.select_worker(&workers, None), Some(1));
}
}
#[test]
fn test_random_no_healthy_workers() {
let policy = RandomPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
))];
workers[0].set_healthy(false);
assert_eq!(policy.select_worker(&workers, None), None);
}
}
sgl-router/src/policies/round_robin.rs 0 → 100644
View file @ c8f31042
//! Round-robin load balancing policy
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
use std::sync::atomic::{AtomicUsize, Ordering};
/// Round-robin selection policy
///
/// Selects workers in sequential order, cycling through all healthy workers.
#[derive(Debug, Default)]
pub struct RoundRobinPolicy {
counter: AtomicUsize,
}
impl RoundRobinPolicy {
pub fn new() -> Self {
Self {
counter: AtomicUsize::new(0),
}
}
}
impl LoadBalancingPolicy for RoundRobinPolicy {
fn select_worker(
&self,
workers: &[Box<dyn Worker>],
_request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
// Get and increment counter atomically
let count = self.counter.fetch_add(1, Ordering::Relaxed);
let selected_idx = count % healthy_indices.len();
Some(healthy_indices[selected_idx])
}
fn name(&self) -> &'static str {
"round_robin"
}
fn reset(&self) {
self.counter.store(0, Ordering::Relaxed);
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
#[test]
fn test_round_robin_selection() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// Should select workers in order: 0, 1, 2, 0, 1, 2, ...
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
assert_eq!(policy.select_worker(&workers, None), Some(2));
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
}
#[test]
fn test_round_robin_with_unhealthy_workers() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
];
// Mark middle worker as unhealthy
workers[1].set_healthy(false);
// Should skip unhealthy worker: 0, 2, 0, 2, ...
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(2));
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(2));
}
#[test]
fn test_round_robin_reset() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Box<dyn Worker>> = vec![
Box::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Box::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
];
// Advance the counter
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
// Reset should start from beginning
policy.reset();
assert_eq!(policy.select_worker(&workers, None), Some(0));
}
}
sgl-router/src/routers/factory.rs 0 → 100644
View file @ c8f31042
//! Factory for creating router instances
use super::{pd_router::PDRouter, router::Router, RouterTrait};
use crate::config::{PolicyConfig, RouterConfig, RoutingMode};
use crate::policies::PolicyFactory;
/// Factory for creating router instances based on configuration
pub struct RouterFactory;
impl RouterFactory {
/// Create a router instance from configuration
pub fn create_router(config: &RouterConfig) -> Result<Box<dyn RouterTrait>, String> {
match &config.mode {
RoutingMode::Regular { worker_urls } => {
Self::create_regular_router(worker_urls, &config.policy, config)
}
RoutingMode::PrefillDecode {
prefill_urls,
decode_urls,
} => Self::create_pd_router(prefill_urls, decode_urls, &config.policy, config),
}
}
/// Create a regular router with injected policy
fn create_regular_router(
worker_urls: &[String],
policy_config: &PolicyConfig,
router_config: &RouterConfig,
) -> Result<Box<dyn RouterTrait>, String> {
// Create policy
let policy = PolicyFactory::create_from_config(policy_config);
// Create regular router with injected policy
let router = Router::new(
worker_urls.to_vec(),
policy,
router_config.worker_startup_timeout_secs,
router_config.worker_startup_check_interval_secs,
)?;
Ok(Box::new(router))
}
/// Create a PD router with injected policy
fn create_pd_router(
prefill_urls: &[(String, Option<u16>)],
decode_urls: &[String],
policy_config: &PolicyConfig,
router_config: &RouterConfig,
) -> Result<Box<dyn RouterTrait>, String> {
// Create policy directly from PolicyConfig
// All policies now support PD mode through the select_worker_pair method
let policy = PolicyFactory::create_from_config(policy_config);
// Create PD router with injected policy
let router = PDRouter::new(
prefill_urls.to_vec(),
decode_urls.to_vec(),
policy,
router_config.worker_startup_timeout_secs,
router_config.worker_startup_check_interval_secs,
)?;
Ok(Box::new(router))
}
}
sgl-router/src/routers/mod.rs 0 → 100644
View file @ c8f31042
//! Router implementations
use actix_web::{HttpRequest, HttpResponse};
use async_trait::async_trait;
use reqwest::Client;
use std::fmt::Debug;
pub mod factory;
pub mod pd_router;
pub mod pd_types;
pub mod request_adapter;
pub mod router;
pub use factory::RouterFactory;
/// Worker management trait for administrative operations
///
/// This trait is separate from RouterTrait to allow Send futures
/// for use in service discovery and other background tasks
#[async_trait]
pub trait WorkerManagement: Send + Sync {
/// Add a worker to the router
async fn add_worker(&self, worker_url: &str) -> Result<String, String>;
/// Remove a worker from the router
fn remove_worker(&self, worker_url: &str);
/// Get all worker URLs
fn get_worker_urls(&self) -> Vec<String>;
}
/// Core trait for all router implementations
///
/// This trait provides a unified interface for routing requests,
/// regardless of whether it's a regular router or PD router.
#[async_trait(?Send)]
pub trait RouterTrait: Send + Sync + Debug + WorkerManagement {
/// Get a reference to self as Any for downcasting
fn as_any(&self) -> &dyn std::any::Any;
/// Route a health check request
async fn health(&self, client: &Client, req: &HttpRequest) -> HttpResponse;
/// Route a health generate request
async fn health_generate(&self, client: &Client, req: &HttpRequest) -> HttpResponse;
/// Get server information
async fn get_server_info(&self, client: &Client, req: &HttpRequest) -> HttpResponse;
/// Get available models
async fn get_models(&self, client: &Client, req: &HttpRequest) -> HttpResponse;
/// Get model information
async fn get_model_info(&self, client: &Client, req: &HttpRequest) -> HttpResponse;
/// Route a generate request
async fn route_generate(
&self,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse;
/// Route a chat completion request
async fn route_chat(
&self,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse;
/// Route a completion request
async fn route_completion(
&self,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse;
/// Flush cache on all workers
async fn flush_cache(&self, client: &Client) -> HttpResponse;
/// Get worker loads (for monitoring)
async fn get_worker_loads(&self, client: &Client) -> HttpResponse;
/// Get router type name
fn router_type(&self) -> &'static str;
/// Check if this is a PD router
fn is_pd_mode(&self) -> bool {
self.router_type() == "pd"
}
/// Server liveness check - is the server process running
fn liveness(&self) -> HttpResponse {
// Simple liveness check - if we can respond, we're alive
HttpResponse::Ok().body("OK")
}
/// Server readiness check - is the server ready to handle requests
fn readiness(&self) -> HttpResponse;
}
sgl-router/src/pd_router.rs → sgl-router/src/routers/pd_router.rs
View file @ c8f31042
// PD (Prefill-Decode) Router Implementation
// This module handles routing for disaggregated prefill-decode systems
use super::pd_types::{api_path, Bootstrap, ChatReqInput, GenerateReqInput, PDRouterError};
use super::request_adapter::ToPdRequest;
use crate::core::{HealthChecker, Worker, WorkerFactory, WorkerLoadGuard};
use crate::pd_types::{
api_path, Bootstrap, ChatReqInput, GenerateReqInput, PDRouterError, PDSelectionPolicy,
};
use crate::openai_api_types::{ChatCompletionRequest, CompletionRequest, GenerateRequest};
use crate::policies::LoadBalancingPolicy;
use crate::tree::Tree;
use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
use actix_web::{HttpRequest, HttpResponse};
......@@ -17,13 +18,11 @@ use std::time::{Duration, Instant};
use tracing::{debug, error, info, warn};
use uuid::Uuid;
// Removed over-engineered ProxyResponse - using HttpResponse directly
#[derive(Debug)]
pub struct PDRouter {
pub prefill_workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
pub decode_workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
pub selection_policy: PDSelectionPolicy,
pub policy: Arc<dyn LoadBalancingPolicy>,
pub prefill_tree: Option<Arc<Mutex<Tree>>>,
pub timeout_secs: u64,
pub interval_secs: u64,
......@@ -42,7 +41,7 @@ impl PDRouter {
bootstrap_port: Option<u16>,
) -> Result<String, PDRouterError> {
// Wait for the new server to be healthy
crate::router::Router::wait_for_healthy_workers(
crate::routers::router::Router::wait_for_healthy_workers(
&[url.clone()],
self.timeout_secs,
self.interval_secs,
......@@ -78,7 +77,7 @@ impl PDRouter {
pub async fn add_decode_server(&self, url: String) -> Result<String, PDRouterError> {
// Wait for the new server to be healthy
crate::router::Router::wait_for_healthy_workers(
crate::routers::router::Router::wait_for_healthy_workers(
&[url.clone()],
self.timeout_secs,
self.interval_secs,
......@@ -103,9 +102,6 @@ impl PDRouter {
workers.push(worker);
// Initialize load tracking
// Worker tracks its own load internally
info!("Added decode server: {}", url);
Ok(format!("Successfully added decode server: {}", url))
}
......@@ -128,9 +124,6 @@ impl PDRouter {
});
}
// Remove from load tracking
// Worker load tracking is internal
// Remove from cache tree if using cache-aware policy
if let Some(ref tree) = self.prefill_tree {
// Note: Tree doesn't have a remove method, so we rebuild it
......@@ -170,7 +163,7 @@ impl PDRouter {
pub fn new(
prefill_urls: Vec<(String, Option<u16>)>,
decode_urls: Vec<String>,
selection_policy: PDSelectionPolicy,
policy: Arc<dyn LoadBalancingPolicy>,
timeout_secs: u64,
interval_secs: u64,
) -> Result<Self, String> {
......@@ -185,25 +178,38 @@ impl PDRouter {
.map(WorkerFactory::create_decode)
.collect();
// Wait for PD workers to be healthy
// Wait for PD workers to be healthy (skip if empty - for service discovery mode)
let all_urls: Vec<String> = prefill_workers
.iter()
.chain(decode_workers.iter())
.map(|worker| worker.url().to_string())
.collect();
crate::router::Router::wait_for_healthy_workers(&all_urls, timeout_secs, interval_secs)?;
if !all_urls.is_empty() {
crate::routers::router::Router::wait_for_healthy_workers(
&all_urls,
timeout_secs,
interval_secs,
)?;
}
// Initialize cache-aware components if needed
let prefill_tree = match &selection_policy {
PDSelectionPolicy::CacheAware { .. } => {
let prefill_tree = if policy.name() == "cache_aware" {
// Initialize the policy's internal tree with prefill workers
if let Some(cache_policy) = policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_policy.init_workers(&prefill_workers);
}
let tree = Arc::new(Mutex::new(Tree::new()));
// Initialize tree with prefill workers
for worker in &prefill_workers {
tree.lock().unwrap().insert("", worker.url());
}
Some(tree)
}
_ => None,
} else {
None
};
// Set up background load monitoring for power-of-two selection
......@@ -216,10 +222,11 @@ impl PDRouter {
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
let load_monitor_handle = if matches!(selection_policy, PDSelectionPolicy::PowerOfTwo) {
let load_monitor_handle = if policy.name() == "power_of_two" {
let monitor_urls = all_urls.clone();
let monitor_interval = interval_secs;
let monitor_client = http_client.clone();
let policy_clone = Arc::clone(&policy);
Some(Arc::new(tokio::spawn(async move {
Self::monitor_worker_loads_with_client(
......@@ -227,6 +234,7 @@ impl PDRouter {
tx,
monitor_interval,
monitor_client,
policy_clone,
)
.await;
})))
......@@ -246,7 +254,7 @@ impl PDRouter {
Ok(PDRouter {
prefill_workers,
decode_workers,
selection_policy,
policy,
prefill_tree,
timeout_secs,
interval_secs,
......@@ -270,15 +278,21 @@ impl PDRouter {
let _request_id = Uuid::new_v4();
// Get stream flag and return_logprob flag before moving the request
let is_stream = typed_req.is_stream();
let is_stream = typed_req.stream;
let return_logprob = typed_req
.other
.get("return_logprob")
.and_then(|v| v.as_bool())
.unwrap_or(false);
// Extract text for cache-aware routing from the typed request
let request_text = typed_req.text.as_ref().and_then(|t| match t {
super::pd_types::InputText::Single(s) => Some(s.as_str()),
super::pd_types::InputText::Batch(v) => v.first().map(|s| s.as_str()),
});
// Select servers
let (prefill, decode) = match self.select_pd_pair(client).await {
let (prefill, decode) = match self.select_pd_pair(client, request_text).await {
Ok(pair) => pair,
Err(e) => {
error!("Failed to select PD pair: {}", e);
......@@ -339,15 +353,24 @@ impl PDRouter {
let start = Instant::now();
// Get stream flag and return_logprob flag before moving the request
let is_stream = typed_req.is_stream();
let is_stream = typed_req.stream;
let return_logprob = typed_req
.other
.get("return_logprob")
.and_then(|v| v.as_bool())
.unwrap_or(false);
// Extract text for cache-aware routing from chat messages
let request_text = typed_req
.other
.get("messages")
.and_then(|messages| messages.as_array())
.and_then(|arr| arr.first())
.and_then(|msg| msg.get("content"))
.and_then(|content| content.as_str());
// Select servers
let (prefill, decode) = match self.select_pd_pair(client).await {
let (prefill, decode) = match self.select_pd_pair(client, request_text).await {
Ok(pair) => pair,
Err(e) => {
error!("Failed to select PD pair: {}", e);
......@@ -424,7 +447,7 @@ impl PDRouter {
.json(&json_request);
// Copy headers from original request
for (name, value) in crate::router::copy_request_headers(req) {
for (name, value) in crate::routers::router::copy_request_headers(req) {
if name.to_lowercase() != "content-type" && name.to_lowercase() != "content-length" {
prefill_request = prefill_request.header(&name, &value);
decode_request = decode_request.header(&name, &value);
......@@ -620,96 +643,38 @@ impl PDRouter {
async fn select_pd_pair(
&self,
_client: &reqwest::Client,
request_text: Option<&str>,
) -> Result<(Box<dyn Worker>, Box<dyn Worker>), String> {
// Check we have workers
if self
// Get read locks for both worker lists
let prefill_workers = self
.prefill_workers
.read()
.map_err(|e| format!("Failed to acquire prefill workers lock: {}", e))?
.is_empty()
{
return Err("No prefill workers available. Please check if prefill servers are configured and healthy.".to_string());
}
if self
.map_err(|e| format!("Failed to acquire prefill workers lock: {}", e))?;
let decode_workers = self
.decode_workers
.read()
.map_err(|e| format!("Failed to acquire decode workers lock: {}", e))?
.is_empty()
{
return Err("No decode workers available. Please check if decode servers are configured and healthy.".to_string());
}
.map_err(|e| format!("Failed to acquire decode workers lock: {}", e))?;
match &self.selection_policy {
PDSelectionPolicy::Random => self.select_random(),
PDSelectionPolicy::PowerOfTwo => self.select_power_of_two().await,
PDSelectionPolicy::CacheAware { .. } => {
// TODO: Implement cache-aware selection
self.select_power_of_two().await
}
// Check we have workers
if prefill_workers.is_empty() {
return Err("No prefill workers available. Please check if prefill servers are configured and healthy.".to_string());
}
if decode_workers.is_empty() {
return Err("No decode workers available. Please check if decode servers are configured and healthy.".to_string());
}
fn select_random(&self) -> Result<(Box<dyn Worker>, Box<dyn Worker>), String> {
let prefill_list = self.prefill_workers.read().map_err(|_| "Lock error")?;
let decode_list = self.decode_workers.read().map_err(|_| "Lock error")?;
let prefill = prefill_list[rand::random::<usize>() % prefill_list.len()].clone_worker();
let decode = decode_list[rand::random::<usize>() % decode_list.len()].clone_worker();
// Use the policy to select worker pair
match self
.policy
.select_worker_pair(&prefill_workers, &decode_workers, request_text)
{
Some((prefill_idx, decode_idx)) => {
let prefill = prefill_workers[prefill_idx].clone_worker();
let decode = decode_workers[decode_idx].clone_worker();
Ok((prefill, decode))
}
async fn select_power_of_two(&self) -> Result<(Box<dyn Worker>, Box<dyn Worker>), String> {
let prefill_list = self.prefill_workers.read().map_err(|_| "Lock error")?;
let decode_list = self.decode_workers.read().map_err(|_| "Lock error")?;
let (p1_idx, p2_idx) = get_two_random_indices(prefill_list.len());
let (d1_idx, d2_idx) = get_two_random_indices(decode_list.len());
let loads = self.worker_loads.borrow();
let p1_load = loads
.get(prefill_list[p1_idx].url())
.copied()
.unwrap_or(isize::MAX);
let p2_load = loads
.get(prefill_list[p2_idx].url())
.copied()
.unwrap_or(isize::MAX);
let d1_load = loads
.get(decode_list[d1_idx].url())
.copied()
.unwrap_or(isize::MAX);
let d2_load = loads
.get(decode_list[d2_idx].url())
.copied()
.unwrap_or(isize::MAX);
info!(
"Power-of-two selection - Prefill: {}={} vs {}={} | Decode: {}={} vs {}={}",
prefill_list[p1_idx].url(),
p1_load,
prefill_list[p2_idx].url(),
p2_load,
decode_list[d1_idx].url(),
d1_load,
decode_list[d2_idx].url(),
d2_load
);
let selected_prefill = if p1_load <= p2_load {
prefill_list[p1_idx].clone_worker()
} else {
prefill_list[p2_idx].clone_worker()
};
let selected_decode = if d1_load <= d2_load {
decode_list[d1_idx].clone_worker()
} else {
decode_list[d2_idx].clone_worker()
};
Ok((selected_prefill, selected_decode))
None => Err("Failed to select worker pair".to_string()),
}
}
// Background task to monitor worker loads with shared client
......@@ -718,6 +683,7 @@ impl PDRouter {
tx: tokio::sync::watch::Sender<HashMap<String, isize>>,
interval_secs: u64,
client: reqwest::Client,
policy: Arc<dyn LoadBalancingPolicy>,
) {
loop {
let mut loads = HashMap::new();
......@@ -742,6 +708,9 @@ impl PDRouter {
debug!("Worker loads updated: {:?}", loads);
// Update the policy with current loads
policy.update_loads(&loads);
// Check if receiver is still active
if tx.send(loads).is_err() {
info!("Load monitor receiver dropped, shutting down monitor task");
......@@ -792,18 +761,6 @@ impl PDRouter {
}
// Helper functions
fn get_two_random_indices(len: usize) -> (usize, usize) {
if len == 1 {
(0, 0)
} else {
let idx1 = rand::random::<usize>() % len;
let mut idx2 = rand::random::<usize>() % len;
while idx2 == idx1 {
idx2 = rand::random::<usize>() % len;
}
(idx1, idx2)
}
}
async fn get_worker_load(client: &reqwest::Client, worker_url: &str) -> Option<isize> {
match client.get(format!("{}/get_load", worker_url)).send().await {
......@@ -841,61 +798,72 @@ async fn get_worker_load(client: &reqwest::Client, worker_url: &str) -> Option<i
// PD-specific endpoints
impl PDRouter {
pub async fn health_generate(&self, client: &reqwest::Client) -> HttpResponse {
let mut all_healthy = true;
let mut unhealthy_servers = Vec::new();
// Collect all worker URLs with their types
let mut worker_infos = Vec::new();
// Test model generation capability by selecting a random pair and testing them
// Note: This endpoint actually causes the model to generate tokens, so we only test one pair
for worker in self.prefill_workers.read().unwrap().iter() {
worker_infos.push((worker.url().to_string(), "prefill"));
// Select a random worker pair using the policy
let (prefill, decode) = match self.select_pd_pair(client, None).await {
Ok(pair) => pair,
Err(e) => {
return HttpResponse::ServiceUnavailable()
.body(format!("No healthy worker pair available: {}", e));
}
};
for worker in self.decode_workers.read().unwrap().iter() {
worker_infos.push((worker.url().to_string(), "decode"));
}
// Test prefill server's health_generate
let prefill_url = format!("{}/health_generate", prefill.url());
let prefill_result = client.get(&prefill_url).send().await;
// Create tasks with URL tracking
let tasks: Vec<_> = worker_infos
.iter()
.map(|(url, _)| {
let health_url = format!("{}/health_generate", url);
client.get(&health_url).send()
})
.collect();
// Test decode server's health_generate
let decode_url = format!("{}/health_generate", decode.url());
let decode_result = client.get(&decode_url).send().await;
let results = futures_util::future::join_all(tasks).await;
// Check results
let mut errors = Vec::new();
for ((url, worker_type), result) in worker_infos.iter().zip(results.into_iter()) {
match result {
match prefill_result {
Ok(res) if res.status().is_success() => {
debug!("Health check passed for {} server: {}", worker_type, url);
debug!(
"Health generate passed for prefill server: {}",
prefill.url()
);
}
Ok(res) => {
all_healthy = false;
let msg = format!(
"{} server {} returned status {}",
worker_type,
url,
errors.push(format!(
"Prefill {} returned status {}",
prefill.url(),
res.status()
);
error!("{}", msg);
unhealthy_servers.push(msg);
));
}
Err(e) => {
all_healthy = false;
let msg = format!("{} server {} error: {}", worker_type, url, e);
error!("{}", msg);
unhealthy_servers.push(msg);
errors.push(format!("Prefill {} error: {}", prefill.url(), e));
}
}
match decode_result {
Ok(res) if res.status().is_success() => {
debug!("Health generate passed for decode server: {}", decode.url());
}
Ok(res) => {
errors.push(format!(
"Decode {} returned status {}",
decode.url(),
res.status()
));
}
Err(e) => {
errors.push(format!("Decode {} error: {}", decode.url(), e));
}
}
if all_healthy {
HttpResponse::Ok().body("Health check passed on all servers")
if errors.is_empty() {
HttpResponse::Ok().body(format!(
"Health generate passed on selected pair: prefill={}, decode={}",
prefill.url(),
decode.url()
))
} else {
HttpResponse::ServiceUnavailable()
.body(format!("Health check failed: {:?}", unhealthy_servers))
HttpResponse::ServiceUnavailable().body(format!("Health generate failed: {:?}", errors))
}
}
......@@ -955,7 +923,7 @@ impl PDRouter {
if let Some(worker_url) = first_worker_url {
// Send request directly without going through Router
let mut request_builder = client.get(format!("{}/v1/models", worker_url));
for (name, value) in crate::router::copy_request_headers(req) {
for (name, value) in crate::routers::router::copy_request_headers(req) {
if name.to_lowercase() != "content-type" && name.to_lowercase() != "content-length"
{
request_builder = request_builder.header(name, value);
......@@ -1035,7 +1003,7 @@ impl PDRouter {
if let Some(worker_url) = first_worker_url {
let mut request_builder = client.get(format!("{}/get_model_info", worker_url));
for (name, value) in crate::router::copy_request_headers(req) {
for (name, value) in crate::routers::router::copy_request_headers(req) {
if name.to_lowercase() != "content-type" && name.to_lowercase() != "content-length"
{
request_builder = request_builder.header(name, value);
......@@ -1102,3 +1070,324 @@ impl PDRouter {
}
}
}
use crate::routers::{RouterTrait, WorkerManagement};
use async_trait::async_trait;
use reqwest::Client;
#[async_trait]
impl WorkerManagement for PDRouter {
async fn add_worker(&self, _worker_url: &str) -> Result<String, String> {
// For PD router, we don't support adding workers via this generic method
Err(
"PD router requires specific add_prefill_server or add_decode_server methods"
.to_string(),
)
}
fn remove_worker(&self, worker_url: &str) {
// For PD router, we would need to know if it's a prefill or decode server
// For now, try both
if let Ok(mut workers) = self.prefill_workers.write() {
if let Some(index) = workers.iter().position(|w| w.url() == worker_url) {
workers.remove(index);
info!("Removed prefill worker: {}", worker_url);
return;
}
}
if let Ok(mut workers) = self.decode_workers.write() {
if let Some(index) = workers.iter().position(|w| w.url() == worker_url) {
workers.remove(index);
info!("Removed decode worker: {}", worker_url);
}
}
}
fn get_worker_urls(&self) -> Vec<String> {
let mut urls = Vec::new();
// Add prefill worker URLs
if let Ok(workers) = self.prefill_workers.read() {
for worker in workers.iter() {
urls.push(worker.url().to_string());
}
}
// Add decode worker URLs
if let Ok(workers) = self.decode_workers.read() {
for worker in workers.iter() {
urls.push(worker.url().to_string());
}
}
urls
}
}
#[async_trait(?Send)]
impl RouterTrait for PDRouter {
fn as_any(&self) -> &dyn std::any::Any {
self
}
async fn health(&self, _client: &Client, _req: &HttpRequest) -> HttpResponse {
// This is a server readiness check - checking if we have healthy workers
// Workers handle their own health checks in the background
let mut all_healthy = true;
let mut unhealthy_servers = Vec::new();
// Check prefill servers
for worker in self.prefill_workers.read().unwrap().iter() {
if !worker.is_healthy() {
all_healthy = false;
unhealthy_servers.push(format!("Prefill: {}", worker.url()));
}
}
// Check decode servers
for worker in self.decode_workers.read().unwrap().iter() {
if !worker.is_healthy() {
all_healthy = false;
unhealthy_servers.push(format!("Decode: {}", worker.url()));
}
}
if all_healthy {
HttpResponse::Ok().body("All servers healthy")
} else {
HttpResponse::ServiceUnavailable()
.body(format!("Unhealthy servers: {:?}", unhealthy_servers))
}
}
async fn health_generate(&self, client: &Client, _req: &HttpRequest) -> HttpResponse {
// Use the existing PDRouter health_generate method
PDRouter::health_generate(self, client).await
}
async fn get_server_info(&self, client: &Client, _req: &HttpRequest) -> HttpResponse {
// Use the existing PDRouter get_server_info method
PDRouter::get_server_info(self, client).await
}
async fn get_models(&self, client: &Client, req: &HttpRequest) -> HttpResponse {
// Get first prefill worker URL to avoid holding lock across await
let first_worker_url = if let Ok(workers) = self.prefill_workers.read() {
workers.first().map(|w| w.url().to_string())
} else {
return HttpResponse::InternalServerError().body("Failed to access prefill workers");
};
if let Some(worker_url) = first_worker_url {
// Send request directly without going through Router
let mut request_builder = client.get(format!("{}/v1/models", worker_url));
for (name, value) in crate::routers::router::copy_request_headers(req) {
if name.to_lowercase() != "content-type" && name.to_lowercase() != "content-length"
{
request_builder = request_builder.header(name, value);
}
}
match request_builder.send().await {
Ok(res) => {
let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
match res.bytes().await {
Ok(body) => HttpResponse::build(status).body(body.to_vec()),
Err(e) => HttpResponse::InternalServerError()
.body(format!("Failed to read response body: {}", e)),
}
}
Err(e) => HttpResponse::InternalServerError()
.body(format!("Failed to send request: {}", e)),
}
} else {
HttpResponse::ServiceUnavailable().body("No prefill servers available")
}
}
async fn get_model_info(&self, client: &Client, req: &HttpRequest) -> HttpResponse {
// For PD router, get model info from the first prefill server
// Get first prefill worker URL to avoid holding lock across await
let first_worker_url = if let Ok(workers) = self.prefill_workers.read() {
workers.first().map(|w| w.url().to_string())
} else {
return HttpResponse::InternalServerError().body("Failed to access prefill workers");
};
if let Some(worker_url) = first_worker_url {
let mut request_builder = client.get(format!("{}/get_model_info", worker_url));
for (name, value) in crate::routers::router::copy_request_headers(req) {
if name.to_lowercase() != "content-type" && name.to_lowercase() != "content-length"
{
request_builder = request_builder.header(name, value);
}
}
match request_builder.send().await {
Ok(res) => {
let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
match res.bytes().await {
Ok(body) => HttpResponse::build(status).body(body.to_vec()),
Err(e) => HttpResponse::InternalServerError()
.body(format!("Failed to read response body: {}", e)),
}
}
Err(e) => HttpResponse::InternalServerError()
.body(format!("Failed to send request: {}", e)),
}
} else {
HttpResponse::ServiceUnavailable().body("No prefill servers available")
}
}
async fn route_generate(
&self,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse {
match serde_json::from_value::<GenerateRequest>(body.clone()) {
Ok(openai_req) => {
// Convert OpenAI format to PD format
let pd_req = openai_req.to_pd_request();
PDRouter::route_generate(self, client, req, pd_req, "/generate").await
}
Err(_) => {
// If that fails, try to deserialize directly as PD format (for backwards compatibility)
match serde_json::from_value::<GenerateReqInput>(body) {
Ok(pd_req) => {
PDRouter::route_generate(self, client, req, pd_req, "/generate").await
}
Err(e) => {
HttpResponse::BadRequest().body(format!("Invalid request format: {}", e))
}
}
}
}
}
async fn route_chat(
&self,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse {
match serde_json::from_value::<ChatCompletionRequest>(body.clone()) {
Ok(openai_req) => {
// Convert OpenAI format to PD format
let pd_req = openai_req.to_pd_request();
PDRouter::route_chat(self, client, req, pd_req, "/v1/chat/completions").await
}
Err(_) => {
// If that fails, try to deserialize directly as PD format (for backwards compatibility)
match serde_json::from_value::<ChatReqInput>(body) {
Ok(pd_req) => {
PDRouter::route_chat(self, client, req, pd_req, "/v1/chat/completions")
.await
}
Err(e) => {
HttpResponse::BadRequest().body(format!("Invalid request format: {}", e))
}
}
}
}
}
async fn route_completion(
&self,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse {
match serde_json::from_value::<CompletionRequest>(body.clone()) {
Ok(openai_req) => {
// Convert OpenAI format to PD format (CompletionRequest -> GenerateReqInput)
let pd_req = openai_req.to_pd_request();
PDRouter::route_generate(self, client, req, pd_req, "/v1/completions").await
}
Err(_) => {
// If that fails, try to deserialize directly as PD format (for backwards compatibility)
match serde_json::from_value::<GenerateReqInput>(body) {
Ok(pd_req) => {
PDRouter::route_generate(self, client, req, pd_req, "/v1/completions").await
}
Err(e) => {
HttpResponse::BadRequest().body(format!("Invalid request format: {}", e))
}
}
}
}
}
async fn flush_cache(&self, client: &Client) -> HttpResponse {
// Use the existing PDRouter flush_cache method
PDRouter::flush_cache(self, client).await
}
async fn get_worker_loads(&self, client: &Client) -> HttpResponse {
// Use the existing PDRouter get_loads method
PDRouter::get_loads(self, client).await
}
fn router_type(&self) -> &'static str {
"pd"
}
fn readiness(&self) -> HttpResponse {
// PD router is ready if it has at least one healthy prefill AND one healthy decode worker
let healthy_prefill_count = self
.prefill_workers
.read()
.unwrap()
.iter()
.filter(|w| w.is_healthy())
.count();
let healthy_decode_count = self
.decode_workers
.read()
.unwrap()
.iter()
.filter(|w| w.is_healthy())
.count();
let total_prefill = self.prefill_workers.read().unwrap().len();
let total_decode = self.decode_workers.read().unwrap().len();
if healthy_prefill_count > 0 && healthy_decode_count > 0 {
HttpResponse::Ok().json(serde_json::json!({
"status": "ready",
"prefill": {
"healthy": healthy_prefill_count,
"total": total_prefill
},
"decode": {
"healthy": healthy_decode_count,
"total": total_decode
}
}))
} else {
let mut reasons = Vec::new();
if healthy_prefill_count == 0 {
reasons.push("no healthy prefill workers");
}
if healthy_decode_count == 0 {
reasons.push("no healthy decode workers");
}
HttpResponse::ServiceUnavailable().json(serde_json::json!({
"status": "not_ready",
"reason": reasons.join(", "),
"prefill": {
"healthy": healthy_prefill_count,
"total": total_prefill
},
"decode": {
"healthy": healthy_decode_count,
"total": total_decode
}
}))
}
}
}
sgl-router/src/request_adapter.rs → sgl-router/src/routers/request_adapter.rs
View file @ c8f31042
// Request adapter to bridge OpenAI API types with PD routing requirements
use super::pd_types::{Bootstrap, ChatReqInput, GenerateReqInput, SingleOrBatch};
use crate::openai_api_types::{
ChatCompletionRequest, CompletionRequest, GenerateRequest, GenerationRequest, StringOrArray,
};
use crate::pd_types::{Bootstrap, ChatReqInput, GenerateReqInput, SingleOrBatch};
use serde_json::Value;
/// Adapter trait to convert OpenAI requests to PD-compatible requests
......
sgl-router/src/router.rs → sgl-router/src/routers/router.rs
View file @ c8f31042
use crate::core::{HealthChecker, Worker, WorkerFactory};
use crate::pd_router::PDRouter;
use crate::pd_types::PDSelectionPolicy;
use crate::tree::Tree;
use crate::policies::LoadBalancingPolicy;
use ::metrics::{counter, gauge, histogram};
use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
use actix_web::{HttpRequest, HttpResponse};
use futures_util::{StreamExt, TryStreamExt};
use std::fmt::Debug;
use std::sync::atomic::AtomicUsize;
use std::sync::{Arc, Mutex, RwLock};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::thread;
use std::time::Duration;
use std::time::Instant;
use tokio;
use std::time::{Duration, Instant};
use tracing::{debug, error, info, warn};
pub fn copy_request_headers(req: &HttpRequest) -> Vec<(String, String)> {
......@@ -27,170 +22,32 @@ pub fn copy_request_headers(req: &HttpRequest) -> Vec<(String, String)> {
.collect()
}
/// Regular router that uses injected load balancing policies
#[derive(Debug)]
pub enum Router {
RoundRobin {
pub struct Router {
workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
current_index: AtomicUsize,
policy: Arc<dyn LoadBalancingPolicy>,
timeout_secs: u64,
interval_secs: u64,
_worker_loads: Arc<tokio::sync::watch::Receiver<HashMap<String, isize>>>,
_load_monitor_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
_health_checker: Option<HealthChecker>,
},
Random {
workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
timeout_secs: u64,
interval_secs: u64,
_health_checker: Option<HealthChecker>,
},
PrefillDecode {
pd_router: Arc<PDRouter>,
},
CacheAware {
/*
Cache-Aware Load Balancing Router
This router combines two strategies to optimize both cache utilization and request distribution:
1. Cache-Aware Routing (Approximate Tree)
2. Load Balancing (Shortest Queue with Balance Thresholds)
The router dynamically switches between these strategies based on load conditions:
- Uses load balancing when the system is imbalanced
- Uses cache-aware routing when the system is balanced
A system is considered imbalanced if both conditions are met:
1. (max - min) > abs_threshold
2. max > rel_threshold * min
Strategy Details:
1. Cache-Aware Routing (Approximate Tree)
-------------------------------------------
This strategy maintains an approximate radix tree for each worker based on request history,
eliminating the need for direct cache state queries. The tree stores raw text characters
instead of token IDs to avoid tokenization overhead.
Process:
a. For each request, find the worker with the highest prefix match
b. If match rate > cache_threshold:
Route to the worker with highest match (likely has relevant data cached)
c. If match rate ≤ cache_threshold:
Route to the worker with smallest tree size (most available cache capacity)
d. Background maintenance:
Periodically evict least recently used leaf nodes to prevent memory overflow
2. Load Balancing (Shortest Queue)
-------------------------------------------
This strategy tracks pending request counts per worker and routes new requests
to the least busy worker when the system is detected to be imbalanced.
Configuration Parameters:
------------------------
1. cache_threshold: (float, 0.0 to 1.0)
Minimum prefix match ratio to use highest-match routing.
Below this threshold, routes to worker with most available cache space.
2. balance_abs_threshold: (integer)
Absolute difference threshold for load imbalance detection.
System is potentially imbalanced if (max_load - min_load) > abs_threshold
3. balance_rel_threshold: (float)
Relative ratio threshold for load imbalance detection.
System is potentially imbalanced if max_load > min_load * rel_threshold
Used in conjunction with abs_threshold to determine final imbalance state.
4. eviction_interval_secs: (integer)
Interval between LRU eviction cycles for the approximate trees.
5. max_tree_size: (integer)
Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted
during the next eviction cycle.
*/
workers: Arc<RwLock<Vec<Box<dyn Worker>>>>,
tree: Arc<Mutex<Tree>>,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
timeout_secs: u64,
interval_secs: u64,
_eviction_thread: Option<thread::JoinHandle<()>>,
_health_checker: Option<HealthChecker>,
},
}
#[derive(Debug, Clone)]
pub enum PolicyConfig {
RandomConfig {
timeout_secs: u64,
interval_secs: u64,
},
RoundRobinConfig {
timeout_secs: u64,
interval_secs: u64,
},
CacheAwareConfig {
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
eviction_interval_secs: u64,
max_tree_size: usize,
timeout_secs: u64,
interval_secs: u64,
},
PrefillDecodeConfig {
selection_policy: PDSelectionPolicy,
prefill_urls: Vec<(String, Option<u16>)>, // (url, bootstrap_port)
decode_urls: Vec<String>,
impl Router {
/// Create a new router with injected policy
pub fn new(
worker_urls: Vec<String>,
policy: Arc<dyn LoadBalancingPolicy>,
timeout_secs: u64,
interval_secs: u64,
},
}
impl Router {
pub fn new(worker_urls: Vec<String>, policy_config: PolicyConfig) -> Result<Self, String> {
) -> Result<Self, String> {
// Update active workers gauge
gauge!("sgl_router_active_workers").set(worker_urls.len() as f64);
// Get timeout and interval from policy config
let (timeout_secs, interval_secs) = match &policy_config {
PolicyConfig::RandomConfig {
timeout_secs,
interval_secs,
} => (*timeout_secs, *interval_secs),
PolicyConfig::RoundRobinConfig {
timeout_secs,
interval_secs,
} => (*timeout_secs, *interval_secs),
PolicyConfig::CacheAwareConfig {
timeout_secs,
interval_secs,
..
} => (*timeout_secs, *interval_secs),
PolicyConfig::PrefillDecodeConfig {
timeout_secs,
interval_secs,
..
} => (*timeout_secs, *interval_secs),
};
// For PrefillDecode, we need to handle workers differently
match &policy_config {
PolicyConfig::PrefillDecodeConfig { .. } => {
// PD mode doesn't use the worker_urls parameter
// We'll validate PD workers separately
}
_ => {
// Wait until all workers are healthy for regular modes
let worker_urls = worker_urls.clone();
std::thread::spawn(move || {
Self::wait_for_healthy_workers(&worker_urls, timeout_secs, interval_secs)
})
.join()
.map_err(|e| {
error!("Health-check thread panicked: {:?}", e);
format!("Health-check thread panicked: {e:?}")
})??;
}
// Wait for workers to be healthy (skip if empty - for service discovery mode)
if !worker_urls.is_empty() {
Self::wait_for_healthy_workers(&worker_urls, timeout_secs, interval_secs)?;
}
// Create Worker trait objects from URLs
......@@ -199,133 +56,52 @@ impl Router {
.map(|url| WorkerFactory::create_regular(url.clone()))
.collect();
// Create router based on policy...
Ok(match policy_config {
PolicyConfig::RandomConfig {
timeout_secs,
interval_secs,
} => {
let workers = Arc::new(RwLock::new(workers));
let health_checker =
crate::core::start_health_checker(Arc::clone(&workers), interval_secs);
Router::Random {
workers,
timeout_secs,
interval_secs,
_health_checker: Some(health_checker),
}
}
PolicyConfig::RoundRobinConfig {
timeout_secs,
interval_secs,
} => {
let workers = Arc::new(RwLock::new(workers));
let health_checker =
crate::core::start_health_checker(Arc::clone(&workers), interval_secs);
Router::RoundRobin {
workers,
current_index: std::sync::atomic::AtomicUsize::new(0),
timeout_secs,
interval_secs,
_health_checker: Some(health_checker),
}
// Initialize policy with workers if needed (e.g., for cache-aware)
if let Some(cache_aware) = policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_aware.init_workers(&workers);
}
PolicyConfig::CacheAwareConfig {
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
eviction_interval_secs,
max_tree_size,
timeout_secs,
interval_secs,
} => {
let tree = Arc::new(Mutex::new(Tree::new()));
// Create background eviction thread
let tree_clone = Arc::clone(&tree);
let workers = Arc::new(RwLock::new(workers));
let workers_clone = Arc::clone(&workers);
let eviction_thread = thread::spawn(move || {
loop {
// Sleep for the specified interval
thread::sleep(Duration::from_secs(eviction_interval_secs));
let locked_tree_clone = tree_clone.lock().unwrap();
// Run eviction
locked_tree_clone.evict_tenant_by_size(max_tree_size);
drop(locked_tree_clone);
let health_checker = crate::core::start_health_checker(Arc::clone(&workers), interval_secs);
// Log worker loads and processed requests
let workers_guard = workers_clone.read().unwrap();
let loads: Vec<(String, usize)> = workers_guard
.iter()
.map(|w| (w.url().to_string(), w.load()))
.collect();
info!("Worker loads: {:?}", loads);
let processed: Vec<(String, usize)> = workers_guard
.iter()
.map(|w| (w.url().to_string(), w.processed_requests()))
.collect();
info!("Processed requests: {:?}", processed);
}
});
// Setup load monitoring for PowerOfTwo policy
let (tx, rx) = tokio::sync::watch::channel(HashMap::new());
let worker_loads = Arc::new(rx);
for worker in workers.read().unwrap().iter() {
tree.lock().unwrap().insert("", worker.url());
}
let load_monitor_handle = if policy.name() == "power_of_two" {
let monitor_urls = worker_urls.clone();
let monitor_interval = interval_secs;
let policy_clone = Arc::clone(&policy);
let health_checker =
crate::core::start_health_checker(Arc::clone(&workers), interval_secs);
Some(Arc::new(tokio::spawn(async move {
Self::monitor_worker_loads(monitor_urls, tx, monitor_interval, policy_clone).await;
})))
} else {
None
};
Router::CacheAware {
Ok(Router {
workers,
tree,
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
policy,
timeout_secs,
interval_secs,
_eviction_thread: Some(eviction_thread),
_worker_loads: worker_loads,
_load_monitor_handle: load_monitor_handle,
_health_checker: Some(health_checker),
}
}
PolicyConfig::PrefillDecodeConfig {
selection_policy,
prefill_urls,
decode_urls,
timeout_secs,
interval_secs,
} => {
// Create PDRouter instance
let pd_router = PDRouter::new(
prefill_urls,
decode_urls,
selection_policy,
timeout_secs,
interval_secs,
)?;
Router::PrefillDecode {
pd_router: Arc::new(pd_router),
}
}
})
}
/// Get the current list of worker URLs
pub fn get_worker_urls(&self) -> Vec<String> {
match self {
Router::RoundRobin { workers, .. }
| Router::Random { workers, .. }
| Router::CacheAware { workers, .. } => workers
self.workers
.read()
.unwrap()
.iter()
.map(|w| w.url().to_string())
.collect(),
Router::PrefillDecode { .. } => Vec::new(),
}
.collect()
}
pub fn wait_for_healthy_workers(
......@@ -390,23 +166,13 @@ impl Router {
}
fn select_first_worker(&self) -> Result<String, String> {
match self {
Router::RoundRobin { workers, .. }
| Router::Random { workers, .. }
| Router::CacheAware { workers, .. } => {
let workers_guard = workers.read().unwrap();
let workers_guard = self.workers.read().unwrap();
if workers_guard.is_empty() {
Err("No workers are available".to_string())
} else {
Ok(workers_guard[0].url().to_string())
}
}
Router::PrefillDecode { .. } => {
// For PD mode, we don't need this method as routing is handled by PDRouter
Err("PrefillDecode mode doesn't use select_first_worker".to_string())
}
}
}
pub async fn send_request(
&self,
......@@ -525,16 +291,8 @@ impl Router {
route: &str,
req: &HttpRequest,
) -> HttpResponse {
// Get all worker URLs based on router type
let worker_urls = match self {
Router::PrefillDecode { .. } => {
// For PD mode, route_to_all is not supported directly
// It should be handled by PDRouter if needed
return HttpResponse::NotImplemented()
.body("route_to_all not implemented for PrefillDecode mode");
}
_ => self.get_worker_urls(),
};
// Get all worker URLs
let worker_urls = self.get_worker_urls();
// Send requests to all workers concurrently
let mut tasks = Vec::new();
......@@ -571,16 +329,6 @@ impl Router {
client: &reqwest::Client,
_req: &HttpRequest,
) -> HttpResponse {
// For PD mode, delegate to PDRouter
match self {
Router::PrefillDecode { pd_router } => {
return pd_router.get_loads(client).await;
}
_ => {
// For non-PD routers, handle normally
}
}
let urls = self.get_worker_urls();
let prefill_urls: Vec<String> = Vec::new();
let decode_urls = urls;
......@@ -623,10 +371,6 @@ impl Router {
typed_req: &T,
route: &str,
) -> HttpResponse {
match self {
Router::PrefillDecode { .. } => HttpResponse::InternalServerError()
.body("PD routing should use specialized typed handlers"),
_ => {
// Handle retries like the original implementation
let start = Instant::now();
const MAX_REQUEST_RETRIES: u32 = 3;
......@@ -646,17 +390,13 @@ impl Router {
while request_retries < MAX_REQUEST_RETRIES {
if total_retries >= 1 {
info!("Retrying request after {} failed attempts", total_retries);
counter!("sgl_router_retries_total", "route" => route.to_string())
.increment(1);
counter!("sgl_router_retries_total", "route" => route.to_string()).increment(1);
}
// For CacheAware router, increment load before request
let load_incremented = match self {
Router::CacheAware { workers, .. } => {
let workers_guard = workers.read().unwrap();
if let Some(worker) =
workers_guard.iter().find(|w| w.url() == &worker_url)
{
// Increment load before request if using RAII load tracking
let load_incremented = if self.policy.name() == "cache_aware" {
let workers_guard = self.workers.read().unwrap();
if let Some(worker) = workers_guard.iter().find(|w| w.url() == &worker_url) {
worker.increment_load();
gauge!("sgl_router_running_requests", "worker" => worker_url.to_string())
.set(worker.load() as f64);
......@@ -664,8 +404,8 @@ impl Router {
} else {
false
}
}
_ => false,
} else {
false
};
// Send typed request directly
......@@ -715,114 +455,19 @@ impl Router {
}
}
counter!("sgl_router_request_errors_total", "route" => route.to_string())
.increment(1);
counter!("sgl_router_request_errors_total", "route" => route.to_string()).increment(1);
HttpResponse::InternalServerError().body("All retry attempts failed")
}
}
}
// Helper method to select worker from text (returns index for RoundRobin/Random, URL for CacheAware)
// Helper method to select worker from text using the policy
fn select_generate_worker_from_text(&self, text: &str) -> String {
match self {
Router::RoundRobin {
workers,
current_index,
..
} => {
let workers_guard = workers.read().unwrap();
let idx = current_index
.fetch_update(
std::sync::atomic::Ordering::SeqCst,
std::sync::atomic::Ordering::SeqCst,
|x| Some((x + 1) % workers_guard.len()),
)
.unwrap();
workers_guard[idx].url().to_string()
}
let workers = self.workers.read().unwrap();
Router::Random { workers, .. } => {
let workers_guard = workers.read().unwrap();
workers_guard[rand::random::<usize>() % workers_guard.len()]
.url()
.to_string()
}
Router::CacheAware {
workers,
tree,
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
..
} => {
let tree = tree.lock().unwrap();
let workers_guard = workers.read().unwrap();
// Get current load statistics from workers
let loads: Vec<usize> = workers_guard.iter().map(|w| w.load()).collect();
let max_load = *loads.iter().max().unwrap_or(&0);
let min_load = *loads.iter().min().unwrap_or(&0);
// Load is considered imbalanced if:
// 1. (max - min) > abs_threshold AND
// 2. max > rel_threshold * min
let is_imbalanced = max_load.saturating_sub(min_load) > *balance_abs_threshold
&& (max_load as f32) > (min_load as f32 * balance_rel_threshold);
let selected_url = if is_imbalanced {
// Log load balancing trigger and current queue state
let worker_loads: Vec<(String, usize)> = workers_guard
.iter()
.map(|w| (w.url().to_string(), w.load()))
.collect();
info!(
"Load balancing triggered due to workload imbalance:\n\
Max load: {}, Min load: {}\n\
Current worker loads: {:?}",
max_load, min_load, worker_loads
);
counter!("sgl_router_load_balancing_events_total").increment(1);
gauge!("sgl_router_max_load").set(max_load as f64);
gauge!("sgl_router_min_load").set(min_load as f64);
// Use shortest queue routing when load is imbalanced
workers_guard
.iter()
.min_by_key(|w| w.load())
.map(|w| w.url().to_string())
.unwrap_or_else(|| workers_guard[0].url().to_string())
} else {
// Use cache-aware routing when load is balanced
let (matched_text, matched_worker) = tree.prefix_match(&text);
let matched_rate =
matched_text.chars().count() as f32 / text.chars().count() as f32;
if matched_rate > *cache_threshold {
counter!("sgl_router_cache_hits_total").increment(1);
matched_worker.to_string()
} else {
counter!("sgl_router_cache_misses_total").increment(1);
tree.get_smallest_tenant()
}
};
// Find the selected worker and increment processed counter only
if let Some(worker) = workers_guard.iter().find(|w| w.url() == &selected_url) {
worker.increment_processed();
counter!("sgl_router_processed_requests_total", "worker" => selected_url.to_string())
.increment(1);
}
tree.insert(&text, &selected_url);
selected_url
}
Router::PrefillDecode { .. } => {
// For PD mode, we don't use this method
return "PD_MODE_ERROR".to_string();
match self.policy.select_worker(&workers, Some(text)) {
Some(idx) => workers[idx].url().to_string(),
None => {
warn!("No healthy workers available");
String::new()
}
}
}
......@@ -862,20 +507,16 @@ impl Router {
Err(e) => {
error!("Failed to send request to {}: {}", worker_url, e);
// Decrement load on error for CacheAware router
// Decrement load on error if it was incremented
if load_incremented {
if let Router::CacheAware { workers, .. } = self {
if let Ok(workers_guard) = workers.read() {
if let Some(worker) =
workers_guard.iter().find(|w| w.url() == worker_url)
{
if let Ok(workers_guard) = self.workers.read() {
if let Some(worker) = workers_guard.iter().find(|w| w.url() == worker_url) {
worker.decrement_load();
gauge!("sgl_router_running_requests", "worker" => worker_url.to_string())
.set(worker.load() as f64);
}
}
}
}
return HttpResponse::InternalServerError().body(format!("Request failed: {}", e));
}
......@@ -894,10 +535,9 @@ impl Router {
}
};
// Decrement load counter for non-streaming CacheAware requests
// Decrement load counter for non-streaming requests if it was incremented
if load_incremented && !is_stream {
if let Router::CacheAware { workers, .. } = self {
if let Ok(workers_guard) = workers.read() {
if let Ok(workers_guard) = self.workers.read() {
if let Some(worker) = workers_guard.iter().find(|w| w.url() == worker_url) {
worker.decrement_load();
gauge!("sgl_router_running_requests", "worker" => worker_url.to_string())
......@@ -905,7 +545,6 @@ impl Router {
}
}
}
}
// Record metrics
let duration = start.elapsed();
......@@ -914,9 +553,9 @@ impl Router {
counter!("sgl_router_requests_total", "route" => route.to_string()).increment(1);
response
} else if let Router::CacheAware { workers, .. } = self {
// For streaming with CacheAware router, we need to manually decrement when done
let workers = Arc::clone(workers);
} else if load_incremented {
// For streaming with load tracking, we need to manually decrement when done
let workers = Arc::clone(&self.workers);
let worker_url = worker_url.to_string();
HttpResponse::build(status)
......@@ -948,7 +587,7 @@ impl Router {
}),
)
} else {
// For non-CacheAware routers, just stream without load tracking
// For requests without load tracking, just stream
HttpResponse::build(status)
.insert_header((CONTENT_TYPE, HeaderValue::from_static("text/event-stream")))
.streaming(res.bytes_stream().map_err(|_| {
......@@ -958,73 +597,47 @@ impl Router {
}
pub async fn add_worker(&self, worker_url: &str) -> Result<String, String> {
let (timeout_secs, interval_secs) = match self {
Router::Random {
timeout_secs,
interval_secs,
..
} => (*timeout_secs, *interval_secs),
Router::RoundRobin {
timeout_secs,
interval_secs,
..
} => (*timeout_secs, *interval_secs),
Router::CacheAware {
timeout_secs,
interval_secs,
..
} => (*timeout_secs, *interval_secs),
Router::PrefillDecode { .. } => {
// For PD mode, we don't support adding workers via this method
return Err("Adding workers to PrefillDecode router not supported via add_worker. Use dedicated PD management methods.".to_string());
}
};
let start_time = std::time::Instant::now();
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(timeout_secs))
.timeout(Duration::from_secs(self.timeout_secs))
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
loop {
if start_time.elapsed() > Duration::from_secs(timeout_secs) {
if start_time.elapsed() > Duration::from_secs(self.timeout_secs) {
error!(
"Timeout {}s waiting for worker {} to become healthy. Please set --router-worker-startup-timeout-secs (sglang_router.launch_server) or --worker-startup-timeout-secs (sglang_worker.router) to a larger value",
timeout_secs, worker_url
self.timeout_secs, worker_url
);
return Err(format!(
"Timeout {}s waiting for worker {} to become healthy. Please set --router-worker-startup-timeout-secs (sglang_router.launch_server) or --worker-startup-timeout-secs (sglang_worker.router) to a larger value",
timeout_secs, worker_url
self.timeout_secs, worker_url
));
}
match client.get(&format!("{}/health", worker_url)).send().await {
Ok(res) => {
if res.status().is_success() {
match self {
Router::RoundRobin { workers, .. }
| Router::Random { workers, .. }
| Router::CacheAware { workers, .. } => {
info!("Worker {} health check passed", worker_url);
let mut workers_guard = workers.write().unwrap();
let mut workers_guard = self.workers.write().unwrap();
if workers_guard.iter().any(|w| w.url() == worker_url) {
return Err(format!("Worker {} already exists", worker_url));
}
info!("Added worker: {}", worker_url);
let new_worker =
WorkerFactory::create_regular(worker_url.to_string());
let new_worker = WorkerFactory::create_regular(worker_url.to_string());
workers_guard.push(new_worker);
gauge!("sgl_router_active_workers").set(workers_guard.len() as f64);
}
Router::PrefillDecode { .. } => {
return Err("Adding workers to PrefillDecode router not supported via add_worker. Use dedicated PD management methods.".to_string());
}
}
// If cache aware, add worker to tree
if let Router::CacheAware { tree, .. } = self {
// Add worker to tree
tree.lock().unwrap().insert("", worker_url);
// If cache aware policy, initialize the worker in the tree
if let Some(cache_aware) =
self.policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
// Get updated workers after adding
drop(workers_guard);
let workers_guard = self.workers.read().unwrap();
cache_aware.init_workers(&workers_guard);
}
return Ok(format!("Successfully added worker: {}", worker_url));
......@@ -1040,7 +653,7 @@ impl Router {
warn!("The worker url {} does not have http or https prefix. Please add the prefix to the url.", worker_url);
}
tokio::time::sleep(Duration::from_secs(interval_secs)).await;
tokio::time::sleep(Duration::from_secs(self.interval_secs)).await;
continue;
}
}
......@@ -1055,7 +668,7 @@ impl Router {
warn!("The worker url {} does not have http or https prefix. Please add the prefix to the url.", worker_url);
}
tokio::time::sleep(Duration::from_secs(interval_secs)).await;
tokio::time::sleep(Duration::from_secs(self.interval_secs)).await;
continue;
}
}
......@@ -1063,11 +676,7 @@ impl Router {
}
pub fn remove_worker(&self, worker_url: &str) {
match self {
Router::RoundRobin { workers, .. }
| Router::Random { workers, .. }
| Router::CacheAware { workers, .. } => {
let mut workers_guard = workers.write().unwrap();
let mut workers_guard = self.workers.write().unwrap();
if let Some(index) = workers_guard.iter().position(|w| w.url() == worker_url) {
workers_guard.remove(index);
info!("Removed worker: {}", worker_url);
......@@ -1076,70 +685,96 @@ impl Router {
warn!("Worker {} not found, skipping removal", worker_url);
return;
}
}
Router::PrefillDecode { .. } => {
warn!("Removing workers from PrefillDecode router not supported via remove_worker. Use dedicated PD management methods.");
return;
}
}
// if cache aware, remove the worker from the tree
if let Router::CacheAware { tree, .. } = self {
tree.lock().unwrap().remove_tenant(&worker_url);
// If cache aware policy, remove the worker from the tree
if let Some(cache_aware) = self
.policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_aware.remove_worker(worker_url);
info!("Removed worker from tree: {}", worker_url);
}
}
/// Add a worker with PD mode support
pub async fn add_pd_worker(
&self,
worker_url: &str,
pod_type: crate::service_discovery::PodType,
bootstrap_port: Option<u16>,
) -> Result<String, String> {
match self {
Router::PrefillDecode { pd_router } => match pod_type {
crate::service_discovery::PodType::Prefill => pd_router
.add_prefill_server(worker_url.to_string(), bootstrap_port)
.await
.map_err(|e| e.to_string()),
crate::service_discovery::PodType::Decode => pd_router
.add_decode_server(worker_url.to_string())
.await
.map_err(|e| e.to_string()),
crate::service_discovery::PodType::Regular => {
Err("Regular pod type not supported in PD mode".to_string())
async fn get_worker_load(&self, client: &reqwest::Client, worker_url: &str) -> Option<isize> {
match client.get(&format!("{}/get_load", worker_url)).send().await {
Ok(res) if res.status().is_success() => match res.bytes().await {
Ok(bytes) => match serde_json::from_slice::<serde_json::Value>(&bytes) {
Ok(data) => data
.get("load")
.and_then(|v| v.as_i64())
.map(|v| v as isize),
Err(e) => {
debug!("Failed to parse load response from {}: {}", worker_url, e);
None
}
},
Err(e) => {
debug!("Failed to read load response from {}: {}", worker_url, e);
None
}
},
_ => Err("add_pd_worker only supported in PD mode".to_string()),
Ok(res) => {
debug!(
"Worker {} returned non-success status: {}",
worker_url,
res.status()
);
None
}
Err(e) => {
debug!("Failed to get load from {}: {}", worker_url, e);
None
}
}
}
/// Remove a worker with PD mode support
pub async fn remove_pd_worker(
&self,
worker_url: &str,
pod_type: crate::service_discovery::PodType,
) -> Result<String, String> {
match self {
Router::PrefillDecode { pd_router } => match pod_type {
crate::service_discovery::PodType::Prefill => pd_router
.remove_prefill_server(worker_url)
.await
.map_err(|e| e.to_string()),
crate::service_discovery::PodType::Decode => pd_router
.remove_decode_server(worker_url)
.await
.map_err(|e| e.to_string()),
crate::service_discovery::PodType::Regular => {
Err("Regular pod type not supported in PD mode".to_string())
// Background task to monitor worker loads
async fn monitor_worker_loads(
worker_urls: Vec<String>,
tx: tokio::sync::watch::Sender<HashMap<String, isize>>,
interval_secs: u64,
policy: Arc<dyn LoadBalancingPolicy>,
) {
let client = match reqwest::Client::builder()
.timeout(Duration::from_secs(5))
.build()
{
Ok(c) => c,
Err(e) => {
error!("Failed to create HTTP client for load monitoring: {}", e);
return;
}
},
_ => Err("remove_pd_worker only supported in PD mode".to_string()),
};
let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
loop {
interval.tick().await;
let mut loads = HashMap::new();
for url in &worker_urls {
if let Some(load) = Self::get_worker_load_static(&client, url).await {
loads.insert(url.clone(), load);
debug!("Worker {} load: {}", url, load);
}
}
async fn get_worker_load(&self, client: &reqwest::Client, worker_url: &str) -> Option<isize> {
if !loads.is_empty() {
// Update policy with new loads
policy.update_loads(&loads);
// Send to watchers
if let Err(e) = tx.send(loads) {
error!("Failed to send load update: {}", e);
}
}
}
}
// Static version of get_worker_load for use in monitoring task
async fn get_worker_load_static(client: &reqwest::Client, worker_url: &str) -> Option<isize> {
match client.get(&format!("{}/get_load", worker_url)).send().await {
Ok(res) if res.status().is_success() => match res.bytes().await {
Ok(bytes) => match serde_json::from_slice::<serde_json::Value>(&bytes) {
......@@ -1171,100 +806,192 @@ impl Router {
}
}
}
}
// PD-specific wrapper methods that delegate to PDRouter
pub async fn route_pd_health_generate(
&self,
_client: &reqwest::Client,
_req: &HttpRequest,
) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router.health_generate(&pd_router.http_client).await
use crate::routers::{RouterTrait, WorkerManagement};
use async_trait::async_trait;
use reqwest::Client;
#[async_trait]
impl WorkerManagement for Router {
async fn add_worker(&self, worker_url: &str) -> Result<String, String> {
Router::add_worker(self, worker_url).await
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
fn remove_worker(&self, worker_url: &str) {
Router::remove_worker(self, worker_url)
}
fn get_worker_urls(&self) -> Vec<String> {
Router::get_worker_urls(self)
}
}
pub async fn route_pd_generate_typed(
&self,
_client: &reqwest::Client,
req: &HttpRequest,
typed_req: crate::pd_types::GenerateReqInput,
route: &str,
) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router
.route_generate(&pd_router.http_client, req, typed_req, route)
.await
#[async_trait(?Send)]
impl RouterTrait for Router {
fn as_any(&self) -> &dyn std::any::Any {
self
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
async fn health(&self, _client: &Client, _req: &HttpRequest) -> HttpResponse {
// Check local health state of all workers (consistent with PD router)
// Note: This uses cached health status from background health checks, not live checks
let mut all_healthy = true;
let mut unhealthy_servers = Vec::new();
for worker in self.workers.read().unwrap().iter() {
if !worker.is_healthy() {
all_healthy = false;
unhealthy_servers.push(worker.url().to_string());
}
}
pub async fn route_pd_chat_typed(
if all_healthy {
HttpResponse::Ok().body("All servers healthy")
} else {
HttpResponse::ServiceUnavailable()
.body(format!("Unhealthy servers: {:?}", unhealthy_servers))
}
}
async fn health_generate(&self, client: &Client, req: &HttpRequest) -> HttpResponse {
// Test model generation capability by sending to first available worker
// Note: This endpoint actually causes the model to generate a token, so we only test one worker
self.route_to_first(client, "/health_generate", req).await
}
async fn get_server_info(&self, client: &Client, req: &HttpRequest) -> HttpResponse {
self.route_to_first(client, "/get_server_info", req).await
}
async fn get_models(&self, client: &Client, req: &HttpRequest) -> HttpResponse {
self.route_to_first(client, "/v1/models", req).await
}
async fn get_model_info(&self, client: &Client, req: &HttpRequest) -> HttpResponse {
self.route_to_first(client, "/get_model_info", req).await
}
async fn route_generate(
&self,
_client: &reqwest::Client,
client: &Client,
req: &HttpRequest,
typed_req: crate::pd_types::ChatReqInput,
route: &str,
body: serde_json::Value,
) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router
.route_chat(&pd_router.http_client, req, typed_req, route)
// Convert JSON to typed request
match serde_json::from_value::<crate::openai_api_types::GenerateRequest>(body) {
Ok(typed_req) => {
self.route_typed_request(client, req, &typed_req, "/generate")
.await
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
Err(e) => HttpResponse::BadRequest().body(format!("Invalid request: {}", e)),
}
}
pub async fn get_pd_server_info(
async fn route_chat(
&self,
_client: &reqwest::Client,
_req: &HttpRequest,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router.get_server_info(&pd_router.http_client).await
// Convert JSON to typed request
match serde_json::from_value::<crate::openai_api_types::ChatCompletionRequest>(body) {
Ok(typed_req) => {
self.route_typed_request(client, req, &typed_req, "/v1/chat/completions")
.await
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
Err(e) => HttpResponse::BadRequest().body(format!("Invalid request: {}", e)),
}
}
pub async fn get_pd_models(
async fn route_completion(
&self,
_client: &reqwest::Client,
client: &Client,
req: &HttpRequest,
body: serde_json::Value,
) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router.get_models(&pd_router.http_client, req).await
// Convert JSON to typed request
match serde_json::from_value::<crate::openai_api_types::CompletionRequest>(body) {
Ok(typed_req) => {
self.route_typed_request(client, req, &typed_req, "/v1/completions")
.await
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
Err(e) => HttpResponse::BadRequest().body(format!("Invalid request: {}", e)),
}
}
pub async fn route_pd_flush_cache(&self, _client: &reqwest::Client) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router.flush_cache(&pd_router.http_client).await
async fn flush_cache(&self, client: &Client) -> HttpResponse {
// Get all worker URLs
let worker_urls = self.get_worker_urls();
// Send requests to all workers concurrently without headers
let mut tasks = Vec::new();
for worker_url in &worker_urls {
let request_builder = client.post(format!("{}/flush_cache", worker_url));
tasks.push(request_builder.send());
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
// Wait for all responses
let results = futures_util::future::join_all(tasks).await;
// Check if all succeeded
let all_success = results.iter().all(|r| {
r.as_ref()
.map(|res| res.status().is_success())
.unwrap_or(false)
});
if all_success {
HttpResponse::Ok().body("Cache flushed on all servers")
} else {
HttpResponse::InternalServerError().body("Cache flush failed on one or more servers")
}
}
pub async fn get_pd_model_info(
&self,
_client: &reqwest::Client,
req: &HttpRequest,
) -> HttpResponse {
match self {
Router::PrefillDecode { pd_router } => {
pd_router.get_model_info(&pd_router.http_client, req).await
async fn get_worker_loads(&self, client: &Client) -> HttpResponse {
let urls = self.get_worker_urls();
let mut loads = Vec::new();
// Get loads from all workers
for url in &urls {
let load = self.get_worker_load(client, url).await.unwrap_or(-1);
loads.push(serde_json::json!({
"worker": url,
"load": load
}));
}
HttpResponse::Ok().json(serde_json::json!({
"workers": loads
}))
}
fn router_type(&self) -> &'static str {
"regular"
}
_ => HttpResponse::InternalServerError().body("Not in PrefillDecode mode"),
fn readiness(&self) -> HttpResponse {
// Regular router is ready if it has at least one healthy worker
let healthy_count = self
.workers
.read()
.unwrap()
.iter()
.filter(|w| w.is_healthy())
.count();
if healthy_count > 0 {
HttpResponse::Ok().json(serde_json::json!({
"status": "ready",
"healthy_workers": healthy_count,
"total_workers": self.workers.read().unwrap().len()
}))
} else {
HttpResponse::ServiceUnavailable().json(serde_json::json!({
"status": "not_ready",
"reason": "no healthy workers available",
"total_workers": self.workers.read().unwrap().len()
}))
}
}
}
......@@ -1272,17 +999,22 @@ impl Router {
#[cfg(test)]
mod tests {
use super::*;
use crate::service_discovery::PodType;
use crate::policies::RandomPolicy;
use std::collections::HashMap;
fn create_test_regular_router() -> Router {
let workers = vec![
WorkerFactory::create_regular("http://worker1:8080".to_string()),
WorkerFactory::create_regular("http://worker2:8080".to_string()),
];
Router::Random {
let (_, rx) = tokio::sync::watch::channel(HashMap::new());
Router {
workers: Arc::new(RwLock::new(workers)),
policy: Arc::new(RandomPolicy::new()),
timeout_secs: 5,
interval_secs: 1,
_worker_loads: Arc::new(rx),
_load_monitor_handle: None,
_health_checker: None,
}
}
......@@ -1297,53 +1029,6 @@ mod tests {
assert!(urls.contains(&"http://worker2:8080".to_string()));
}
// #[test]
// fn test_router_get_worker_urls_pd_mode() {
// // For PD mode, get_worker_urls returns empty list
// // Note: PDRouter::new requires health checks which fail in tests
// // This test would need a mock server or different test setup
// }
#[tokio::test]
async fn test_add_pd_worker_with_regular_router() {
let router = create_test_regular_router();
let result = router
.add_pd_worker("http://new-worker:8080", PodType::Prefill, Some(8081))
.await;
assert!(result.is_err());
assert!(result
.unwrap_err()
.contains("add_pd_worker only supported in PD mode"));
}
#[tokio::test]
async fn test_remove_pd_worker_with_regular_router() {
let router = create_test_regular_router();
let result = router
.remove_pd_worker("http://worker:8080", PodType::Decode)
.await;
assert!(result.is_err());
assert!(result
.unwrap_err()
.contains("remove_pd_worker only supported in PD mode"));
}
// #[tokio::test]
// async fn test_add_pd_worker_with_pd_router_regular_type() {
// // Note: PDRouter::new requires health checks which fail in tests
// // This test would need a mock server or different test setup
// }
// #[tokio::test]
// async fn test_remove_pd_worker_with_pd_router_regular_type() {
// // Note: PDRouter::new requires health checks which fail in tests
// // This test would need a mock server or different test setup
// }
#[test]
fn test_select_first_worker_regular() {
let router = create_test_regular_router();
......@@ -1353,12 +1038,6 @@ mod tests {
assert_eq!(result.unwrap(), "http://worker1:8080");
}
// #[test]
// fn test_select_first_worker_pd_mode() {
// // Note: PDRouter::new requires health checks which fail in tests
// // This test would need a mock server or different test setup
// }
#[test]
fn test_wait_for_healthy_workers_empty_list() {
let result = Router::wait_for_healthy_workers(&[], 1, 1);
......
sgl-router/src/server.rs
View file @ c8f31042
use crate::config::RouterConfig;
use crate::logging::{self, LoggingConfig};
use crate::openai_api_types::{ChatCompletionRequest, CompletionRequest, GenerateRequest};
use crate::prometheus::{self, PrometheusConfig};
use crate::request_adapter::ToPdRequest;
use crate::router::PolicyConfig;
use crate::router::Router;
use crate::routers::{RouterFactory, RouterTrait};
use crate::service_discovery::{start_service_discovery, ServiceDiscoveryConfig};
use actix_web::{
error, get, post, web, App, Error, HttpRequest, HttpResponse, HttpServer, Responder,
......@@ -19,27 +18,19 @@ use tracing::{error, info, warn, Level};
#[derive(Debug)]
pub struct AppState {
router: Arc<Router>,
router: Arc<dyn RouterTrait>,
client: Client,
is_pd_mode: bool, // Add flag to track PD mode
}
impl AppState {
pub fn new(
worker_urls: Vec<String>,
client: Client,
policy_config: PolicyConfig,
) -> Result<Self, String> {
// Check if this is PD mode from policy config
let is_pd_mode = matches!(policy_config, PolicyConfig::PrefillDecodeConfig { .. });
// Create router based on policy
let router = Arc::new(Router::new(worker_urls, policy_config)?);
Ok(Self {
router,
client,
is_pd_mode,
})
pub fn new(router_config: RouterConfig, client: Client) -> Result<Self, String> {
// Use RouterFactory to create the appropriate router type
let router = RouterFactory::create_router(&router_config)?;
// Convert Box<dyn RouterTrait> to Arc<dyn RouterTrait>
let router = Arc::from(router);
Ok(Self { router, client })
}
}
......@@ -76,65 +67,39 @@ fn json_error_handler(err: error::JsonPayloadError, _req: &HttpRequest) -> Error
}
}
#[get("/liveness")]
async fn liveness(_req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
data.router.liveness()
}
#[get("/readiness")]
async fn readiness(_req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
data.router.readiness()
}
#[get("/health")]
async fn health(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
data.router
.route_to_first(&data.client, "/health", &req)
.await
data.router.health(&data.client, &req).await
}
#[get("/health_generate")]
async fn health_generate(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
// Check if we're in PD mode
if data.is_pd_mode {
// For PD mode, check health on all servers
data.router
.route_pd_health_generate(&data.client, &req)
.await
} else {
// Regular mode
data.router
.route_to_first(&data.client, "/health_generate", &req)
.await
}
data.router.health_generate(&data.client, &req).await
}
#[get("/get_server_info")]
async fn get_server_info(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
if data.is_pd_mode {
// For PD mode, aggregate info from both prefill and decode servers
data.router.get_pd_server_info(&data.client, &req).await
} else {
// Regular mode - return first server's info
data.router
.route_to_first(&data.client, "/get_server_info", &req)
.await
}
data.router.get_server_info(&data.client, &req).await
}
#[get("/v1/models")]
async fn v1_models(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
if data.is_pd_mode {
// For PD mode, return models from the first prefill server
data.router.get_pd_models(&data.client, &req).await
} else {
// Regular mode
data.router
.route_to_first(&data.client, "/v1/models", &req)
.await
}
data.router.get_models(&data.client, &req).await
}
#[get("/get_model_info")]
async fn get_model_info(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
if data.is_pd_mode {
// For PD mode, get model info from the first prefill server
data.router.get_pd_model_info(&data.client, &req).await
} else {
data.router
.route_to_first(&data.client, "/get_model_info", &req)
.await
}
data.router.get_model_info(&data.client, &req).await
}
#[post("/generate")]
......@@ -143,24 +108,12 @@ async fn generate(
body: web::Json<GenerateRequest>,
state: web::Data<AppState>,
) -> Result<HttpResponse, Error> {
let client = &state.client;
let router = &state.router;
// Use typed request directly for both PD and regular routing
if state.is_pd_mode {
// For PD mode, convert to PD request with bootstrap
let pd_request = body.into_inner().to_pd_request();
Ok(router
.route_pd_generate_typed(&client, &req, pd_request, "/generate")
.await)
} else {
// For regular mode, use typed request directly
let request = body.into_inner();
Ok(router
.route_typed_request(&client, &req, &request, "/generate")
let json_body = serde_json::to_value(body.into_inner())
.map_err(|e| error::ErrorBadRequest(format!("Invalid JSON: {}", e)))?;
Ok(state
.router
.route_generate(&state.client, &req, json_body)
.await)
}
}
#[post("/v1/chat/completions")]
......@@ -169,24 +122,12 @@ async fn v1_chat_completions(
body: web::Json<ChatCompletionRequest>,
state: web::Data<AppState>,
) -> Result<HttpResponse, Error> {
let client = &state.client;
let router = &state.router;
// Use typed request directly for both PD and regular routing
if state.is_pd_mode {
// For PD mode, convert to PD request with bootstrap
let pd_request = body.into_inner().to_pd_request();
Ok(router
.route_pd_chat_typed(&client, &req, pd_request, "/v1/chat/completions")
.await)
} else {
// For regular mode, use typed request directly
let request = body.into_inner();
Ok(router
.route_typed_request(&client, &req, &request, "/v1/chat/completions")
let json_body = serde_json::to_value(body.into_inner())
.map_err(|e| error::ErrorBadRequest(format!("Invalid JSON: {}", e)))?;
Ok(state
.router
.route_chat(&state.client, &req, json_body)
.await)
}
}
#[post("/v1/completions")]
......@@ -195,24 +136,12 @@ async fn v1_completions(
body: web::Json<CompletionRequest>,
state: web::Data<AppState>,
) -> Result<HttpResponse, Error> {
let client = &state.client;
let router = &state.router;
// Use typed request directly for both PD and regular routing
if state.is_pd_mode {
// For PD mode, convert to PD request with bootstrap
let pd_request = body.into_inner().to_pd_request();
Ok(router
.route_pd_generate_typed(&client, &req, pd_request, "/v1/completions")
.await)
} else {
// For regular mode, use typed request directly
let request = body.into_inner();
Ok(router
.route_typed_request(&client, &req, &request, "/v1/completions")
let json_body = serde_json::to_value(body.into_inner())
.map_err(|e| error::ErrorBadRequest(format!("Invalid JSON: {}", e)))?;
Ok(state
.router
.route_completion(&state.client, &req, json_body)
.await)
}
}
#[post("/add_worker")]
......@@ -254,29 +183,19 @@ async fn remove_worker(
}
#[post("/flush_cache")]
async fn flush_cache(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
if data.is_pd_mode {
// For PD mode, flush cache on both prefill and decode servers
data.router.route_pd_flush_cache(&data.client).await
} else {
// Route to all workers for cache flushing
data.router
.route_to_all(&data.client, "/flush_cache", &req)
.await
}
async fn flush_cache(_req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
data.router.flush_cache(&data.client).await
}
#[get("/get_loads")]
async fn get_loads(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
// Get loads from all workers
data.router.get_all_loads(&data.client, &req).await
async fn get_loads(_req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
data.router.get_worker_loads(&data.client).await
}
pub struct ServerConfig {
pub host: String,
pub port: u16,
pub worker_urls: Vec<String>,
pub policy_config: PolicyConfig,
pub router_config: RouterConfig,
pub max_payload_size: usize,
pub log_dir: Option<String>,
pub log_level: Option<String>,
......@@ -324,8 +243,8 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
}
info!("🚧 Initializing router on {}:{}", config.host, config.port);
info!("🚧 Initializing workers on {:?}", config.worker_urls);
info!("🚧 Policy Config: {:?}", config.policy_config);
info!("🚧 Router mode: {:?}", config.router_config.mode);
info!("🚧 Policy: {:?}", config.router_config.policy);
info!(
"🚧 Max payload size: {} MB",
config.max_payload_size / (1024 * 1024)
......@@ -345,11 +264,7 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
.build()
.expect("Failed to create HTTP client");
let app_state_init = AppState::new(
config.worker_urls.clone(),
client.clone(),
config.policy_config.clone(),
)
let app_state_init = AppState::new(config.router_config.clone(), client.clone())
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
let router_arc = Arc::clone(&app_state_init.router);
let app_state = web::Data::new(app_state_init);
......@@ -397,6 +312,8 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
.service(v1_completions)
.service(v1_models)
.service(get_model_info)
.service(liveness)
.service(readiness)
.service(health)
.service(health_generate)
.service(get_server_info)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment