[pd-router] Add Configurable Retry Logic for reduce backend pressure (sgl-project#8744)

Author: slin1237

sgl-router/src/config/types.rs

@@ -39,6 +39,8 @@ pub struct RouterConfig {
     pub max_concurrent_requests: usize,
     /// CORS allowed origins
     pub cors_allowed_origins: Vec<String>,
+    /// Retry configuration
+    pub retry: RetryConfig,
 }
 
 /// Routing mode configuration
@@ -182,6 +184,30 @@ impl Default for DiscoveryConfig {
     }
 }
 
+/// Retry configuration for request handling
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RetryConfig {
+    /// Maximum number of retry attempts
+    pub max_retries: u32,
+    /// Initial backoff delay in milliseconds
+    pub initial_backoff_ms: u64,
+    /// Maximum backoff delay in milliseconds
+    pub max_backoff_ms: u64,
+    /// Backoff multiplier for exponential backoff
+    pub backoff_multiplier: f32,
+}
+
+impl Default for RetryConfig {
+    fn default() -> Self {
+        Self {
+            max_retries: 3,
+            initial_backoff_ms: 100,
+            max_backoff_ms: 10000,
+            backoff_multiplier: 2.0,
+        }
+    }
+}
+
 /// Metrics configuration
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MetricsConfig {
@@ -210,7 +236,7 @@ impl Default for RouterConfig {
             host: "127.0.0.1".to_string(),
             port: 3001,
             max_payload_size: 268_435_456, // 256MB
-            request_timeout_secs: 600,
+            request_timeout_secs: 3600,    // 1 hour to match Python mini LB
             worker_startup_timeout_secs: 300,
             worker_startup_check_interval_secs: 10,
             dp_aware: false,
@@ -222,6 +248,7 @@ impl Default for RouterConfig {
             request_id_headers: None,
             max_concurrent_requests: 64,
             cors_allowed_origins: vec![],
+            retry: RetryConfig::default(),
         }
     }
 }
@@ -277,7 +304,7 @@ mod tests {
         assert_eq!(config.host, "127.0.0.1");
         assert_eq!(config.port, 3001);
         assert_eq!(config.max_payload_size, 268_435_456);
-        assert_eq!(config.request_timeout_secs, 600);
+        assert_eq!(config.request_timeout_secs, 3600);
         assert_eq!(config.worker_startup_timeout_secs, 300);
         assert_eq!(config.worker_startup_check_interval_secs, 10);
         assert!(config.discovery.is_none());
@@ -332,6 +359,7 @@ mod tests {
             request_id_headers: None,
             max_concurrent_requests: 64,
             cors_allowed_origins: vec![],
+            retry: RetryConfig::default(),
         };
 
         let json = serde_json::to_string(&config).unwrap();
@@ -759,6 +787,7 @@ mod tests {
             request_id_headers: None,
             max_concurrent_requests: 64,
             cors_allowed_origins: vec![],
+            retry: RetryConfig::default(),
         };
 
         assert!(config.mode.is_pd_mode());
@@ -810,6 +839,7 @@ mod tests {
             request_id_headers: None,
             max_concurrent_requests: 64,
             cors_allowed_origins: vec![],
+            retry: RetryConfig::default(),
         };
 
         assert!(!config.mode.is_pd_mode());
@@ -857,6 +887,7 @@ mod tests {
             request_id_headers: None,
             max_concurrent_requests: 64,
             cors_allowed_origins: vec![],
+            retry: RetryConfig::default(),
         };
 
         assert!(config.has_service_discovery());

sgl-router/src/lib.rs

@@ -19,7 +19,7 @@ pub enum PolicyType {
     Random,
     RoundRobin,
     CacheAware,
-    PowerOfTwo, // Moved from PD-specific, now shared
+    PowerOfTwo,
 }
 
 #[pyclass]
@@ -45,22 +45,18 @@ struct Router {
     selector: HashMap<String, String>,
     service_discovery_port: u16,
     service_discovery_namespace: Option<String>,
-    // PD service discovery fields
     prefill_selector: HashMap<String, String>,
     decode_selector: HashMap<String, String>,
     bootstrap_port_annotation: String,
     prometheus_port: Option<u16>,
     prometheus_host: Option<String>,
     request_timeout_secs: u64,
     request_id_headers: Option<Vec<String>>,
-    // PD mode flag
     pd_disaggregation: bool,
-    // PD-specific fields (only used when pd_disaggregation is true)
     prefill_urls: Option<Vec<(String, Option<u16>)>>,
     decode_urls: Option<Vec<String>>,
     prefill_policy: Option<PolicyType>,
     decode_policy: Option<PolicyType>,
-    // Additional server config fields
     max_concurrent_requests: usize,
     cors_allowed_origins: Vec<String>,
 }
@@ -150,6 +146,7 @@ impl Router {
             request_id_headers: self.request_id_headers.clone(),
             max_concurrent_requests: self.max_concurrent_requests,
             cors_allowed_origins: self.cors_allowed_origins.clone(),
+            retry: config::RetryConfig::default(),
         })
     }
 }
@@ -289,7 +286,6 @@ impl Router {
                 check_interval: std::time::Duration::from_secs(60),
                 port: self.service_discovery_port,
                 namespace: self.service_discovery_namespace.clone(),
-                // PD mode configuration
                 pd_mode: self.pd_disaggregation,
                 prefill_selector: self.prefill_selector.clone(),
                 decode_selector: self.decode_selector.clone(),

sgl-router/src/routers/factory.rs

@@ -50,6 +50,7 @@ impl RouterFactory {
             ctx.router_config.worker_startup_check_interval_secs,
             ctx.router_config.dp_aware,
             ctx.router_config.api_key.clone(),
+            ctx.router_config.retry.clone(),
         )?;
 
         Ok(Box::new(router))
@@ -79,6 +80,7 @@ impl RouterFactory {
             ctx.client.clone(),
             ctx.router_config.worker_startup_timeout_secs,
             ctx.router_config.worker_startup_check_interval_secs,
+            ctx.router_config.retry.clone(),
         )?;
 
         Ok(Box::new(router))