router.rs 62 KB
Newer Older
1
use crate::config::types::RetryConfig;
2
use crate::core::{
3
    is_retryable_status, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, RetryExecutor,
4
    Worker, WorkerRegistry, WorkerType,
5
};
6
use crate::metrics::RouterMetrics;
7
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
8
use crate::protocols::spec::{
9
10
    ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, GenerationRequest,
    RerankRequest, RerankResponse, RerankResult, ResponsesRequest,
11
};
12
use crate::routers::header_utils;
13
use crate::routers::{RouterTrait, WorkerManagement};
14
use axum::body::to_bytes;
15
16
17
use axum::{
    body::Body,
    extract::Request,
18
19
20
    http::{
        header::CONTENT_LENGTH, header::CONTENT_TYPE, HeaderMap, HeaderValue, Method, StatusCode,
    },
21
22
23
24
    response::{IntoResponse, Response},
    Json,
};
use futures_util::StreamExt;
25
use reqwest::Client;
26
use std::collections::HashMap;
27
use std::sync::Arc;
28
use std::time::{Duration, Instant};
29
use tokio_stream::wrappers::UnboundedReceiverStream;
30
31
32
33
34
use tracing::{debug, error, info, warn};

/// Regular router that uses injected load balancing policies
#[derive(Debug)]
pub struct Router {
35
36
    worker_registry: Arc<WorkerRegistry>,
    policy_registry: Arc<PolicyRegistry>,
37
    client: Client,
38
39
    worker_startup_timeout_secs: u64,
    worker_startup_check_interval_secs: u64,
40
    dp_aware: bool,
41
    #[allow(dead_code)]
42
    api_key: Option<String>,
43
    retry_config: RetryConfig,
44
    circuit_breaker_config: CircuitBreakerConfig,
45
46
47
48
49
    _worker_loads: Arc<tokio::sync::watch::Receiver<HashMap<String, isize>>>,
    _load_monitor_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
}

impl Router {
50
    /// Create a new router with injected policy and client
51
52
53
54
55
56
57
    pub async fn new(ctx: &Arc<crate::server::AppContext>) -> Result<Self, String> {
        let workers = ctx.worker_registry.get_workers_filtered(
            None, // any model
            Some(WorkerType::Regular),
            Some(ConnectionMode::Http),
            false, // include all workers
        );
58

59
60
        // Update active workers gauge
        RouterMetrics::set_active_workers(workers.len());
61

62
63
        // Get worker URLs for monitoring
        let worker_urls: Vec<String> = workers.iter().map(|w| w.url().to_string()).collect();
64

65
        // Convert config CircuitBreakerConfig to core CircuitBreakerConfig
66
        let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config();
67
68
69
        let core_cb_config = CircuitBreakerConfig {
            failure_threshold: circuit_breaker_config.failure_threshold,
            success_threshold: circuit_breaker_config.success_threshold,
70
71
            timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs),
            window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs),
72
73
        };

74
        // Cache-aware policies are initialized in WorkerInitializer
75
76
77
78
        // Setup load monitoring for PowerOfTwo policy
        let (tx, rx) = tokio::sync::watch::channel(HashMap::new());
        let worker_loads = Arc::new(rx);

79
80
81
        // Get default policy to check if we need load monitoring
        let default_policy = ctx.policy_registry.get_default_policy();

82
83
        // Check if default policy is power_of_two for load monitoring
        let load_monitor_handle = if default_policy.name() == "power_of_two" {
84
            let monitor_urls = worker_urls.clone();
85
86
87
88
89
90
91
92
            let monitor_api_keys = monitor_urls
                .iter()
                .map(|url| {
                    ctx.worker_registry
                        .get_by_url(url)
                        .and_then(|w| w.api_key().clone())
                })
                .collect::<Vec<Option<String>>>();
93
            let monitor_interval = ctx.router_config.worker_startup_check_interval_secs;
94
            let policy_clone = default_policy.clone();
95
            let client_clone = ctx.client.clone();
96
97

            Some(Arc::new(tokio::spawn(async move {
98
99
                Self::monitor_worker_loads(
                    monitor_urls,
100
                    monitor_api_keys,
101
102
103
104
105
106
                    tx,
                    monitor_interval,
                    policy_clone,
                    client_clone,
                )
                .await;
107
108
109
110
111
112
            })))
        } else {
            None
        };

        Ok(Router {
113
114
            worker_registry: ctx.worker_registry.clone(),
            policy_registry: ctx.policy_registry.clone(),
115
116
117
118
119
120
121
122
            client: ctx.client.clone(),
            worker_startup_timeout_secs: ctx.router_config.worker_startup_timeout_secs,
            worker_startup_check_interval_secs: ctx
                .router_config
                .worker_startup_check_interval_secs,
            dp_aware: ctx.router_config.dp_aware,
            api_key: ctx.router_config.api_key.clone(),
            retry_config: ctx.router_config.effective_retry_config(),
123
            circuit_breaker_config: core_cb_config,
124
125
126
127
128
129
130
            _worker_loads: worker_loads,
            _load_monitor_handle: load_monitor_handle,
        })
    }

    /// Get the current list of worker URLs
    pub fn get_worker_urls(&self) -> Vec<String> {
131
132
133
134
135
        self.worker_registry.get_all_urls()
    }

    /// Get worker URLs for a specific model
    pub fn get_worker_urls_for_model(&self, model_id: Option<&str>) -> Vec<String> {
136
137
138
139
140
141
        let workers = self.worker_registry.get_workers_filtered(
            model_id,
            Some(WorkerType::Regular),
            Some(ConnectionMode::Http),
            false, // get all workers
        );
142
        workers.iter().map(|w| w.url().to_string()).collect()
143
144
    }

145
    pub async fn wait_for_healthy_workers(
146
        worker_urls: &[String],
147
148
        worker_startup_timeout_secs: u64,
        worker_startup_check_interval_secs: u64,
149
    ) -> Result<(), String> {
150
151
152
153
154
155
        if worker_urls.is_empty() {
            return Err(
                "Timeout waiting for workers to become healthy: no workers provided".to_string(),
            );
        }

156
        // Perform health check asynchronously
157
158
159
160
161
162
        Self::wait_for_healthy_workers_async(
            worker_urls,
            worker_startup_timeout_secs,
            worker_startup_check_interval_secs,
        )
        .await
163
164
165
166
    }

    async fn wait_for_healthy_workers_async(
        worker_urls: &[String],
167
168
        worker_startup_timeout_secs: u64,
        worker_startup_check_interval_secs: u64,
169
170
171
172
    ) -> Result<(), String> {
        info!(
            "Waiting for {} workers to become healthy (timeout: {}s)",
            worker_urls.len(),
173
            worker_startup_timeout_secs
174
175
        );

176
        let start_time = std::time::Instant::now();
177
178
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(2))
179
180
181
182
            .build()
            .map_err(|e| format!("Failed to create HTTP client: {}", e))?;

        loop {
183
            if start_time.elapsed() > Duration::from_secs(worker_startup_timeout_secs) {
184
185
                error!(
                    "Timeout {}s waiting for workers {:?} to become healthy. Please set --router-worker-startup-timeout-secs (sglang_router.launch_server) or --worker-startup-timeout-secs (sglang_worker.router) to a larger value",
186
                    worker_startup_timeout_secs, worker_urls
187
188
189
                );
                return Err(format!(
                    "Timeout {}s waiting for workers {:?} to become healthy. Please set --router-worker-startup-timeout-secs (sglang_router.launch_server) or --worker-startup-timeout-secs (sglang_worker.router) to a larger value",
190
                    worker_startup_timeout_secs, worker_urls
191
192
193
                ));
            }

194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
            // Perform all health checks concurrently
            let mut health_checks = Vec::new();
            for url in worker_urls {
                let client_clone = client.clone();
                let url_clone = url.clone();

                let check_health = tokio::spawn(async move {
                    let health_url = format!("{}/health", url_clone);
                    match client_clone.get(&health_url).send().await {
                        Ok(res) => {
                            if res.status().is_success() {
                                None
                            } else {
                                Some((url_clone, format!("status: {}", res.status())))
                            }
                        }
                        Err(_) => Some((url_clone, "not ready".to_string())),
                    }
                });

                health_checks.push(check_health);
            }

            // Wait for all health checks to complete
            let results = futures::future::join_all(health_checks).await;

220
221
222
            let mut all_healthy = true;
            let mut unhealthy_workers = Vec::new();

223
224
225
226
            for result in results {
                match result {
                    Ok(None) => {
                        // Worker is healthy
227
                    }
228
                    Ok(Some((url, reason))) => {
229
                        all_healthy = false;
230
231
232
233
234
235
                        unhealthy_workers.push((url, reason));
                    }
                    Err(e) => {
                        all_healthy = false;
                        unhealthy_workers
                            .push(("unknown".to_string(), format!("task error: {}", e)));
236
237
238
239
240
                    }
                }
            }

            if all_healthy {
241
                info!("All {} workers are healthy", worker_urls.len());
242
243
                return Ok(());
            } else {
244
                debug!(
245
                    "Waiting for {} workers to become healthy ({} unhealthy: {:?})",
246
                    worker_urls.len(),
247
248
                    unhealthy_workers.len(),
                    unhealthy_workers
249
                );
250
                tokio::time::sleep(Duration::from_secs(worker_startup_check_interval_secs)).await;
251
252
253
254
            }
        }
    }

255
256
    fn get_worker_dp_size(worker_url: &str, api_key: &Option<String>) -> Result<usize, String> {
        let sync_client = reqwest::blocking::Client::new();
257
        let mut req_builder = sync_client.get(format!("{}/get_server_info", worker_url));
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
        if let Some(key) = api_key {
            req_builder = req_builder.bearer_auth(key);
        }

        match req_builder.send() {
            Ok(res) => {
                if res.status().is_success() {
                    let server_info = res
                        .text()
                        .map_err(|e| format!("failed to read text from response: {}", e))?;

                    let server_info: serde_json::Value = serde_json::from_str(&server_info)
                        .map_err(|e| format!("failed to decode JSON: {}", e))?;

                    let dp_size = server_info
                        .get("dp_size")
                        .and_then(|v| v.as_u64())
                        .ok_or_else(|| String::from("dp_size not found or not an u64"))?;

                    Ok(if dp_size > usize::MAX as u64 {
                        return Err(format!("dp_size is too large: {}", dp_size));
                    } else {
                        dp_size as usize
                    })
                } else {
                    Err(format!("unexpected status code: {}", res.status()))
                }
            }
            Err(e) => Err(format!("error response: {}", e)),
        }
    }

    // Given a list of workers, return a list of workers with dp_rank as suffix
    fn get_dp_aware_workers(
        worker_urls: &[String],
        api_key: &Option<String>,
    ) -> Result<Vec<String>, String> {
        let mut dp_aware_workers: Vec<String> = Vec::new();

        for url in worker_urls {
            match Self::get_worker_dp_size(url, api_key) {
                Ok(dp_size) => {
                    for i in 0..dp_size {
                        dp_aware_workers.push(format!("{}@{}", url, i));
                    }
                }
                Err(e) => return Err(format!("Failed to get DP size for {}: {}", url, e)),
            }
        }

        Ok(dp_aware_workers)
    }

311
    fn select_first_worker(&self) -> Result<String, String> {
312
313
        let workers = self.worker_registry.get_all();
        if workers.is_empty() {
314
315
            Err("No workers are available".to_string())
        } else {
316
317
318
319
            Ok(workers[0].url().to_string())
        }
    }

320
    pub async fn send_health_check(&self, worker_url: &str) -> Response {
321
        let health_url = if self.dp_aware {
322
            // Need to extract the URL from "http://host:port@dp_rank"
323
324
            match Self::extract_dp_rank(worker_url) {
                Ok((worker_url_prefix, _dp_rank)) => worker_url_prefix,
325
                Err(e) => {
326
327
328
329
330
331
                    error!("Failed to extract dp_rank for health check: {}", e);
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Failed to extract dp_rank: {}", e),
                    )
                        .into_response();
332
                }
333
            }
334
335
336
337
        } else {
            worker_url
        };

338
        let request_builder = self.client.get(format!("{}/health", health_url));
339
340
341

        let response = match request_builder.send().await {
            Ok(res) => {
342
343
                let status = StatusCode::from_u16(res.status().as_u16())
                    .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
344
345

                match res.bytes().await {
346
                    Ok(body) => (status, body).into_response(),
347
348
                    Err(e) => {
                        error!(
349
                            worker_url = %health_url,
350
                            error = %e,
351
                            "Failed to read health response body"
352
                        );
353
354
355
356
357
                        (
                            StatusCode::INTERNAL_SERVER_ERROR,
                            format!("Failed to read response body: {}", e),
                        )
                            .into_response()
358
                    }
359
360
                }
            }
361
362
            Err(e) => {
                error!(
363
                    worker_url = %health_url,
364
                    error = %e,
365
                    "Failed to send health request to worker"
366
                );
367
368
369
370
371
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Failed to send request to worker {}: {}", health_url, e),
                )
                    .into_response()
372
            }
373
374
        };

375
        // Don't record metrics for health checks
376
377
378
        response
    }

379
    // Helper method to proxy GET requests to the first available worker
380
    async fn proxy_get_request(&self, req: Request<Body>, endpoint: &str) -> Response {
381
        let headers = header_utils::copy_request_headers(&req);
382
383
384

        match self.select_first_worker() {
            Ok(worker_url) => {
385
                let mut request_builder = self.client.get(format!("{}/{}", worker_url, endpoint));
386
                for (name, value) in headers {
387
388
                    let name_lc = name.to_lowercase();
                    if name_lc != "content-type" && name_lc != "content-length" {
389
390
391
                        request_builder = request_builder.header(name, value);
                    }
                }
392

393
394
395
396
                match request_builder.send().await {
                    Ok(res) => {
                        let status = StatusCode::from_u16(res.status().as_u16())
                            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
397
398
399
400
401

                        // Preserve headers from backend
                        let response_headers =
                            header_utils::preserve_response_headers(res.headers());

402
                        match res.bytes().await {
403
404
405
406
407
408
                            Ok(body) => {
                                let mut response = Response::new(axum::body::Body::from(body));
                                *response.status_mut() = status;
                                *response.headers_mut() = response_headers;
                                response
                            }
409
410
411
412
413
                            Err(e) => (
                                StatusCode::INTERNAL_SERVER_ERROR,
                                format!("Failed to read response: {}", e),
                            )
                                .into_response(),
414
415
                        }
                    }
416
417
418
419
420
                    Err(e) => (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Request failed: {}", e),
                    )
                        .into_response(),
421
422
                }
            }
423
            Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(),
424
425
426
        }
    }

427
428
429
430
431
432
    /// Select worker for a specific model considering circuit breaker state
    fn select_worker_for_model(
        &self,
        model_id: Option<&str>,
        text: Option<&str>,
    ) -> Option<Arc<dyn Worker>> {
433
434
435
436
437
438
439
        // Get workers for the specified model O(1), filtered by connection mode
        let workers = self.worker_registry.get_workers_filtered(
            model_id,
            Some(WorkerType::Regular),
            Some(ConnectionMode::Http),
            false, // get all workers, we'll filter by is_available() next
        );
440
441

        let available: Vec<Arc<dyn Worker>> = workers
442
443
            .iter()
            .filter(|w| w.is_available())
444
            .cloned()
445
446
447
448
            .collect();
        if available.is_empty() {
            return None;
        }
449
450
451
452
453
454
455
456
457

        // Get the appropriate policy for this model
        let policy = match model_id {
            Some(model) => self.policy_registry.get_policy_or_default(model),
            None => self.policy_registry.get_default_policy(),
        };

        let idx = policy.select_worker(&available, text)?;
        Some(available[idx].clone())
458
459
    }

460
    pub async fn route_typed_request<T: GenerationRequest + serde::Serialize + Clone>(
461
        &self,
462
        headers: Option<&HeaderMap>,
463
464
        typed_req: &T,
        route: &str,
465
        model_id: Option<&str>,
466
    ) -> Response {
467
        let start = Instant::now();
468
469
470
471
472
473
474
        let is_stream = typed_req.is_stream();
        let text = typed_req.extract_text_for_routing();

        let response = RetryExecutor::execute_response_with_retry(
            &self.retry_config,
            // operation per attempt
            |_: u32| async {
475
                let worker = match self.select_worker_for_model(model_id, Some(&text)) {
476
477
478
479
480
481
482
483
484
485
                    Some(w) => w,
                    None => {
                        RouterMetrics::record_request_error(route, "no_available_workers");
                        return (
                            StatusCode::SERVICE_UNAVAILABLE,
                            "No available workers (all circuits open or unhealthy)",
                        )
                            .into_response();
                    }
                };
486

487
                // Optional load tracking for cache-aware policy
488
489
490
491
492
493
494
                // Get the policy for this model to check if it's cache-aware
                let policy = match model_id {
                    Some(model) => self.policy_registry.get_policy_or_default(model),
                    None => self.policy_registry.get_default_policy(),
                };

                let load_incremented = if policy.name() == "cache_aware" {
495
496
497
                    worker.increment_load();
                    RouterMetrics::set_running_requests(worker.url(), worker.load());
                    true
498
499
500
501
                } else {
                    false
                };

502
503
                // Keep a clone for potential cleanup on retry
                let worker_for_cleanup = if load_incremented {
504
                    Some(worker.clone())
505
506
507
508
                } else {
                    None
                };

509
510
                let response = self
                    .send_typed_request(
511
                        headers,
512
513
                        typed_req,
                        route,
514
                        worker.url(),
515
516
517
518
519
                        is_stream,
                        load_incremented,
                    )
                    .await;

520
                worker.record_outcome(response.status().is_success());
521
522
523
524
525
526
527
528
529
530
531
532
533

                // For retryable failures, we need to decrement load since send_typed_request
                // won't have done it (it only decrements on success or non-retryable failures)
                if is_retryable_status(response.status()) && load_incremented {
                    if let Some(cleanup_worker) = worker_for_cleanup {
                        cleanup_worker.decrement_load();
                        RouterMetrics::set_running_requests(
                            cleanup_worker.url(),
                            cleanup_worker.load(),
                        );
                    }
                }

534
535
536
                response
            },
            // should_retry predicate
537
            |res, _attempt| is_retryable_status(res.status()),
538
539
540
541
542
543
544
            // on_backoff hook
            |delay, attempt| {
                RouterMetrics::record_retry(route);
                RouterMetrics::record_retry_backoff_duration(delay, attempt);
            },
            // on_exhausted hook
            || RouterMetrics::record_retries_exhausted(route),
545
        )
546
547
548
549
550
551
        .await;

        if response.status().is_success() {
            let duration = start.elapsed();
            RouterMetrics::record_request(route);
            RouterMetrics::record_generate_duration(duration);
552
        } else if !is_retryable_status(response.status()) {
553
            RouterMetrics::record_request_error(route, "non_retryable_error");
554
        }
555
556

        response
557
558
    }

559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
    // Helper: return base worker URL (strips DP suffix when enabled)
    fn worker_base_url(&self, worker_url: &str) -> String {
        if self.dp_aware {
            if let Ok((prefix, _)) = Self::extract_dp_rank(worker_url) {
                return prefix.to_string();
            }
        }
        worker_url.to_string()
    }

    // Generic simple routing for GET/POST without JSON body
    async fn route_simple_request(
        &self,
        headers: Option<&HeaderMap>,
        endpoint: &str,
        method: Method,
    ) -> Response {
        // TODO: currently the sglang worker is using in-memory state management, so this implementation has to fan out to all workers.
        // Eventually, we need to have router to manage the chat history with a proper database, will update this implementation accordingly.
        let worker_urls = self.get_worker_urls();
        if worker_urls.is_empty() {
            return (StatusCode::SERVICE_UNAVAILABLE, "No available workers").into_response();
        }

        let mut last_response: Option<Response> = None;
        for worker_url in worker_urls {
            let base = self.worker_base_url(&worker_url);

            let url = format!("{}/{}", base, endpoint);
            let mut request_builder = match method {
                Method::GET => self.client.get(url),
                Method::POST => self.client.post(url),
                _ => {
                    return (
                        StatusCode::METHOD_NOT_ALLOWED,
                        "Unsupported method for simple routing",
                    )
                        .into_response()
                }
            };

            if let Some(hdrs) = headers {
                for (name, value) in hdrs {
                    let name_lc = name.as_str().to_lowercase();
                    if name_lc != "content-type" && name_lc != "content-length" {
                        request_builder = request_builder.header(name, value);
                    }
                }
            }

            match request_builder.send().await {
                Ok(res) => {
                    let status = StatusCode::from_u16(res.status().as_u16())
                        .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                    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() = status;
                            *response.headers_mut() = response_headers;
                            if status.is_success() {
                                return response;
                            }
                            last_response = Some(response);
                        }
                        Err(e) => {
                            last_response = Some(
                                (
                                    StatusCode::INTERNAL_SERVER_ERROR,
                                    format!("Failed to read response: {}", e),
                                )
                                    .into_response(),
                            );
                        }
                    }
                }
                Err(e) => {
                    last_response = Some(
                        (
                            StatusCode::INTERNAL_SERVER_ERROR,
                            format!("Request failed: {}", e),
                        )
                            .into_response(),
                    );
                }
            }
        }

        last_response
            .unwrap_or_else(|| (StatusCode::BAD_GATEWAY, "No worker response").into_response())
    }

    // Route a GET request with provided headers to a specific endpoint
    async fn route_get_request(&self, headers: Option<&HeaderMap>, endpoint: &str) -> Response {
        self.route_simple_request(headers, endpoint, Method::GET)
            .await
    }

    // Route a POST request with empty body to a specific endpoint
    async fn route_post_empty_request(
        &self,
        headers: Option<&HeaderMap>,
        endpoint: &str,
    ) -> Response {
        self.route_simple_request(headers, endpoint, Method::POST)
            .await
    }

667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
    // TODO (rui): Better accommodate to the Worker abstraction
    fn extract_dp_rank(worker_url: &str) -> Result<(&str, usize), String> {
        let parts: Vec<&str> = worker_url.split('@').collect();
        if parts.len() != 2 {
            return Err(format!("invalid worker_url format: {}", worker_url));
        }

        // Parse the second part (dp_rank) into an integer
        match parts[1].parse::<usize>() {
            Ok(dp_rank) => Ok((parts[0], dp_rank)),
            Err(_) => Err(format!(
                "failed to parse dp_rank from worker_url: {}",
                worker_url
            )),
        }
    }

684
685
686
    // Send typed request directly without conversion
    async fn send_typed_request<T: serde::Serialize>(
        &self,
687
        headers: Option<&HeaderMap>,
688
689
690
691
692
        typed_req: &T,
        route: &str,
        worker_url: &str,
        is_stream: bool,
        load_incremented: bool, // Whether load was incremented for this request
693
    ) -> Response {
694
695
696
697
698
        let mut request_builder = if self.dp_aware {
            let (worker_url_prefix, dp_rank) = match Self::extract_dp_rank(worker_url) {
                Ok(tup) => tup,
                Err(e) => {
                    error!("Failed to extract dp_rank: {}", e);
699
700
701
702
703
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Failed to extract dp_rank: {}", e),
                    )
                        .into_response();
704
705
706
707
708
709
710
                }
            };

            // Parse the request body
            let mut json_val = match serde_json::to_value(typed_req) {
                Ok(j) => j,
                Err(e) => {
711
712
713
714
715
                    return (
                        StatusCode::BAD_REQUEST,
                        format!("Convert into serde_json::Value failed: {}", e),
                    )
                        .into_response();
716
717
718
719
720
721
722
723
724
725
726
727
728
729
                }
            };

            // Insert the data_parallel_rank field
            if let Some(map) = json_val.as_object_mut() {
                map.insert(
                    String::from("data_parallel_rank"),
                    serde_json::json!(dp_rank),
                );
                debug!(
                    "Modified request body: {}",
                    serde_json::to_string(&json_val).unwrap_or(String::from("ERR"))
                );
            } else {
730
731
732
733
734
                return (
                    StatusCode::BAD_REQUEST,
                    "Failed to insert the data_parallel_rank field into the request body",
                )
                    .into_response();
735
736
            }

737
            self.client
738
739
740
                .post(format!("{}{}", worker_url_prefix, route))
                .json(&json_val)
        } else {
741
            self.client
742
743
744
                .post(format!("{}{}", worker_url, route))
                .json(typed_req) // Use json() directly with typed request
        };
745

746
747
748
749
        // Copy all headers from original request if provided
        if let Some(headers) = headers {
            for (name, value) in headers {
                // Skip Content-Type and Content-Length as .json() sets them
750
                if *name != CONTENT_TYPE && *name != CONTENT_LENGTH {
751
752
                    request_builder = request_builder.header(name, value);
                }
753
754
755
756
757
758
            }
        }

        let res = match request_builder.send().await {
            Ok(res) => res,
            Err(e) => {
759
760
761
762
                error!(
                    "Failed to send typed request worker_url={} route={} error={}",
                    worker_url, route, e
                );
763
764
765

                // Decrement load on error if it was incremented
                if load_incremented {
766
767
768
                    if let Some(worker) = self.worker_registry.get_by_url(worker_url) {
                        worker.decrement_load();
                        RouterMetrics::set_running_requests(worker_url, worker.load());
769
770
771
                    }
                }

772
773
774
775
776
                return (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Request failed: {}", e),
                )
                    .into_response();
777
778
779
            }
        };

780
781
        let status = StatusCode::from_u16(res.status().as_u16())
            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
782
783

        if !is_stream {
784
            // For non-streaming requests, preserve headers
785
            let response_headers = header_utils::preserve_response_headers(res.headers());
786

787
            let response = match res.bytes().await {
788
789
790
791
792
793
                Ok(body) => {
                    let mut response = Response::new(axum::body::Body::from(body));
                    *response.status_mut() = status;
                    *response.headers_mut() = response_headers;
                    response
                }
794
                Err(e) => {
795
796
                    // IMPORTANT: Decrement load on error before returning
                    if load_incremented {
797
798
799
                        if let Some(worker) = self.worker_registry.get_by_url(worker_url) {
                            worker.decrement_load();
                            RouterMetrics::set_running_requests(worker_url, worker.load());
800
801
802
                        }
                    }

803
                    let error_msg = format!("Failed to get response body: {}", e);
804
                    (StatusCode::INTERNAL_SERVER_ERROR, error_msg).into_response()
805
806
807
808
                }
            };

            // Decrement load counter for non-streaming requests if it was incremented
809
            if load_incremented {
810
811
812
                if let Some(worker) = self.worker_registry.get_by_url(worker_url) {
                    worker.decrement_load();
                    RouterMetrics::set_running_requests(worker_url, worker.load());
813
814
815
816
817
818
                }
            }

            response
        } else if load_incremented {
            // For streaming with load tracking, we need to manually decrement when done
819
            let registry = Arc::clone(&self.worker_registry);
820
821
            let worker_url = worker_url.to_string();

822
823
824
825
826
            // Preserve headers for streaming response
            let mut response_headers = header_utils::preserve_response_headers(res.headers());
            // Ensure we set the correct content-type for SSE
            response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));

827
828
829
830
831
832
            let stream = res.bytes_stream();
            let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

            // Spawn task to forward stream and detect completion
            tokio::spawn(async move {
                let mut stream = stream;
833
                let mut decremented = false;
834
835
836
837
838
839
840
841
842
                while let Some(chunk) = stream.next().await {
                    match chunk {
                        Ok(bytes) => {
                            // Check for stream end marker
                            if bytes
                                .as_ref()
                                .windows(12)
                                .any(|window| window == b"data: [DONE]")
                            {
843
844
845
846
                                if let Some(worker) = registry.get_by_url(&worker_url) {
                                    worker.decrement_load();
                                    RouterMetrics::set_running_requests(&worker_url, worker.load());
                                    decremented = true;
847
848
                                }
                            }
849
850
851
852
853
854
855
856
857
858
                            if tx.send(Ok(bytes)).is_err() {
                                break;
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(format!("Stream error: {}", e)));
                            break;
                        }
                    }
                }
859
                if !decremented {
860
861
862
                    if let Some(worker) = registry.get_by_url(&worker_url) {
                        worker.decrement_load();
                        RouterMetrics::set_running_requests(&worker_url, worker.load());
863
864
                    }
                }
865
866
867
868
869
870
871
            });

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

            let mut response = Response::new(body);
            *response.status_mut() = status;
872
            *response.headers_mut() = response_headers;
873
            response
874
875
        } else {
            // For requests without load tracking, just stream
876
877
878
879
880
            // Preserve headers for streaming response
            let mut response_headers = header_utils::preserve_response_headers(res.headers());
            // Ensure we set the correct content-type for SSE
            response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));

881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
            let stream = res.bytes_stream();
            let (tx, rx) = tokio::sync::mpsc::unbounded_channel();

            // Spawn task to forward stream
            tokio::spawn(async move {
                let mut stream = stream;
                while let Some(chunk) = stream.next().await {
                    match chunk {
                        Ok(bytes) => {
                            if tx.send(Ok(bytes)).is_err() {
                                break;
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(format!("Stream error: {}", e)));
                            break;
                        }
                    }
                }
            });

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

            let mut response = Response::new(body);
            *response.status_mut() = status;
907
            *response.headers_mut() = response_headers;
908
            response
909
910
911
        }
    }

912
913
914
915
916
    pub async fn add_worker(
        &self,
        worker_url: &str,
        api_key: &Option<String>,
    ) -> Result<String, String> {
917
918
        let start_time = std::time::Instant::now();
        let client = reqwest::Client::builder()
919
            .timeout(Duration::from_secs(self.worker_startup_timeout_secs))
920
921
922
923
            .build()
            .map_err(|e| format!("Failed to create HTTP client: {}", e))?;

        loop {
924
            if start_time.elapsed() > Duration::from_secs(self.worker_startup_timeout_secs) {
925
926
                error!(
                    "Timeout {}s waiting for worker {} to become healthy. Please set --router-worker-startup-timeout-secs (sglang_router.launch_server) or --worker-startup-timeout-secs (sglang_worker.router) to a larger value",
927
                    self.worker_startup_timeout_secs, worker_url
928
929
930
                );
                return Err(format!(
                    "Timeout {}s waiting for worker {} to become healthy. Please set --router-worker-startup-timeout-secs (sglang_router.launch_server) or --worker-startup-timeout-secs (sglang_worker.router) to a larger value",
931
                    self.worker_startup_timeout_secs, worker_url
932
933
934
                ));
            }

935
            match client.get(format!("{}/health", worker_url)).send().await {
936
937
                Ok(res) => {
                    if res.status().is_success() {
938
939
940
941
                        if self.dp_aware {
                            // Need to contact the worker to extract the dp_size,
                            // and add them as multiple workers
                            let url_vec = vec![String::from(worker_url)];
942
                            let dp_url_vec = Self::get_dp_aware_workers(&url_vec, api_key)
943
944
945
                                .map_err(|e| format!("Failed to get dp-aware workers: {}", e))?;
                            let mut worker_added: bool = false;
                            for dp_url in &dp_url_vec {
946
                                if self.worker_registry.get_by_url(dp_url).is_some() {
947
948
949
950
                                    warn!("Worker {} already exists", dp_url);
                                    continue;
                                }
                                info!("Added worker: {}", dp_url);
951
                                // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
952
953
954
955
956
957
958
959
960
961
962
963
                                let new_worker_builder =
                                    BasicWorkerBuilder::new(dp_url.to_string())
                                        .worker_type(WorkerType::Regular)
                                        .circuit_breaker_config(
                                            self.circuit_breaker_config.clone(),
                                        );

                                let new_worker = if let Some(api_key) = api_key {
                                    new_worker_builder.api_key(api_key).build()
                                } else {
                                    new_worker_builder.build()
                                };
964
965
966
967
968
969

                                let worker_arc = Arc::new(new_worker);
                                self.worker_registry.register(worker_arc.clone());

                                // Notify PolicyRegistry about the new worker
                                let model_id = worker_arc.model_id();
970
971
972
                                self.policy_registry.on_worker_added(model_id, None);

                                // Initialize cache-aware policy if applicable
973
974
975
976
977
978
                                let model_workers = self.worker_registry.get_workers_filtered(
                                    Some(model_id),
                                    Some(WorkerType::Regular),
                                    Some(ConnectionMode::Http),
                                    false,
                                );
979
980
                                self.policy_registry
                                    .init_cache_aware_policy(model_id, &model_workers);
981

982
983
984
985
986
987
                                worker_added = true;
                            }
                            if !worker_added {
                                return Err(format!("No worker added for {}", worker_url));
                            }
                        } else {
988
                            if self.worker_registry.get_by_url(worker_url).is_some() {
989
990
991
                                return Err(format!("Worker {} already exists", worker_url));
                            }
                            info!("Added worker: {}", worker_url);
992

993
                            // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
994
995
996
997
998
999
1000
1001
1002
1003
                            let new_worker_builder =
                                BasicWorkerBuilder::new(worker_url.to_string())
                                    .worker_type(WorkerType::Regular)
                                    .circuit_breaker_config(self.circuit_breaker_config.clone());

                            let new_worker = if let Some(api_key) = api_key {
                                new_worker_builder.api_key(api_key).build()
                            } else {
                                new_worker_builder.build()
                            };
1004
1005
1006
1007
1008
1009

                            let worker_arc = Arc::new(new_worker);
                            self.worker_registry.register(worker_arc.clone());

                            // Notify PolicyRegistry about the new worker
                            let model_id = worker_arc.model_id();
1010
1011
1012
                            self.policy_registry.on_worker_added(model_id, None);

                            // Initialize cache-aware policy if applicable
1013
1014
1015
1016
1017
1018
                            let model_workers = self.worker_registry.get_workers_filtered(
                                Some(model_id),
                                Some(WorkerType::Regular),
                                Some(ConnectionMode::Http),
                                false,
                            );
1019
1020
                            self.policy_registry
                                .init_cache_aware_policy(model_id, &model_workers);
1021
1022
                        }

1023
1024
                        RouterMetrics::set_active_workers(self.worker_registry.get_all().len());

1025
1026
                        return Ok(format!("Successfully added worker: {}", worker_url));
                    } else {
1027
1028
                        debug!(
                            "Worker {} health check pending - status: {}",
1029
1030
1031
1032
1033
1034
1035
1036
1037
                            worker_url,
                            res.status()
                        );
                        // if the url does not have http or https prefix, warn users
                        if !worker_url.starts_with("http://") && !worker_url.starts_with("https://")
                        {
                            warn!("The worker url {} does not have http or https prefix. Please add the prefix to the url.", worker_url);
                        }

1038
1039
1040
1041
                        tokio::time::sleep(Duration::from_secs(
                            self.worker_startup_check_interval_secs,
                        ))
                        .await;
1042
1043
1044
1045
                        continue;
                    }
                }
                Err(e) => {
1046
                    debug!("Worker {} health check pending - error: {}", worker_url, e);
1047
1048
1049
1050
1051
1052

                    // if the url does not have http or https prefix, warn users
                    if !worker_url.starts_with("http://") && !worker_url.starts_with("https://") {
                        warn!("The worker url {} does not have http or https prefix. Please add the prefix to the url.", worker_url);
                    }

1053
1054
1055
1056
                    tokio::time::sleep(Duration::from_secs(
                        self.worker_startup_check_interval_secs,
                    ))
                    .await;
1057
1058
1059
1060
1061
1062
1063
                    continue;
                }
            }
        }
    }

    pub fn remove_worker(&self, worker_url: &str) {
1064
1065
1066
1067
1068
1069
        if self.dp_aware {
            // remove dp-aware workers in a prefix-matching fashion
            // without contacting the remote worker
            let mut removed_workers: Vec<String> = Vec::new();
            let worker_url_prefix = format!("{}@", worker_url);

1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
            // Find and remove all workers with matching prefix
            let all_workers = self.worker_registry.get_all();
            for w in all_workers.iter() {
                if w.url().starts_with(&worker_url_prefix) {
                    // Get model_id before removing
                    let model_id = w.model_id().to_string();

                    if self.worker_registry.remove_by_url(w.url()).is_some() {
                        info!("Removed worker: {}", w.url());
                        removed_workers.push(w.url().to_string());
1080

1081
1082
                        // Notify PolicyRegistry about the removed worker
                        self.policy_registry.on_worker_removed(&model_id);
1083
                    } else {
1084
                        warn!("Worker {} not found, skipping removal", w.url());
1085
1086
1087
1088
                    }
                }
            }

1089
1090
1091
1092
1093
            RouterMetrics::set_active_workers(self.worker_registry.get_all().len());

            for dp_url in removed_workers.iter() {
                if let Some(worker) = self.worker_registry.get_by_url(dp_url) {
                    let model_id = worker.model_id();
1094
1095
                    self.policy_registry
                        .remove_worker_from_cache_aware(model_id, dp_url);
1096
1097
1098
                }
            }
        } else {
1099
1100
1101
            // Get the worker first to extract model_id
            let model_id = if let Some(worker) = self.worker_registry.get_by_url(worker_url) {
                worker.model_id().to_string()
1102
1103
1104
            } else {
                warn!("Worker {} not found, skipping removal", worker_url);
                return;
1105
1106
1107
1108
1109
1110
1111
1112
1113
            };

            if self.worker_registry.remove_by_url(worker_url).is_some() {
                info!("Removed worker: {}", worker_url);

                // Notify PolicyRegistry about the removed worker
                self.policy_registry.on_worker_removed(&model_id);

                RouterMetrics::set_active_workers(self.worker_registry.get_all().len());
1114
1115
            }

1116
1117
            self.policy_registry
                .remove_worker_from_cache_aware(&model_id, worker_url);
1118
1119
1120
        }
    }

1121
    async fn get_worker_load(&self, worker_url: &str, api_key: &Option<String>) -> Option<isize> {
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
        let worker_url = if self.dp_aware {
            // Need to extract the URL from "http://host:port@dp_rank"
            let (worker_url_prefix, _dp_rank) = match Self::extract_dp_rank(worker_url) {
                Ok(tup) => tup,
                Err(e) => {
                    error!("Failed to extract dp_rank: {}", e);
                    return None;
                }
            };
            worker_url_prefix
        } else {
            worker_url
        };

1136
1137
1138
1139
1140
1141
        let mut req_builder = self.client.get(format!("{}/get_load", worker_url));
        if let Some(key) = api_key {
            req_builder = req_builder.bearer_auth(key);
        }

        match req_builder.send().await {
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
            Ok(res) if res.status().is_success() => match res.bytes().await {
                Ok(bytes) => match serde_json::from_slice::<serde_json::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
            }
        }
    }

    // Background task to monitor worker loads
    async fn monitor_worker_loads(
        worker_urls: Vec<String>,
1176
        worker_api_keys: Vec<Option<String>>,
1177
1178
1179
        tx: tokio::sync::watch::Sender<HashMap<String, isize>>,
        interval_secs: u64,
        policy: Arc<dyn LoadBalancingPolicy>,
1180
        client: Client,
1181
1182
1183
1184
1185
1186
1187
    ) {
        let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));

        loop {
            interval.tick().await;

            let mut loads = HashMap::new();
1188
1189
            for (url, api_key) in worker_urls.iter().zip(worker_api_keys.iter()) {
                if let Some(load) = Self::get_worker_load_static(&client, url, api_key).await {
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
                    loads.insert(url.clone(), load);
                }
            }

            if !loads.is_empty() {
                // Update policy with new loads
                policy.update_loads(&loads);

                // Send to watchers
                if let Err(e) = tx.send(loads) {
                    error!("Failed to send load update: {}", e);
                }
            }
        }
    }

    // Static version of get_worker_load for use in monitoring task
1207
1208
1209
1210
1211
    async fn get_worker_load_static(
        client: &reqwest::Client,
        worker_url: &str,
        api_key: &Option<String>,
    ) -> Option<isize> {
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
        let worker_url = if worker_url.contains("@") {
            // Need to extract the URL from "http://host:port@dp_rank"
            let (worker_url_prefix, _dp_rank) = match Self::extract_dp_rank(worker_url) {
                Ok(tup) => tup,
                Err(e) => {
                    debug!("Failed to extract dp_rank: {}", e);
                    return None;
                }
            };
            worker_url_prefix
        } else {
            worker_url
        };

1226
1227
1228
1229
1230
        let mut req_builder = client.get(format!("{}/get_load", worker_url));
        if let Some(key) = api_key {
            req_builder = req_builder.bearer_auth(key);
        }
        match req_builder.send().await {
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
            Ok(res) if res.status().is_success() => match res.bytes().await {
                Ok(bytes) => match serde_json::from_slice::<serde_json::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
            }
        }
    }
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279

    async fn build_rerank_response(
        req: &RerankRequest,
        response: Response,
    ) -> anyhow::Result<Response> {
        let (_, response_body) = response.into_parts();
        let body_bytes = to_bytes(response_body, usize::MAX).await?;
        let rerank_results = serde_json::from_slice::<Vec<RerankResult>>(&body_bytes)?;
        let mut rerank_response =
            RerankResponse::new(rerank_results, req.model.clone(), req.rid.clone());
        rerank_response.sort_by_score();
        if let Some(top_k) = req.top_k {
            rerank_response.apply_top_k(top_k);
        }
        if !req.return_documents {
            rerank_response.drop_documents();
        }
        Ok(Json(rerank_response).into_response())
    }
1280
1281
1282
1283
1284
1285
}

use async_trait::async_trait;

#[async_trait]
impl WorkerManagement for Router {
1286
1287
1288
1289
1290
1291
    async fn add_worker(
        &self,
        worker_url: &str,
        api_key: &Option<String>,
    ) -> Result<String, String> {
        Router::add_worker(self, worker_url, api_key).await
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
    }

    fn remove_worker(&self, worker_url: &str) {
        Router::remove_worker(self, worker_url)
    }

    fn get_worker_urls(&self) -> Vec<String> {
        Router::get_worker_urls(self)
    }
}

1303
#[async_trait]
1304
1305
1306
1307
1308
impl RouterTrait for Router {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

1309
    async fn health(&self, _req: Request<Body>) -> Response {
1310
        let workers = self.worker_registry.get_all();
1311
1312
1313
1314
1315
        let unhealthy_servers: Vec<_> = workers
            .iter()
            .filter(|w| !w.is_healthy())
            .map(|w| w.url().to_string())
            .collect();
1316

1317
1318
        if unhealthy_servers.is_empty() {
            (StatusCode::OK, "All servers healthy").into_response()
1319
        } else {
1320
1321
1322
1323
1324
            (
                StatusCode::SERVICE_UNAVAILABLE,
                format!("Unhealthy servers: {:?}", unhealthy_servers),
            )
                .into_response()
1325
1326
1327
        }
    }

1328
1329
    async fn health_generate(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "health_generate").await
1330
1331
    }

1332
1333
    async fn get_server_info(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "get_server_info").await
1334
1335
    }

1336
1337
    async fn get_models(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "v1/models").await
1338
1339
    }

1340
1341
    async fn get_model_info(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "get_model_info").await
1342
1343
1344
1345
    }

    async fn route_generate(
        &self,
1346
1347
        headers: Option<&HeaderMap>,
        body: &GenerateRequest,
1348
        model_id: Option<&str>,
1349
    ) -> Response {
1350
1351
        self.route_typed_request(headers, body, "/generate", model_id)
            .await
1352
1353
1354
1355
    }

    async fn route_chat(
        &self,
1356
1357
        headers: Option<&HeaderMap>,
        body: &ChatCompletionRequest,
1358
        model_id: Option<&str>,
1359
    ) -> Response {
1360
        self.route_typed_request(headers, body, "/v1/chat/completions", model_id)
1361
            .await
1362
1363
1364
1365
    }

    async fn route_completion(
        &self,
1366
1367
        headers: Option<&HeaderMap>,
        body: &CompletionRequest,
1368
        model_id: Option<&str>,
1369
    ) -> Response {
1370
        self.route_typed_request(headers, body, "/v1/completions", model_id)
1371
            .await
1372
1373
    }

1374
1375
1376
1377
    async fn route_responses(
        &self,
        headers: Option<&HeaderMap>,
        body: &ResponsesRequest,
1378
        model_id: Option<&str>,
1379
    ) -> Response {
1380
        self.route_typed_request(headers, body, "/v1/responses", model_id)
1381
1382
1383
            .await
    }

1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
    async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
        let endpoint = format!("v1/responses/{}", response_id);
        self.route_get_request(headers, &endpoint).await
    }

    async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response {
        let endpoint = format!("v1/responses/{}/cancel", response_id);
        self.route_post_empty_request(headers, &endpoint).await
    }

1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
    async fn route_embeddings(
        &self,
        headers: Option<&HeaderMap>,
        body: &EmbeddingRequest,
        model_id: Option<&str>,
    ) -> Response {
        // Record embeddings-specific metrics in addition to general request metrics
        let start = Instant::now();
        let res = self
            .route_typed_request(headers, body, "/v1/embeddings", model_id)
            .await;

        // Embedding specific metrics
        if res.status().is_success() {
            RouterMetrics::record_embeddings_request();
            RouterMetrics::record_embeddings_duration(start.elapsed());
        } else {
            let error_type = format!("http_{}", res.status().as_u16());
            RouterMetrics::record_embeddings_error(&error_type);
        }

        res
1416
1417
    }

1418
1419
1420
1421
1422
1423
    async fn route_rerank(
        &self,
        headers: Option<&HeaderMap>,
        body: &RerankRequest,
        model_id: Option<&str>,
    ) -> Response {
1424
1425
1426
        if let Err(e) = body.validate() {
            return (StatusCode::BAD_REQUEST, e).into_response();
        }
1427
1428
1429
        let response = self
            .route_typed_request(headers, body, "/v1/rerank", model_id)
            .await;
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
        if response.status().is_success() {
            match Self::build_rerank_response(body, response).await {
                Ok(rerank_response) => rerank_response,
                Err(e) => {
                    error!("Failed to build rerank response: {}", e);
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        "Failed to build rerank response".to_string(),
                    )
                        .into_response();
                }
            }
        } else {
            response
        }
1445
1446
    }

1447
    async fn flush_cache(&self) -> Response {
1448
1449
1450
1451
1452
1453
        // Get all worker URLs
        let worker_urls = self.get_worker_urls();

        // Send requests to all workers concurrently without headers
        let mut tasks = Vec::new();
        for worker_url in &worker_urls {
1454
1455
1456
1457
1458
1459
            let worker_url = if self.dp_aware {
                // Need to extract the URL from "http://host:port@dp_rank"
                let (worker_url_prefix, _dp_rank) = match Self::extract_dp_rank(worker_url) {
                    Ok(tup) => tup,
                    Err(e) => {
                        error!("Failed to extract dp_rank: {}", e);
1460
1461
1462
1463
1464
                        return (
                            StatusCode::INTERNAL_SERVER_ERROR,
                            format!("Failed to extract dp_rank: {}", e),
                        )
                            .into_response();
1465
1466
1467
1468
1469
1470
                    }
                };
                worker_url_prefix
            } else {
                worker_url
            };
1471
            let request_builder = self.client.post(format!("{}/flush_cache", worker_url));
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
            tasks.push(request_builder.send());
        }

        // Wait for all responses
        let results = futures_util::future::join_all(tasks).await;

        // Check if all succeeded
        let all_success = results.iter().all(|r| {
            r.as_ref()
                .map(|res| res.status().is_success())
                .unwrap_or(false)
        });

        if all_success {
1486
            (StatusCode::OK, "Cache flushed on all servers").into_response()
1487
        } else {
1488
1489
1490
1491
1492
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "Cache flush failed on one or more servers",
            )
                .into_response()
1493
1494
1495
        }
    }

1496
    async fn get_worker_loads(&self) -> Response {
1497
        let urls_with_key = self.worker_registry.get_all_urls_with_api_key();
1498
1499
1500
        let mut loads = Vec::new();

        // Get loads from all workers
1501
1502
        for (url, api_key) in &urls_with_key {
            let load = self.get_worker_load(url, api_key).await.unwrap_or(-1);
1503
1504
1505
1506
1507
1508
            loads.push(serde_json::json!({
                "worker": url,
                "load": load
            }));
        }

1509
        Json(serde_json::json!({
1510
1511
            "workers": loads
        }))
1512
        .into_response()
1513
1514
1515
1516
1517
1518
    }

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

1519
    fn readiness(&self) -> Response {
1520
        // Regular router is ready if it has at least one healthy worker
1521
1522
1523
        let workers = self.worker_registry.get_all();
        let healthy_count = workers.iter().filter(|w| w.is_healthy()).count();
        let total_workers = workers.len();
1524
1525

        if healthy_count > 0 {
1526
            Json(serde_json::json!({
1527
1528
                "status": "ready",
                "healthy_workers": healthy_count,
1529
                "total_workers": total_workers
1530
            }))
1531
            .into_response()
1532
        } else {
1533
1534
1535
1536
1537
            (
                StatusCode::SERVICE_UNAVAILABLE,
                Json(serde_json::json!({
                    "status": "not_ready",
                    "reason": "no healthy workers available",
1538
                    "total_workers": total_workers
1539
1540
1541
                })),
            )
                .into_response()
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn create_test_regular_router() -> Router {
1552
1553
1554
1555
1556
1557
1558
        // Create registries
        let worker_registry = Arc::new(WorkerRegistry::new());
        let policy_registry = Arc::new(PolicyRegistry::new(
            crate::config::types::PolicyConfig::RoundRobin,
        ));

        // Register test workers
1559
1560
        let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
            .worker_type(WorkerType::Regular)
1561
            .api_key("test_api_key")
1562
1563
1564
            .build();
        let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
            .worker_type(WorkerType::Regular)
1565
            .api_key("test_api_key")
1566
            .build();
1567
1568
1569
        worker_registry.register(Arc::new(worker1));
        worker_registry.register(Arc::new(worker2));

1570
1571
        let (_, rx) = tokio::sync::watch::channel(HashMap::new());
        Router {
1572
1573
            worker_registry,
            policy_registry,
1574
1575
            worker_startup_timeout_secs: 5,
            worker_startup_check_interval_secs: 1,
1576
1577
            dp_aware: false,
            api_key: None,
1578
            client: Client::new(),
1579
            retry_config: RetryConfig::default(),
1580
            circuit_breaker_config: CircuitBreakerConfig::default(),
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
            _worker_loads: Arc::new(rx),
            _load_monitor_handle: None,
        }
    }

    #[test]
    fn test_router_get_worker_urls_regular() {
        let router = create_test_regular_router();
        let urls = router.get_worker_urls();

        assert_eq!(urls.len(), 2);
        assert!(urls.contains(&"http://worker1:8080".to_string()));
        assert!(urls.contains(&"http://worker2:8080".to_string()));
    }

    #[test]
    fn test_select_first_worker_regular() {
        let router = create_test_regular_router();
        let result = router.select_first_worker();

        assert!(result.is_ok());
1602
1603
1604
        let url = result.unwrap();
        // DashMap doesn't guarantee order, so just check we get one of the workers
        assert!(url == "http://worker1:8080" || url == "http://worker2:8080");
1605
1606
    }

1607
1608
1609
1610
    #[tokio::test]
    async fn test_wait_for_healthy_workers_empty_list() {
        // Empty list will return error immediately
        let result = Router::wait_for_healthy_workers(&[], 1, 1).await;
1611
        assert!(result.is_err());
1612
        assert!(result.unwrap_err().contains("no workers provided"));
1613
1614
    }

1615
1616
    #[tokio::test]
    async fn test_wait_for_healthy_workers_invalid_urls() {
1617
1618
        // This test will timeout quickly since the URLs are invalid
        let result =
1619
            Router::wait_for_healthy_workers(&["http://nonexistent:8080".to_string()], 1, 1).await;
1620
1621
1622
1623
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Timeout"));
    }
}