router.rs 60.3 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
41
    dp_aware: bool,
    api_key: Option<String>,
42
    retry_config: RetryConfig,
43
    circuit_breaker_config: CircuitBreakerConfig,
44
45
46
47
48
    _worker_loads: Arc<tokio::sync::watch::Receiver<HashMap<String, isize>>>,
    _load_monitor_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
}

impl Router {
49
    /// Create a new router with injected policy and client
50
51
52
53
54
55
56
    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
        );
57

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

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

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

73
        // Cache-aware policies are initialized in WorkerInitializer
74
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
            let monitor_interval = ctx.router_config.worker_startup_check_interval_secs;
86
            let policy_clone = default_policy.clone();
87
            let client_clone = ctx.client.clone();
88
89

            Some(Arc::new(tokio::spawn(async move {
90
91
92
93
94
95
96
97
                Self::monitor_worker_loads(
                    monitor_urls,
                    tx,
                    monitor_interval,
                    policy_clone,
                    client_clone,
                )
                .await;
98
99
100
101
102
103
            })))
        } else {
            None
        };

        Ok(Router {
104
105
            worker_registry: ctx.worker_registry.clone(),
            policy_registry: ctx.policy_registry.clone(),
106
107
108
109
110
111
112
113
            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(),
114
            circuit_breaker_config: core_cb_config,
115
116
117
118
119
120
121
            _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> {
122
123
124
125
126
127
128
129
130
131
        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> {
        let workers = match model_id {
            Some(model) => self.worker_registry.get_by_model_fast(model),
            None => self.worker_registry.get_all(),
        };
        workers.iter().map(|w| w.url().to_string()).collect()
132
133
    }

134
    pub async fn wait_for_healthy_workers(
135
        worker_urls: &[String],
136
137
        worker_startup_timeout_secs: u64,
        worker_startup_check_interval_secs: u64,
138
    ) -> Result<(), String> {
139
140
141
142
143
144
        if worker_urls.is_empty() {
            return Err(
                "Timeout waiting for workers to become healthy: no workers provided".to_string(),
            );
        }

145
        // Perform health check asynchronously
146
147
148
149
150
151
        Self::wait_for_healthy_workers_async(
            worker_urls,
            worker_startup_timeout_secs,
            worker_startup_check_interval_secs,
        )
        .await
152
153
154
155
    }

    async fn wait_for_healthy_workers_async(
        worker_urls: &[String],
156
157
        worker_startup_timeout_secs: u64,
        worker_startup_check_interval_secs: u64,
158
159
160
161
    ) -> Result<(), String> {
        info!(
            "Waiting for {} workers to become healthy (timeout: {}s)",
            worker_urls.len(),
162
            worker_startup_timeout_secs
163
164
        );

165
        let start_time = std::time::Instant::now();
166
167
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(2))
168
169
170
171
            .build()
            .map_err(|e| format!("Failed to create HTTP client: {}", e))?;

        loop {
172
            if start_time.elapsed() > Duration::from_secs(worker_startup_timeout_secs) {
173
174
                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",
175
                    worker_startup_timeout_secs, worker_urls
176
177
178
                );
                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",
179
                    worker_startup_timeout_secs, worker_urls
180
181
182
                ));
            }

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
            // 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;

209
210
211
            let mut all_healthy = true;
            let mut unhealthy_workers = Vec::new();

212
213
214
215
            for result in results {
                match result {
                    Ok(None) => {
                        // Worker is healthy
216
                    }
217
                    Ok(Some((url, reason))) => {
218
                        all_healthy = false;
219
220
221
222
223
224
                        unhealthy_workers.push((url, reason));
                    }
                    Err(e) => {
                        all_healthy = false;
                        unhealthy_workers
                            .push(("unknown".to_string(), format!("task error: {}", e)));
225
226
227
228
229
                    }
                }
            }

            if all_healthy {
230
                info!("All {} workers are healthy", worker_urls.len());
231
232
                return Ok(());
            } else {
233
                debug!(
234
                    "Waiting for {} workers to become healthy ({} unhealthy: {:?})",
235
                    worker_urls.len(),
236
237
                    unhealthy_workers.len(),
                    unhealthy_workers
238
                );
239
                tokio::time::sleep(Duration::from_secs(worker_startup_check_interval_secs)).await;
240
241
242
243
            }
        }
    }

244
245
    fn get_worker_dp_size(worker_url: &str, api_key: &Option<String>) -> Result<usize, String> {
        let sync_client = reqwest::blocking::Client::new();
246
        let mut req_builder = sync_client.get(format!("{}/get_server_info", worker_url));
247
248
249
250
251
252
253
254
255
256
257
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
        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)
    }

300
    fn select_first_worker(&self) -> Result<String, String> {
301
302
        let workers = self.worker_registry.get_all();
        if workers.is_empty() {
303
304
            Err("No workers are available".to_string())
        } else {
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
            Ok(workers[0].url().to_string())
        }
    }

    #[allow(dead_code)]
    fn select_first_worker_for_model(&self, model_id: Option<&str>) -> Result<String, String> {
        let workers = match model_id {
            Some(model) => self.worker_registry.get_by_model_fast(model),
            None => self.worker_registry.get_all(),
        };
        if workers.is_empty() {
            Err(format!(
                "No workers are available for model: {:?}",
                model_id
            ))
        } else {
            Ok(workers[0].url().to_string())
322
323
324
        }
    }

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

343
        let request_builder = self.client.get(format!("{}/health", health_url));
344
345
346

        let response = match request_builder.send().await {
            Ok(res) => {
347
348
                let status = StatusCode::from_u16(res.status().as_u16())
                    .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
349
350

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

380
        // Don't record metrics for health checks
381
382
383
        response
    }

384
    // Helper method to proxy GET requests to the first available worker
385
    async fn proxy_get_request(&self, req: Request<Body>, endpoint: &str) -> Response {
386
        let headers = header_utils::copy_request_headers(&req);
387
388
389

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

398
399
400
401
                match request_builder.send().await {
                    Ok(res) => {
                        let status = StatusCode::from_u16(res.status().as_u16())
                            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
402
403
404
405
406

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

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

432
433
434
435
436
437
438
439
440
441
442
443
444
    /// 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>> {
        // Get workers for the specified model (O(1) lookup if model_id is provided)
        let workers = match model_id {
            Some(model) => self.worker_registry.get_by_model_fast(model),
            None => self.worker_registry.get_all(),
        };

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

        // 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())
461
462
    }

463
    pub async fn route_typed_request<T: GenerationRequest + serde::Serialize + Clone>(
464
        &self,
465
        headers: Option<&HeaderMap>,
466
467
        typed_req: &T,
        route: &str,
468
        model_id: Option<&str>,
469
    ) -> Response {
470
        let start = Instant::now();
471
472
473
474
475
476
477
        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 {
478
                let worker = match self.select_worker_for_model(model_id, Some(&text)) {
479
480
481
482
483
484
485
486
487
488
                    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();
                    }
                };
489

490
                // Optional load tracking for cache-aware policy
491
492
493
494
495
496
497
                // 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" {
498
499
500
                    worker.increment_load();
                    RouterMetrics::set_running_requests(worker.url(), worker.load());
                    true
501
502
503
504
                } else {
                    false
                };

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

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

523
                worker.record_outcome(response.status().is_success());
524
525
526
527
528
529
530
531
532
533
534
535
536

                // 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(),
                        );
                    }
                }

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

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

        response
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
667
668
669
    // 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
    }

670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
    // 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
            )),
        }
    }

687
688
689
    // Send typed request directly without conversion
    async fn send_typed_request<T: serde::Serialize>(
        &self,
690
        headers: Option<&HeaderMap>,
691
692
693
694
695
        typed_req: &T,
        route: &str,
        worker_url: &str,
        is_stream: bool,
        load_incremented: bool, // Whether load was incremented for this request
696
    ) -> Response {
697
698
699
700
701
        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);
702
703
704
705
706
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Failed to extract dp_rank: {}", e),
                    )
                        .into_response();
707
708
709
710
711
712
713
                }
            };

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

            // 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 {
733
734
735
736
737
                return (
                    StatusCode::BAD_REQUEST,
                    "Failed to insert the data_parallel_rank field into the request body",
                )
                    .into_response();
738
739
            }

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

749
750
751
752
        // 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
753
                if *name != CONTENT_TYPE && *name != CONTENT_LENGTH {
754
755
                    request_builder = request_builder.header(name, value);
                }
756
757
758
759
760
761
            }
        }

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

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

775
776
777
778
779
                return (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Request failed: {}", e),
                )
                    .into_response();
780
781
782
            }
        };

783
784
        let status = StatusCode::from_u16(res.status().as_u16())
            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
785
786

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

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

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

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

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

825
826
827
828
829
            // 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"));

830
831
832
833
834
835
            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;
836
                let mut decremented = false;
837
838
839
840
841
842
843
844
845
                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]")
                            {
846
847
848
849
                                if let Some(worker) = registry.get_by_url(&worker_url) {
                                    worker.decrement_load();
                                    RouterMetrics::set_running_requests(&worker_url, worker.load());
                                    decremented = true;
850
851
                                }
                            }
852
853
854
855
856
857
858
859
860
861
                            if tx.send(Ok(bytes)).is_err() {
                                break;
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(format!("Stream error: {}", e)));
                            break;
                        }
                    }
                }
862
                if !decremented {
863
864
865
                    if let Some(worker) = registry.get_by_url(&worker_url) {
                        worker.decrement_load();
                        RouterMetrics::set_running_requests(&worker_url, worker.load());
866
867
                    }
                }
868
869
870
871
872
873
874
            });

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

            let mut response = Response::new(body);
            *response.status_mut() = status;
875
            *response.headers_mut() = response_headers;
876
            response
877
878
        } else {
            // For requests without load tracking, just stream
879
880
881
882
883
            // 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"));

884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
            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;
910
            *response.headers_mut() = response_headers;
911
            response
912
913
914
915
916
917
        }
    }

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

        loop {
923
            if start_time.elapsed() > Duration::from_secs(self.worker_startup_timeout_secs) {
924
925
                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",
926
                    self.worker_startup_timeout_secs, worker_url
927
928
929
                );
                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",
930
                    self.worker_startup_timeout_secs, worker_url
931
932
933
                ));
            }

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

                                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();
961
962
963
964
965
966
967
                                self.policy_registry.on_worker_added(model_id, None);

                                // Initialize cache-aware policy if applicable
                                let model_workers =
                                    self.worker_registry.get_by_model_fast(model_id);
                                self.policy_registry
                                    .init_cache_aware_policy(model_id, &model_workers);
968

969
970
971
972
973
974
                                worker_added = true;
                            }
                            if !worker_added {
                                return Err(format!("No worker added for {}", worker_url));
                            }
                        } else {
975
                            if self.worker_registry.get_by_url(worker_url).is_some() {
976
977
978
                                return Err(format!("Worker {} already exists", worker_url));
                            }
                            info!("Added worker: {}", worker_url);
979

980
                            // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
981
982
983
984
                            let new_worker = BasicWorkerBuilder::new(worker_url.to_string())
                                .worker_type(WorkerType::Regular)
                                .circuit_breaker_config(self.circuit_breaker_config.clone())
                                .build();
985
986
987
988
989
990

                            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();
991
992
993
994
995
996
                            self.policy_registry.on_worker_added(model_id, None);

                            // Initialize cache-aware policy if applicable
                            let model_workers = self.worker_registry.get_by_model_fast(model_id);
                            self.policy_registry
                                .init_cache_aware_policy(model_id, &model_workers);
997
998
                        }

999
1000
                        RouterMetrics::set_active_workers(self.worker_registry.get_all().len());

1001
1002
                        return Ok(format!("Successfully added worker: {}", worker_url));
                    } else {
1003
1004
                        debug!(
                            "Worker {} health check pending - status: {}",
1005
1006
1007
1008
1009
1010
1011
1012
1013
                            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);
                        }

1014
1015
1016
1017
                        tokio::time::sleep(Duration::from_secs(
                            self.worker_startup_check_interval_secs,
                        ))
                        .await;
1018
1019
1020
1021
                        continue;
                    }
                }
                Err(e) => {
1022
                    debug!("Worker {} health check pending - error: {}", worker_url, e);
1023
1024
1025
1026
1027
1028

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

1029
1030
1031
1032
                    tokio::time::sleep(Duration::from_secs(
                        self.worker_startup_check_interval_secs,
                    ))
                    .await;
1033
1034
1035
1036
1037
1038
1039
                    continue;
                }
            }
        }
    }

    pub fn remove_worker(&self, worker_url: &str) {
1040
1041
1042
1043
1044
1045
        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);

1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
            // 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());
1056

1057
1058
                        // Notify PolicyRegistry about the removed worker
                        self.policy_registry.on_worker_removed(&model_id);
1059
                    } else {
1060
                        warn!("Worker {} not found, skipping removal", w.url());
1061
1062
1063
1064
                    }
                }
            }

1065
1066
1067
1068
1069
            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();
1070
1071
                    self.policy_registry
                        .remove_worker_from_cache_aware(model_id, dp_url);
1072
1073
1074
                }
            }
        } else {
1075
1076
1077
            // 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()
1078
1079
1080
            } else {
                warn!("Worker {} not found, skipping removal", worker_url);
                return;
1081
1082
1083
1084
1085
1086
1087
1088
1089
            };

            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());
1090
1091
            }

1092
1093
            self.policy_registry
                .remove_worker_from_cache_aware(&model_id, worker_url);
1094
1095
1096
        }
    }

1097
    async fn get_worker_load(&self, worker_url: &str) -> Option<isize> {
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
        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
        };

1112
1113
        match self
            .client
1114
            .get(format!("{}/get_load", worker_url))
1115
1116
1117
            .send()
            .await
        {
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
            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>,
        tx: tokio::sync::watch::Sender<HashMap<String, isize>>,
        interval_secs: u64,
        policy: Arc<dyn LoadBalancingPolicy>,
1155
        client: Client,
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
    ) {
        let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));

        loop {
            interval.tick().await;

            let mut loads = HashMap::new();
            for url in &worker_urls {
                if let Some(load) = Self::get_worker_load_static(&client, url).await {
                    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
    async fn get_worker_load_static(client: &reqwest::Client, worker_url: &str) -> Option<isize> {
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
        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
        };

1197
        match client.get(format!("{}/get_load", worker_url)).send().await {
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
            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
            }
        }
    }
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246

    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())
    }
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
}

use async_trait::async_trait;

#[async_trait]
impl WorkerManagement for Router {
    async fn add_worker(&self, worker_url: &str) -> Result<String, String> {
        Router::add_worker(self, worker_url).await
    }

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

1266
#[async_trait]
1267
1268
1269
1270
1271
impl RouterTrait for Router {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

1272
    async fn health(&self, _req: Request<Body>) -> Response {
1273
        let workers = self.worker_registry.get_all();
1274
1275
1276
1277
1278
        let unhealthy_servers: Vec<_> = workers
            .iter()
            .filter(|w| !w.is_healthy())
            .map(|w| w.url().to_string())
            .collect();
1279

1280
1281
        if unhealthy_servers.is_empty() {
            (StatusCode::OK, "All servers healthy").into_response()
1282
        } else {
1283
1284
1285
1286
1287
            (
                StatusCode::SERVICE_UNAVAILABLE,
                format!("Unhealthy servers: {:?}", unhealthy_servers),
            )
                .into_response()
1288
1289
1290
        }
    }

1291
1292
    async fn health_generate(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "health_generate").await
1293
1294
    }

1295
1296
    async fn get_server_info(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "get_server_info").await
1297
1298
    }

1299
1300
    async fn get_models(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "v1/models").await
1301
1302
    }

1303
1304
    async fn get_model_info(&self, req: Request<Body>) -> Response {
        self.proxy_get_request(req, "get_model_info").await
1305
1306
1307
1308
    }

    async fn route_generate(
        &self,
1309
1310
        headers: Option<&HeaderMap>,
        body: &GenerateRequest,
1311
        model_id: Option<&str>,
1312
    ) -> Response {
1313
1314
        self.route_typed_request(headers, body, "/generate", model_id)
            .await
1315
1316
1317
1318
    }

    async fn route_chat(
        &self,
1319
1320
        headers: Option<&HeaderMap>,
        body: &ChatCompletionRequest,
1321
        model_id: Option<&str>,
1322
    ) -> Response {
1323
        self.route_typed_request(headers, body, "/v1/chat/completions", model_id)
1324
            .await
1325
1326
1327
1328
    }

    async fn route_completion(
        &self,
1329
1330
        headers: Option<&HeaderMap>,
        body: &CompletionRequest,
1331
        model_id: Option<&str>,
1332
    ) -> Response {
1333
        self.route_typed_request(headers, body, "/v1/completions", model_id)
1334
            .await
1335
1336
    }

1337
1338
1339
1340
    async fn route_responses(
        &self,
        headers: Option<&HeaderMap>,
        body: &ResponsesRequest,
1341
        model_id: Option<&str>,
1342
    ) -> Response {
1343
        self.route_typed_request(headers, body, "/v1/responses", model_id)
1344
1345
1346
            .await
    }

1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
    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
    }

1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
    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
1379
1380
    }

1381
1382
1383
1384
1385
1386
    async fn route_rerank(
        &self,
        headers: Option<&HeaderMap>,
        body: &RerankRequest,
        model_id: Option<&str>,
    ) -> Response {
1387
1388
1389
        if let Err(e) = body.validate() {
            return (StatusCode::BAD_REQUEST, e).into_response();
        }
1390
1391
1392
        let response = self
            .route_typed_request(headers, body, "/v1/rerank", model_id)
            .await;
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
        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
        }
1408
1409
    }

1410
    async fn flush_cache(&self) -> Response {
1411
1412
1413
1414
1415
1416
        // 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 {
1417
1418
1419
1420
1421
1422
            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);
1423
1424
1425
1426
1427
                        return (
                            StatusCode::INTERNAL_SERVER_ERROR,
                            format!("Failed to extract dp_rank: {}", e),
                        )
                            .into_response();
1428
1429
1430
1431
1432
1433
                    }
                };
                worker_url_prefix
            } else {
                worker_url
            };
1434
            let request_builder = self.client.post(format!("{}/flush_cache", worker_url));
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
            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 {
1449
            (StatusCode::OK, "Cache flushed on all servers").into_response()
1450
        } else {
1451
1452
1453
1454
1455
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "Cache flush failed on one or more servers",
            )
                .into_response()
1456
1457
1458
        }
    }

1459
    async fn get_worker_loads(&self) -> Response {
1460
1461
1462
1463
1464
        let urls = self.get_worker_urls();
        let mut loads = Vec::new();

        // Get loads from all workers
        for url in &urls {
1465
            let load = self.get_worker_load(url).await.unwrap_or(-1);
1466
1467
1468
1469
1470
1471
            loads.push(serde_json::json!({
                "worker": url,
                "load": load
            }));
        }

1472
        Json(serde_json::json!({
1473
1474
            "workers": loads
        }))
1475
        .into_response()
1476
1477
1478
1479
1480
1481
    }

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

1482
    fn readiness(&self) -> Response {
1483
        // Regular router is ready if it has at least one healthy worker
1484
1485
1486
        let workers = self.worker_registry.get_all();
        let healthy_count = workers.iter().filter(|w| w.is_healthy()).count();
        let total_workers = workers.len();
1487
1488

        if healthy_count > 0 {
1489
            Json(serde_json::json!({
1490
1491
                "status": "ready",
                "healthy_workers": healthy_count,
1492
                "total_workers": total_workers
1493
            }))
1494
            .into_response()
1495
        } else {
1496
1497
1498
1499
1500
            (
                StatusCode::SERVICE_UNAVAILABLE,
                Json(serde_json::json!({
                    "status": "not_ready",
                    "reason": "no healthy workers available",
1501
                    "total_workers": total_workers
1502
1503
1504
                })),
            )
                .into_response()
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
        }
    }
}

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

    fn create_test_regular_router() -> Router {
1515
1516
1517
1518
1519
1520
1521
        // Create registries
        let worker_registry = Arc::new(WorkerRegistry::new());
        let policy_registry = Arc::new(PolicyRegistry::new(
            crate::config::types::PolicyConfig::RoundRobin,
        ));

        // Register test workers
1522
1523
1524
1525
1526
1527
        let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
            .worker_type(WorkerType::Regular)
            .build();
        let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
            .worker_type(WorkerType::Regular)
            .build();
1528
1529
1530
        worker_registry.register(Arc::new(worker1));
        worker_registry.register(Arc::new(worker2));

1531
1532
        let (_, rx) = tokio::sync::watch::channel(HashMap::new());
        Router {
1533
1534
            worker_registry,
            policy_registry,
1535
1536
            worker_startup_timeout_secs: 5,
            worker_startup_check_interval_secs: 1,
1537
1538
            dp_aware: false,
            api_key: None,
1539
            client: Client::new(),
1540
            retry_config: RetryConfig::default(),
1541
            circuit_breaker_config: CircuitBreakerConfig::default(),
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
            _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());
1563
1564
1565
        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");
1566
1567
    }

1568
1569
1570
1571
    #[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;
1572
        assert!(result.is_err());
1573
        assert!(result.unwrap_err().contains("no workers provided"));
1574
1575
    }

1576
1577
    #[tokio::test]
    async fn test_wait_for_healthy_workers_invalid_urls() {
1578
1579
        // This test will timeout quickly since the URLs are invalid
        let result =
1580
            Router::wait_for_healthy_workers(&["http://nonexistent:8080".to_string()], 1, 1).await;
1581
1582
1583
1584
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Timeout"));
    }
}