pd_router.rs 47.2 KB
Newer Older
1
use super::pd_types::api_path;
2
use crate::config::types::RetryConfig;
3
use crate::core::{
4
    is_retryable_status, RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType,
5
};
6
use crate::metrics::RouterMetrics;
7
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
8
use crate::protocols::spec::{
9
    ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, RerankRequest,
10
    ResponsesGetParams, ResponsesRequest, StringOrArray, UserMessageContent,
11
};
12
use crate::routers::header_utils;
13
use crate::routers::RouterTrait;
14
use async_trait::async_trait;
15
16
17
18
19
20
21
use axum::{
    body::Body,
    extract::Request,
    http::{header::CONTENT_TYPE, HeaderMap, HeaderValue, StatusCode},
    response::{IntoResponse, Response},
};
use futures_util::StreamExt;
22
use reqwest::Client;
23
use serde::Serialize;
24
use serde_json::{json, Value};
25
use std::sync::Arc;
26
use std::time::Instant;
27
use tokio_stream::wrappers::UnboundedReceiverStream;
28
use tracing::{debug, error, warn};
29
30
31

#[derive(Debug)]
pub struct PDRouter {
32
33
    pub worker_registry: Arc<WorkerRegistry>,
    pub policy_registry: Arc<PolicyRegistry>,
34
    pub client: Client,
35
    pub retry_config: RetryConfig,
36
    pub api_key: Option<String>,
37
    pub enable_igw: bool,
38
39
}

40
#[derive(Clone)]
41
struct PDRequestContext<'a> {
42
43
44
45
46
    route: &'static str,
    batch_size: Option<usize>,
    is_stream: bool,
    return_logprob: bool,
    request_text: Option<String>,
47
    model_id: Option<&'a str>,
48
49
}

50
impl PDRouter {
51
    async fn proxy_to_first_prefill_worker(
52
53
54
55
        &self,
        endpoint: &str,
        headers: Option<Vec<(String, String)>>,
    ) -> Response {
56
57
        let workers = self.worker_registry.get_prefill_workers();
        let first_worker_url = workers.first().map(|w| w.url().to_string());
58
59

        if let Some(worker_url) = first_worker_url {
60
61
62
63
64
65
66
67
68
            self.proxy_to_worker(worker_url, endpoint, headers).await
        } else {
            (
                StatusCode::SERVICE_UNAVAILABLE,
                "No prefill servers available".to_string(),
            )
                .into_response()
        }
    }
69

70
71
72
73
74
75
76
77
78
79
80
81
    async fn proxy_to_worker(
        &self,
        worker_url: String,
        endpoint: &str,
        headers: Option<Vec<(String, String)>>,
    ) -> Response {
        let url = format!("{}/{}", worker_url, endpoint);
        let mut request_builder = self.client.get(&url);

        if let Some(headers) = headers {
            for (name, value) in headers {
                request_builder = request_builder.header(name, value);
82
            }
83
        }
84

85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        match request_builder.send().await {
            Ok(res) if res.status().is_success() => {
                let response_headers = header_utils::preserve_response_headers(res.headers());

                match res.bytes().await {
                    Ok(body) => {
                        let mut response = Response::new(axum::body::Body::from(body));
                        *response.status_mut() = StatusCode::OK;
                        *response.headers_mut() = response_headers;
                        response
                    }
                    Err(e) => {
                        error!("Failed to read response body: {}", e);
                        (
                            StatusCode::INTERNAL_SERVER_ERROR,
                            format!("Failed to read response body: {}", e),
                        )
                            .into_response()
103
104
105
                    }
                }
            }
106
107
108
109
110
111
112
113
114
115
116
117
118
            Ok(res) => {
                let status = StatusCode::from_u16(res.status().as_u16())
                    .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                (status, format!("{} server returned status: ", res.status())).into_response()
            }
            Err(e) => {
                error!("Failed to proxy request server: {}", e);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Failed to proxy request: {}", e),
                )
                    .into_response()
            }
119
120
121
        }
    }

122
    pub async fn new(ctx: &Arc<crate::server::AppContext>) -> Result<Self, String> {
123
        Ok(PDRouter {
124
125
            worker_registry: Arc::clone(&ctx.worker_registry),
            policy_registry: Arc::clone(&ctx.policy_registry),
126
127
            client: ctx.client.clone(),
            retry_config: ctx.router_config.effective_retry_config(),
128
            api_key: ctx.router_config.api_key.clone(),
129
            enable_igw: ctx.router_config.enable_igw,
130
131
132
        })
    }

133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    fn handle_server_selection_error(error: String) -> Response {
        error!("Failed to select PD pair error={}", error);
        RouterMetrics::record_pd_error("server_selection");
        (
            StatusCode::SERVICE_UNAVAILABLE,
            format!("No available servers: {}", error),
        )
            .into_response()
    }

    fn handle_serialization_error(error: impl std::fmt::Display) -> Response {
        error!("Failed to serialize request error={}", error);
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            "Failed to serialize request",
        )
            .into_response()
    }

152
    fn get_generate_batch_size(req: &GenerateRequest) -> Option<usize> {
153
        if let Some(StringOrArray::Array(arr)) = &req.prompt {
154
155
            if !arr.is_empty() {
                return Some(arr.len());
156
            }
157
158
159
        }
        if let Some(text) = &req.text {
            if text.contains("[") && text.contains("]") {
160
                return None;
161
            }
162
163
164
        }
        None
    }
165

166
167
168
169
    fn get_chat_batch_size(req: &ChatCompletionRequest) -> Option<usize> {
        if let Some(n) = req.n {
            if n > 1 {
                return Some(n as usize);
170
            }
171
172
173
        }
        None
    }
174

175
    fn get_completion_batch_size(req: &CompletionRequest) -> Option<usize> {
176
        if let StringOrArray::Array(arr) = &req.prompt {
177
178
            if !arr.is_empty() {
                return Some(arr.len());
179
180
            }
        }
181
182
        None
    }
183

184
185
    fn inject_bootstrap_into_value(
        mut original: Value,
186
187
        prefill_worker: &dyn Worker,
        batch_size: Option<usize>,
188
    ) -> Result<Value, String> {
189
190
191
192
193
194
        let bootstrap_port = match prefill_worker.worker_type() {
            crate::core::WorkerType::Prefill { bootstrap_port } => bootstrap_port,
            _ => None,
        };
        let hostname = super::pd_types::get_hostname(prefill_worker.url());

195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
        let obj = original
            .as_object_mut()
            .ok_or_else(|| "Request must be a JSON object".to_string())?;

        if let Some(n) = batch_size {
            let mut hosts = Vec::with_capacity(n);
            let mut ports = Vec::with_capacity(n);
            let mut rooms = Vec::with_capacity(n);
            for _ in 0..n {
                hosts.push(hostname.clone());
                ports.push(bootstrap_port);
                rooms.push(super::pd_types::generate_room_id());
            }
            obj.insert(
                "bootstrap_host".to_string(),
                Value::Array(hosts.into_iter().map(serde_json::Value::from).collect()),
            );
            obj.insert(
                "bootstrap_port".to_string(),
                Value::Array(
                    ports
                        .into_iter()
                        .map(|p| match p {
                            Some(v) => serde_json::Value::from(v),
                            None => Value::Null,
                        })
                        .collect(),
                ),
            );
            obj.insert(
                "bootstrap_room".to_string(),
                Value::Array(rooms.into_iter().map(serde_json::Value::from).collect()),
            );
228
        } else {
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
            obj.insert(
                "bootstrap_host".to_string(),
                serde_json::Value::from(hostname),
            );
            obj.insert(
                "bootstrap_port".to_string(),
                match bootstrap_port {
                    Some(v) => serde_json::Value::from(v),
                    None => Value::Null,
                },
            );
            obj.insert(
                "bootstrap_room".to_string(),
                serde_json::Value::from(super::pd_types::generate_room_id()),
            );
244
        }
245
        Ok(original)
246
247
    }

248
249
250
251
    async fn execute_dual_dispatch<T: Serialize + Clone>(
        &self,
        headers: Option<&HeaderMap>,
        original_request: &T,
252
        context: PDRequestContext<'_>,
253
254
255
256
257
258
259
260
261
262
263
264
    ) -> Response {
        let start_time = Instant::now();

        let route = context.route;
        RetryExecutor::execute_response_with_retry(
            &self.retry_config,
            {
                let original_request = original_request.clone();
                move |attempt: u32| {
                    let original_request = original_request.clone();
                    let context = context.clone();
                    async move {
265
266
267
268
269
270
271
272
273
274
                        let (prefill, decode) = match self
                            .select_pd_pair(context.request_text.as_deref(), context.model_id)
                            .await
                        {
                            Ok(pair) => pair,
                            Err(e) => {
                                RouterMetrics::record_pd_error("server_selection");
                                return Self::handle_server_selection_error(e);
                            }
                        };
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300

                        debug!(
                            "PD retry attempt {} using prefill={} decode={}",
                            attempt,
                            prefill.url(),
                            decode.url()
                        );

                        let mut json_request = match serde_json::to_value(&original_request) {
                            Ok(v) => v,
                            Err(e) => return Self::handle_serialization_error(e),
                        };

                        json_request = match Self::inject_bootstrap_into_value(
                            json_request,
                            prefill.as_ref(),
                            context.batch_size,
                        ) {
                            Ok(v) => v,
                            Err(e) => return Self::handle_serialization_error(e),
                        };

                        let response = self
                            .execute_dual_dispatch_internal(
                                headers,
                                json_request,
301
                                context,
302
303
304
305
306
307
                                prefill.as_ref(),
                                decode.as_ref(),
                                start_time,
                            )
                            .await;

308
309
310
311
                        let _status = response.status();
                        let not_error = _status.is_success() || _status.is_client_error();
                        prefill.record_outcome(not_error);
                        decode.record_outcome(not_error);
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326

                        response
                    }
                }
            },
            |res, _attempt| is_retryable_status(res.status()),
            |delay, attempt| {
                RouterMetrics::record_retry(route);
                RouterMetrics::record_retry_backoff_duration(delay, attempt);
            },
            || RouterMetrics::record_retries_exhausted(route),
        )
        .await
    }

327
328
329
    async fn handle_decode_error_response(
        &self,
        res: reqwest::Response,
330
        context: &PDRequestContext<'_>,
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        prefill: &dyn Worker,
        decode: &dyn Worker,
    ) -> Response {
        let status = res.status();

        if context.is_stream {
            // Handle streaming error response
            let response_headers = header_utils::preserve_response_headers(res.headers());
            let error_payload = match res.bytes().await {
                Ok(error_body) => {
                    if let Ok(error_json) = serde_json::from_slice::<Value>(&error_body) {
                        json!({ "message": error_json, "status": status.as_u16() })
                    } else {
                        json!({ "message": String::from_utf8_lossy(&error_body).to_string(), "status": status.as_u16() })
                    }
                }
                Err(e) => {
                    json!({ "message": format!("Decode server error: {}", e), "status": status.as_u16() })
                }
            };

            let sse_data = format!(
                "data: {{'error': {}}}",
                serde_json::to_string(&error_payload).unwrap_or_default()
            );
            let error_stream = tokio_stream::once(Ok(axum::body::Bytes::from(sse_data)));

            let decode_url = decode.url().to_string();
            self.create_streaming_response(
                error_stream,
                status,
                None,
                context.return_logprob,
                Some(decode_url),
                Some(response_headers),
                prefill,
                decode,
            )
        } else {
            // Handle non-streaming error response
            match res.bytes().await {
                Ok(error_body) => (status, error_body).into_response(),
                Err(e) => (status, format!("Decode server error: {}", e)).into_response(),
            }
        }
    }

378
379
    // Internal method that performs the actual dual dispatch (without retry logic)
    async fn execute_dual_dispatch_internal(
380
381
382
        &self,
        headers: Option<&HeaderMap>,
        json_request: Value,
383
        context: PDRequestContext<'_>,
384
385
386
387
        prefill: &dyn Worker,
        decode: &dyn Worker,
        start_time: Instant,
    ) -> Response {
388
389
390
391
392
393
394
        // For non-streaming: use guard for automatic load management
        // For streaming: load will be managed in create_streaming_response
        let _guard = if !context.is_stream {
            Some(WorkerLoadGuard::new_multi(vec![prefill, decode]))
        } else {
            None
        };
395

396
397
398
399
400
401
402
403
404
        // Build both requests
        let prefill_request = self.build_post_with_headers(
            &self.client,
            prefill.url(),
            context.route,
            &json_request,
            headers,
            false,
        );
405
406
407
        let decode_request = self.build_post_with_headers(
            &self.client,
            decode.url(),
408
            context.route,
409
410
411
412
            &json_request,
            headers,
            false,
        );
413

414
        // Send both requests concurrently and wait for both
415
416
417
418
419
420
        debug!(
            "Sending concurrent requests to prefill={} decode={}",
            prefill.url(),
            decode.url()
        );

421
422
423
        let (prefill_result, decode_result) =
            tokio::join!(prefill_request.send(), decode_request.send());
        debug!("Received responses from both servers");
424

425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
        let duration = start_time.elapsed();
        RouterMetrics::record_pd_request_duration(context.route, duration);
        RouterMetrics::record_pd_request(context.route);
        RouterMetrics::record_pd_prefill_request(prefill.url());
        RouterMetrics::record_pd_decode_request(decode.url());

        // Process decode response
        match decode_result {
            Ok(res) => {
                let status = StatusCode::from_u16(res.status().as_u16())
                    .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                debug!("Decode response status: {}", status);

                if !status.is_success() {
                    RouterMetrics::record_pd_decode_error(decode.url());
                    error!(
                        "Decode server returned error status decode_url={} status={}",
                        decode.url(),
                        status
                    );

                    return self
                        .handle_decode_error_response(res, &context, prefill, decode)
                        .await;
                }
450

451
452
453
                // Process prefill response
                let prefill_body = if context.return_logprob {
                    match self
454
455
456
457
458
                        .process_prefill_response(
                            prefill_result,
                            prefill.url(),
                            context.return_logprob,
                        )
459
460
461
462
                        .await
                    {
                        Ok((_, body)) => body,
                        Err(error_response) => return error_response,
463
464
465
466
467
468
469
470
471
472
473
                    }
                } else {
                    // Even if we don't need logprobs, we should check prefill status
                    match self
                        .process_prefill_response(prefill_result, prefill.url(), false)
                        .await
                    {
                        Ok((_, body)) => body,
                        Err(error_response) => return error_response,
                    }
                };
474

475
476
477
478
                if context.is_stream {
                    // Streaming response
                    let prefill_logprobs = if context.return_logprob {
                        prefill_body
479
480
481
482
                            .as_ref()
                            .and_then(|body| serde_json::from_slice::<Value>(body).ok())
                            .and_then(|json| {
                                json.pointer("/meta_info/input_token_logprobs").cloned()
483
                            })
484
                    } else {
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
                        None
                    };

                    let response_headers = header_utils::preserve_response_headers(res.headers());

                    self.create_streaming_response(
                        res.bytes_stream(),
                        status,
                        prefill_logprobs,
                        context.return_logprob,
                        None,
                        Some(response_headers),
                        prefill,
                        decode,
                    )
                } else {
                    // Non-streaming response
                    if context.return_logprob {
503
504
505
                        self.process_non_streaming_response(
                            res,
                            status,
506
                            context.return_logprob,
507
508
509
510
                            prefill_body,
                        )
                        .await
                    } else {
511
                        // Direct passthrough when no logprobs needed
512
513
514
                        let response_headers =
                            header_utils::preserve_response_headers(res.headers());

515
                        match res.bytes().await {
516
517
518
519
520
521
522
                            Ok(decode_body) => {
                                let mut response =
                                    Response::new(axum::body::Body::from(decode_body));
                                *response.status_mut() = status;
                                *response.headers_mut() = response_headers;
                                response
                            }
523
524
525
526
527
528
529
530
                            Err(e) => {
                                error!("Failed to read decode response: {}", e);
                                (StatusCode::INTERNAL_SERVER_ERROR, "Failed to read response")
                                    .into_response()
                            }
                        }
                    }
                }
531
532
533
534
535
536
537
538
539
540
541
542
543
            }
            Err(e) => {
                error!(
                    decode_url = %decode.url(),
                    error = %e,
                    "Decode request failed"
                );
                RouterMetrics::record_pd_decode_error(decode.url());
                (
                    StatusCode::BAD_GATEWAY,
                    format!("Decode server error: {}", e),
                )
                    .into_response()
544
545
546
547
            }
        }
    }

548
    fn policies_need_request_text(&self) -> bool {
549
550
551
        let prefill_policy = self.policy_registry.get_prefill_policy();
        let decode_policy = self.policy_registry.get_decode_policy();
        prefill_policy.needs_request_text() || decode_policy.needs_request_text()
552
553
    }

554
555
    async fn select_pd_pair(
        &self,
556
        request_text: Option<&str>,
557
558
        model_id: Option<&str>,
    ) -> Result<(Arc<dyn Worker>, Arc<dyn Worker>), String> {
559
560
561
562
563
564
565
566
        let effective_model_id = if !self.enable_igw { None } else { model_id };

        debug!(
            "Selecting PD pair: enable_igw={}, model_id={:?}, effective_model_id={:?}",
            self.enable_igw, model_id, effective_model_id
        );

        let prefill_workers = if let Some(model) = effective_model_id {
567
568
569
570
571
572
573
574
575
            self.worker_registry
                .get_by_model_fast(model)
                .into_iter()
                .filter(|w| matches!(w.worker_type(), WorkerType::Prefill { .. }))
                .collect()
        } else {
            self.worker_registry.get_prefill_workers()
        };

576
        let decode_workers = if let Some(model) = effective_model_id {
577
578
579
580
581
582
583
584
            self.worker_registry
                .get_by_model_fast(model)
                .into_iter()
                .filter(|w| matches!(w.worker_type(), WorkerType::Decode))
                .collect()
        } else {
            self.worker_registry.get_decode_workers()
        };
585

586
587
588
589
        let prefill_policy = self.policy_registry.get_prefill_policy();
        let decode_policy = self.policy_registry.get_decode_policy();

        let prefill = Self::pick_worker_by_policy_arc(
590
            &prefill_workers,
591
            &*prefill_policy,
592
593
594
595
            request_text,
            "prefill",
        )?;

596
        let decode = Self::pick_worker_by_policy_arc(
597
            &decode_workers,
598
            &*decode_policy,
599
600
601
602
603
604
605
            request_text,
            "decode",
        )?;

        Ok((prefill, decode))
    }

606
607
    fn pick_worker_by_policy_arc(
        workers: &[Arc<dyn Worker>],
608
609
610
        policy: &dyn LoadBalancingPolicy,
        request_text: Option<&str>,
        worker_type: &str,
611
    ) -> Result<Arc<dyn Worker>, String> {
612
613
614
615
616
        if workers.is_empty() {
            return Err(format!(
                "No {} workers available. Please check if {} servers are configured and healthy.",
                worker_type, worker_type
            ));
617
618
        }

619
        let available_workers: Vec<Arc<dyn Worker>> = workers
620
621
            .iter()
            .filter(|w| w.is_available())
622
            .cloned()
623
            .collect();
624

625
626
627
628
629
630
        if available_workers.is_empty() {
            return Err(format!(
                "No available {} workers (all circuits open or unhealthy)",
                worker_type
            ));
        }
631

632
633
634
635
636
637
638
639
640
641
642
        let selected_idx = policy
            .select_worker(&available_workers, request_text)
            .ok_or_else(|| {
                format!(
                    "Policy {} failed to select a {} worker",
                    policy.name(),
                    worker_type
                )
            })?;

        Ok(available_workers[selected_idx].clone())
643
644
    }

645
    #[allow(clippy::too_many_arguments)]
646
    fn create_streaming_response(
647
        &self,
648
649
650
651
652
        stream: impl futures_util::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send + 'static,
        status: StatusCode,
        prefill_logprobs: Option<Value>,
        return_logprob: bool,
        decode_url: Option<String>,
653
        headers: Option<HeaderMap>,
654
655
        prefill: &dyn Worker,
        decode: &dyn Worker,
656
    ) -> Response {
657
658
659
660
661
662
        prefill.increment_load();
        decode.increment_load();

        let prefill_url = prefill.url().to_string();
        let decode_url_str = decode.url().to_string();

663
664
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

665
        let registry = self.worker_registry.clone();
666

667
        tokio::spawn(async move {
668
669
            let mut stream_completed = false;

670
671
672
673
            futures_util::pin_mut!(stream);
            while let Some(chunk_result) = stream.next().await {
                match chunk_result {
                    Ok(chunk) => {
674
675
676
677
678
                        let is_done = chunk
                            .as_ref()
                            .windows(12)
                            .any(|window| window == b"data: [DONE]");

679
680
681
682
683
684
685
686
687
688
                        let result = if return_logprob && prefill_logprobs.is_some() {
                            Self::merge_streaming_logprobs(prefill_logprobs.clone(), &chunk)
                                .unwrap_or(chunk)
                        } else {
                            chunk
                        };

                        if tx.send(Ok(result)).is_err() {
                            break;
                        }
689
690
691
692
693

                        if is_done {
                            stream_completed = true;
                            break;
                        }
694
695
696
697
698
699
700
701
702
703
704
                    }
                    Err(e) => {
                        if let Some(ref url) = decode_url {
                            error!("Stream error from decode server {}: {}", url, e);
                            RouterMetrics::record_pd_stream_error(url);
                        }
                        let _ = tx.send(Err(format!("Stream error: {}", e)));
                        break;
                    }
                }
            }
705

706
707
708
709
710
711
            if let Some(worker) = registry.get_by_url(&prefill_url) {
                worker.decrement_load();
                debug!(
                    "Decremented load for prefill worker: {} (stream_completed: {})",
                    prefill_url, stream_completed
                );
712
713
            }

714
715
716
717
718
719
            if let Some(worker) = registry.get_by_url(&decode_url_str) {
                worker.decrement_load();
                debug!(
                    "Decremented load for decode worker: {} (stream_completed: {})",
                    decode_url_str, stream_completed
                );
720
            }
721
722
723
724
725
726
727
        });

        let stream = UnboundedReceiverStream::new(rx);
        let body = Body::from_stream(stream);

        let mut response = Response::new(body);
        *response.status_mut() = status;
728

729
        let mut headers = headers.unwrap_or_default();
730
731
732
        headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
        *response.headers_mut() = headers;

733
734
735
736
737
738
739
740
741
742
743
        response
    }

    // Helper to process non-streaming decode response with logprob merging
    async fn process_non_streaming_response(
        &self,
        res: reqwest::Response,
        status: StatusCode,
        return_logprob: bool,
        prefill_body: Option<bytes::Bytes>,
    ) -> Response {
744
745
746
        let response = res.bytes().await;
        let decode_body = match response {
            Ok(decode_body) => decode_body,
747
748
            Err(e) => {
                error!("Failed to read decode response: {}", e);
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
                return (StatusCode::INTERNAL_SERVER_ERROR, "Failed to read response")
                    .into_response();
            }
        };

        if !return_logprob {
            return (status, decode_body).into_response();
        }

        let Some(prefill_body) = prefill_body else {
            return (status, decode_body).into_response();
        };

        // Merge logprobs from prefill and decode
        let (Ok(prefill_json), Ok(mut decode_json)) = (
            serde_json::from_slice::<Value>(&prefill_body),
            serde_json::from_slice::<Value>(&decode_body),
        ) else {
            warn!("Failed to parse responses for logprob merging");
            return (status, decode_body).into_response();
        };

        Self::merge_logprobs_in_json(&prefill_json, &mut decode_json);

        // Return merged response
        match serde_json::to_vec(&decode_json) {
            Ok(body) => (status, body).into_response(),
            Err(e) => {
                error!("Failed to serialize merged response: {}", e);
                (status, decode_body).into_response()
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
            }
        }
    }

    // Helper to process prefill response and extract body if needed for logprobs
    async fn process_prefill_response(
        &self,
        prefill_result: Result<reqwest::Response, reqwest::Error>,
        prefill_url: &str,
        return_logprob: bool,
    ) -> Result<(StatusCode, Option<bytes::Bytes>), Response> {
        // Check prefill result first - it's critical for disaggregated mode
        let prefill_response = match prefill_result {
            Ok(response) => response,
            Err(e) => {
                RouterMetrics::record_pd_prefill_error(prefill_url);
                error!(
                    "Prefill server failed (CRITICAL) prefill_url={} error={}. Decode will timeout without prefill KV cache.",
                    prefill_url,
                    e
                );

                // Return error immediately - don't wait for decode to timeout
                return Err((
                    StatusCode::BAD_GATEWAY,
                    format!(
                        "Prefill server error: {}. This will cause decode timeout.",
                        e
                    ),
                )
                    .into_response());
            }
        };

        let prefill_status = StatusCode::from_u16(prefill_response.status().as_u16())
            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

        // Check if prefill succeeded
        if !prefill_status.is_success() {
            RouterMetrics::record_pd_prefill_error(prefill_url);

            // Get error body from prefill
            let error_msg = prefill_response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown prefill error".to_string());

            error!(
                "Prefill server returned error status prefill_url={} status={} body={}",
                prefill_url, prefill_status, error_msg
            );

            return Err((
                prefill_status,
                format!("Prefill server error ({}): {}", prefill_status, error_msg),
            )
                .into_response());
        }

        // Read prefill body if needed for logprob merging
        let prefill_body = if return_logprob {
            match prefill_response.bytes().await {
                Ok(body) => Some(body),
                Err(e) => {
                    warn!("Failed to read prefill response body for logprobs: {}", e);
                    None
                }
            }
        } else {
            // For non-logprob requests, just consume the response without storing
            debug!("Consuming prefill response body (non-logprob request)");
            match prefill_response.bytes().await {
                Ok(_) => debug!("Prefill response consumed successfully"),
                Err(e) => warn!("Error consuming prefill response: {}", e),
            }
            None
        };

        Ok((prefill_status, prefill_body))
    }

860
    fn build_post_with_headers(
861
        &self,
862
        client: &Client,
863
864
        url: &str,
        route: &str,
865
        json_request: &Value,
866
        headers: Option<&HeaderMap>,
867
        connection_close: bool,
868
    ) -> reqwest::RequestBuilder {
869
870
871
872
        let mut request = client.post(api_path(url, route)).json(json_request);
        if connection_close {
            request = request.header("Connection", "close");
        }
873
874
        if let Some(headers) = headers {
            for (name, value) in headers.iter() {
875
876
877
878
879
880
881
882
883
                let name_lc = name.as_str().to_ascii_lowercase();
                // Whitelist important end-to-end headers, skip hop-by-hop
                let forward = matches!(
                    name_lc.as_str(),
                    "authorization" | "x-request-id" | "x-correlation-id"
                ) || name_lc.starts_with("x-request-id-");
                if forward {
                    if let Ok(val) = value.to_str() {
                        request = request.header(name, val);
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
                    }
                }
            }
        }
        request
    }

    // Helper to merge logprobs from prefill and decode responses
    fn merge_logprobs_in_json(prefill_json: &Value, decode_json: &mut Value) -> bool {
        if let (Some(prefill_meta), Some(decode_meta)) = (
            prefill_json.get("meta_info"),
            decode_json.get_mut("meta_info"),
        ) {
            if let (Some(prefill_logprobs), Some(decode_logprobs)) = (
                prefill_meta.get("input_token_logprobs"),
                decode_meta.get_mut("input_token_logprobs"),
            ) {
                if let (Some(prefill_arr), Some(decode_arr)) =
                    (prefill_logprobs.as_array(), decode_logprobs.as_array_mut())
                {
                    let mut merged = prefill_arr.clone();
                    merged.extend(decode_arr.clone());
                    decode_meta["input_token_logprobs"] = Value::Array(merged);
                    return true;
                }
            }
        }
        false
    }

914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
    // Simple helper to merge logprobs in streaming responses
    fn merge_streaming_logprobs(
        prefill_logprobs: Option<Value>,
        decode_chunk: &[u8],
    ) -> Result<bytes::Bytes, ()> {
        // Skip non-data chunks
        let chunk_str = std::str::from_utf8(decode_chunk).map_err(|_| ())?;
        if !chunk_str.starts_with("data: ") || chunk_str.contains("[DONE]") {
            return Err(());
        }

        // Parse JSON from chunk
        let json_str = chunk_str.trim_start_matches("data: ").trim();
        let mut decode_json: Value = serde_json::from_str(json_str).map_err(|_| ())?;

        // Merge prefill logprobs if available
        if let Some(ref p_logprobs) = prefill_logprobs {
            if let Some(meta) = decode_json.get_mut("meta_info") {
                if let Some(d_logprobs) = meta.get_mut("input_token_logprobs") {
                    if let (Some(p_arr), Some(d_arr)) =
                        (p_logprobs.as_array(), d_logprobs.as_array())
                    {
                        let mut merged = p_arr.clone();
                        merged.extend(d_arr.clone());
                        *d_logprobs = Value::Array(merged);
                    }
                }
            }
        }

        // Re-serialize
        let merged_str = format!(
            "data: {}\n\n",
            serde_json::to_string(&decode_json).unwrap_or_default()
        );
        Ok(bytes::Bytes::from(merged_str))
    }
}

953
954
955
956
957
958
959
#[async_trait]
impl RouterTrait for PDRouter {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    async fn health_generate(&self, _req: Request<Body>) -> Response {
960
        // Note: This endpoint actually causes the model to generate tokens, so we only test one pair
961

962
        // Select a random worker pair using the policy
963
        let (prefill, decode) = match self.select_pd_pair(None, None).await {
964
965
            Ok(pair) => pair,
            Err(e) => {
966
967
968
969
970
                return (
                    StatusCode::SERVICE_UNAVAILABLE,
                    format!("No healthy worker pair available: {}", e),
                )
                    .into_response();
971
972
            }
        };
973

974
        let prefill_url = format!("{}/health_generate", prefill.url());
975
976
977
        let (prefill_result, decode_result) = tokio::join!(
            self.client.get(&prefill_url).send(),
            self.client
978
                .get(format!("{}/health_generate", decode.url()))
979
980
                .send()
        );
981

982
983
        // Check results
        let mut errors = Vec::new();
984

985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
        match prefill_result {
            Ok(res) if res.status().is_success() => {
                debug!(
                    "Health generate passed for prefill server: {}",
                    prefill.url()
                );
            }
            Ok(res) => {
                errors.push(format!(
                    "Prefill {} returned status {}",
                    prefill.url(),
                    res.status()
                ));
            }
            Err(e) => {
                errors.push(format!("Prefill {} error: {}", prefill.url(), e));
            }
        }
1003

1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
        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));
1017
1018
1019
            }
        }

1020
        if errors.is_empty() {
1021
1022
1023
1024
1025
1026
1027
1028
1029
            (
                StatusCode::OK,
                format!(
                    "Health generate passed on selected pair: prefill={}, decode={}",
                    prefill.url(),
                    decode.url()
                ),
            )
                .into_response()
1030
        } else {
1031
1032
1033
1034
1035
            (
                StatusCode::SERVICE_UNAVAILABLE,
                format!("Health generate failed: {:?}", errors),
            )
                .into_response()
1036
1037
1038
        }
    }

1039
    async fn get_server_info(&self, _req: Request<Body>) -> Response {
1040
        // Get info from the first decode server to match sglang's server info format
1041
        // Note: We use decode workers for server info to match expected format
1042
        self.proxy_to_first_prefill_worker("get_server_info", None)
1043
            .await
1044
1045
    }

1046
    async fn get_models(&self, req: Request<Body>) -> Response {
1047
        // Extract headers first to avoid Send issues
1048
        let headers = header_utils::copy_request_headers(&req);
1049

1050
        // Proxy to first prefill worker
1051
        self.proxy_to_first_prefill_worker("v1/models", Some(headers))
1052
            .await
1053
1054
    }

1055
    async fn get_model_info(&self, req: Request<Body>) -> Response {
1056
        // Extract headers first to avoid Send issues
1057
        let headers = header_utils::copy_request_headers(&req);
1058

1059
        // Proxy to first prefill worker
1060
1061
        self.proxy_to_first_prefill_worker("get_model_info", Some(headers))
            .await
1062
1063
    }

1064
1065
1066
1067
    async fn route_generate(
        &self,
        headers: Option<&HeaderMap>,
        body: &GenerateRequest,
1068
        model_id: Option<&str>,
1069
1070
1071
    ) -> Response {
        let is_stream = body.stream;
        let return_logprob = body.return_logprob;
1072

1073
        let request_text = if self.policies_need_request_text() {
1074
1075
1076
1077
            body.text
                .as_deref()
                .or_else(|| {
                    body.prompt.as_ref().and_then(|p| match p {
1078
1079
                        StringOrArray::String(s) => Some(s.as_str()),
                        StringOrArray::Array(v) => v.first().map(|s| s.as_str()),
1080
                    })
1081
                })
1082
                .map(|s| s.to_string())
1083
1084
1085
        } else {
            None
        };
1086

1087
        let batch_size = Self::get_generate_batch_size(body);
1088

1089
1090
1091
        let context = PDRequestContext {
            route: "/generate",
            batch_size,
1092
1093
            is_stream,
            return_logprob,
1094
            request_text,
1095
            model_id,
1096
1097
1098
        };

        self.execute_dual_dispatch(headers, body, context).await
1099
1100
    }

1101
1102
1103
1104
    async fn route_chat(
        &self,
        headers: Option<&HeaderMap>,
        body: &ChatCompletionRequest,
1105
        model_id: Option<&str>,
1106
1107
1108
1109
    ) -> Response {
        let is_stream = body.stream;
        let return_logprob = body.logprobs;

1110
1111
        let request_text = if self.policies_need_request_text() {
            body.messages.first().and_then(|msg| match msg {
1112
1113
1114
                ChatMessage::User { content, .. } => match content {
                    UserMessageContent::Text(text) => Some(text.clone()),
                    UserMessageContent::Parts(_) => None,
1115
                },
1116
                ChatMessage::System { content, .. } => Some(content.clone()),
1117
1118
1119
1120
1121
                _ => None,
            })
        } else {
            None
        };
1122

1123
        // Calculate batch size
1124
        let batch_size = Self::get_chat_batch_size(body);
1125

1126
1127
1128
        let context = PDRequestContext {
            route: "/v1/chat/completions",
            batch_size,
1129
1130
            is_stream,
            return_logprob,
1131
            request_text,
1132
            model_id,
1133
1134
1135
        };

        self.execute_dual_dispatch(headers, body, context).await
1136
1137
    }

1138
1139
1140
1141
    async fn route_completion(
        &self,
        headers: Option<&HeaderMap>,
        body: &CompletionRequest,
1142
        model_id: Option<&str>,
1143
1144
1145
1146
    ) -> Response {
        let is_stream = body.stream;
        let return_logprob = body.logprobs.is_some();

1147
1148
        let request_text = if self.policies_need_request_text() {
            match &body.prompt {
1149
1150
                StringOrArray::String(s) => Some(s.clone()),
                StringOrArray::Array(v) => v.first().map(|s| s.to_string()),
1151
1152
1153
            }
        } else {
            None
1154
1155
        };

1156
        // Calculate batch size
1157
        let batch_size = Self::get_completion_batch_size(body);
1158

1159
1160
1161
        let context = PDRequestContext {
            route: "/v1/completions",
            batch_size,
1162
1163
            is_stream,
            return_logprob,
1164
            request_text,
1165
            model_id,
1166
1167
1168
        };

        self.execute_dual_dispatch(headers, body, context).await
1169
1170
    }

1171
1172
1173
1174
    async fn route_responses(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &ResponsesRequest,
1175
        _model_id: Option<&str>,
1176
1177
1178
1179
1180
1181
1182
1183
    ) -> Response {
        (
            StatusCode::NOT_IMPLEMENTED,
            "Responses endpoint not implemented for PD router",
        )
            .into_response()
    }

1184
1185
1186
1187
1188
1189
    async fn get_response(
        &self,
        _headers: Option<&HeaderMap>,
        _response_id: &str,
        _params: &ResponsesGetParams,
    ) -> Response {
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
        (
            StatusCode::NOT_IMPLEMENTED,
            "Responses retrieve endpoint not implemented for PD router",
        )
            .into_response()
    }

    async fn cancel_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
        (
            StatusCode::NOT_IMPLEMENTED,
            "Responses cancel endpoint not implemented for PD router",
        )
            .into_response()
    }

1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
    async fn route_embeddings(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::EmbeddingRequest,
        _model_id: Option<&str>,
    ) -> Response {
        (
            StatusCode::NOT_IMPLEMENTED,
            "Embeddings endpoint not implemented for PD router",
        )
            .into_response()
1216
1217
    }

1218
1219
1220
1221
1222
1223
    async fn route_rerank(
        &self,
        headers: Option<&HeaderMap>,
        body: &RerankRequest,
        model_id: Option<&str>,
    ) -> Response {
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
        // Extract text for cache-aware routing
        let req_text = if self.policies_need_request_text() {
            Some(body.query.clone())
        } else {
            None
        };

        let context = PDRequestContext {
            route: "/v1/rerank",
            batch_size: None,
            is_stream: false,
            return_logprob: false,
            request_text: req_text,
1237
            model_id,
1238
1239
1240
        };

        self.execute_dual_dispatch(headers, body, context).await
1241
1242
    }

1243
1244
1245
1246
    fn router_type(&self) -> &'static str {
        "pd"
    }
}
1247
1248
1249
1250

#[cfg(test)]
mod tests {
    use super::*;
1251
    use crate::core::{BasicWorkerBuilder, WorkerType};
1252
1253

    fn create_test_pd_router() -> PDRouter {
1254
1255
1256
        let worker_registry = Arc::new(WorkerRegistry::new());
        let policy_registry =
            Arc::new(PolicyRegistry::new(crate::config::PolicyConfig::RoundRobin));
1257
1258

        PDRouter {
1259
1260
            worker_registry,
            policy_registry,
1261
            client: Client::new(),
1262
            retry_config: RetryConfig::default(),
1263
            api_key: Some("test_api_key".to_string()),
1264
            enable_igw: false,
1265
1266
1267
1268
        }
    }

    fn create_test_worker(url: String, worker_type: WorkerType, healthy: bool) -> Box<dyn Worker> {
1269
1270
1271
        let worker = BasicWorkerBuilder::new(url)
            .worker_type(worker_type)
            .build();
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
        worker.set_healthy(healthy);
        Box::new(worker)
    }

    #[tokio::test]
    async fn test_select_healthy_prefill_worker() {
        let router = create_test_pd_router();

        let healthy_worker = create_test_worker(
            "http://healthy".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            true,
        );
        let unhealthy_worker = create_test_worker(
            "http://unhealthy".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            false,
        );
        let decode_worker =
            create_test_worker("http://decode".to_string(), WorkerType::Decode, true);

1297
1298
1299
        router.worker_registry.register(Arc::from(unhealthy_worker));
        router.worker_registry.register(Arc::from(healthy_worker));
        router.worker_registry.register(Arc::from(decode_worker));
1300

1301
        let result = router.select_pd_pair(None, None).await;
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313

        assert!(result.is_ok());
        let (prefill, _decode) = result.unwrap();

        assert_eq!(prefill.url(), "http://healthy");
        assert!(prefill.is_healthy());
    }

    #[tokio::test]
    async fn test_empty_worker_lists() {
        let router = create_test_pd_router();

1314
        let result = router.select_pd_pair(None, None).await;
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("No prefill workers available"));
    }

    #[test]
    fn test_worker_load_metrics() {
        let prefill_worker = create_test_worker(
            "http://prefill".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            true,
        );
        let decode_worker =
            create_test_worker("http://decode".to_string(), WorkerType::Decode, true);

        let _guard =
            WorkerLoadGuard::new_multi(vec![prefill_worker.as_ref(), decode_worker.as_ref()]);

        assert_eq!(prefill_worker.load(), 1);
        assert_eq!(decode_worker.load(), 1);

        drop(_guard);

        assert_eq!(prefill_worker.load(), 0);
        assert_eq!(decode_worker.load(), 0);
    }

1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
    #[tokio::test]
    async fn test_streaming_load_tracking() {
        use futures_util::StreamExt;
        use tokio::time::{sleep, Duration};

        let router = create_test_pd_router();

        let prefill_worker = create_test_worker(
            "http://prefill".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            true,
        );
        let decode_worker =
            create_test_worker("http://decode".to_string(), WorkerType::Decode, true);

1361
1362
1363
1364
1365
1366
1367
1368
        router.worker_registry.register(Arc::from(prefill_worker));
        router.worker_registry.register(Arc::from(decode_worker));

        let prefill_workers = router.worker_registry.get_prefill_workers();
        let decode_workers = router.worker_registry.get_decode_workers();

        let prefill_ref = prefill_workers[0].clone();
        let decode_ref = decode_workers[0].clone();
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403

        assert_eq!(prefill_ref.load(), 0);
        assert_eq!(decode_ref.load(), 0);

        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx);

        let _response = router.create_streaming_response(
            stream.map(Ok),
            StatusCode::OK,
            None,
            false,
            None,
            None,
            prefill_ref.as_ref(),
            decode_ref.as_ref(),
        );

        assert_eq!(prefill_ref.load(), 1);
        assert_eq!(decode_ref.load(), 1);

        tx.send(bytes::Bytes::from("test data")).unwrap();

        sleep(Duration::from_millis(10)).await;

        assert_eq!(prefill_ref.load(), 1);
        assert_eq!(decode_ref.load(), 1);

        drop(tx);

        sleep(Duration::from_millis(100)).await;

        assert_eq!(prefill_ref.load(), 0);
        assert_eq!(decode_ref.load(), 0);
    }
1404
}