[router] fix pd model completion request

Author: slin1237

sgl-router/benches/request_processing.rs

@@ -97,6 +97,7 @@ fn create_sample_completion_request() -> CompletionRequest {
         logit_bias: None,
         user: None,
         seed: None,
+        other: serde_json::Map::new(),
     }
 }
 

sgl-router/src/openai_api_types.rs

@@ -91,6 +91,10 @@ pub struct CompletionRequest {
     /// If specified, our system will make a best effort to sample deterministically
     #[serde(skip_serializing_if = "Option::is_none")]
     pub seed: Option<i64>,
+
+    /// Additional fields including bootstrap info for PD routing
+    #[serde(flatten)]
+    pub other: serde_json::Map<String, serde_json::Value>,
 }
 
 impl GenerationRequest for CompletionRequest {

sgl-router/src/routers/pd_router.rs

@@ -420,6 +420,77 @@ impl PDRouter {
         .await
     }
 
+    // Route a completion request while preserving OpenAI format
+    pub async fn route_completion(
+        &self,
+        client: &reqwest::Client,
+        req: &HttpRequest,
+        mut typed_req: CompletionRequest,
+        route: &str,
+    ) -> HttpResponse {
+        let start = Instant::now();
+
+        // Get stream flag and return_logprob flag before moving the request
+        let is_stream = typed_req.stream;
+        let return_logprob = typed_req.logprobs.is_some();
+
+        // Extract text for cache-aware routing from the typed request
+        let request_text = match &typed_req.prompt {
+            crate::openai_api_types::StringOrArray::String(s) => Some(s.as_str()),
+            crate::openai_api_types::StringOrArray::Array(arr) => arr.first().map(|s| s.as_str()),
+        };
+
+        // Select servers
+        let (prefill, decode) = match self.select_pd_pair(client, request_text).await {
+            Ok(pair) => pair,
+            Err(e) => {
+                error!("Failed to select PD pair: {}", e);
+                RouterMetrics::record_pd_error("server_selection");
+                return HttpResponse::ServiceUnavailable()
+                    .body(format!("No available servers: {}", e));
+            }
+        };
+
+        // Log routing decision
+        info!(
+            "PD routing: {} -> prefill={}, decode={}",
+            route,
+            prefill.url(),
+            decode.url()
+        );
+
+        // Add bootstrap info using the trait method
+        if let Err(e) = typed_req.add_bootstrap_info(prefill.as_ref()) {
+            error!("Failed to add bootstrap info: {}", e);
+            RouterMetrics::record_pd_error("bootstrap_injection");
+            return HttpResponse::InternalServerError()
+                .body(format!("Bootstrap injection failed: {}", e));
+        }
+
+        // Convert to JSON after bootstrap injection
+        let json_with_bootstrap = match serde_json::to_value(&typed_req) {
+            Ok(json) => json,
+            Err(e) => {
+                error!("Failed to serialize request: {}", e);
+                return HttpResponse::InternalServerError().body("Failed to serialize request");
+            }
+        };
+
+        // Execute dual dispatch
+        self.execute_dual_dispatch(
+            client,
+            req,
+            json_with_bootstrap,
+            route,
+            prefill.as_ref(),
+            decode.as_ref(),
+            is_stream,
+            return_logprob,
+            start,
+        )
+        .await
+    }
+
     // Execute the dual dispatch to prefill and decode servers
     #[allow(clippy::too_many_arguments)]
     async fn execute_dual_dispatch(
@@ -1302,23 +1373,12 @@ impl RouterTrait for PDRouter {
         req: &HttpRequest,
         body: serde_json::Value,
     ) -> HttpResponse {
-        match serde_json::from_value::<CompletionRequest>(body.clone()) {
+        match serde_json::from_value::<CompletionRequest>(body) {
             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))
-                    }
-                }
+                // Use the new method that preserves OpenAI format
+                PDRouter::route_completion(self, client, req, openai_req, "/v1/completions").await
             }
+            Err(e) => HttpResponse::BadRequest().body(format!("Invalid request format: {}", e)),
         }
     }
 

sgl-router/src/routers/pd_types.rs

@@ -1,6 +1,7 @@
 // Essential PDLB types extracted for PD routing
 
 use crate::core::{Worker, WorkerType};
+use crate::openai_api_types::CompletionRequest;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
@@ -233,3 +234,39 @@ impl Bootstrap for ChatReqInput {
         self.bootstrap_room = Some(bootstrap_room);
     }
 }
+
+// Bootstrap implementation for CompletionRequest to preserve OpenAI format
+impl Bootstrap for CompletionRequest {
+    fn is_stream(&self) -> bool {
+        self.stream
+    }
+
+    fn get_batch_size(&self) -> Result<Option<usize>, String> {
+        // Check if 'n' parameter is present and > 1
+        if let Some(n) = self.n {
+            if n > 1 {
+                return Ok(Some(n as usize));
+            }
+        }
+        Ok(None)
+    }
+
+    fn set_bootstrap_info(
+        &mut self,
+        bootstrap_host: BootstrapHost,
+        bootstrap_port: BootstrapPort,
+        bootstrap_room: BootstrapRoom,
+    ) {
+        // Add bootstrap info to the 'other' field to preserve OpenAI format
+        // This follows the same pattern as ChatReqInput
+        if let Ok(host_value) = serde_json::to_value(bootstrap_host) {
+            self.other.insert("bootstrap_host".to_string(), host_value);
+        }
+        if let Ok(port_value) = serde_json::to_value(bootstrap_port) {
+            self.other.insert("bootstrap_port".to_string(), port_value);
+        }
+        if let Ok(room_value) = serde_json::to_value(bootstrap_room) {
+            self.other.insert("bootstrap_room".to_string(), room_value);
+        }
+    }
+}

sgl-router/src/routers/request_adapter.rs

@@ -648,6 +648,7 @@ mod tests {
             user: None,
             seed: None,
             suffix: None,
+            other: serde_json::Map::new(),
         };
 
         let pd_req = req.to_pd_request();
@@ -687,6 +688,7 @@ mod tests {
             user: None,
             seed: None,
             suffix: None,
+            other: serde_json::Map::new(),
         };
 
         let pd_req = req.to_pd_request();
@@ -725,6 +727,7 @@ mod tests {
             user: Some("user123".to_string()),
             seed: Some(42),
             suffix: Some("...".to_string()),
+            other: serde_json::Map::new(),
         };
 
         let pd_req = req.to_pd_request();
@@ -768,6 +771,7 @@ mod tests {
             user: None,
             seed: None,
             suffix: None,
+            other: serde_json::Map::new(),
         };
 
         let pd_req = req.to_pd_request();
@@ -799,6 +803,7 @@ mod tests {
             user: None,
             seed: None,
             suffix: None,
+            other: serde_json::Map::new(),
         };
 
         let pd_req = req.to_pd_request();

sgl-router/tests/benchmark_integration.rs

@@ -86,6 +86,7 @@ fn test_benchmark_request_creation() {
         logit_bias: None,
         user: None,
         seed: None,
+        other: serde_json::Map::new(),
     };
 
     // Test serialization works
@@ -181,6 +182,7 @@ fn test_benchmark_request_adaptation() {
         logit_bias: None,
         user: None,
         seed: None,
+        other: serde_json::Map::new(),
     };
 
     // Test PD adaptation (should not panic)