From 2c7fbb85f575bfdc79eb6984f284ebb0d4297bc3 Mon Sep 17 00:00:00 2001 From: WU Hang Date: Thu, 20 Aug 2026 18:10:19 +0800 Subject: [PATCH 1/3] fix: strip @rank suffix in DP intra-node GET proxy URLs proxy_get_request forwarded worker URLs including the @rank suffix (e.g. http://127.0.0.1:18100@0/v1/models) to reqwest, which parses @ as RFC 3986 userinfo and sends the request to the rank digit as host instead of the worker. GET /v1/models, /health_generate, /get_server_info, and /get_model_info returned 500 with DP > 1 (error sending request for url (http://0.0.0.0/v1/models)). Parse the worker URL with dp_utils::parse_worker_url and add the X-data-parallel-rank header via dp_utils::add_dp_rank_header, matching the chat POST path. DP=1 is unchanged (no @rank, no extra header). Fixes #220 Co-authored-by: herotai214 Signed-off-by: WU Hang --- src/routers/http/router.rs | 6 +++-- tests/test_dp_routing.rs | 47 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/src/routers/http/router.rs b/src/routers/http/router.rs index f68e4b08..c94aced4 100644 --- a/src/routers/http/router.rs +++ b/src/routers/http/router.rs @@ -428,9 +428,11 @@ impl Router { match self.select_first_worker() { Ok(worker_url) => { - let url = format!("{}/{}", worker_url, endpoint); + let (base_url, dp_rank) = dp_utils::parse_worker_url(&worker_url); + let url = format!("{}/{}", base_url, endpoint); let route_name = format!("/{}", endpoint); - let mut request_builder = self.client.get(&url); + let mut request_builder = + dp_utils::add_dp_rank_header(self.client.get(&url), dp_rank); for (name, value) in headers { let name_lc = name.to_lowercase(); if name_lc != "content-type" diff --git a/tests/test_dp_routing.rs b/tests/test_dp_routing.rs index 81bd9145..36b2b15b 100644 --- a/tests/test_dp_routing.rs +++ b/tests/test_dp_routing.rs @@ -449,6 +449,53 @@ mod dp_e2e_tests { worker.stop().await; } + // GET /v1/models used to send hostname:port@rank to reqwest as userinfo, + // routing the request to the rank digit as host instead of the worker. + #[tokio::test] + async fn test_regular_router_dp2_get_v1_models() { + let mut worker = MockWorker::new(MockWorkerConfig::default()); + let worker_url = worker.start().await.unwrap(); + let port: u16 = worker_url.split(':').next_back().unwrap().parse().unwrap(); + clear_captured_requests(port); + + let config = make_regular_config(vec![worker_url.clone()], 2); + let app_context = common::create_test_context(config.clone()); + let router = RouterFactory::create_router(&app_context).await.unwrap(); + let router = Arc::from(router); + + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + + let app = common::test_app::create_test_app(Arc::clone(&router), Client::new(), &config); + + let req = Request::builder() + .uri("/v1/models") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!( + resp.status().as_u16(), + 200, + "GET /v1/models must succeed when DP > 1 (got {}). @rank in the worker URL was misinterpreted as userinfo.", + resp.status() + ); + + // The 200 alone could be misleading if a fallback were added later; + // assert the body is the mock worker's model list. + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + body["data"][0]["id"].as_str(), + Some("mock-model"), + "GET /v1/models should return the worker's model list. Body: {:?}", + body + ); + + worker.stop().await; + } + // ----------------------------------------------------------------- // Regular Router + DP > 1: worker registry verification // ----------------------------------------------------------------- From 6c6dcf1f67648999973d217697ca1d7c7165de9d Mon Sep 17 00:00:00 2001 From: WU Hang Date: Wed, 9 Sep 2026 21:38:36 +0800 Subject: [PATCH 2/3] upate test comment Signed-off-by: WU Hang --- src/routers/http/router.rs | 1 + tests/test_dp_routing.rs | 2 -- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/routers/http/router.rs b/src/routers/http/router.rs index c94aced4..c0d394ad 100644 --- a/src/routers/http/router.rs +++ b/src/routers/http/router.rs @@ -433,6 +433,7 @@ impl Router { let route_name = format!("/{}", endpoint); let mut request_builder = dp_utils::add_dp_rank_header(self.client.get(&url), dp_rank); + for (name, value) in headers { let name_lc = name.to_lowercase(); if name_lc != "content-type" diff --git a/tests/test_dp_routing.rs b/tests/test_dp_routing.rs index 36b2b15b..b46075bd 100644 --- a/tests/test_dp_routing.rs +++ b/tests/test_dp_routing.rs @@ -449,8 +449,6 @@ mod dp_e2e_tests { worker.stop().await; } - // GET /v1/models used to send hostname:port@rank to reqwest as userinfo, - // routing the request to the rank digit as host instead of the worker. #[tokio::test] async fn test_regular_router_dp2_get_v1_models() { let mut worker = MockWorker::new(MockWorkerConfig::default()); From 9480f7e66d46f930e80de762baab9dbff36f5ee8 Mon Sep 17 00:00:00 2001 From: WU Hang Date: Thu, 10 Sep 2026 18:09:51 +0800 Subject: [PATCH 3/3] fix: drop client X-data-parallel-rank when router selects DP rank Address PR review feedback: - proxy_get_request: skip forwarding a client-supplied x-data-parallel-rank when the router injects its selected rank, so workers see exactly one value. DP=1 passthrough is unchanged. - Mock worker /v1/models now captures request headers; CapturedRequest stores all header values (HashMap>) to preserve duplicates. - test_regular_router_dp2_get_v1_models now asserts the forwarded rank matches the selected worker's rank. - New test_regular_router_dp2_get_v1_models_overrides_client_dp_rank sends a conflicting client rank (99) and asserts the worker receives exactly one router-selected rank value. Signed-off-by: WU Hang --- src/routers/http/router.rs | 4 ++ tests/common/mock_worker.rs | 33 ++++++++---- tests/otel_integration_test.rs | 2 + tests/test_dp_routing.rs | 91 ++++++++++++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 11 deletions(-) diff --git a/src/routers/http/router.rs b/src/routers/http/router.rs index c0d394ad..e8e1259e 100644 --- a/src/routers/http/router.rs +++ b/src/routers/http/router.rs @@ -436,8 +436,12 @@ impl Router { for (name, value) in headers { let name_lc = name.to_lowercase(); + // When the router selects a DP rank, it owns the + // X-data-parallel-rank header: skip any client-supplied + // value so the worker sees exactly one rank (ours). if name_lc != "content-type" && name_lc != "content-length" + && !(dp_rank.is_some() && name_lc == "x-data-parallel-rank") && !header_utils::TRACE_HEADER_NAMES.contains(&name_lc.as_str()) { request_builder = request_builder.header(name, value); diff --git a/tests/common/mock_worker.rs b/tests/common/mock_worker.rs index 10abcc40..24521e29 100644 --- a/tests/common/mock_worker.rs +++ b/tests/common/mock_worker.rs @@ -684,9 +684,15 @@ async fn flush_cache_handler(State(config): State>> .into_response() } -async fn v1_models_handler(State(config): State>>) -> Response { +async fn v1_models_handler( + State(config): State>>, + headers: axum::http::HeaderMap, +) -> Response { let config = config.read().await; + // Capture request for test inspection (e.g. X-data-parallel-rank) + capture_request(config.port, "/v1/models", &headers); + if should_fail(&config).await { return ( StatusCode::INTERNAL_SERVER_ERROR, @@ -872,10 +878,14 @@ impl Default for MockWorkerConfig { // --- Request header capture for verifying router behavior (e.g., X-data-parallel-rank) --- /// A captured request with headers and path +/// +/// `headers` maps each header name to *all* values received for that name, +/// so duplicate headers (e.g. a client-supplied and a router-injected +/// X-data-parallel-rank) are preserved for inspection. #[derive(Debug, Clone)] pub struct CapturedRequest { pub path: String, - pub headers: HashMap, + pub headers: HashMap>, } static REQ_CAPTURE_STORE: OnceLock>>> = OnceLock::new(); @@ -886,17 +896,18 @@ fn get_capture_store() -> &'static Mutex>> { /// Record a request for a given worker port pub fn capture_request(port: u16, path: &str, headers: &axum::http::HeaderMap) { + let mut captured_headers: HashMap> = HashMap::new(); + for (name, value) in headers.iter() { + if let Ok(v) = value.to_str() { + captured_headers + .entry(name.as_str().to_string()) + .or_default() + .push(v.to_string()); + } + } let captured = CapturedRequest { path: path.to_string(), - headers: headers - .iter() - .filter_map(|(name, value)| { - value - .to_str() - .ok() - .map(|v| (name.as_str().to_string(), v.to_string())) - }) - .collect(), + headers: captured_headers, }; let mut store = get_capture_store().lock().unwrap(); store.entry(port).or_default().push(captured); diff --git a/tests/otel_integration_test.rs b/tests/otel_integration_test.rs index 2d84f179..cb052d3b 100644 --- a/tests/otel_integration_test.rs +++ b/tests/otel_integration_test.rs @@ -207,6 +207,7 @@ async fn test_otel_integration() { let backend_tp = backend_req .headers .get("traceparent") + .and_then(|values| values.first()) .expect("Backend should receive traceparent header"); let parts: Vec<&str> = backend_tp.split('-').collect(); assert_eq!(parts.len(), 4, "traceparent should have 4 parts"); @@ -223,6 +224,7 @@ async fn test_otel_integration() { let backend_baggage = backend_req .headers .get("baggage") + .and_then(|values| values.first()) .expect("Backend should receive baggage header"); assert!( backend_baggage.contains("userId=alice"), diff --git a/tests/test_dp_routing.rs b/tests/test_dp_routing.rs index b46075bd..5d8db08b 100644 --- a/tests/test_dp_routing.rs +++ b/tests/test_dp_routing.rs @@ -491,6 +491,97 @@ mod dp_e2e_tests { body ); + // Verify the forwarded X-data-parallel-rank matches the selected + // worker's rank (0 or 1 here; both DP ranks share this mock worker, + // and worker selection order is not deterministic). + let captured = get_captured_requests(port); + let models_req = captured + .iter() + .find(|r| r.path == "/v1/models") + .expect("Mock worker should have received the GET /v1/models request"); + let rank_values = models_req + .headers + .get("x-data-parallel-rank") + .expect("GET proxy must forward X-data-parallel-rank when DP > 1"); + assert_eq!( + rank_values.len(), + 1, + "exactly one X-data-parallel-rank value expected, got {:?}", + rank_values + ); + let rank: usize = rank_values[0] + .parse() + .expect("X-data-parallel-rank must be numeric"); + assert!( + rank < 2, + "forwarded rank must be one of the router's DP ranks (0 or 1), got {}", + rank + ); + + worker.stop().await; + } + + #[tokio::test] + async fn test_regular_router_dp2_get_v1_models_overrides_client_dp_rank() { + // A client-supplied X-data-parallel-rank must not leak through or + // duplicate the router-selected rank: the worker should receive + // exactly one value, matching the router's selection. + let mut worker = MockWorker::new(MockWorkerConfig::default()); + let worker_url = worker.start().await.unwrap(); + let port: u16 = worker_url.split(':').next_back().unwrap().parse().unwrap(); + clear_captured_requests(port); + + let config = make_regular_config(vec![worker_url.clone()], 2); + let app_context = common::create_test_context(config.clone()); + let router = RouterFactory::create_router(&app_context).await.unwrap(); + let router = Arc::from(router); + + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + + let app = common::test_app::create_test_app(Arc::clone(&router), Client::new(), &config); + + let req = Request::builder() + .uri("/v1/models") + .header("x-data-parallel-rank", "99") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!( + resp.status().as_u16(), + 200, + "GET /v1/models must succeed even when the client sends a conflicting rank (got {})", + resp.status() + ); + + let captured = get_captured_requests(port); + let models_req = captured + .iter() + .find(|r| r.path == "/v1/models") + .expect("Mock worker should have received the GET /v1/models request"); + let rank_values = models_req + .headers + .get("x-data-parallel-rank") + .expect("router-selected X-data-parallel-rank must be forwarded"); + assert_eq!( + rank_values.len(), + 1, + "client-supplied rank must be filtered out; worker received {:?}", + rank_values + ); + assert_ne!( + rank_values[0], "99", + "worker must receive the router's selection, not the client's value" + ); + let rank: usize = rank_values[0] + .parse() + .expect("X-data-parallel-rank must be numeric"); + assert!( + rank < 2, + "forwarded rank must be one of the router's DP ranks (0 or 1), got {}", + rank + ); + worker.stop().await; }