chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/tests/routing/load_balancing_test.rs

@@ -0,0 +1,373 @@
+//! Load balancing integration tests
+//!
+//! Tests for various load balancing policies: round_robin, random, cache_aware, etc.
+
+use axum::{
+    body::Body,
+    extract::Request,
+    http::{header::CONTENT_TYPE, StatusCode},
+};
+use serde_json::json;
+use smg::config::{CircuitBreakerConfig, RetryConfig};
+use tower::ServiceExt;
+
+use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
+
+#[cfg(test)]
+mod round_robin_tests {
+    use super::*;
+
+    /// Test that round robin distributes requests evenly across workers
+    #[tokio::test]
+    async fn test_round_robin_distribution() {
+        let config = TestRouterConfig::round_robin(3100);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19001, 3))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 30;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed with round robin across 3 healthy workers
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with round robin"
+        );
+
+        ctx.shutdown().await;
+    }
+
+    /// Test round robin with one worker failing
+    #[tokio::test]
+    async fn test_round_robin_with_failing_worker() {
+        let config = TestRouterConfig::round_robin_with_retry(
+            3101,
+            RetryConfig {
+                max_retries: 2,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+        );
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19004, 1.0), // Always fail
+                TestWorkerConfig::healthy(19005),    // Always succeed
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // With retries enabled, requests should eventually succeed
+        // by being retried on the healthy worker
+        let payload = json!({
+            "text": "Test with retries",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        assert_eq!(
+            resp.status(),
+            StatusCode::OK,
+            "Request should succeed via retry on healthy worker"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod random_tests {
+    use super::*;
+
+    /// Test that random policy distributes requests across workers
+    #[tokio::test]
+    async fn test_random_distribution() {
+        let config = TestRouterConfig::random(3102);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19010, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let num_requests = 20;
+        let mut success_count = 0;
+
+        for i in 0..num_requests {
+            let payload = json!({
+                "text": format!("Random request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed with random policy
+        assert_eq!(
+            success_count, num_requests,
+            "All requests should succeed with random policy"
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod cache_aware_tests {
+    use super::*;
+
+    /// Test cache-aware routing uses consistent hashing
+    #[tokio::test]
+    async fn test_cache_aware_consistent_routing() {
+        let config = TestRouterConfig::cache_aware(3103);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19020, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+
+        // Same prompt should route to same worker (consistent hashing)
+        let same_prompt = "Hello, cache-aware routing test!";
+        let mut worker_ids: Vec<Option<String>> = Vec::new();
+
+        for _ in 0..5 {
+            let payload = json!({
+                "text": same_prompt,
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+
+            // Extract worker ID from response header if available
+            let worker_id = resp
+                .headers()
+                .get("x-worker-id")
+                .map(|v| v.to_str().unwrap().to_string());
+            worker_ids.push(worker_id);
+        }
+
+        // All requests should succeed
+        assert_eq!(worker_ids.len(), 5);
+
+        ctx.shutdown().await;
+    }
+
+    /// Test cache-aware routing with different prompts
+    #[tokio::test]
+    async fn test_cache_aware_different_prompts() {
+        let config = TestRouterConfig::cache_aware(3104);
+        let ctx =
+            AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19022, 2))
+                .await;
+
+        let app = ctx.create_app().await;
+        let prompts = vec![
+            "First unique prompt",
+            "Second unique prompt",
+            "Third unique prompt",
+            "Fourth unique prompt",
+        ];
+
+        let mut success_count = 0;
+        for prompt in prompts {
+            let payload = json!({
+                "text": prompt,
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            if resp.status() == StatusCode::OK {
+                success_count += 1;
+            }
+        }
+
+        // All requests should succeed
+        assert_eq!(success_count, 4);
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod worker_health_tests {
+    use super::*;
+
+    /// Test that failing workers are avoided via circuit breaker/retry
+    #[tokio::test]
+    async fn test_skip_failing_workers() {
+        let config = TestRouterConfig::round_robin_with_reliability(
+            3105,
+            RetryConfig {
+                max_retries: 3,
+                initial_backoff_ms: 10,
+                max_backoff_ms: 50,
+                ..Default::default()
+            },
+            CircuitBreakerConfig {
+                failure_threshold: 2,
+                success_threshold: 1,
+                timeout_duration_secs: 2,
+                window_duration_secs: 10,
+            },
+        );
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![
+                TestWorkerConfig::flaky(19030, 1.0), // Always fails
+                TestWorkerConfig::healthy(19031),    // Always succeeds
+            ],
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        // Requests should succeed by retrying on healthy worker
+        // or by circuit breaker opening on failing worker
+        for i in 0..10 {
+            let payload = json!({
+                "text": format!("Request {}", i),
+                "stream": false
+            });
+
+            let req = Request::builder()
+                .method("POST")
+                .uri("/generate")
+                .header(CONTENT_TYPE, "application/json")
+                .body(Body::from(serde_json::to_string(&payload).unwrap()))
+                .unwrap();
+
+            let resp = app.clone().oneshot(req).await.unwrap();
+            assert_eq!(resp.status(), StatusCode::OK);
+        }
+
+        ctx.shutdown().await;
+    }
+
+    /// Test behavior when all workers are unhealthy
+    #[tokio::test]
+    async fn test_all_workers_unhealthy() {
+        let ctx = AppTestContext::new(vec![]).await;
+
+        let app = ctx.create_app().await;
+
+        let payload = json!({
+            "text": "Test with no workers",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        // Should return service unavailable when no workers available
+        assert!(
+            resp.status() == StatusCode::SERVICE_UNAVAILABLE
+                || resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
+            "Expected 503 or 500 when no workers available, got {}",
+            resp.status()
+        );
+
+        ctx.shutdown().await;
+    }
+}
+
+#[cfg(test)]
+mod worker_response_delay_tests {
+    use super::*;
+
+    /// Test that response delays are handled correctly
+    #[tokio::test]
+    async fn test_worker_with_delay() {
+        let config = TestRouterConfig::random(3106);
+        let ctx = AppTestContext::new_with_config(
+            config,
+            vec![TestWorkerConfig::slow(19040, 100)], // 100ms delay
+        )
+        .await;
+
+        let app = ctx.create_app().await;
+
+        let start = std::time::Instant::now();
+        let payload = json!({
+            "text": "Test with delay",
+            "stream": false
+        });
+
+        let req = Request::builder()
+            .method("POST")
+            .uri("/generate")
+            .header(CONTENT_TYPE, "application/json")
+            .body(Body::from(serde_json::to_string(&payload).unwrap()))
+            .unwrap();
+
+        let resp = app.oneshot(req).await.unwrap();
+        let elapsed = start.elapsed();
+
+        assert_eq!(resp.status(), StatusCode::OK);
+        // Response should take at least 100ms due to configured delay
+        assert!(
+            elapsed.as_millis() >= 100,
+            "Response should be delayed by at least 100ms, got {}ms",
+            elapsed.as_millis()
+        );
+
+        ctx.shutdown().await;
+    }
+}