pd_router.rs 91.1 KB
Newer Older
1
2
// PD (Prefill-Decode) Router Implementation
// This module handles routing for disaggregated prefill-decode systems
3
use super::pd_types::{api_path, PDRouterError};
4
use crate::config::types::RetryConfig;
5
use crate::core::{
6
7
    is_retryable_status, BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, RetryExecutor,
    Worker, WorkerFactory, WorkerLoadGuard, WorkerRegistry, WorkerType,
8
};
9
use crate::metrics::RouterMetrics;
10
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
11
use crate::protocols::spec::{
12
13
    ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, RerankRequest,
    ResponsesRequest, StringOrArray, UserMessageContent,
14
};
15
use crate::routers::header_utils;
16
17
use crate::routers::{RouterTrait, WorkerManagement};
use async_trait::async_trait;
18
19
20
21
22
23
24
25
use axum::{
    body::Body,
    extract::Request,
    http::{header::CONTENT_TYPE, HeaderMap, HeaderValue, StatusCode},
    response::{IntoResponse, Response},
    Json,
};
use futures_util::StreamExt;
26
use reqwest::Client;
27
use serde::Serialize;
28
use serde_json::{json, Value};
29
use std::collections::HashMap;
30
use std::sync::Arc;
31
use std::time::{Duration, Instant};
32
use tokio::sync::mpsc;
33
use tokio_stream::wrappers::UnboundedReceiverStream;
34
35
36
37
use tracing::{debug, error, info, warn};

#[derive(Debug)]
pub struct PDRouter {
38
39
    pub worker_registry: Arc<WorkerRegistry>,
    pub policy_registry: Arc<PolicyRegistry>,
40
41
    pub worker_startup_timeout_secs: u64,
    pub worker_startup_check_interval_secs: u64,
42
43
    pub worker_loads: Arc<tokio::sync::watch::Receiver<HashMap<String, isize>>>,
    pub load_monitor_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
44
    pub client: Client,
45
46
    // Dedicated client for prefill fire-and-forget (non-logprob) requests
    pub prefill_client: Client,
47
    pub retry_config: RetryConfig,
48
    pub circuit_breaker_config: CircuitBreakerConfig,
49
50
    // Channel for sending prefill responses to background workers for draining
    prefill_drain_tx: mpsc::Sender<reqwest::Response>,
51
52
}

53
54
// Request context for PD router operations
#[derive(Clone)]
55
struct PDRequestContext<'a> {
56
57
58
59
60
    route: &'static str,
    batch_size: Option<usize>,
    is_stream: bool,
    return_logprob: bool,
    request_text: Option<String>,
61
    model_id: Option<&'a str>,
62
63
}

64
impl PDRouter {
65
66
    // Private helper method to perform health check on a new server
    async fn wait_for_server_health(&self, url: &str) -> Result<(), PDRouterError> {
67
        crate::routers::http::router::Router::wait_for_healthy_workers(
68
            &[url.to_string()],
69
70
            self.worker_startup_timeout_secs,
            self.worker_startup_check_interval_secs,
71
        )
72
        .await
73
74
75
76
77
        .map_err(|_| PDRouterError::HealthCheckFailed {
            url: url.to_string(),
        })
    }

78
79
80
    // Generic helper for processing all workers with an endpoint
    async fn process_workers(
        &self,
81
        worker_type_enum: WorkerType,
82
83
84
85
86
87
        worker_type: &str,
        endpoint: &str,
    ) -> (Vec<String>, Vec<String>) {
        let mut results = Vec::new();
        let mut errors = Vec::new();

88
89
90
        // Get workers from registry based on type
        let workers = self.worker_registry.get_by_type(&worker_type_enum);
        let urls: Vec<String> = workers.iter().map(|w| w.url().to_string()).collect();
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

        // Process each worker
        for worker_url in urls {
            let url = format!("{}/{}", worker_url, endpoint);
            match self.client.post(&url).send().await {
                Ok(res) if res.status().is_success() => {
                    results.push(format!("{} {}: OK", worker_type, worker_url));
                }
                Ok(res) => {
                    errors.push(format!(
                        "{} {} returned status: {}",
                        worker_type,
                        worker_url,
                        res.status()
                    ));
                }
                Err(e) => {
                    errors.push(format!("{} {} error: {}", worker_type, worker_url, e));
                }
            }
        }

        (results, errors)
    }

116
117
118
119
120
121
122
    // Helper to get prefill worker URLs
    fn get_prefill_worker_urls(&self) -> Vec<String> {
        self.worker_registry
            .get_prefill_workers()
            .iter()
            .map(|w| w.url().to_string())
            .collect()
123
124
    }

125
126
127
128
129
130
131
132
133
134
135
    // Helper to get decode worker URLs
    fn get_decode_worker_urls(&self) -> Vec<String> {
        self.worker_registry
            .get_decode_workers()
            .iter()
            .map(|w| w.url().to_string())
            .collect()
    }

    // Helper for proxying requests to the first prefill worker
    async fn proxy_to_first_prefill_worker(
136
137
138
139
        &self,
        endpoint: &str,
        headers: Option<Vec<(String, String)>>,
    ) -> Response {
140
141
        let workers = self.worker_registry.get_prefill_workers();
        let first_worker_url = workers.first().map(|w| w.url().to_string());
142
143

        if let Some(worker_url) = first_worker_url {
144
145
146
147
148
149
150
151
152
            self.proxy_to_worker(worker_url, endpoint, headers).await
        } else {
            (
                StatusCode::SERVICE_UNAVAILABLE,
                "No prefill servers available".to_string(),
            )
                .into_response()
        }
    }
153

154
155
156
157
158
159
160
161
162
163
164
165
166
167
    // Generic helper for proxying to a specific worker
    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);

        // Add headers if provided
        if let Some(headers) = headers {
            for (name, value) in headers {
                request_builder = request_builder.header(name, value);
168
            }
169
        }
170

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
        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()
189
190
191
                    }
                }
            }
192
193
194
195
196
197
198
199
200
201
202
203
204
            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()
            }
205
206
207
        }
    }

208
209
210
211
212
213
    pub async fn add_prefill_server(
        &self,
        url: String,
        bootstrap_port: Option<u16>,
    ) -> Result<String, PDRouterError> {
        // Wait for the new server to be healthy
214
        self.wait_for_server_health(&url).await?;
215

216
217
218
219
220
        // Check if already exists
        if self.worker_registry.get_by_url(&url).is_some() {
            return Err(PDRouterError::WorkerAlreadyExists { url: url.clone() });
        }

221
        // Create Worker for the new prefill server with circuit breaker configuration
222
        // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
223
224
225
226
227
        let worker = WorkerFactory::create_prefill_with_config(
            url.clone(),
            bootstrap_port,
            self.circuit_breaker_config.clone(),
        );
228

229
        let worker_arc: Arc<dyn Worker> = Arc::from(worker);
230

231
232
        // Register the worker in the registry
        self.worker_registry.register(worker_arc.clone());
233

234
235
236
        // Notify PolicyRegistry about the new worker
        let model_id = worker_arc.model_id();
        let policy = self.policy_registry.on_worker_added(model_id, None);
237

238
239
240
241
242
243
244
245
246
        // If this is a cache-aware policy, update it with all workers for this model
        if policy.name() == "cache_aware" {
            if let Some(cache_aware) = policy
                .as_any()
                .downcast_ref::<crate::policies::CacheAwarePolicy>()
            {
                let model_workers = self.worker_registry.get_by_model_fast(model_id);
                cache_aware.init_workers(&model_workers);
            }
247
248
249
250
251
252
253
254
        }

        info!("Added prefill server: {}", url);
        Ok(format!("Successfully added prefill server: {}", url))
    }

    pub async fn add_decode_server(&self, url: String) -> Result<String, PDRouterError> {
        // Wait for the new server to be healthy
255
        self.wait_for_server_health(&url).await?;
256

257
258
259
260
261
        // Check if already exists
        if self.worker_registry.get_by_url(&url).is_some() {
            return Err(PDRouterError::WorkerAlreadyExists { url: url.clone() });
        }

262
        // Create Worker for the new decode server with circuit breaker configuration
263
        // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
264
265
266
267
        let worker = WorkerFactory::create_decode_with_config(
            url.clone(),
            self.circuit_breaker_config.clone(),
        );
268

269
        let worker_arc: Arc<dyn Worker> = Arc::from(worker);
270

271
272
        // Register the worker in the registry
        self.worker_registry.register(worker_arc.clone());
273

274
275
276
        // Notify PolicyRegistry about the new worker
        let model_id = worker_arc.model_id();
        let policy = self.policy_registry.on_worker_added(model_id, None);
277

278
279
280
281
282
283
284
285
286
        // If this is a cache-aware policy, update it with all workers for this model
        if policy.name() == "cache_aware" {
            if let Some(cache_aware) = policy
                .as_any()
                .downcast_ref::<crate::policies::CacheAwarePolicy>()
            {
                let model_workers = self.worker_registry.get_by_model_fast(model_id);
                cache_aware.init_workers(&model_workers);
            }
287
288
        }

289
290
291
292
293
        info!("Added decode server: {}", url);
        Ok(format!("Successfully added decode server: {}", url))
    }

    pub async fn remove_prefill_server(&self, url: &str) -> Result<String, PDRouterError> {
294
295
296
297
298
299
300
301
302
        // Check if worker exists and get model_id
        let model_id = match self.worker_registry.get_by_url(url) {
            Some(worker) => worker.model_id().to_string(),
            None => {
                return Err(PDRouterError::WorkerNotFound {
                    url: url.to_string(),
                });
            }
        };
303

304
305
        // Remove from registry
        let removed = self.worker_registry.remove_by_url(url);
306

307
308
309
        if removed.is_some() {
            // Notify PolicyRegistry about the removed worker
            self.policy_registry.on_worker_removed(&model_id);
310

311
312
313
314
315
316
317
318
319
320
321
            // Get the policy for this model to update cache-aware if needed
            if let Some(policy) = self.policy_registry.get_policy(&model_id) {
                if policy.name() == "cache_aware" {
                    if let Some(cache_aware) = policy
                        .as_any()
                        .downcast_ref::<crate::policies::CacheAwarePolicy>()
                    {
                        cache_aware.remove_worker_by_url(url);
                    }
                }
            }
322
323
        }

324
325
326
327
328
329
330
331
        if removed.is_some() {
            info!("Removed prefill server: {}", url);
            Ok(format!("Successfully removed prefill server: {}", url))
        } else {
            Err(PDRouterError::WorkerNotFound {
                url: url.to_string(),
            })
        }
332
333
334
    }

    pub async fn remove_decode_server(&self, url: &str) -> Result<String, PDRouterError> {
335
336
337
338
339
340
341
342
343
        // Check if worker exists and get model_id
        let model_id = match self.worker_registry.get_by_url(url) {
            Some(worker) => worker.model_id().to_string(),
            None => {
                return Err(PDRouterError::WorkerNotFound {
                    url: url.to_string(),
                });
            }
        };
344

345
346
        // Remove from registry
        let removed = self.worker_registry.remove_by_url(url);
347

348
349
350
        if removed.is_some() {
            // Notify PolicyRegistry about the removed worker
            self.policy_registry.on_worker_removed(&model_id);
351

352
353
354
355
356
357
358
359
360
361
362
            // Get the policy for this model to update cache-aware if needed
            if let Some(policy) = self.policy_registry.get_policy(&model_id) {
                if policy.name() == "cache_aware" {
                    if let Some(cache_aware) = policy
                        .as_any()
                        .downcast_ref::<crate::policies::CacheAwarePolicy>()
                    {
                        cache_aware.remove_worker_by_url(url);
                    }
                }
            }
363
364
        }

365
366
367
368
369
370
371
372
        if removed.is_some() {
            info!("Removed decode server: {}", url);
            Ok(format!("Successfully removed decode server: {}", url))
        } else {
            Err(PDRouterError::WorkerNotFound {
                url: url.to_string(),
            })
        }
373
    }
374

375
    #[allow(clippy::too_many_arguments)]
376
    pub async fn new(
377
378
        prefill_urls: Vec<(String, Option<u16>)>,
        decode_urls: Vec<String>,
379
        ctx: &Arc<crate::server::AppContext>,
380
    ) -> Result<Self, String> {
381
        // Convert config CircuitBreakerConfig to core CircuitBreakerConfig
382
        let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config();
383
384
385
        let core_cb_config = CircuitBreakerConfig {
            failure_threshold: circuit_breaker_config.failure_threshold,
            success_threshold: circuit_breaker_config.success_threshold,
386
387
            timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs),
            window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs),
388
389
        };

390
391
        // Register prefill workers in the registry
        for (url, port) in prefill_urls {
392
393
            let worker = BasicWorkerBuilder::new(url)
                .worker_type(WorkerType::Prefill {
394
                    bootstrap_port: port,
395
396
397
398
399
400
401
402
403
404
                })
                .circuit_breaker_config(core_cb_config.clone())
                .health_config(HealthConfig {
                    timeout_secs: ctx.router_config.health_check.timeout_secs,
                    check_interval_secs: ctx.router_config.health_check.check_interval_secs,
                    endpoint: ctx.router_config.health_check.endpoint.clone(),
                    failure_threshold: ctx.router_config.health_check.failure_threshold,
                    success_threshold: ctx.router_config.health_check.success_threshold,
                })
                .build();
405
406
407
408
409
            ctx.worker_registry.register(Arc::new(worker));
        }

        // Register decode workers in the registry
        for url in decode_urls {
410
411
412
413
            let worker = BasicWorkerBuilder::new(url)
                .worker_type(WorkerType::Decode)
                .circuit_breaker_config(core_cb_config.clone())
                .health_config(HealthConfig {
414
415
416
417
418
                    timeout_secs: ctx.router_config.health_check.timeout_secs,
                    check_interval_secs: ctx.router_config.health_check.check_interval_secs,
                    endpoint: ctx.router_config.health_check.endpoint.clone(),
                    failure_threshold: ctx.router_config.health_check.failure_threshold,
                    success_threshold: ctx.router_config.health_check.success_threshold,
419
420
                })
                .build();
421
422
            ctx.worker_registry.register(Arc::new(worker));
        }
423

424
425
426
        // Get all workers from registry for health check
        let all_workers = ctx.worker_registry.get_all();
        let all_urls: Vec<String> = all_workers
427
            .iter()
428
            .map(|worker| worker.url().to_string())
429
            .collect();
430
        if !all_urls.is_empty() {
431
            crate::routers::http::router::Router::wait_for_healthy_workers(
432
                &all_urls,
433
434
                ctx.router_config.worker_startup_timeout_secs,
                ctx.router_config.worker_startup_check_interval_secs,
435
436
            )
            .await?;
437
        }
438

439
440
441
442
        // Initialize cache-aware policies with workers from registry
        // Note: We need to get workers by type and convert to Box<dyn Worker> for CacheAwarePolicy
        // This is a temporary workaround until CacheAwarePolicy is updated to work with Arc<dyn Worker>
        // TODO: Update CacheAwarePolicy to accept Arc<dyn Worker> instead of Box<dyn Worker>
443

444
445
446
447
        // Set up background load monitoring for power-of-two selection
        let (tx, rx) = tokio::sync::watch::channel(HashMap::new());
        let worker_loads = Arc::new(rx);

448
449
450
451
        // Get policies from registry to check if we need load monitoring
        let prefill_policy = ctx.policy_registry.get_prefill_policy();
        let decode_policy = ctx.policy_registry.get_decode_policy();

452
453
454
        let load_monitor_handle =
            if prefill_policy.name() == "power_of_two" || decode_policy.name() == "power_of_two" {
                let monitor_urls = all_urls.clone();
455
456
                let monitor_interval = ctx.router_config.worker_startup_check_interval_secs;
                let monitor_client = ctx.client.clone();
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
                let prefill_policy_clone = Arc::clone(&prefill_policy);
                let decode_policy_clone = Arc::clone(&decode_policy);

                Some(Arc::new(tokio::spawn(async move {
                    Self::monitor_worker_loads_with_client(
                        monitor_urls,
                        tx,
                        monitor_interval,
                        monitor_client,
                        prefill_policy_clone,
                        decode_policy_clone,
                    )
                    .await;
                })))
            } else {
                None
            };
474

475
476
        // Note: Health checking is now handled centrally by RouterManager
        // Individual routers no longer need to manage health checkers
477

478
479
480
481
482
        // Build a dedicated prefill client for fire-and-forget semantics
        let prefill_client = reqwest::Client::builder()
            .pool_max_idle_per_host(0)
            .http1_only()
            .connect_timeout(Duration::from_millis(300))
483
            .timeout(Duration::from_secs(ctx.router_config.request_timeout_secs))
484
485
486
            .build()
            .map_err(|e| format!("Failed to build prefill client: {}", e))?;

487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
        // Create bounded channel for prefill response draining
        // Larger buffer for high concurrency scenarios
        let (prefill_drain_tx, mut prefill_drain_rx) = mpsc::channel::<reqwest::Response>(2000);

        // Spawn a coordinator with limited concurrent drain tasks
        // This prevents unbounded task spawning under extreme load
        tokio::spawn(async move {
            info!("Prefill drain coordinator started");

            // Use a semaphore to limit concurrent drain operations
            let max_concurrent_drains = 100;
            let semaphore = Arc::new(tokio::sync::Semaphore::new(max_concurrent_drains));

            while let Some(response) = prefill_drain_rx.recv().await {
                let permit = semaphore.clone().acquire_owned().await;

                match permit {
                    Ok(permit) => {
                        // Spawn a task to drain this response
                        tokio::spawn(async move {
                            let url = response.url().to_string();
                            let status = response.status();

                            if !status.is_success() {
                                error!("Prefill drain: error status={} url={}", status, url);
                                RouterMetrics::record_pd_prefill_error(&url);
                            }

                            // Drain the response body efficiently
                            // Use streaming to avoid loading entire body into memory
                            let start = std::time::Instant::now();
                            let mut stream = response.bytes_stream();
                            let mut bytes_drained = 0;

                            while let Some(chunk_result) = stream.next().await {
                                match chunk_result {
                                    Ok(chunk) => bytes_drained += chunk.len(),
                                    Err(e) => {
                                        debug!(
                                            "Prefill drain: error streaming url={} error={}",
                                            url, e
                                        );
                                        break;
                                    }
                                }
                            }

                            let elapsed = start.elapsed();
                            if elapsed > Duration::from_millis(100) {
                                // Only log slow drains
                                debug!(
                                    "Prefill drain: slow drain {} bytes from {} in {:?}",
                                    bytes_drained, url, elapsed
                                );
                            }

                            // Permit is automatically released when dropped
                            drop(permit);
                        });
                    }
                    Err(_) => {
                        // Semaphore closed, shutting down
                        break;
                    }
                }
            }
            info!("Prefill drain coordinator shutting down");
        });

556
        Ok(PDRouter {
557
558
            worker_registry: Arc::clone(&ctx.worker_registry),
            policy_registry: Arc::clone(&ctx.policy_registry),
559
560
561
562
            worker_startup_timeout_secs: ctx.router_config.worker_startup_timeout_secs,
            worker_startup_check_interval_secs: ctx
                .router_config
                .worker_startup_check_interval_secs,
563
564
            worker_loads,
            load_monitor_handle,
565
            client: ctx.client.clone(),
566
            prefill_client,
567
            prefill_drain_tx,
568
            retry_config: ctx.router_config.effective_retry_config(),
569
            circuit_breaker_config: core_cb_config,
570
571
572
        })
    }

573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
    // Helper to handle server selection errors
    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()
    }

    // Helper to handle serialization errors
    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()
    }

594
595
596
    // Helper to determine batch size from a GenerateRequest
    fn get_generate_batch_size(req: &GenerateRequest) -> Option<usize> {
        // Check prompt array
597
        if let Some(StringOrArray::Array(arr)) = &req.prompt {
598
599
            if !arr.is_empty() {
                return Some(arr.len());
600
            }
601
602
603
604
605
606
        }
        // Check text array
        if let Some(text) = &req.text {
            if text.contains("[") && text.contains("]") {
                // This is a simplified check - in reality we'd need to parse JSON
                return None; // For now, fall back to non-batch
607
            }
608
609
610
        }
        None
    }
611

612
613
614
615
616
617
    // Helper to determine batch size from a ChatCompletionRequest
    fn get_chat_batch_size(req: &ChatCompletionRequest) -> Option<usize> {
        // Check 'n' parameter for multiple responses
        if let Some(n) = req.n {
            if n > 1 {
                return Some(n as usize);
618
            }
619
620
621
        }
        None
    }
622

623
624
625
    // Helper to determine batch size from a CompletionRequest
    fn get_completion_batch_size(req: &CompletionRequest) -> Option<usize> {
        // Check prompt array
626
        if let StringOrArray::Array(arr) = &req.prompt {
627
628
            if !arr.is_empty() {
                return Some(arr.len());
629
630
            }
        }
631
632
        None
    }
633

634
635
636
    // Helper to inject bootstrap fields into an existing JSON request value
    fn inject_bootstrap_into_value(
        mut original: Value,
637
638
        prefill_worker: &dyn Worker,
        batch_size: Option<usize>,
639
    ) -> Result<Value, String> {
640
641
642
643
644
645
        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());

646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
        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()),
            );
679
        } else {
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
            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()),
            );
695
        }
696
        Ok(original)
697
698
    }

699
700
701
702
703
    // Execute the dual dispatch to prefill and decode servers with retries and bootstrap injection
    async fn execute_dual_dispatch<T: Serialize + Clone>(
        &self,
        headers: Option<&HeaderMap>,
        original_request: &T,
704
        context: PDRequestContext<'_>,
705
706
707
708
709
710
711
712
713
714
715
716
717
718
    ) -> Response {
        let start_time = Instant::now();

        let route = context.route;
        RetryExecutor::execute_response_with_retry(
            &self.retry_config,
            // Operation per attempt
            {
                let original_request = original_request.clone();
                move |attempt: u32| {
                    let original_request = original_request.clone();
                    let context = context.clone();
                    async move {
                        // Select workers fresh for each attempt
719
720
721
722
723
724
725
726
727
728
                        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);
                            }
                        };
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757

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

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

                        // Inject bootstrap based on current prefill worker
                        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),
                        };

                        // Execute the actual dual dispatch
                        let response = self
                            .execute_dual_dispatch_internal(
                                headers,
                                json_request,
758
                                context,
759
760
761
762
763
764
765
                                prefill.as_ref(),
                                decode.as_ref(),
                                start_time,
                            )
                            .await;

                        // Record outcomes for circuit breakers
766
767
768
769
                        let _status = response.status();
                        let not_error = _status.is_success() || _status.is_client_error();
                        prefill.record_outcome(not_error);
                        decode.record_outcome(not_error);
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787

                        response
                    }
                }
            },
            // Should retry predicate
            |res, _attempt| is_retryable_status(res.status()),
            // On backoff hook
            |delay, attempt| {
                RouterMetrics::record_retry(route);
                RouterMetrics::record_retry_backoff_duration(delay, attempt);
            },
            // On exhausted hook
            || RouterMetrics::record_retries_exhausted(route),
        )
        .await
    }

788
789
790
    async fn handle_decode_error_response(
        &self,
        res: reqwest::Response,
791
        context: &PDRequestContext<'_>,
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
        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(),
            }
        }
    }

839
840
    // Internal method that performs the actual dual dispatch (without retry logic)
    async fn execute_dual_dispatch_internal(
841
842
843
        &self,
        headers: Option<&HeaderMap>,
        json_request: Value,
844
        context: PDRequestContext<'_>,
845
846
847
848
        prefill: &dyn Worker,
        decode: &dyn Worker,
        start_time: Instant,
    ) -> Response {
849
850
851
852
853
854
855
        // 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
        };
856

857
858
859
860
        // Build decode request with shared client
        let decode_request = self.build_post_with_headers(
            &self.client,
            decode.url(),
861
            context.route,
862
863
864
865
            &json_request,
            headers,
            false,
        );
866

867
        // Send both requests concurrently
868
869
870
871
872
873
        debug!(
            "Sending concurrent requests to prefill={} decode={}",
            prefill.url(),
            decode.url()
        );

874
        if context.return_logprob {
875
876
877
878
            // Build prefill request with shared client when we need response body
            let prefill_request = self.build_post_with_headers(
                &self.client,
                prefill.url(),
879
                context.route,
880
881
882
883
                &json_request,
                headers,
                false,
            );
884
885
886
887
888
889
890
            // When we need logprobs, wait for both responses
            let (prefill_result, decode_result) =
                tokio::join!(prefill_request.send(), decode_request.send());
            debug!("Received responses from both servers");

            // Update metrics
            let duration = start_time.elapsed();
891
892
            RouterMetrics::record_pd_request_duration(context.route, duration);
            RouterMetrics::record_pd_request(context.route);
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
            RouterMetrics::record_pd_prefill_request(prefill.url());
            RouterMetrics::record_pd_decode_request(decode.url());

            // Process decode response with prefill for logprobs
            debug!("Processing decode response with logprobs");
            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
                        );

912
913
914
                        return self
                            .handle_decode_error_response(res, &context, prefill, decode)
                            .await;
915
916
                    }

917
918
                    // Process prefill response for logprobs
                    let prefill_body = match self
919
920
921
922
923
                        .process_prefill_response(
                            prefill_result,
                            prefill.url(),
                            context.return_logprob,
                        )
924
925
926
927
928
929
                        .await
                    {
                        Ok((_, body)) => body,
                        Err(error_response) => return error_response,
                    };

930
                    if context.is_stream {
931
932
                        // Streaming response with logprobs
                        let prefill_logprobs = prefill_body
933
934
935
936
                            .as_ref()
                            .and_then(|body| serde_json::from_slice::<Value>(body).ok())
                            .and_then(|json| {
                                json.pointer("/meta_info/input_token_logprobs").cloned()
937
938
                            });

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

942
                        self.create_streaming_response(
943
944
945
                            res.bytes_stream(),
                            status,
                            prefill_logprobs,
946
                            context.return_logprob,
947
                            None,
948
                            Some(response_headers),
949
950
                            prefill,
                            decode,
951
                        )
952
                    } else {
953
954
955
956
                        // Non-streaming response with logprobs
                        self.process_non_streaming_response(
                            res,
                            status,
957
                            context.return_logprob,
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
                            prefill_body,
                        )
                        .await
                    }
                }
                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()
                }
            }
        } else {
            // When we don't need logprobs, only wait for decode response
            // Send both requests concurrently but don't wait for prefill
980
981
982
983
984
            // Use dedicated prefill client with Connection: close
            let prefill_future = self
                .build_post_with_headers(
                    &self.prefill_client,
                    prefill.url(),
985
                    context.route,
986
987
988
989
990
                    &json_request,
                    headers,
                    true,
                )
                .send();
991
992
            let decode_future = decode_request.send();

993
994
995
996
            // Send prefill response to background worker for draining
            // This ensures HTTP compliance without blocking
            let drain_tx = self.prefill_drain_tx.clone();
            let prefill_url = prefill.url().to_string();
997
998
            tokio::spawn(async move {
                if let Ok(response) = prefill_future.await {
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
                    // Try to send to drain worker
                    // If channel is full (under extreme load), drain inline as fallback
                    match drain_tx.try_send(response) {
                        Ok(_) => {
                            // Successfully queued for draining
                            debug!("Prefill response queued for draining");
                        }
                        Err(mpsc::error::TrySendError::Full(response)) => {
                            // Channel full - drain inline as fallback
                            warn!("Prefill drain channel full (capacity exceeded), draining inline for {}", prefill_url);
                            RouterMetrics::record_pd_prefill_error(&prefill_url);

                            // Drain inline with timeout to prevent blocking too long
                            let drain_future = async {
                                let mut stream = response.bytes_stream();
                                while stream.next().await.is_some() {
                                    // Just drain
                                }
                            };

                            match tokio::time::timeout(Duration::from_secs(1), drain_future).await {
                                Ok(_) => debug!("Inline drain completed for {}", prefill_url),
                                Err(_) => error!("Inline drain timeout for {}", prefill_url),
                            }
                        }
                        Err(mpsc::error::TrySendError::Closed(_)) => {
                            error!("Prefill drain channel closed!");
                        }
                    }
1028
1029
                }
            });
1030

1031
1032
1033
1034
1035
1036
            // Wait only for decode response
            let decode_result = decode_future.await;
            debug!("Received decode response");

            // Update metrics
            let duration = start_time.elapsed();
1037
1038
            RouterMetrics::record_pd_request_duration(context.route, duration);
            RouterMetrics::record_pd_request(context.route);
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
            RouterMetrics::record_pd_prefill_request(prefill.url());
            RouterMetrics::record_pd_decode_request(decode.url());

            // Process decode response immediately
            debug!("Processing decode response (no logprobs)");
            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
                        );

1058
1059
                        self.handle_decode_error_response(res, &context, prefill, decode)
                            .await
1060
                    } else if context.is_stream {
1061
1062
                        // Streaming response without logprobs - direct passthrough
                        let decode_url = decode.url().to_string();
1063
1064
1065
                        let response_headers =
                            header_utils::preserve_response_headers(res.headers());

1066
                        self.create_streaming_response(
1067
1068
1069
1070
1071
                            res.bytes_stream(),
                            status,
                            None,
                            false,
                            Some(decode_url),
1072
                            Some(response_headers),
1073
1074
                            prefill,
                            decode,
1075
1076
1077
                        )
                    } else {
                        // Non-streaming response without logprobs - direct passthrough like fast version
1078
1079
1080
                        let response_headers =
                            header_utils::preserve_response_headers(res.headers());

1081
                        match res.bytes().await {
1082
1083
1084
1085
1086
1087
1088
                            Ok(decode_body) => {
                                let mut response =
                                    Response::new(axum::body::Body::from(decode_body));
                                *response.status_mut() = status;
                                *response.headers_mut() = response_headers;
                                response
                            }
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
                            Err(e) => {
                                error!("Failed to read decode response: {}", e);
                                (StatusCode::INTERNAL_SERVER_ERROR, "Failed to read response")
                                    .into_response()
                            }
                        }
                    }
                }
                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),
1107
                    )
1108
                        .into_response()
1109
1110
1111
1112
1113
                }
            }
        }
    }

1114
1115
    // Check if either prefill or decode policy needs request text
    fn policies_need_request_text(&self) -> bool {
1116
1117
1118
1119
        // Check both prefill and decode policies
        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()
1120
1121
    }

1122
    // Select a pair of prefill and decode servers considering circuit breaker state
1123
1124
    async fn select_pd_pair(
        &self,
1125
        request_text: Option<&str>,
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
        model_id: Option<&str>,
    ) -> Result<(Arc<dyn Worker>, Arc<dyn Worker>), String> {
        // Get workers from registry - filter by model if provided
        let prefill_workers = if let Some(model) = model_id {
            // Get model-specific workers and filter for prefill type
            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()
        };

        let decode_workers = if let Some(model) = model_id {
            // Get model-specific workers and filter for decode type
            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()
        };
1150

1151
        // Select workers using helper function
1152
1153
1154
1155
1156
        // Use separate policies for prefill and decode to avoid counter conflicts
        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(
1157
            &prefill_workers,
1158
            &*prefill_policy,
1159
1160
1161
1162
            request_text,
            "prefill",
        )?;

1163
        let decode = Self::pick_worker_by_policy_arc(
1164
            &decode_workers,
1165
            &*decode_policy,
1166
1167
1168
1169
1170
1171
1172
            request_text,
            "decode",
        )?;

        Ok((prefill, decode))
    }

1173
1174
1175
    // Helper function to select a worker using the policy (Arc version)
    fn pick_worker_by_policy_arc(
        workers: &[Arc<dyn Worker>],
1176
1177
1178
        policy: &dyn LoadBalancingPolicy,
        request_text: Option<&str>,
        worker_type: &str,
1179
    ) -> Result<Arc<dyn Worker>, String> {
1180
1181
1182
1183
1184
1185
        // Check if we have any workers
        if workers.is_empty() {
            return Err(format!(
                "No {} workers available. Please check if {} servers are configured and healthy.",
                worker_type, worker_type
            ));
1186
1187
        }

1188
        // Filter available workers (healthy + circuit breaker not open)
1189
        let available_workers: Vec<Arc<dyn Worker>> = workers
1190
1191
            .iter()
            .filter(|w| w.is_available())
1192
            .cloned()
1193
            .collect();
1194

1195
1196
1197
1198
1199
1200
        if available_workers.is_empty() {
            return Err(format!(
                "No available {} workers (all circuits open or unhealthy)",
                worker_type
            ));
        }
1201

1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
        // Let policy select from available workers (no conversion needed now!)
        let selected_idx = policy
            .select_worker(&available_workers, request_text)
            .ok_or_else(|| {
                format!(
                    "Policy {} failed to select a {} worker",
                    policy.name(),
                    worker_type
                )
            })?;

        // Return the selected Arc worker
        Ok(available_workers[selected_idx].clone())
1215
1216
1217
1218
1219
1220
1221
    }

    // Background task to monitor worker loads with shared client
    async fn monitor_worker_loads_with_client(
        worker_urls: Vec<String>,
        tx: tokio::sync::watch::Sender<HashMap<String, isize>>,
        interval_secs: u64,
1222
        client: Client,
1223
1224
        prefill_policy: Arc<dyn LoadBalancingPolicy>,
        decode_policy: Arc<dyn LoadBalancingPolicy>,
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
    ) {
        loop {
            let mut loads = HashMap::new();

            let futures: Vec<_> = worker_urls
                .iter()
                .map(|url| {
                    let client = client.clone();
                    let url = url.clone();
                    async move {
                        let load = get_worker_load(&client, &url).await.unwrap_or(0);
                        (url, load)
                    }
                })
                .collect();

            let results = futures_util::future::join_all(futures).await;

            for (url, load) in results {
                loads.insert(url, load);
            }

            debug!("Worker loads updated: {:?}", loads);

1249
1250
1251
            // Update both policies with current loads
            prefill_policy.update_loads(&loads);
            decode_policy.update_loads(&loads);
1252

1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
            // Check if receiver is still active
            if tx.send(loads).is_err() {
                info!("Load monitor receiver dropped, shutting down monitor task");
                break;
            }

            tokio::time::sleep(Duration::from_secs(interval_secs)).await;
        }
    }

1263
    // Helper to create a streaming response
1264
    #[allow(clippy::too_many_arguments)]
1265
    fn create_streaming_response(
1266
        &self,
1267
1268
1269
1270
1271
        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>,
1272
        headers: Option<HeaderMap>,
1273
1274
        prefill: &dyn Worker,
        decode: &dyn Worker,
1275
    ) -> Response {
1276
1277
1278
1279
1280
1281
1282
1283
        // For streaming, increment load now - will be decremented when streaming completes
        prefill.increment_load();
        decode.increment_load();

        // Store URLs to find workers later for decrementing
        let prefill_url = prefill.url().to_string();
        let decode_url_str = decode.url().to_string();

1284
1285
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

1286
1287
        // Clone the registry for the spawned task
        let registry = self.worker_registry.clone();
1288

1289
        tokio::spawn(async move {
1290
1291
1292
            // Use a flag to track whether stream completed successfully
            let mut stream_completed = false;

1293
1294
1295
1296
            futures_util::pin_mut!(stream);
            while let Some(chunk_result) = stream.next().await {
                match chunk_result {
                    Ok(chunk) => {
1297
1298
1299
1300
1301
1302
                        // Check for stream end marker to decrement load early
                        let is_done = chunk
                            .as_ref()
                            .windows(12)
                            .any(|window| window == b"data: [DONE]");

1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
                        let result = if return_logprob && prefill_logprobs.is_some() {
                            // Try to merge logprobs
                            Self::merge_streaming_logprobs(prefill_logprobs.clone(), &chunk)
                                .unwrap_or(chunk)
                        } else {
                            chunk
                        };

                        if tx.send(Ok(result)).is_err() {
                            break;
                        }
1314
1315
1316
1317
1318
1319

                        // If we see the done marker, decrement load immediately
                        if is_done {
                            stream_completed = true;
                            break;
                        }
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
                    }
                    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;
                    }
                }
            }
1331

1332
1333
            // Always decrement load after streaming (either completes or errors)
            // Find and decrement prefill worker
1334
1335
1336
1337
1338
1339
            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
                );
1340
1341
            }

1342
            // Find and decrement decode worker
1343
1344
1345
1346
1347
1348
            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
                );
1349
            }
1350
1351
1352
1353
1354
1355
1356
        });

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

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

        // Use provided headers or create new ones, then ensure content-type is set for streaming
1359
        let mut headers = headers.unwrap_or_default();
1360
1361
1362
        headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
        *response.headers_mut() = headers;

1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
        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 {
1374
1375
1376
        let response = res.bytes().await;
        let decode_body = match response {
            Ok(decode_body) => decode_body,
1377
1378
            Err(e) => {
                error!("Failed to read decode response: {}", e);
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
1404
1405
1406
1407
1408
                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()
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
            }
        }
    }

    // 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))
    }

1490
    fn build_post_with_headers(
1491
        &self,
1492
        client: &Client,
1493
1494
        url: &str,
        route: &str,
1495
        json_request: &Value,
1496
        headers: Option<&HeaderMap>,
1497
        connection_close: bool,
1498
    ) -> reqwest::RequestBuilder {
1499
1500
1501
1502
        let mut request = client.post(api_path(url, route)).json(json_request);
        if connection_close {
            request = request.header("Connection", "close");
        }
1503
1504
        if let Some(headers) = headers {
            for (name, value) in headers.iter() {
1505
1506
1507
1508
1509
1510
1511
1512
1513
                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);
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
                    }
                }
            }
        }
        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
    }

1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
    // 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))
    }
}

// Helper functions

1585
async fn get_worker_load(client: &Client, worker_url: &str) -> Option<isize> {
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
    match client.get(format!("{}/get_load", worker_url)).send().await {
        Ok(res) if res.status().is_success() => match res.bytes().await {
            Ok(bytes) => match serde_json::from_slice::<Value>(&bytes) {
                Ok(data) => data
                    .get("load")
                    .and_then(|v| v.as_i64())
                    .map(|v| v as isize),
                Err(e) => {
                    debug!("Failed to parse load response from {}: {}", worker_url, e);
                    None
                }
            },
            Err(e) => {
                debug!("Failed to read load response from {}: {}", worker_url, e);
                None
            }
        },
        Ok(res) => {
            debug!(
                "Worker {} returned non-success status: {}",
                worker_url,
                res.status()
            );
            None
        }
        Err(e) => {
            debug!("Failed to get load from {}: {}", worker_url, e);
            None
        }
    }
}

1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
#[async_trait]
impl WorkerManagement for PDRouter {
    async fn add_worker(&self, _worker_url: &str) -> Result<String, String> {
        // For PD router, we don't support adding workers via this generic method
        Err(
            "PD router requires specific add_prefill_server or add_decode_server methods"
                .to_string(),
        )
    }

    fn remove_worker(&self, worker_url: &str) {
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
        // Remove from registry
        if let Some(worker) = self.worker_registry.remove_by_url(worker_url) {
            match worker.worker_type() {
                WorkerType::Prefill { .. } => {
                    info!("Removed prefill worker: {}", worker_url);
                }
                WorkerType::Decode => {
                    info!("Removed decode worker: {}", worker_url);
                }
                _ => {
                    info!("Removed worker: {}", worker_url);
                }
1641
1642
1643
1644
1645
            }
        }
    }

    fn get_worker_urls(&self) -> Vec<String> {
1646
        self.worker_registry.get_all_urls()
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
    }
}

#[async_trait]
impl RouterTrait for PDRouter {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    async fn health(&self, _req: Request<Body>) -> Response {
        // This is a server readiness check - checking if we have healthy workers
        // Workers handle their own health checks in the background
        let mut all_healthy = true;
        let mut unhealthy_servers = Vec::new();

1662
1663
        // Check all workers
        for worker in self.worker_registry.get_all() {
1664
1665
            if !worker.is_healthy() {
                all_healthy = false;
1666
1667
1668
1669
1670
1671
                let worker_type = match worker.worker_type() {
                    WorkerType::Prefill { .. } => "Prefill",
                    WorkerType::Decode => "Decode",
                    _ => "Worker",
                };
                unhealthy_servers.push(format!("{}: {}", worker_type, worker.url()));
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
            }
        }

        if all_healthy {
            (StatusCode::OK, "All servers healthy").into_response()
        } else {
            (
                StatusCode::SERVICE_UNAVAILABLE,
                format!("Unhealthy servers: {:?}", unhealthy_servers),
            )
                .into_response()
        }
    }

    async fn health_generate(&self, _req: Request<Body>) -> Response {
1687
1688
        // Test model generation capability by selecting a random pair and testing them
        // Note: This endpoint actually causes the model to generate tokens, so we only test one pair
1689

1690
        // Select a random worker pair using the policy
1691
        let (prefill, decode) = match self.select_pd_pair(None, None).await {
1692
1693
            Ok(pair) => pair,
            Err(e) => {
1694
1695
1696
1697
1698
                return (
                    StatusCode::SERVICE_UNAVAILABLE,
                    format!("No healthy worker pair available: {}", e),
                )
                    .into_response();
1699
1700
            }
        };
1701

1702
1703
        // Test prefill server's health_generate
        let prefill_url = format!("{}/health_generate", prefill.url());
1704
1705
1706
        let (prefill_result, decode_result) = tokio::join!(
            self.client.get(&prefill_url).send(),
            self.client
1707
                .get(format!("{}/health_generate", decode.url()))
1708
1709
                .send()
        );
1710

1711
1712
        // Check results
        let mut errors = Vec::new();
1713

1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
        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));
            }
        }
1732

1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
        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));
1746
1747
1748
            }
        }

1749
        if errors.is_empty() {
1750
1751
1752
1753
1754
1755
1756
1757
1758
            (
                StatusCode::OK,
                format!(
                    "Health generate passed on selected pair: prefill={}, decode={}",
                    prefill.url(),
                    decode.url()
                ),
            )
                .into_response()
1759
        } else {
1760
1761
1762
1763
1764
            (
                StatusCode::SERVICE_UNAVAILABLE,
                format!("Health generate failed: {:?}", errors),
            )
                .into_response()
1765
1766
1767
        }
    }

1768
    async fn get_server_info(&self, _req: Request<Body>) -> Response {
1769
        // Get info from the first decode server to match sglang's server info format
1770
        // Note: We use decode workers for server info to match expected format
1771
        self.proxy_to_first_prefill_worker("get_server_info", None)
1772
            .await
1773
1774
    }

1775
    async fn get_models(&self, req: Request<Body>) -> Response {
1776
        // Extract headers first to avoid Send issues
1777
        let headers = header_utils::copy_request_headers(&req);
1778

1779
        // Proxy to first prefill worker
1780
        self.proxy_to_first_prefill_worker("v1/models", Some(headers))
1781
            .await
1782
1783
    }

1784
    async fn get_model_info(&self, req: Request<Body>) -> Response {
1785
        // Extract headers first to avoid Send issues
1786
        let headers = header_utils::copy_request_headers(&req);
1787

1788
        // Proxy to first prefill worker
1789
1790
        self.proxy_to_first_prefill_worker("get_model_info", Some(headers))
            .await
1791
1792
    }

1793
1794
1795
1796
    async fn route_generate(
        &self,
        headers: Option<&HeaderMap>,
        body: &GenerateRequest,
1797
        model_id: Option<&str>,
1798
    ) -> Response {
1799
        // Extract parameters
1800
1801
        let is_stream = body.stream;
        let return_logprob = body.return_logprob;
1802

1803
        // Extract text for cache-aware routing
1804
        let request_text = if self.policies_need_request_text() {
1805
1806
1807
1808
            body.text
                .as_deref()
                .or_else(|| {
                    body.prompt.as_ref().and_then(|p| match p {
1809
1810
                        StringOrArray::String(s) => Some(s.as_str()),
                        StringOrArray::Array(v) => v.first().map(|s| s.as_str()),
1811
                    })
1812
                })
1813
                .map(|s| s.to_string())
1814
1815
1816
        } else {
            None
        };
1817

1818
        // Calculate batch size
1819
        let batch_size = Self::get_generate_batch_size(body);
1820

1821
1822
1823
1824
        // Create context
        let context = PDRequestContext {
            route: "/generate",
            batch_size,
1825
1826
            is_stream,
            return_logprob,
1827
            request_text,
1828
            model_id,
1829
1830
1831
1832
        };

        // Execute with retry and bootstrap injection
        self.execute_dual_dispatch(headers, body, context).await
1833
1834
    }

1835
1836
1837
1838
    async fn route_chat(
        &self,
        headers: Option<&HeaderMap>,
        body: &ChatCompletionRequest,
1839
        model_id: Option<&str>,
1840
    ) -> Response {
1841
        // Extract parameters
1842
1843
1844
        let is_stream = body.stream;
        let return_logprob = body.logprobs;

1845
        // Extract text for cache-aware routing
1846
1847
        let request_text = if self.policies_need_request_text() {
            body.messages.first().and_then(|msg| match msg {
1848
1849
1850
                ChatMessage::User { content, .. } => match content {
                    UserMessageContent::Text(text) => Some(text.clone()),
                    UserMessageContent::Parts(_) => None,
1851
                },
1852
                ChatMessage::System { content, .. } => Some(content.clone()),
1853
1854
1855
1856
1857
                _ => None,
            })
        } else {
            None
        };
1858

1859
        // Calculate batch size
1860
        let batch_size = Self::get_chat_batch_size(body);
1861

1862
1863
1864
1865
        // Create context
        let context = PDRequestContext {
            route: "/v1/chat/completions",
            batch_size,
1866
1867
            is_stream,
            return_logprob,
1868
            request_text,
1869
            model_id,
1870
1871
1872
1873
        };

        // Execute with retry and bootstrap injection
        self.execute_dual_dispatch(headers, body, context).await
1874
1875
    }

1876
1877
1878
1879
    async fn route_completion(
        &self,
        headers: Option<&HeaderMap>,
        body: &CompletionRequest,
1880
        model_id: Option<&str>,
1881
    ) -> Response {
1882
        // Extract parameters
1883
1884
1885
        let is_stream = body.stream;
        let return_logprob = body.logprobs.is_some();

1886
        // Extract text for cache-aware routing
1887
1888
        let request_text = if self.policies_need_request_text() {
            match &body.prompt {
1889
1890
                StringOrArray::String(s) => Some(s.clone()),
                StringOrArray::Array(v) => v.first().map(|s| s.to_string()),
1891
1892
1893
            }
        } else {
            None
1894
1895
        };

1896
        // Calculate batch size
1897
        let batch_size = Self::get_completion_batch_size(body);
1898

1899
1900
1901
1902
        // Create context
        let context = PDRequestContext {
            route: "/v1/completions",
            batch_size,
1903
1904
            is_stream,
            return_logprob,
1905
            request_text,
1906
            model_id,
1907
1908
1909
1910
        };

        // Execute with retry and bootstrap injection
        self.execute_dual_dispatch(headers, body, context).await
1911
1912
    }

1913
1914
1915
1916
    async fn route_responses(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &ResponsesRequest,
1917
        _model_id: Option<&str>,
1918
1919
1920
1921
1922
1923
1924
1925
    ) -> Response {
        (
            StatusCode::NOT_IMPLEMENTED,
            "Responses endpoint not implemented for PD router",
        )
            .into_response()
    }

1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
    async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
        (
            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()
    }

1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
    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()
1953
1954
    }

1955
1956
1957
1958
1959
1960
    async fn route_rerank(
        &self,
        headers: Option<&HeaderMap>,
        body: &RerankRequest,
        model_id: Option<&str>,
    ) -> Response {
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
        // Extract text for cache-aware routing
        let req_text = if self.policies_need_request_text() {
            Some(body.query.clone())
        } else {
            None
        };

        // Create context
        let context = PDRequestContext {
            route: "/v1/rerank",
            batch_size: None,
            is_stream: false,
            return_logprob: false,
            request_text: req_text,
1975
            model_id,
1976
1977
1978
1979
        };

        // Execute with retry and bootstrap injection
        self.execute_dual_dispatch(headers, body, context).await
1980
1981
    }

1982
    async fn flush_cache(&self) -> Response {
1983
1984
        // Process both prefill and decode workers
        let (prefill_results, prefill_errors) = self
1985
1986
1987
1988
1989
1990
1991
            .process_workers(
                WorkerType::Prefill {
                    bootstrap_port: None,
                },
                "Prefill",
                "flush_cache",
            )
1992
1993
            .await;
        let (decode_results, decode_errors) = self
1994
            .process_workers(WorkerType::Decode, "Decode", "flush_cache")
1995
1996
1997
1998
1999
2000
2001
            .await;

        // Combine results and errors
        let mut results = prefill_results;
        results.extend(decode_results);
        let mut errors = prefill_errors;
        errors.extend(decode_errors);
2002

2003
2004
2005
2006
2007
2008
        if errors.is_empty() {
            (
                StatusCode::OK,
                format!("Cache flushed successfully: {:?}", results),
            )
                .into_response()
2009
        } else {
2010
            (
2011
2012
2013
2014
2015
                StatusCode::PARTIAL_CONTENT,
                format!(
                    "Partial success. Results: {:?}, Errors: {:?}",
                    results, errors
                ),
2016
2017
            )
                .into_response()
2018
2019
2020
        }
    }

2021
2022
2023
    async fn get_worker_loads(&self) -> Response {
        let mut loads = HashMap::new();
        let mut errors = Vec::new();
2024

2025
        // Process prefill workers
2026
2027
2028
2029
2030
2031
2032
2033
        let prefill_urls = self.get_prefill_worker_urls();
        for worker_url in prefill_urls {
            match get_worker_load(&self.client, &worker_url).await {
                Some(load) => {
                    loads.insert(format!("prefill_{}", worker_url), load);
                }
                None => {
                    errors.push(format!("Failed to get load from prefill {}", worker_url));
2034
2035
2036
                }
            }
        }
2037

2038
        // Process decode workers
2039
2040
2041
2042
2043
2044
2045
2046
        let decode_urls = self.get_decode_worker_urls();
        for worker_url in decode_urls {
            match get_worker_load(&self.client, &worker_url).await {
                Some(load) => {
                    loads.insert(format!("decode_{}", worker_url), load);
                }
                None => {
                    errors.push(format!("Failed to get load from decode {}", worker_url));
2047
2048
2049
                }
            }
        }
2050

2051
2052
2053
2054
        let response_data = serde_json::json!({
            "loads": loads,
            "errors": errors
        });
2055

2056
        (StatusCode::OK, Json(response_data)).into_response()
2057
2058
2059
2060
2061
2062
    }

    fn router_type(&self) -> &'static str {
        "pd"
    }

2063
    fn readiness(&self) -> Response {
2064
        // PD router is ready if it has at least one healthy prefill AND one healthy decode worker
2065
2066
        let prefill_workers = self.worker_registry.get_prefill_workers();
        let decode_workers = self.worker_registry.get_decode_workers();
2067

2068
        let healthy_prefill_count = prefill_workers.iter().filter(|w| w.is_healthy()).count();
2069

2070
2071
2072
2073
        let healthy_decode_count = decode_workers.iter().filter(|w| w.is_healthy()).count();

        let total_prefill = prefill_workers.len();
        let total_decode = decode_workers.len();
2074
2075

        if healthy_prefill_count > 0 && healthy_decode_count > 0 {
2076
            Json(json!({
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
                "status": "ready",
                "prefill": {
                    "healthy": healthy_prefill_count,
                    "total": total_prefill
                },
                "decode": {
                    "healthy": healthy_decode_count,
                    "total": total_decode
                }
            }))
2087
            .into_response()
2088
2089
2090
2091
2092
2093
2094
2095
2096
        } else {
            let mut reasons = Vec::new();
            if healthy_prefill_count == 0 {
                reasons.push("no healthy prefill workers");
            }
            if healthy_decode_count == 0 {
                reasons.push("no healthy decode workers");
            }

2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
            (
                StatusCode::SERVICE_UNAVAILABLE,
                Json(serde_json::json!({
                    "status": "not_ready",
                    "reason": reasons.join(", "),
                    "prefill": {
                        "healthy": healthy_prefill_count,
                        "total": total_prefill
                    },
                    "decode": {
                        "healthy": healthy_decode_count,
                        "total": total_decode
                    }
                })),
            )
                .into_response()
2113
2114
2115
        }
    }
}
2116
2117
2118
2119

#[cfg(test)]
mod tests {
    use super::*;
2120
    use crate::core::WorkerType;
2121
2122

    fn create_test_pd_router() -> PDRouter {
2123
2124
2125
        let worker_registry = Arc::new(WorkerRegistry::new());
        let policy_registry =
            Arc::new(PolicyRegistry::new(crate::config::PolicyConfig::RoundRobin));
2126
2127

        PDRouter {
2128
2129
            worker_registry,
            policy_registry,
2130
2131
            worker_startup_timeout_secs: 5,
            worker_startup_check_interval_secs: 1,
2132
2133
            worker_loads: Arc::new(tokio::sync::watch::channel(HashMap::new()).1),
            load_monitor_handle: None,
2134
            client: Client::new(),
2135
            prefill_client: Client::new(),
2136
            prefill_drain_tx: mpsc::channel(100).0,
2137
            retry_config: RetryConfig::default(),
2138
            circuit_breaker_config: CircuitBreakerConfig::default(),
2139
2140
2141
2142
        }
    }

    fn create_test_worker(url: String, worker_type: WorkerType, healthy: bool) -> Box<dyn Worker> {
2143
2144
2145
        let worker = BasicWorkerBuilder::new(url)
            .worker_type(worker_type)
            .build();
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
        worker.set_healthy(healthy);
        Box::new(worker)
    }

    // ============= Worker Management Tests =============

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

        // Add a worker first
        let worker = create_test_worker(
            "http://localhost:8000".to_string(),
            WorkerType::Prefill {
                bootstrap_port: Some(8080),
            },
            true,
        );
2164
        router.worker_registry.register(Arc::from(worker));
2165
2166
2167

        // Try to add the same URL again - this would fail during health check in real scenario
        // For unit test, we test the duplicate check logic
2168
2169
2170
2171
        let exists = router
            .worker_registry
            .get_by_url("http://localhost:8000")
            .is_some();
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
        assert!(exists);
    }

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

        // Add servers first
        let worker1 = create_test_worker(
            "http://worker1".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            true,
        );
        let worker2 = create_test_worker(
            "http://worker2".to_string(),
            WorkerType::Prefill {
                bootstrap_port: Some(8080),
            },
            true,
        );

2195
2196
        router.worker_registry.register(Arc::from(worker1));
        router.worker_registry.register(Arc::from(worker2));
2197
2198
2199
2200
2201
2202
2203

        // Remove one
        let result = router.remove_prefill_server("http://worker1").await;

        assert!(result.is_ok());
        assert!(result.unwrap().contains("Successfully removed"));

2204
        let workers = router.worker_registry.get_prefill_workers();
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
        assert_eq!(workers.len(), 1);
        assert_eq!(workers[0].url(), "http://worker2");
    }

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

        let result = router.remove_prefill_server("http://nonexistent").await;

        assert!(result.is_err());
        match result.unwrap_err() {
            PDRouterError::WorkerNotFound { url } => {
                assert_eq!(url, "http://nonexistent");
            }
            _ => panic!("Expected WorkerNotFound error"),
        }
    }

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

        // Add server first
        let worker = create_test_worker("http://decode1".to_string(), WorkerType::Decode, true);
2230
        router.worker_registry.register(Arc::from(worker));
2231
2232
2233
2234
2235
2236

        let result = router.remove_decode_server("http://decode1").await;

        assert!(result.is_ok());
        assert!(result.unwrap().contains("Successfully removed"));

2237
        let workers = router.worker_registry.get_decode_workers();
2238
2239
2240
2241
2242
2243
        assert_eq!(workers.len(), 0);
    }

    // ============= Lock Error Handling Tests =============

    #[test]
2244
    fn test_registry_operations() {
2245
2246
        let router = create_test_pd_router();

2247
2248
2249
        // Test registry operations
        let workers = router.worker_registry.get_all();
        assert_eq!(workers.len(), 0);
2250

2251
2252
2253
2254
2255
2256
2257
2258
2259
        // Add a worker
        let worker = create_test_worker(
            "http://test".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            true,
        );
        router.worker_registry.register(Arc::from(worker));
2260

2261
2262
2263
2264
2265
        let workers = router.worker_registry.get_all();
        assert_eq!(workers.len(), 1);

        let prefill_workers = router.worker_registry.get_prefill_workers();
        assert_eq!(prefill_workers.len(), 1);
2266
2267
    }

2268
2269
2270
2271
2272
2273
2274
    // ============= Bootstrap Injection Tests =============
    // Note: These tests are commented out as we've moved to the optimized bootstrap injection
    // approach that doesn't use the Bootstrap trait on GenerateReqInput anymore.

    // TODO: Add new tests for the optimized bootstrap injection approach using
    // RequestWithBootstrap and BatchRequestWithBootstrap wrappers

2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
    // ============= Worker Selection Tests =============

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

        // Add mix of healthy and unhealthy workers
        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);

2299
2300
2301
        router.worker_registry.register(Arc::from(unhealthy_worker));
        router.worker_registry.register(Arc::from(healthy_worker));
        router.worker_registry.register(Arc::from(decode_worker));
2302

2303
        let result = router.select_pd_pair(None, None).await;
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316

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

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

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

2317
        let result = router.select_pd_pair(None, None).await;
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328

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

    // ============= Health Endpoints Tests =============

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

2329
        // Add healthy workers - create_test_worker returns Box<dyn Worker>, convert to Arc
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
        let prefill_worker = create_test_worker(
            "http://localhost:8000".to_string(),
            WorkerType::Prefill {
                bootstrap_port: None,
            },
            true,
        );
        let decode_worker = create_test_worker(
            "http://localhost:8001".to_string(),
            WorkerType::Decode,
            true,
        );

2343
2344
        router.worker_registry.register(Arc::from(prefill_worker));
        router.worker_registry.register(Arc::from(decode_worker));
2345
2346

        // Test health endpoint
2347
2348
2349
        let http_req = axum::http::Request::builder()
            .body(axum::body::Body::empty())
            .unwrap();
2350
        let response = router.health(http_req).await;
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362

        assert_eq!(response.status(), 200);

        // Test readiness endpoint
        let response = router.readiness();
        assert_eq!(response.status(), 200);
    }

    // ============= Load Monitoring Tests =============

    #[tokio::test]
    async fn test_load_monitor_updates() {
2363
        let power_of_two_policy = Arc::new(crate::policies::PowerOfTwoPolicy::new());
2364
        let mut router = create_test_pd_router();
2365
2366
2367
2368
2369
2370
2371
        // Set power_of_two policies in the registry
        router
            .policy_registry
            .set_prefill_policy(power_of_two_policy.clone());
        router
            .policy_registry
            .set_decode_policy(power_of_two_policy);
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418

        // Create load channel
        let (tx, rx) = tokio::sync::watch::channel(HashMap::new());
        router.worker_loads = Arc::new(rx);

        // Simulate load updates
        let mut loads = HashMap::new();
        loads.insert("http://worker1".to_string(), 10);
        loads.insert("http://worker2".to_string(), 5);

        let _ = tx.send(loads.clone());

        // Router should receive updates
        let received = router.worker_loads.borrow().clone();
        assert_eq!(received.get("http://worker1"), Some(&10));
        assert_eq!(received.get("http://worker2"), Some(&5));
    }

    // ============= Worker Load Tests =============

    #[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);

        // Create load guard for both workers
        let _guard =
            WorkerLoadGuard::new_multi(vec![prefill_worker.as_ref(), decode_worker.as_ref()]);

        // Load should be incremented
        assert_eq!(prefill_worker.load(), 1);
        assert_eq!(decode_worker.load(), 1);

        // Drop guard - load should decrement
        drop(_guard);

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

2419
2420
2421
2422
2423
2424
2425
    #[tokio::test]
    async fn test_streaming_load_tracking() {
        use futures_util::StreamExt;
        use tokio::time::{sleep, Duration};

        let router = create_test_pd_router();

2426
        // Add workers - create_test_worker returns Box<dyn Worker>, convert to Arc
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
        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);

2437
2438
2439
2440
2441
2442
2443
2444
2445
        router.worker_registry.register(Arc::from(prefill_worker));
        router.worker_registry.register(Arc::from(decode_worker));

        // Get references to the workers from registry
        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();
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491

        // Initially load should be 0
        assert_eq!(prefill_ref.load(), 0);
        assert_eq!(decode_ref.load(), 0);

        // Create a mock streaming response
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx);

        // Call create_streaming_response which should increment load
        let _response = router.create_streaming_response(
            stream.map(Ok),
            StatusCode::OK,
            None,
            false,
            None,
            None,
            prefill_ref.as_ref(),
            decode_ref.as_ref(),
        );

        // Load should be incremented immediately
        assert_eq!(prefill_ref.load(), 1);
        assert_eq!(decode_ref.load(), 1);

        // Send some data through the stream
        tx.send(bytes::Bytes::from("test data")).unwrap();

        // Give time for the spawned task to process
        sleep(Duration::from_millis(10)).await;

        // Load should still be 1 (streaming in progress)
        assert_eq!(prefill_ref.load(), 1);
        assert_eq!(decode_ref.load(), 1);

        // Close the stream
        drop(tx);

        // Give time for cleanup
        sleep(Duration::from_millis(100)).await;

        // Load should be decremented after streaming completes
        assert_eq!(prefill_ref.load(), 0);
        assert_eq!(decode_ref.load(), 0);
    }

2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
    // ============= Concurrent Operations Tests =============

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

        let mut handles = vec![];

        // Spawn tasks to add workers
        for i in 0..5 {
            let router_clone = Arc::clone(&router);
            let url = format!("http://worker{}", i);
            let handle = tokio::spawn(async move {
                let worker = create_test_worker(
                    url,
                    WorkerType::Prefill {
                        bootstrap_port: None,
                    },
                    true,
                );
2512
                router_clone.worker_registry.register(Arc::from(worker));
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
            });
            handles.push(handle);
        }

        // Wait for all tasks
        for handle in handles {
            let _ = handle.await;
        }

        // Check final state
2523
        let workers = router.worker_registry.get_prefill_workers();
2524
2525
2526
        assert_eq!(workers.len(), 5);
    }
}