service_v2.rs 27.5 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
// SPDX-License-Identifier: Apache-2.0
3

4
use std::collections::HashMap;
5
use std::env::var;
Graham King's avatar
Graham King committed
6
use std::path::PathBuf;
7
use std::sync::Arc;
8
9
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
10
use std::time::Duration;
11

12
13
14
use axum::body::Body;
use axum::http::Response;

15
use super::Metrics;
16
use super::RouteDoc;
17
use super::metrics;
18
use super::metrics::register_worker_timing_metrics;
19
use crate::discovery::ModelManager;
20
use crate::endpoint_type::EndpointType;
21
22
23
use crate::kv_router::metrics::{
    RoutingOverheadMetrics, register_router_queue_metrics, register_worker_load_metrics,
};
24
use crate::request_template::RequestTemplate;
25
use anyhow::Result;
Graham King's avatar
Graham King committed
26
use axum_server::tls_rustls::RustlsConfig;
27
use derive_builder::Builder;
28
use dynamo_runtime::config::env_is_truthy;
29
use dynamo_runtime::config::environment_names::llm as env_llm;
30
use dynamo_runtime::discovery::Discovery;
31
use dynamo_runtime::logging::{make_inference_request_span, make_system_request_span};
32
33
34
35
36
37
use dynamo_runtime::metrics::{
    frontend_perf::ensure_frontend_perf_metrics_registered_prometheus,
    request_plane::ensure_request_plane_metrics_registered_prometheus,
    tokio_perf::{ensure_tokio_perf_metrics_registered_prometheus, tokio_metrics_and_canary_loop},
    transport_metrics::ensure_transport_metrics_registered_prometheus,
};
Graham King's avatar
Graham King committed
38
use std::net::SocketAddr;
39
use tokio::task::JoinHandle;
40
use tokio_util::sync::CancellationToken;
41
use tower_http::trace::TraceLayer;
42

43
44
45
46
47
48
49
50
51
52
53
54
55
/// Middleware that echoes `x-request-id` from request to response headers.
async fn echo_request_id_header(
    request: axum::extract::Request,
    next: axum::middleware::Next,
) -> axum::response::Response {
    let x_request_id = request.headers().get("x-request-id").cloned();
    let mut response = next.run(request).await;
    if let Some(value) = x_request_id {
        response.headers_mut().insert("x-request-id", value);
    }
    response
}

56
57
58
59
/// HTTP service shared state
pub struct State {
    metrics: Arc<Metrics>,
    manager: Arc<ModelManager>,
60
    discovery_client: Arc<dyn Discovery>,
61
    flags: StateFlags,
62
    cancel_token: CancellationToken,
63
64
65
66
67
68
69
}

#[derive(Default, Debug)]
struct StateFlags {
    chat_endpoints_enabled: AtomicBool,
    cmpl_endpoints_enabled: AtomicBool,
    embeddings_endpoints_enabled: AtomicBool,
70
    images_endpoints_enabled: AtomicBool,
71
    videos_endpoints_enabled: AtomicBool,
72
    audios_endpoints_enabled: AtomicBool,
73
    responses_endpoints_enabled: AtomicBool,
74
    anthropic_endpoints_enabled: AtomicBool,
75
76
77
78
79
80
81
82
}

impl StateFlags {
    pub fn get(&self, endpoint_type: &EndpointType) -> bool {
        match endpoint_type {
            EndpointType::Chat => self.chat_endpoints_enabled.load(Ordering::Relaxed),
            EndpointType::Completion => self.cmpl_endpoints_enabled.load(Ordering::Relaxed),
            EndpointType::Embedding => self.embeddings_endpoints_enabled.load(Ordering::Relaxed),
83
            EndpointType::Images => self.images_endpoints_enabled.load(Ordering::Relaxed),
84
            EndpointType::Videos => self.videos_endpoints_enabled.load(Ordering::Relaxed),
85
            EndpointType::Audios => self.audios_endpoints_enabled.load(Ordering::Relaxed),
86
            EndpointType::Responses => self.responses_endpoints_enabled.load(Ordering::Relaxed),
87
88
89
            EndpointType::AnthropicMessages => {
                self.anthropic_endpoints_enabled.load(Ordering::Relaxed)
            }
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        }
    }

    pub fn set(&self, endpoint_type: &EndpointType, enabled: bool) {
        match endpoint_type {
            EndpointType::Chat => self
                .chat_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
            EndpointType::Completion => self
                .cmpl_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
            EndpointType::Embedding => self
                .embeddings_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
104
105
106
            EndpointType::Images => self
                .images_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
107
108
109
            EndpointType::Videos => self
                .videos_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
110
111
112
            EndpointType::Audios => self
                .audios_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
113
114
115
            EndpointType::Responses => self
                .responses_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
116
117
118
            EndpointType::AnthropicMessages => self
                .anthropic_endpoints_enabled
                .store(enabled, Ordering::Relaxed),
119
120
        }
    }
121
122
123
}

impl State {
124
125
    pub fn new(
        manager: Arc<ModelManager>,
126
        discovery_client: Arc<dyn Discovery>,
127
128
        cancel_token: CancellationToken,
    ) -> Self {
129
130
131
        Self {
            manager,
            metrics: Arc::new(Metrics::default()),
132
            discovery_client,
133
134
135
136
            flags: StateFlags {
                chat_endpoints_enabled: AtomicBool::new(false),
                cmpl_endpoints_enabled: AtomicBool::new(false),
                embeddings_endpoints_enabled: AtomicBool::new(false),
137
                images_endpoints_enabled: AtomicBool::new(false),
138
                videos_endpoints_enabled: AtomicBool::new(false),
139
                audios_endpoints_enabled: AtomicBool::new(false),
140
                responses_endpoints_enabled: AtomicBool::new(false),
141
                anthropic_endpoints_enabled: AtomicBool::new(false),
142
            },
143
            cancel_token,
144
145
146
        }
    }

147
148
149
150
151
152
153
154
155
156
157
158
159
    /// Get the Prometheus [`Metrics`] object which tracks request counts and inflight requests
    pub fn metrics_clone(&self) -> Arc<Metrics> {
        self.metrics.clone()
    }

    pub fn manager(&self) -> &ModelManager {
        Arc::as_ref(&self.manager)
    }

    pub fn manager_clone(&self) -> Arc<ModelManager> {
        self.manager.clone()
    }

160
161
162
163
    pub fn discovery(&self) -> Arc<dyn Discovery> {
        self.discovery_client.clone()
    }

164
165
166
167
168
169
170
171
172
173
    /// Check if the service is shutting down
    pub fn is_cancelled(&self) -> bool {
        self.cancel_token.is_cancelled()
    }

    /// Get the cancellation token
    pub fn cancel_token(&self) -> &CancellationToken {
        &self.cancel_token
    }

174
175
176
177
    // TODO
    pub fn sse_keep_alive(&self) -> Option<Duration> {
        None
    }
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197

    /// Returns true if streaming tool call dispatch is enabled via
    /// [`env_llm::DYN_ENABLE_STREAMING_TOOL_DISPATCH`].
    ///
    /// When enabled, the chat completions streaming path emits `event: tool_call_dispatch`
    /// SSE events for each complete tool call, letting clients start processing tool calls
    /// before `finish_reason="tool_calls"` arrives.
    pub fn streaming_tool_dispatch_enabled(&self) -> bool {
        env_is_truthy(env_llm::DYN_ENABLE_STREAMING_TOOL_DISPATCH)
    }

    /// Returns true if streaming reasoning dispatch is enabled via
    /// [`env_llm::DYN_ENABLE_STREAMING_REASONING_DISPATCH`].
    ///
    /// When enabled, the chat completions streaming path accumulates reasoning tokens and
    /// emits a single `event: reasoning_dispatch` SSE event with the complete reasoning
    /// block once thinking ends (DeepSeek-R1, Qwen3, etc.).
    pub fn streaming_reasoning_dispatch_enabled(&self) -> bool {
        env_is_truthy(env_llm::DYN_ENABLE_STREAMING_REASONING_DISPATCH)
    }
198
199
}

200
201
#[derive(Clone)]
pub struct HttpService {
202
203
204
    // The state we share with every request handler
    state: Arc<State>,

205
206
    router: axum::Router,
    port: u16,
207
    host: String,
Graham King's avatar
Graham King committed
208
209
210
    enable_tls: bool,
    tls_cert_path: Option<PathBuf>,
    tls_key_path: Option<PathBuf>,
211
    route_docs: Vec<RouteDoc>,
212
213
214
}

#[derive(Clone, Builder)]
215
#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
216
217
218
219
pub struct HttpServiceConfig {
    #[builder(default = "8787")]
    port: u16,

220
221
222
    #[builder(setter(into), default = "String::from(\"0.0.0.0\")")]
    host: String,

Graham King's avatar
Graham King committed
223
224
225
226
227
228
229
230
231
    #[builder(default = "false")]
    enable_tls: bool,

    #[builder(default = "None")]
    tls_cert_path: Option<PathBuf>,

    #[builder(default = "None")]
    tls_key_path: Option<PathBuf>,

232
233
    // #[builder(default)]
    // custom: Vec<axum::Router>
234
    #[builder(default = "false")]
235
236
    enable_chat_endpoints: bool,

237
    #[builder(default = "false")]
238
    enable_cmpl_endpoints: bool,
239

240
    #[builder(default = "true")]
241
242
    enable_embeddings_endpoints: bool,

243
244
245
    #[builder(default = "true")]
    enable_responses_endpoints: bool,

246
247
248
    #[builder(default = "false")]
    enable_anthropic_endpoints: bool,

249
250
    #[builder(default = "None")]
    request_template: Option<RequestTemplate>,
251

252
253
    #[builder(default = "None")]
    discovery: Option<Arc<dyn Discovery>>,
254
255
256

    #[builder(default = "None")]
    cancel_token: Option<CancellationToken>,
257
258
259
260
261
262
263
264
265
266

    /// When set, the `/metrics` endpoint will also expose metrics from the
    /// DRT's registry tree (anything created via `metrics().create*()`).
    #[builder(default = "None")]
    drt_metrics: Option<dynamo_runtime::metrics::MetricsRegistry>,

    /// When set (e.g. DRT discovery), router metrics (dynamo_router_* with router_id label)
    /// are registered using discovery.instance_id() and exposed on /metrics.
    #[builder(default = "None")]
    drt_discovery: Option<Arc<dyn Discovery>>,
267
268
269
270
271
272
273
}

impl HttpService {
    pub fn builder() -> HttpServiceConfigBuilder {
        HttpServiceConfigBuilder::default()
    }

274
275
276
277
278
279
280
281
    pub fn state_clone(&self) -> Arc<State> {
        self.state.clone()
    }

    pub fn state(&self) -> &State {
        Arc::as_ref(&self.state)
    }

282
    pub fn model_manager(&self) -> &ModelManager {
283
        self.state().manager()
284
285
    }

286
287
288
289
290
291
    pub async fn spawn(&self, cancel_token: CancellationToken) -> JoinHandle<Result<()>> {
        let this = self.clone();
        tokio::spawn(async move { this.run(cancel_token).await })
    }

    pub async fn run(&self, cancel_token: CancellationToken) -> Result<()> {
292
        let address = format!("{}:{}", self.host, self.port);
Graham King's avatar
Graham King committed
293
294
        let protocol = if self.enable_tls { "HTTPS" } else { "HTTP" };
        tracing::info!(protocol, address, "Starting HTTP(S) service");
295
296
297
298

        let router = self.router.clone();
        let observer = cancel_token.child_token();

299
300
        let state_cancel = self.state.cancel_token().clone();

Graham King's avatar
Graham King committed
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
        let addr: SocketAddr = address
            .parse()
            .map_err(|e| anyhow::anyhow!("Invalid address '{}': {}", address, e))?;

        if self.enable_tls {
            let cert_path = self
                .tls_cert_path
                .as_ref()
                .ok_or_else(|| anyhow::anyhow!("TLS certificate path not provided"))?;
            let key_path = self
                .tls_key_path
                .as_ref()
                .ok_or_else(|| anyhow::anyhow!("TLS private key path not provided"))?;

            // aws_lc_rs is the default but other crates pull in `ring` also,
            // so rustls doesn't know which one to use. Tell it.
            if let Err(e) = rustls::crypto::aws_lc_rs::default_provider().install_default() {
                tracing::debug!("TLS crypto provider already installed: {e:?}");
            }

            let config = RustlsConfig::from_pem_file(cert_path, key_path)
                .await
                .map_err(|e| anyhow::anyhow!("Failed to create TLS config: {}", e))?;

            let handle = axum_server::Handle::new();
            let server = axum_server::bind_rustls(addr, config)
                .handle(handle.clone())
                .serve(router.into_make_service());

330
331
332
            // Spawn canary after all fallible startup so it won't leak on early errors
            tokio::spawn(tokio_metrics_and_canary_loop(cancel_token.clone()));

Graham King's avatar
Graham King committed
333
334
            tokio::select! {
                result = server => {
335
336
337
                    let result = result.map_err(|e| anyhow::anyhow!("HTTPS server error: {}", e));
                    cancel_token.cancel();
                    result?;
Graham King's avatar
Graham King committed
338
339
                }
                _ = observer.cancelled() => {
340
                    state_cancel.cancel();
Graham King's avatar
Graham King committed
341
                    tracing::info!("HTTPS server shutdown requested");
342
343
344
                    // accepting requests for 5 more seconds, to allow incorrectly routed requests to arrive
                    handle.graceful_shutdown(Some(Duration::from_secs(get_graceful_shutdown_timeout() as u64)));
                    // no longer accepting requests, draining all existing connections
Graham King's avatar
Graham King committed
345
346
347
                }
            }
        } else {
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
            let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
                tracing::error!(
                    protocol = %protocol,
                    address = %address,
                    error = %e,
                    "Failed to bind server to address"
                );
                match e.kind() {
                    std::io::ErrorKind::AddrInUse => anyhow::anyhow!(
                        "Failed to start {} server: port {} already in use. Use --http-port to specify a different port.",
                        protocol,
                        self.port
                    ),
                    _ => anyhow::anyhow!(
                        "Failed to start {} server on {}: {}",
                        protocol,
                        address,
                        e
                    ),
                }
            })?;
Graham King's avatar
Graham King committed
369

370
371
372
            // Spawn canary after all fallible startup so it won't leak on early errors
            tokio::spawn(tokio_metrics_and_canary_loop(cancel_token.clone()));

Graham King's avatar
Graham King committed
373
            axum::serve(listener, router)
374
375
376
377
378
379
380
381
382
                .with_graceful_shutdown(async move {
                    observer.cancelled_owned().await;
                    state_cancel.cancel();
                    tracing::info!("HTTP server shutdown requested");
                    // accepting requests for 5 more seconds, to allow incorrectly routed requests to arrive
                    tokio::time::sleep(Duration::from_secs(get_graceful_shutdown_timeout() as u64))
                        .await;
                    // no longer accepting requests, draining all existing connections
                })
Graham King's avatar
Graham King committed
383
384
                .await
                .inspect_err(|_| cancel_token.cancel())?;
385
            cancel_token.cancel();
Graham King's avatar
Graham King committed
386
        }
387
388

        Ok(())
389
    }
390
391
392
393
394

    /// Documentation of exposed HTTP endpoints
    pub fn route_docs(&self) -> &[RouteDoc] {
        &self.route_docs
    }
395

396
    pub fn enable_model_endpoint(&self, endpoint_type: EndpointType, enable: bool) {
397
398
399
400
401
402
403
        self.state.flags.set(&endpoint_type, enable);
        tracing::info!(
            "{} endpoints {}",
            endpoint_type.as_str(),
            if enable { "enabled" } else { "disabled" }
        );
    }
404
405
}

406
407
408
409
410
411
412
fn get_graceful_shutdown_timeout() -> usize {
    std::env::var(env_llm::DYN_HTTP_GRACEFUL_SHUTDOWN_TIMEOUT_SECS)
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(5)
}

413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
/// Environment variable to set the metrics endpoint path (default: `/metrics`)
static HTTP_SVC_METRICS_PATH_ENV: &str = "DYN_HTTP_SVC_METRICS_PATH";
/// Environment variable to set the models endpoint path (default: `/v1/models`)
static HTTP_SVC_MODELS_PATH_ENV: &str = "DYN_HTTP_SVC_MODELS_PATH";
/// Environment variable to set the health endpoint path (default: `/health`)
static HTTP_SVC_HEALTH_PATH_ENV: &str = "DYN_HTTP_SVC_HEALTH_PATH";
/// Environment variable to set the live endpoint path (default: `/live`)
static HTTP_SVC_LIVE_PATH_ENV: &str = "DYN_HTTP_SVC_LIVE_PATH";
/// Environment variable to set the chat completions endpoint path (default: `/v1/chat/completions`)
static HTTP_SVC_CHAT_PATH_ENV: &str = "DYN_HTTP_SVC_CHAT_PATH";
/// Environment variable to set the completions endpoint path (default: `/v1/completions`)
static HTTP_SVC_CMP_PATH_ENV: &str = "DYN_HTTP_SVC_CMP_PATH";
/// Environment variable to set the embeddings endpoint path (default: `/v1/embeddings`)
static HTTP_SVC_EMB_PATH_ENV: &str = "DYN_HTTP_SVC_EMB_PATH";
/// Environment variable to set the responses endpoint path (default: `/v1/responses`)
static HTTP_SVC_RESPONSES_PATH_ENV: &str = "DYN_HTTP_SVC_RESPONSES_PATH";
429
430
/// Environment variable to set the anthropic messages endpoint path (default: `/v1/messages`)
static HTTP_SVC_ANTHROPIC_PATH_ENV: &str = "DYN_HTTP_SVC_ANTHROPIC_PATH";
431

432
433
impl HttpServiceConfigBuilder {
    pub fn build(self) -> Result<HttpService, anyhow::Error> {
434
        let config: HttpServiceConfig = self.build_internal()?;
435

436
        let model_manager = Arc::new(ModelManager::new());
437
        let cancel_token = config.cancel_token.unwrap_or_default();
438
439
440
441
442
443
        // Use the provided discovery client, or fall back to a no-op memory-backed one
        // (for in-process modes that don't need discovery)
        let discovery_client = config.discovery.unwrap_or_else(|| {
            use dynamo_runtime::discovery::KVStoreDiscovery;
            Arc::new(KVStoreDiscovery::new(
                dynamo_runtime::storage::kv::Manager::memory(),
444
                cancel_token.child_token(),
445
446
            )) as Arc<dyn Discovery>
        });
447
        let state = Arc::new(State::new(model_manager, discovery_client, cancel_token));
448
449
450
451
452
453
454
455
456
457
458
459
        state
            .flags
            .set(&EndpointType::Chat, config.enable_chat_endpoints);
        state
            .flags
            .set(&EndpointType::Completion, config.enable_cmpl_endpoints);
        state
            .flags
            .set(&EndpointType::Embedding, config.enable_embeddings_endpoints);
        state
            .flags
            .set(&EndpointType::Responses, config.enable_responses_endpoints);
460
461
462
463
        state.flags.set(
            &EndpointType::AnthropicMessages,
            config.enable_anthropic_endpoints,
        );
464

465
466
        // enable prometheus metrics
        let registry = metrics::Registry::new();
467
        state.metrics_clone().register(&registry)?;
468

469
470
471
472
473
474
475
476
477
478
479
480
        // Register worker load metrics (active_decode_blocks, active_prefill_tokens per worker)
        // These are updated by KvWorkerMonitor when receiving ActiveLoad events
        if let Err(e) = register_worker_load_metrics(&registry) {
            tracing::warn!("Failed to register worker load metrics: {}", e);
        }

        // Register worker timing metrics (last_ttft, last_itl per worker)
        // These are updated by ResponseMetricCollector when observing TTFT/ITL
        if let Err(e) = register_worker_timing_metrics(&registry) {
            tracing::warn!("Failed to register worker timing metrics: {}", e);
        }

481
482
483
484
485
486
        // Register router queue metrics (pending requests per worker_type)
        // These are updated by KvScheduler on enqueue/update/free
        if let Err(e) = register_router_queue_metrics(&registry) {
            tracing::warn!("Failed to register router queue metrics: {}", e);
        }

487
488
489
490
491
        if let Some(ref discovery) = config.drt_discovery {
            let instance_id = discovery.instance_id();
            if let Err(e) = RoutingOverheadMetrics::register(&registry, instance_id) {
                tracing::warn!("Failed to register routing overhead metrics: {}", e);
            }
492
493
        }

494
495
496
497
498
499
500
501
502
503
504
505
506
        if let Err(e) = ensure_request_plane_metrics_registered_prometheus(&registry) {
            tracing::warn!("Failed to register request-plane metrics: {}", e);
        }
        if let Err(e) = ensure_frontend_perf_metrics_registered_prometheus(&registry) {
            tracing::warn!("Failed to register frontend perf metrics: {}", e);
        }
        if let Err(e) = ensure_tokio_perf_metrics_registered_prometheus(&registry) {
            tracing::warn!("Failed to register tokio perf metrics: {}", e);
        }
        if let Err(e) = ensure_transport_metrics_registered_prometheus(&registry) {
            tracing::warn!("Failed to register transport metrics: {}", e);
        }

507
508
        let mut all_docs = Vec::new();

509
510
511
512
513
514
515
516
517
518
519
520
521
        // Shared on_response callback for both system and inference routes
        let on_response = |response: &Response<Body>, latency: Duration, _span: &tracing::Span| {
            let status = response.status();
            let latency_ms = latency.as_millis();
            if status.is_server_error() || status.is_client_error() {
                tracing::error!(status = %status.as_u16(), latency_ms = %latency_ms, "http response sent");
            } else {
                tracing::info!(status = %status.as_u16(), latency_ms = %latency_ms, "http response sent");
            }
        };

        // System routes (health, metrics, models) — debug-level spans
        let system_routes = vec![
522
523
524
525
526
            metrics::router(
                registry,
                var(HTTP_SVC_METRICS_PATH_ENV).ok(),
                config.drt_metrics,
            ),
527
528
529
            super::openai::list_models_router(state.clone(), var(HTTP_SVC_MODELS_PATH_ENV).ok()),
            super::health::health_check_router(state.clone(), var(HTTP_SVC_HEALTH_PATH_ENV).ok()),
            super::health::live_check_router(state.clone(), var(HTTP_SVC_LIVE_PATH_ENV).ok()),
530
            super::busy_threshold::busy_threshold_router(state.clone(), None),
531
        ];
532
533
534
535
536
537
        let mut system_router = axum::Router::new();
        for (route_docs, route) in system_routes {
            system_router = system_router.merge(route);
            all_docs.extend(route_docs);
        }
        // Inference routes (completions, chat, embeddings, etc.) — info-level spans
538
539
        let endpoint_routes =
            HttpServiceConfigBuilder::get_endpoints_router(state.clone(), &config.request_template);
540
541
542
        let mut inference_router = axum::Router::new();
        for (route_docs, route) in endpoint_routes {
            inference_router = inference_router.merge(route);
543
544
            all_docs.extend(route_docs);
        }
545
546
547
548
549
        inference_router = inference_router.layer(
            TraceLayer::new_for_http()
                .make_span_with(make_inference_request_span)
                .on_response(on_response),
        );
550

551
        // OpenAPI documentation routes (system)
552
553
        let (openapi_docs, openapi_route) =
            super::openapi_docs::openapi_router(all_docs.clone(), None);
554
        system_router = system_router.merge(openapi_route);
555
556
        all_docs.extend(openapi_docs);

557
        system_router = system_router.layer(
558
            TraceLayer::new_for_http()
559
560
                .make_span_with(make_system_request_span)
                .on_response(on_response),
561
        );
562

563
564
565
566
567
        let router = system_router.merge(inference_router);

        // Echo x-request-id from request to response headers for client correlation
        let router = router.layer(axum::middleware::from_fn(echo_request_id_header));

568
        Ok(HttpService {
569
            state,
570
571
            router,
            port: config.port,
572
            host: config.host,
Graham King's avatar
Graham King committed
573
574
575
            enable_tls: config.enable_tls,
            tls_cert_path: config.tls_cert_path,
            tls_key_path: config.tls_key_path,
576
            route_docs: all_docs,
577
578
        })
    }
579
580
581
582
583

    pub fn with_request_template(mut self, request_template: Option<RequestTemplate>) -> Self {
        self.request_template = Some(request_template);
        self
    }
584

585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
    fn get_endpoints_router(
        state: Arc<State>,
        request_template: &Option<RequestTemplate>,
    ) -> Vec<(Vec<RouteDoc>, axum::Router)> {
        let mut routes = Vec::new();
        // Add chat completions route with conditional middleware
        let (chat_docs, chat_route) = super::openai::chat_completions_router(
            state.clone(),
            request_template.clone(),
            var(HTTP_SVC_CHAT_PATH_ENV).ok(),
        );
        let (cmpl_docs, cmpl_route) =
            super::openai::completions_router(state.clone(), var(HTTP_SVC_CMP_PATH_ENV).ok());
        let (embed_docs, embed_route) =
            super::openai::embeddings_router(state.clone(), var(HTTP_SVC_EMB_PATH_ENV).ok());
600
        let (images_docs, images_route) = super::openai::images_router(state.clone(), None);
601
        let (videos_docs, videos_route) = super::openai::videos_router(state.clone(), None);
602
        let (audios_docs, audios_route) = super::openai::audios_router(state.clone(), None);
603
604
605
606
607
608
609
610
611
        let (responses_docs, responses_route) = super::openai::responses_router(
            state.clone(),
            request_template.clone(),
            var(HTTP_SVC_RESPONSES_PATH_ENV).ok(),
        );
        let mut endpoint_routes = HashMap::new();
        endpoint_routes.insert(EndpointType::Chat, (chat_docs, chat_route));
        endpoint_routes.insert(EndpointType::Completion, (cmpl_docs, cmpl_route));
        endpoint_routes.insert(EndpointType::Embedding, (embed_docs, embed_route));
612
        endpoint_routes.insert(EndpointType::Images, (images_docs, images_route));
613
        endpoint_routes.insert(EndpointType::Videos, (videos_docs, videos_route));
614
        endpoint_routes.insert(EndpointType::Audios, (audios_docs, audios_route));
615
616
        endpoint_routes.insert(EndpointType::Responses, (responses_docs, responses_route));

617
618
619
620
621
622
623
624
625
626
627
628
629
        if env_is_truthy(env_llm::DYN_ENABLE_ANTHROPIC_API) {
            tracing::warn!("Anthropic Messages API (/v1/messages) is experimental.");
            let (anthropic_docs, anthropic_route) = super::anthropic::anthropic_messages_router(
                state.clone(),
                request_template.clone(),
                var(HTTP_SVC_ANTHROPIC_PATH_ENV).ok(),
            );
            endpoint_routes.insert(
                EndpointType::AnthropicMessages,
                (anthropic_docs, anthropic_route),
            );
        }

630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
        for endpoint_type in EndpointType::all() {
            let state_route = state.clone();
            if !endpoint_routes.contains_key(&endpoint_type) {
                tracing::debug!("{} endpoints are disabled", endpoint_type.as_str());
                continue;
            }
            let (docs, route) = endpoint_routes.get(&endpoint_type).cloned().unwrap();
            let route = route.route_layer(axum::middleware::from_fn(
                move |req: axum::http::Request<axum::body::Body>, next: axum::middleware::Next| {
                    let state: Arc<State> = state_route.clone();
                    async move {
                        // Check if the endpoint is enabled
                        let enabled = state.flags.get(&endpoint_type);
                        if enabled {
                            Ok(next.run(req).await)
                        } else {
                            tracing::debug!("{} endpoints are disabled", endpoint_type.as_str());
647
                            Err(axum::http::StatusCode::NOT_FOUND)
648
649
650
651
652
653
654
655
                        }
                    }
                },
            ));
            routes.push((docs, route));
        }
        routes
    }
656
}
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701

#[cfg(test)]
mod tests {
    use super::*;
    use serial_test::serial;
    use std::sync::Arc;
    use tokio_util::sync::CancellationToken;

    #[tokio::test]
    #[serial]
    async fn test_liveness_endpoint_reflects_cancellation() {
        // 1. Setup service & token
        let cancel_token = Arc::new(CancellationToken::new());
        let service = HttpService::builder().build().unwrap();
        let port = service.port;

        // 2. Spawn service with shared token
        let service_token = cancel_token.clone();
        let handle = tokio::spawn(async move {
            service.run((*service_token).clone()).await.unwrap();
        });

        tokio::time::sleep(std::time::Duration::from_millis(1)).await;

        // 3. Cancel the token
        cancel_token.cancel();

        // 4. Wait a tiny bit for propagation
        tokio::time::sleep(std::time::Duration::from_millis(20)).await;

        // 5. Hit the endpoint
        let client = reqwest::Client::new();
        let resp = client
            .get(format!("http://localhost:{}/live", port))
            .send()
            .await
            .expect("Request failed");

        // 6. ASSERTION: Should be 503 Service Unavailable
        assert_eq!(resp.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE);

        // Clean up
        handle.abort();
    }
}