dynamo_runtime/
system_status_server.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4// TODO: (DEP-635) this file should be renamed to system_http_server.rs
5//  it is being used not just for status, health, but others like loras management.
6
7use crate::config::HealthStatus;
8use crate::config::environment_names::logging as env_logging;
9use crate::config::environment_names::runtime::canary as env_canary;
10use crate::config::environment_names::runtime::system as env_system;
11use crate::logging::make_system_request_span;
12use crate::metrics::MetricsHierarchy;
13use crate::traits::DistributedRuntimeProvider;
14use axum::{
15    Router,
16    body::Bytes,
17    extract::{Json, Path, State},
18    http::StatusCode,
19    response::IntoResponse,
20    routing::{any, delete, get, post},
21};
22use futures::StreamExt;
23use serde::{Deserialize, Serialize};
24use serde_json::json;
25use std::collections::HashMap;
26use std::sync::{Arc, OnceLock};
27use std::time::Instant;
28use tokio::{net::TcpListener, task::JoinHandle};
29use tokio_util::sync::CancellationToken;
30use tower_http::trace::TraceLayer;
31
32/// System status server information containing socket address and handle
33#[derive(Debug)]
34pub struct SystemStatusServerInfo {
35    pub socket_addr: std::net::SocketAddr,
36    pub handle: Option<Arc<JoinHandle<()>>>,
37}
38
39impl SystemStatusServerInfo {
40    pub fn new(socket_addr: std::net::SocketAddr, handle: Option<JoinHandle<()>>) -> Self {
41        Self {
42            socket_addr,
43            handle: handle.map(Arc::new),
44        }
45    }
46
47    pub fn address(&self) -> String {
48        self.socket_addr.to_string()
49    }
50
51    pub fn hostname(&self) -> String {
52        self.socket_addr.ip().to_string()
53    }
54
55    pub fn port(&self) -> u16 {
56        self.socket_addr.port()
57    }
58}
59
60impl Clone for SystemStatusServerInfo {
61    fn clone(&self) -> Self {
62        Self {
63            socket_addr: self.socket_addr,
64            handle: self.handle.clone(),
65        }
66    }
67}
68
69/// System status server state containing the distributed runtime reference
70pub struct SystemStatusState {
71    // global drt registry is for printing out the entire Prometheus format output
72    root_drt: Arc<crate::DistributedRuntime>,
73    // Discovery metadata (only for Kubernetes backend)
74    discovery_metadata: Option<Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>>,
75}
76
77impl SystemStatusState {
78    /// Create new system status server state with the provided distributed runtime
79    pub fn new(
80        drt: Arc<crate::DistributedRuntime>,
81        discovery_metadata: Option<Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>>,
82    ) -> anyhow::Result<Self> {
83        Ok(Self {
84            root_drt: drt,
85            discovery_metadata,
86        })
87    }
88
89    /// Get a reference to the distributed runtime
90    pub fn drt(&self) -> &crate::DistributedRuntime {
91        &self.root_drt
92    }
93
94    /// Get a reference to the discovery metadata if available
95    pub fn discovery_metadata(
96        &self,
97    ) -> Option<&Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>> {
98        self.discovery_metadata.as_ref()
99    }
100}
101
102/// Request body for POST /v1/loras
103#[derive(Debug, Clone, Deserialize, Serialize)]
104pub struct LoadLoraRequest {
105    pub lora_name: String,
106    pub source: LoraSource,
107}
108
109/// Source information for loading a LoRA
110#[derive(Debug, Clone, Deserialize, Serialize)]
111pub struct LoraSource {
112    pub uri: String,
113}
114
115/// Response body for LoRA operations
116#[derive(Debug, Clone, Deserialize, Serialize)]
117pub struct LoraResponse {
118    pub status: String,
119    #[serde(skip_serializing_if = "Option::is_none")]
120    pub message: Option<String>,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub lora_name: Option<String>,
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub lora_id: Option<u64>,
125    #[serde(skip_serializing_if = "Option::is_none")]
126    pub loras: Option<serde_json::Value>,
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub count: Option<usize>,
129}
130
131/// Start system status server with metrics support
132pub async fn spawn_system_status_server(
133    host: &str,
134    port: u16,
135    cancel_token: CancellationToken,
136    drt: Arc<crate::DistributedRuntime>,
137    discovery_metadata: Option<Arc<tokio::sync::RwLock<crate::discovery::DiscoveryMetadata>>>,
138) -> anyhow::Result<(std::net::SocketAddr, tokio::task::JoinHandle<()>)> {
139    // Create system status server state with the provided distributed runtime
140    let server_state = Arc::new(SystemStatusState::new(drt, discovery_metadata)?);
141    let health_path = server_state
142        .drt()
143        .system_health()
144        .lock()
145        .health_path()
146        .to_string();
147    let live_path = server_state
148        .drt()
149        .system_health()
150        .lock()
151        .live_path()
152        .to_string();
153
154    // Check if LoRA feature is enabled
155    let lora_enabled =
156        crate::config::env_is_truthy(crate::config::environment_names::llm::DYN_LORA_ENABLED);
157
158    let mut app = Router::new()
159        .route(
160            &health_path,
161            get({
162                let state = Arc::clone(&server_state);
163                move || health_handler(state)
164            }),
165        )
166        .route(
167            &live_path,
168            get({
169                let state = Arc::clone(&server_state);
170                move || health_handler(state)
171            }),
172        )
173        .route(
174            "/metrics",
175            get({
176                let state = Arc::clone(&server_state);
177                move || metrics_handler(state)
178            }),
179        )
180        .route(
181            "/metadata",
182            get({
183                let state = Arc::clone(&server_state);
184                move || metadata_handler(state)
185            }),
186        )
187        .route(
188            "/engine/{*path}",
189            any({
190                let state = Arc::clone(&server_state);
191                move |path, body| engine_route_handler(state, path, body)
192            }),
193        );
194
195    // Add LoRA routes only if DYN_LORA_ENABLED is set to true
196    if lora_enabled {
197        app = app
198            .route(
199                "/v1/loras",
200                get({
201                    let state = Arc::clone(&server_state);
202                    move || list_loras_handler(State(state))
203                })
204                .post({
205                    let state = Arc::clone(&server_state);
206                    move |body| load_lora_handler(State(state), body)
207                }),
208            )
209            .route(
210                "/v1/loras/{*lora_name}",
211                delete({
212                    let state = Arc::clone(&server_state);
213                    move |path| unload_lora_handler(State(state), path)
214                }),
215            );
216    }
217
218    // Self-hosted MDC files. Always mounted; empty registry → 404.
219    // The endpoint triple disambiguates multi-LocalModel-per-DRT; the
220    // suffix segment (LoRA slug or `_base`) scopes per-registration so
221    // detaching one doesn't wipe another's entries.
222    app = app.route(
223        "/v1/metadata/{namespace}/{component}/{endpoint}/{model_slug}/{model_suffix}/{*filename}",
224        get({
225            let state = Arc::clone(&server_state);
226            move |path| metadata_file_handler(State(state), path)
227        }),
228    );
229
230    let app = app
231        .fallback(|| async {
232            tracing::info!("[fallback handler] called");
233            (StatusCode::NOT_FOUND, "Route not found").into_response()
234        })
235        .layer(TraceLayer::new_for_http().make_span_with(make_system_request_span));
236
237    let address = format!("{}:{}", host, port);
238    tracing::info!("[spawn_system_status_server] binding to: {address}");
239
240    let listener = match TcpListener::bind(&address).await {
241        Ok(listener) => {
242            // get the actual address and port, print in debug level
243            let actual_address = listener.local_addr()?;
244            tracing::info!(
245                "[spawn_system_status_server] system status server bound to: {}",
246                actual_address
247            );
248            (listener, actual_address)
249        }
250        Err(e) => {
251            tracing::error!("Failed to bind to address {}: {}", address, e);
252            return Err(anyhow::anyhow!("Failed to bind to address: {}", e));
253        }
254    };
255    let (listener, actual_address) = listener;
256
257    let observer = cancel_token.child_token();
258    // Spawn the server in the background and return the handle
259    let handle = tokio::spawn(async move {
260        if let Err(e) = axum::serve(listener, app)
261            .with_graceful_shutdown(observer.cancelled_owned())
262            .await
263        {
264            tracing::error!("System status server error: {e}");
265        }
266    });
267
268    Ok((actual_address, handle))
269}
270
271/// Health handler with optional active health checking
272#[tracing::instrument(skip_all, level = "trace")]
273async fn health_handler(state: Arc<SystemStatusState>) -> impl IntoResponse {
274    // Get basic health status
275    let system_health = state.drt().system_health();
276    let system_health_lock = system_health.lock();
277    let (healthy, endpoints) = system_health_lock.get_health_status();
278    let uptime = Some(system_health_lock.uptime());
279    drop(system_health_lock);
280
281    let healthy_string = if healthy { "ready" } else { "notready" };
282    let status_code = if healthy {
283        StatusCode::OK
284    } else {
285        StatusCode::SERVICE_UNAVAILABLE
286    };
287
288    let response = json!({
289        "status": healthy_string,
290        "uptime": uptime,
291        "endpoints": endpoints,
292    });
293
294    tracing::trace!("Response {}", response.to_string());
295
296    (status_code, response.to_string())
297}
298
299/// Metrics handler with DistributedRuntime uptime
300#[tracing::instrument(skip_all, level = "trace")]
301async fn metrics_handler(state: Arc<SystemStatusState>) -> impl IntoResponse {
302    // Get all metrics from the DistributedRuntime.
303    // The uptime gauge is updated automatically via a PrometheusUpdateCallback
304    // registered in DistributedRuntime::new(), so it is always fresh before scrape.
305    //
306    // NOTE: We use a multi-registry model (e.g. one registry per endpoint) and merge at scrape time,
307    // so /metrics traverses registered child registries and produces a single combined output.
308    let response = match state.drt().metrics().prometheus_expfmt() {
309        Ok(r) => r,
310        Err(e) => {
311            tracing::error!("Failed to get metrics from registry: {e}");
312            return (
313                StatusCode::INTERNAL_SERVER_ERROR,
314                "Failed to get metrics".to_string(),
315            );
316        }
317    };
318
319    (StatusCode::OK, response)
320}
321
322/// Metadata handler
323#[tracing::instrument(skip_all, level = "trace")]
324async fn metadata_handler(state: Arc<SystemStatusState>) -> impl IntoResponse {
325    // Check if discovery metadata is available
326    let metadata = match state.discovery_metadata() {
327        Some(metadata) => metadata,
328        None => {
329            tracing::debug!("Metadata endpoint called but no discovery metadata available");
330            return (
331                StatusCode::NOT_FOUND,
332                "Discovery metadata not available".to_string(),
333            )
334                .into_response();
335        }
336    };
337
338    // Read the metadata
339    let metadata_guard = metadata.read().await;
340
341    // Serialize to JSON
342    match serde_json::to_string(&*metadata_guard) {
343        Ok(json) => {
344            tracing::trace!("Returning metadata: {} bytes", json.len());
345            (StatusCode::OK, json).into_response()
346        }
347        Err(e) => {
348            tracing::error!("Failed to serialize metadata: {e}");
349            (
350                StatusCode::INTERNAL_SERVER_ERROR,
351                "Failed to serialize metadata".to_string(),
352            )
353                .into_response()
354        }
355    }
356}
357
358/// Handler for POST /v1/loras - Load a LoRA adapter
359#[tracing::instrument(skip_all, level = "debug")]
360async fn load_lora_handler(
361    State(state): State<Arc<SystemStatusState>>,
362    Json(request): Json<LoadLoraRequest>,
363) -> impl IntoResponse {
364    tracing::info!("Loading LoRA: {}", request.lora_name);
365
366    // Call the load_lora endpoint for each available backend
367    match call_lora_endpoint(
368        state.drt(),
369        "load_lora",
370        json!({
371            "lora_name": request.lora_name,
372            "source": {
373                "uri": request.source.uri
374            },
375        }),
376    )
377    .await
378    {
379        Ok(response) => {
380            if response.status == "error" {
381                tracing::error!(
382                    "Failed to load LoRA {}: {}",
383                    request.lora_name,
384                    response.message.as_deref().unwrap_or("Unknown error")
385                );
386                (StatusCode::INTERNAL_SERVER_ERROR, Json(response))
387            } else {
388                tracing::info!("LoRA loaded successfully: {}", request.lora_name);
389                (StatusCode::OK, Json(response))
390            }
391        }
392        Err(e) => {
393            tracing::error!("Failed to load LoRA {}: {}", request.lora_name, e);
394            (
395                StatusCode::INTERNAL_SERVER_ERROR,
396                Json(LoraResponse {
397                    status: "error".to_string(),
398                    message: Some(e.to_string()),
399                    lora_name: Some(request.lora_name),
400                    lora_id: None,
401                    loras: None,
402                    count: None,
403                }),
404            )
405        }
406    }
407}
408
409/// Handler for DELETE /v1/loras/*lora_name - Unload a LoRA adapter
410#[tracing::instrument(skip_all, level = "debug")]
411async fn unload_lora_handler(
412    State(state): State<Arc<SystemStatusState>>,
413    Path(lora_name): Path<String>,
414) -> impl IntoResponse {
415    // Strip the leading slash from the wildcard capture
416    let lora_name = lora_name
417        .strip_prefix('/')
418        .unwrap_or(&lora_name)
419        .to_string();
420    tracing::info!("Unloading LoRA: {lora_name}");
421
422    // Call the unload_lora endpoint for each available backend
423    match call_lora_endpoint(
424        state.drt(),
425        "unload_lora",
426        json!({
427            "lora_name": lora_name.clone(),
428        }),
429    )
430    .await
431    {
432        Ok(response) => {
433            if response.status == "error" {
434                tracing::error!(
435                    "Failed to unload LoRA {}: {}",
436                    lora_name,
437                    response.message.as_deref().unwrap_or("Unknown error")
438                );
439                (StatusCode::INTERNAL_SERVER_ERROR, Json(response))
440            } else {
441                tracing::info!("LoRA unloaded successfully: {lora_name}");
442                (StatusCode::OK, Json(response))
443            }
444        }
445        Err(e) => {
446            tracing::error!("Failed to unload LoRA {}: {}", lora_name, e);
447            (
448                StatusCode::INTERNAL_SERVER_ERROR,
449                Json(LoraResponse {
450                    status: "error".to_string(),
451                    message: Some(e.to_string()),
452                    lora_name: Some(lora_name),
453                    lora_id: None,
454                    loras: None,
455                    count: None,
456                }),
457            )
458        }
459    }
460}
461
462/// Handler for GET /v1/loras - List all LoRA adapters
463#[tracing::instrument(skip_all, level = "debug")]
464async fn list_loras_handler(State(state): State<Arc<SystemStatusState>>) -> impl IntoResponse {
465    tracing::info!("Listing all LoRAs");
466
467    // Call the list_loras endpoint for each available backend
468    match call_lora_endpoint(state.drt(), "list_loras", json!({})).await {
469        Ok(response) => {
470            tracing::info!("Successfully retrieved LoRA list");
471            (StatusCode::OK, Json(response))
472        }
473        Err(e) => {
474            tracing::error!("Failed to list LoRAs: {e}");
475            (
476                StatusCode::INTERNAL_SERVER_ERROR,
477                Json(LoraResponse {
478                    status: "error".to_string(),
479                    message: Some(e.to_string()),
480                    lora_name: None,
481                    lora_id: None,
482                    loras: None,
483                    count: None,
484                }),
485            )
486        }
487    }
488}
489
490/// `GET /v1/metadata/{namespace}/{component}/{endpoint}/{slug}/{suffix}/{filename}`
491/// — 404 on miss, 500 on read error, raw bytes on hit. Consumer blake3-verifies.
492async fn metadata_file_handler(
493    State(state): State<Arc<SystemStatusState>>,
494    Path((namespace, component, endpoint, model_slug, model_suffix, filename)): Path<(
495        String,
496        String,
497        String,
498        String,
499        String,
500        String,
501    )>,
502) -> impl IntoResponse {
503    let path = match state.drt().metadata_artifacts().get(
504        &namespace,
505        &component,
506        &endpoint,
507        &model_slug,
508        &model_suffix,
509        &filename,
510    ) {
511        Some(p) => p,
512        None => {
513            tracing::debug!(
514                namespace,
515                component,
516                endpoint,
517                model_slug,
518                model_suffix,
519                filename,
520                "metadata artifact not registered for self-host"
521            );
522            return (StatusCode::NOT_FOUND, "Not found").into_response();
523        }
524    };
525
526    match tokio::fs::read(&path).await {
527        Ok(bytes) => (StatusCode::OK, bytes).into_response(),
528        Err(err) => {
529            tracing::error!(
530                namespace,
531                component,
532                endpoint,
533                model_slug,
534                model_suffix,
535                filename,
536                path = %path.display(),
537                %err,
538                "failed to read self-hosted metadata file"
539            );
540            (StatusCode::INTERNAL_SERVER_ERROR, "Failed to read file").into_response()
541        }
542    }
543}
544
545/// Helper function to call a LoRA management endpoint for the local worker.
546///
547/// Resolution order (both are in-process, never network discovery):
548/// 1. The legacy local endpoint registry, populated by non-unified workers via
549///    `.register_local_engine()`.
550/// 2. The generic engine-route registry (`/engine/*`). Unified-backend workers
551///    advertise LoRA lifecycle ops (`load_lora`/`unload_lora`/`list_loras`) as
552///    engine *updates*, registered under `update/<name>`; this fallback maps the
553///    bare LoRA name onto that key so the legacy `/v1/loras` surface forwards to
554///    them (the `/v1/loras` compatibility shim).
555///
556/// Because legacy workers populate the local registry, they never reach the
557/// fallback — their `/v1/loras` behavior is unchanged. If neither registry
558/// holds the name, returns an explicit "LoRA management not available" error
559/// rather than an opaque "endpoint not found".
560async fn call_lora_endpoint(
561    drt: &crate::DistributedRuntime,
562    endpoint_name: &str,
563    request_body: serde_json::Value,
564) -> anyhow::Result<LoraResponse> {
565    use crate::engine::AsyncEngine;
566
567    tracing::debug!("Calling LoRA endpoint: '{endpoint_name}'");
568
569    // 1. Legacy local registry (in-process call only).
570    if let Some(engine) = drt.local_endpoint_registry().get(endpoint_name) {
571        tracing::debug!(
572            "Found endpoint '{}' in local registry, calling directly",
573            endpoint_name
574        );
575
576        let request = crate::pipeline::SingleIn::new(request_body);
577        let mut stream = engine.generate(request).await?;
578
579        if let Some(response) = stream.next().await {
580            let response_data = response.data.unwrap_or_default();
581            let lora_response = serde_json::from_value::<LoraResponse>(response_data.clone())
582                .unwrap_or_else(|_| parse_lora_response(&response_data));
583            return Ok(lora_response);
584        }
585
586        anyhow::bail!("No response received from endpoint '{}'", endpoint_name)
587    }
588
589    // 2. Unified-backend engine-update registry fallback. The unified Worker
590    //    registers LoRA ops as engine updates under `update/<name>`, so map the
591    //    bare LoRA endpoint name onto that namespaced key.
592    let update_key = format!("update/{endpoint_name}");
593    if let Some(callback) = drt.engine_routes().get(&update_key) {
594        tracing::debug!(
595            "Found '{}' in engine routes registry, invoking update callback",
596            update_key
597        );
598        let response_data = callback(request_body).await?;
599        let lora_response = serde_json::from_value::<LoraResponse>(response_data.clone())
600            .unwrap_or_else(|_| parse_lora_response(&response_data));
601        return Ok(lora_response);
602    }
603
604    anyhow::bail!(
605        "LoRA management not available: no '{}' handler is registered \
606         (neither a local LoRA endpoint nor an engine update). This worker \
607         either has LoRA disabled or its backend does not support LoRA \
608         management.",
609        endpoint_name
610    )
611}
612
613/// Helper to parse response data into LoraResponse
614fn parse_lora_response(response_data: &serde_json::Value) -> LoraResponse {
615    LoraResponse {
616        status: response_data
617            .get("status")
618            .and_then(|s| s.as_str())
619            .unwrap_or("success")
620            .to_string(),
621        message: response_data
622            .get("message")
623            .and_then(|m| m.as_str())
624            .map(|s| s.to_string()),
625        lora_name: response_data
626            .get("lora_name")
627            .and_then(|n| n.as_str())
628            .map(|s| s.to_string()),
629        lora_id: response_data.get("lora_id").and_then(|id| id.as_u64()),
630        loras: response_data.get("loras").cloned(),
631        count: response_data
632            .get("count")
633            .and_then(|c| c.as_u64())
634            .map(|c| c as usize),
635    }
636}
637
638/// Engine route handler for /engine/* routes
639///
640/// This handler looks up registered callbacks in the engine routes registry
641/// and invokes them with the request body, returning the response as JSON.
642#[tracing::instrument(skip_all, level = "trace", fields(path = %path))]
643async fn engine_route_handler(
644    state: Arc<SystemStatusState>,
645    Path(path): Path<String>,
646    body: Bytes,
647) -> impl IntoResponse {
648    tracing::trace!("Engine route request to /engine/{path}");
649
650    // Parse body as JSON (empty object for GET/empty body)
651    let body_json: serde_json::Value = if body.is_empty() {
652        serde_json::json!({})
653    } else {
654        match serde_json::from_slice(&body) {
655            Ok(json) => json,
656            Err(e) => {
657                tracing::warn!("Invalid JSON in request body: {e}");
658                return (
659                    StatusCode::BAD_REQUEST,
660                    json!({
661                        "error": "Invalid JSON",
662                        "message": format!("{}", e)
663                    })
664                    .to_string(),
665                )
666                    .into_response();
667            }
668        }
669    };
670
671    // Look up callback
672    let callback = match state.drt().engine_routes().get(&path) {
673        Some(cb) => cb,
674        None => {
675            tracing::debug!("Route /engine/{path} not found");
676            return (
677                StatusCode::NOT_FOUND,
678                json!({
679                    "error": "Route not found",
680                    "message": format!("Route /engine/{} not found", path)
681                })
682                .to_string(),
683            )
684                .into_response();
685        }
686    };
687
688    // Call callback (it's async, so await it)
689    match callback(body_json).await {
690        Ok(response) => {
691            tracing::trace!("Engine route handler succeeded for /engine/{path}");
692            (StatusCode::OK, response.to_string()).into_response()
693        }
694        Err(e) => {
695            tracing::error!("Engine route handler error for /engine/{}: {}", path, e);
696            (
697                StatusCode::INTERNAL_SERVER_ERROR,
698                json!({
699                    "error": "Handler error",
700                    "message": format!("{}", e)
701                })
702                .to_string(),
703            )
704                .into_response()
705        }
706    }
707}
708
709// Regular tests: cargo test system_status_server --lib
710#[cfg(test)]
711mod tests {
712    use super::*;
713    use tokio::time::Duration;
714
715    // This is a basic test to verify the HTTP server is working before testing other more complicated tests
716    #[tokio::test]
717    async fn test_http_server_lifecycle() {
718        let cancel_token = CancellationToken::new();
719        let cancel_token_for_server = cancel_token.clone();
720
721        // Test basic HTTP server lifecycle without DistributedRuntime
722        let app = Router::new().route("/test", get(|| async { (StatusCode::OK, "test") }));
723
724        // start HTTP server
725        let server_handle = tokio::spawn(async move {
726            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
727            let _ = axum::serve(listener, app)
728                .with_graceful_shutdown(cancel_token_for_server.cancelled_owned())
729                .await;
730        });
731
732        // server starts immediately, no need to wait
733
734        // cancel token
735        cancel_token.cancel();
736
737        // wait for the server to shut down
738        let result = tokio::time::timeout(Duration::from_secs(5), server_handle).await;
739        assert!(
740            result.is_ok(),
741            "HTTP server should shut down when cancel token is cancelled"
742        );
743    }
744}
745
746// Integration tests: cargo test system_status_server --lib --features integration
747#[cfg(all(test, feature = "integration"))]
748mod integration_tests {
749    use super::*;
750    use crate::config::environment_names::logging as env_logging;
751    use crate::config::environment_names::runtime::canary as env_canary;
752    use crate::distributed::distributed_test_utils::create_test_drt_async;
753    use crate::metrics::MetricsHierarchy;
754    use anyhow::Result;
755    use rstest::rstest;
756    use std::sync::Arc;
757    use tokio::time::Duration;
758
759    #[tokio::test]
760    async fn test_uptime_from_system_health() {
761        // Test that uptime is available from SystemHealth
762        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
763            let drt = create_test_drt_async().await;
764
765            // Get uptime from SystemHealth
766            let uptime = drt.system_health().lock().uptime();
767            // Uptime should exist (even if close to zero)
768            assert!(uptime.as_nanos() > 0 || uptime.is_zero());
769
770            // Sleep briefly and check uptime increases
771            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
772            let uptime_after = drt.system_health().lock().uptime();
773            assert!(uptime_after > uptime);
774        })
775        .await;
776    }
777
778    #[tokio::test]
779    async fn test_runtime_metrics_initialization_and_namespace() {
780        // Test that metrics have correct namespace
781        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
782            let drt = create_test_drt_async().await;
783            // SystemStatusState is already created in distributed.rs
784            // so we don't need to create it again here
785
786            // The uptime_seconds metric should already be registered and available
787            let response = drt.metrics().prometheus_expfmt().unwrap();
788            println!("Full metrics response:\n{}", response);
789
790            // Check that uptime_seconds metric is present with correct namespace
791            assert!(
792                response.contains("# HELP dynamo_component_uptime_seconds"),
793                "Should contain uptime_seconds help text"
794            );
795            assert!(
796                response.contains("# TYPE dynamo_component_uptime_seconds gauge"),
797                "Should contain uptime_seconds type"
798            );
799            assert!(
800                response.contains("dynamo_component_uptime_seconds"),
801                "Should contain uptime_seconds metric with correct namespace"
802            );
803        })
804        .await;
805    }
806
807    #[tokio::test]
808    async fn test_uptime_gauge_updates() {
809        // Test that the uptime gauge is properly updated and increases over time
810        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
811            let drt = create_test_drt_async().await;
812
813            // Get initial uptime
814            let initial_uptime = drt.system_health().lock().uptime();
815
816            // Update the gauge with initial value
817            drt.system_health().lock().update_uptime_gauge();
818
819            // Sleep for 100ms
820            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
821
822            // Get uptime after sleep
823            let uptime_after_sleep = drt.system_health().lock().uptime();
824
825            // Update the gauge again
826            drt.system_health().lock().update_uptime_gauge();
827
828            // Verify uptime increased by at least 100ms
829            let elapsed = uptime_after_sleep - initial_uptime;
830            assert!(
831                elapsed >= std::time::Duration::from_millis(100),
832                "Uptime should have increased by at least 100ms after sleep, but only increased by {:?}",
833                elapsed
834            );
835        })
836        .await;
837    }
838
839    #[tokio::test]
840    async fn test_http_requests_fail_when_system_disabled() {
841        // Test that system status server is not running when disabled
842        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
843            let drt = create_test_drt_async().await;
844
845            // Verify that system status server info is None when disabled
846            let system_info = drt.system_status_server_info();
847            assert!(
848                system_info.is_none(),
849                "System status server should not be running when disabled"
850            );
851
852            println!("✓ System status server correctly disabled when not enabled");
853        })
854        .await;
855    }
856
857    /// This test verifies the health and liveness endpoints of the system status server.
858    /// It checks that the endpoints respond with the correct HTTP status codes and bodies
859    /// based on the initial health status and any custom endpoint paths provided via environment variables.
860    /// The test is parameterized using multiple #[case] attributes to cover various scenarios,
861    /// including different initial health states ("ready" and "notready"), default and custom endpoint paths,
862    /// and expected response codes and bodies.
863    #[rstest]
864    #[case("ready", 200, "ready", None, None, 3)]
865    #[case("notready", 503, "notready", None, None, 3)]
866    #[case("ready", 200, "ready", Some("/custom/health"), Some("/custom/live"), 5)]
867    #[case(
868        "notready",
869        503,
870        "notready",
871        Some("/custom/health"),
872        Some("/custom/live"),
873        5
874    )]
875    #[tokio::test]
876    #[cfg(feature = "integration")]
877    async fn test_health_endpoints(
878        #[case] starting_health_status: &'static str,
879        #[case] expected_status: u16,
880        #[case] expected_body: &'static str,
881        #[case] custom_health_path: Option<&'static str>,
882        #[case] custom_live_path: Option<&'static str>,
883        #[case] expected_num_tests: usize,
884    ) {
885        use std::sync::Arc;
886        // use tokio::io::{AsyncReadExt, AsyncWriteExt};
887        // use reqwest for HTTP requests
888
889        // Closure call is needed here to satisfy async_with_vars
890
891        crate::logging::init();
892
893        #[allow(clippy::redundant_closure_call)]
894        temp_env::async_with_vars(
895            [
896                (env_system::DYN_SYSTEM_PORT, Some("0")),
897                (
898                    env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS,
899                    Some(starting_health_status),
900                ),
901                (env_system::DYN_SYSTEM_HEALTH_PATH, custom_health_path),
902                (env_system::DYN_SYSTEM_LIVE_PATH, custom_live_path),
903            ],
904            (async || {
905                let drt = Arc::new(create_test_drt_async().await);
906
907                // Get system status server info from DRT (instead of manually spawning)
908                let system_info = drt
909                    .system_status_server_info()
910                    .expect("System status server should be started by DRT");
911                let addr = system_info.socket_addr;
912
913                let client = reqwest::Client::new();
914
915                // Prepare test cases
916                let mut test_cases = vec![];
917                match custom_health_path {
918                    None => {
919                        // When using default paths, test the default paths
920                        test_cases.push(("/health", expected_status, expected_body));
921                    }
922                    Some(chp) => {
923                        // When using custom paths, default paths should not exist
924                        test_cases.push(("/health", 404, "Route not found"));
925                        test_cases.push((chp, expected_status, expected_body));
926                    }
927                }
928                match custom_live_path {
929                    None => {
930                        // When using default paths, test the default paths
931                        test_cases.push(("/live", expected_status, expected_body));
932                    }
933                    Some(clp) => {
934                        // When using custom paths, default paths should not exist
935                        test_cases.push(("/live", 404, "Route not found"));
936                        test_cases.push((clp, expected_status, expected_body));
937                    }
938                }
939                test_cases.push(("/someRandomPathNotFoundHere", 404, "Route not found"));
940                assert_eq!(test_cases.len(), expected_num_tests);
941
942                for (path, expect_status, expect_body) in test_cases {
943                    println!("[test] Sending request to {}", path);
944                    let url = format!("http://{}{}", addr, path);
945                    let response = client.get(&url).send().await.unwrap();
946                    let status = response.status();
947                    let body = response.text().await.unwrap();
948                    println!(
949                        "[test] Response for {}: status={}, body={:?}",
950                        path, status, body
951                    );
952                    assert_eq!(
953                        status, expect_status,
954                        "Response: status={}, body={:?}",
955                        status, body
956                    );
957                    assert!(
958                        body.contains(expect_body),
959                        "Response: status={}, body={:?}",
960                        status,
961                        body
962                    );
963                }
964            })(),
965        )
966        .await;
967    }
968
969    #[tokio::test]
970    async fn test_health_endpoint_tracing() -> Result<()> {
971        use std::sync::Arc;
972
973        // Closure call is needed here to satisfy async_with_vars
974
975        #[allow(clippy::redundant_closure_call)]
976        let _ = temp_env::async_with_vars(
977            [
978                (env_system::DYN_SYSTEM_PORT, Some("0")),
979                (env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS, Some("ready")),
980                (env_logging::DYN_LOGGING_JSONL, Some("1")),
981                (env_logging::DYN_LOG, Some("trace")),
982            ],
983            (async || {
984                // TODO Add proper testing for
985                // trace id and parent id
986
987                crate::logging::init();
988
989                let drt = Arc::new(create_test_drt_async().await);
990
991                // Get system status server info from DRT (instead of manually spawning)
992                let system_info = drt
993                    .system_status_server_info()
994                    .expect("System status server should be started by DRT");
995                let addr = system_info.socket_addr;
996                let client = reqwest::Client::new();
997                for path in [("/health"), ("/live"), ("/someRandomPathNotFoundHere")] {
998                    let traceparent_value =
999                        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
1000                    let tracestate_value = "vendor1=opaqueValue1,vendor2=opaqueValue2";
1001                    let mut headers = reqwest::header::HeaderMap::new();
1002                    headers.insert(
1003                        reqwest::header::HeaderName::from_static("traceparent"),
1004                        reqwest::header::HeaderValue::from_str(traceparent_value)?,
1005                    );
1006                    headers.insert(
1007                        reqwest::header::HeaderName::from_static("tracestate"),
1008                        reqwest::header::HeaderValue::from_str(tracestate_value)?,
1009                    );
1010                    let url = format!("http://{}{}", addr, path);
1011                    let response = client.get(&url).headers(headers).send().await.unwrap();
1012                    let status = response.status();
1013                    let body = response.text().await.unwrap();
1014                    tracing::info!(body = body, status = status.to_string());
1015                }
1016
1017                Ok::<(), anyhow::Error>(())
1018            })(),
1019        )
1020        .await;
1021        Ok(())
1022    }
1023
1024    #[tokio::test]
1025    async fn test_health_endpoint_with_changing_health_status() {
1026        // Test health endpoint starts in not ready status, then becomes ready
1027        // when endpoints are created (DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS=generate)
1028        const ENDPOINT_NAME: &str = "generate";
1029        const ENDPOINT_HEALTH_CONFIG: &str = "[\"generate\"]";
1030        temp_env::async_with_vars(
1031            [
1032                (env_system::DYN_SYSTEM_PORT, Some("0")),
1033                (env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS, Some("notready")),
1034                (env_system::DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS, Some(ENDPOINT_HEALTH_CONFIG)),
1035            ],
1036            async {
1037                let drt = Arc::new(create_test_drt_async().await);
1038
1039                // Check if system status server was started
1040                let system_info_opt = drt.system_status_server_info();
1041
1042                // Ensure system status server was spawned by DRT
1043                assert!(
1044                    system_info_opt.is_some(),
1045                    "System status server was not spawned by DRT. Expected DRT to spawn server when DYN_SYSTEM_PORT is set to a positive value, but system_status_server_info() returned None. Environment: DYN_SYSTEM_PORT={:?}",
1046                    std::env::var(env_system::DYN_SYSTEM_PORT)
1047                );
1048
1049                // Get the system status server info from DRT - this should never fail now due to above check
1050                let system_info = system_info_opt.unwrap();
1051                let addr = system_info.socket_addr;
1052
1053                // Initially check health - should be not ready
1054                let client = reqwest::Client::new();
1055                let health_url = format!("http://{}/health", addr);
1056
1057                let response = client.get(&health_url).send().await.unwrap();
1058                let status = response.status();
1059                let body = response.text().await.unwrap();
1060
1061                // Health should be not ready (503) initially
1062                assert_eq!(status, 503, "Health should be 503 (not ready) initially, got: {}", status);
1063                assert!(body.contains("\"status\":\"notready\""), "Health should contain status notready");
1064
1065                // Now create a namespace, component, and endpoint to make the system healthy
1066                let namespace = drt.namespace("ns1234").unwrap();
1067                let component = namespace.component("comp1234").unwrap();
1068
1069                // Create a simple test handler
1070                use crate::pipeline::{async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, SingleIn};
1071                use crate::protocols::annotated::Annotated;
1072
1073                struct TestHandler;
1074
1075                #[async_trait]
1076                impl AsyncEngine<SingleIn<String>, ManyOut<Annotated<String>>, anyhow::Error> for TestHandler {
1077                    async fn generate(&self, input: SingleIn<String>) -> anyhow::Result<ManyOut<Annotated<String>>> {
1078                        let (data, ctx) = input.into_parts();
1079                        let response = Annotated::from_data(format!("You responded: {}", data));
1080                        Ok(crate::pipeline::ResponseStream::new(
1081                            Box::pin(crate::stream::iter(vec![response])),
1082                            ctx.context()
1083                        ))
1084                    }
1085                }
1086
1087                // Create the ingress and start the endpoint service
1088                let ingress = Ingress::for_engine(std::sync::Arc::new(TestHandler)).unwrap();
1089
1090                // Start the service and endpoint with a health check payload
1091                // This will automatically register the endpoint for health monitoring
1092                tokio::spawn(async move {
1093                    let _ = component.endpoint(ENDPOINT_NAME)
1094                        .endpoint_builder()
1095                        .handler(ingress)
1096                        .health_check_payload(serde_json::json!({
1097                            "test": "health_check"
1098                        }))
1099                        .start()
1100                        .await;
1101                });
1102
1103                // Hit health endpoint 200 times to verify consistency
1104                let mut success_count = 0;
1105                let mut failures = Vec::new();
1106
1107                for i in 1..=200 {
1108                    let response = client.get(&health_url).send().await.unwrap();
1109                    let status = response.status();
1110                    let body = response.text().await.unwrap();
1111
1112                    if status == 200 && body.contains("\"status\":\"ready\"") {
1113                        success_count += 1;
1114                    } else {
1115                        failures.push((i, status.as_u16(), body.clone()));
1116                        if failures.len() <= 5 {  // Only log first 5 failures
1117                            tracing::warn!("Request {}: status={}, body={}", i, status, body);
1118                        }
1119                    }
1120                }
1121
1122                tracing::info!("Health endpoint test results: {success_count}/200 requests succeeded");
1123                if !failures.is_empty() {
1124                    tracing::warn!("Failed requests: {}", failures.len());
1125                }
1126
1127                // Expect at least 150 out of 200 requests to be successful
1128                assert!(success_count >= 150, "Expected at least 150 out of 200 requests to succeed, but only {} succeeded", success_count);
1129            },
1130        )
1131        .await;
1132    }
1133
1134    #[tokio::test]
1135    async fn test_spawn_system_status_server_endpoints() {
1136        // use reqwest for HTTP requests
1137        temp_env::async_with_vars(
1138            [
1139                (env_system::DYN_SYSTEM_PORT, Some("0")),
1140                (env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS, Some("ready")),
1141            ],
1142            async {
1143                let drt = Arc::new(create_test_drt_async().await);
1144
1145                // Get system status server info from DRT (instead of manually spawning)
1146                let system_info = drt
1147                    .system_status_server_info()
1148                    .expect("System status server should be started by DRT");
1149                let addr = system_info.socket_addr;
1150                let client = reqwest::Client::new();
1151                for (path, expect_200, expect_body) in [
1152                    ("/health", true, "ready"),
1153                    ("/live", true, "ready"),
1154                    ("/someRandomPathNotFoundHere", false, "Route not found"),
1155                ] {
1156                    println!("[test] Sending request to {}", path);
1157                    let url = format!("http://{}{}", addr, path);
1158                    let response = client.get(&url).send().await.unwrap();
1159                    let status = response.status();
1160                    let body = response.text().await.unwrap();
1161                    println!(
1162                        "[test] Response for {}: status={}, body={:?}",
1163                        path, status, body
1164                    );
1165                    if expect_200 {
1166                        assert_eq!(status, 200, "Response: status={}, body={:?}", status, body);
1167                    } else {
1168                        assert_eq!(status, 404, "Response: status={}, body={:?}", status, body);
1169                    }
1170                    assert!(
1171                        body.contains(expect_body),
1172                        "Response: status={}, body={:?}",
1173                        status,
1174                        body
1175                    );
1176                }
1177                // DRT handles server cleanup automatically
1178            },
1179        )
1180        .await;
1181    }
1182
1183    #[cfg(feature = "integration")]
1184    #[tokio::test]
1185    async fn test_health_check_with_payload_and_timeout() {
1186        // Test the complete health check flow with the new canary-based system:
1187        crate::logging::init();
1188
1189        temp_env::async_with_vars(
1190            [
1191                (env_system::DYN_SYSTEM_PORT, Some("0")),
1192                (
1193                    env_system::DYN_SYSTEM_STARTING_HEALTH_STATUS,
1194                    Some("notready"),
1195                ),
1196                (
1197                    env_system::DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS,
1198                    Some("[\"test.endpoint\"]"),
1199                ),
1200                // Enable health check with short intervals for testing
1201                ("DYN_HEALTH_CHECK_ENABLED", Some("true")),
1202                (env_canary::DYN_CANARY_WAIT_TIME, Some("1")), // Send canary after 1 second of inactivity
1203                ("DYN_HEALTH_CHECK_REQUEST_TIMEOUT", Some("1")), // Immediately timeout to mimic unresponsiveness
1204                ("RUST_LOG", Some("info")),                      // Enable logging for test
1205            ],
1206            async {
1207                let drt = Arc::new(create_test_drt_async().await);
1208
1209                // Get system status server info
1210                let system_info = drt
1211                    .system_status_server_info()
1212                    .expect("System status server should be started");
1213                let addr = system_info.socket_addr;
1214
1215                let client = reqwest::Client::new();
1216                let health_url = format!("http://{}/health", addr);
1217
1218                // Register an endpoint with health check payload
1219                let endpoint = "test.endpoint";
1220                let health_check_payload = serde_json::json!({
1221                    "prompt": "health check test",
1222                    "_health_check": true
1223                });
1224
1225                // Register the endpoint and its health check payload
1226                {
1227                    let system_health = drt.system_health();
1228                    let system_health_lock = system_health.lock();
1229                    system_health_lock.register_health_check_target(
1230                        endpoint,
1231                        crate::component::Instance {
1232                            component: "test_component".to_string(),
1233                            endpoint: "health".to_string(),
1234                            namespace: "test_namespace".to_string(),
1235                            instance_id: 1,
1236                            transport: crate::component::TransportType::Nats(endpoint.to_string()),
1237                            device_type: None,
1238                            request_plane_codec: None,
1239                        },
1240                        health_check_payload.clone(),
1241                    );
1242                }
1243
1244                // Check initial health - should be ready (default state)
1245                let response = client.get(&health_url).send().await.unwrap();
1246                let status = response.status();
1247                let body = response.text().await.unwrap();
1248                assert_eq!(status, 503, "Should be unhealthy initially (default state)");
1249                assert!(
1250                    body.contains("\"status\":\"notready\""),
1251                    "Should show notready status initially"
1252                );
1253
1254                // Set endpoint to healthy state
1255                drt.system_health()
1256                    .lock()
1257                    .set_endpoint_health_status(endpoint, HealthStatus::Ready);
1258
1259                // Check health again - should now be healthy
1260                let response = client.get(&health_url).send().await.unwrap();
1261                let status = response.status();
1262                let body = response.text().await.unwrap();
1263
1264                assert_eq!(status, 200, "Should be healthy due to recent response");
1265                assert!(
1266                    body.contains("\"status\":\"ready\""),
1267                    "Should show ready status after response"
1268                );
1269
1270                // Verify the endpoint status in SystemHealth directly
1271                let endpoint_status = drt
1272                    .system_health()
1273                    .lock()
1274                    .get_endpoint_health_status(endpoint);
1275                assert_eq!(
1276                    endpoint_status,
1277                    Some(HealthStatus::Ready),
1278                    "SystemHealth should show endpoint as Ready after response"
1279                );
1280            },
1281        )
1282        .await;
1283    }
1284
1285    /// `/v1/loras` compat shim: with the legacy local registry empty, a LoRA
1286    /// update registered in `engine_routes()` under `update/<name>` resolves via
1287    /// the fallback and its JSON response is parsed into a `LoraResponse`. This
1288    /// is the path unified-backend workers take for `/v1/loras`.
1289    #[tokio::test]
1290    async fn test_call_lora_endpoint_resolves_via_engine_routes() {
1291        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1292            let drt = create_test_drt_async().await;
1293
1294            let callback: crate::engine_routes::EngineRouteCallback = Arc::new(|_body| {
1295                Box::pin(async move {
1296                    Ok(serde_json::json!({
1297                        "status": "success",
1298                        "lora_name": "adapterA",
1299                        "lora_id": 42,
1300                    }))
1301                })
1302            });
1303            // Unified Worker registers LoRA ops under the `update/` namespace.
1304            drt.engine_routes().register("update/load_lora", callback);
1305
1306            // Local registry is empty, so resolution must fall through to
1307            // engine_routes.
1308            assert!(drt.local_endpoint_registry().get("load_lora").is_none());
1309
1310            let response = call_lora_endpoint(
1311                &drt,
1312                "load_lora",
1313                serde_json::json!({"lora_name": "adapterA"}),
1314            )
1315            .await
1316            .expect("engine_routes fallback should resolve the control");
1317
1318            assert_eq!(response.status, "success");
1319            assert_eq!(response.lora_name.as_deref(), Some("adapterA"));
1320            assert_eq!(response.lora_id, Some(42));
1321        })
1322        .await;
1323    }
1324
1325    /// When neither the local registry nor `engine_routes()` holds the name,
1326    /// the caller gets an explicit "LoRA management not available" error
1327    /// rather than an opaque "endpoint not found".
1328    #[tokio::test]
1329    async fn test_call_lora_endpoint_missing_returns_clean_error() {
1330        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1331            let drt = create_test_drt_async().await;
1332
1333            let err = call_lora_endpoint(&drt, "load_lora", serde_json::json!({}))
1334                .await
1335                .expect_err("missing handler must error");
1336
1337            assert!(
1338                err.to_string().contains("LoRA management not available"),
1339                "expected explicit unavailable message, got: {err}"
1340            );
1341        })
1342        .await;
1343    }
1344
1345    /// An update callback that returns `{"status":"error",...}` (rather than
1346    /// raising) surfaces as a `LoraResponse{status:"error"}`. The `/v1/loras`
1347    /// load/unload handlers map this to HTTP 500, preserving legacy semantics;
1348    /// the direct `/engine/*` route would instead return HTTP 200 + this JSON.
1349    #[tokio::test]
1350    async fn test_call_lora_endpoint_propagates_error_status() {
1351        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1352            let drt = create_test_drt_async().await;
1353
1354            let callback: crate::engine_routes::EngineRouteCallback = Arc::new(|_body| {
1355                Box::pin(async move {
1356                    Ok(serde_json::json!({
1357                        "status": "error",
1358                        "message": "adapter not found",
1359                    }))
1360                })
1361            });
1362            // Unified Worker registers LoRA ops under the `update/` namespace.
1363            drt.engine_routes().register("update/unload_lora", callback);
1364
1365            let response = call_lora_endpoint(&drt, "unload_lora", serde_json::json!({}))
1366                .await
1367                .expect("a non-raising callback returns Ok even on logical error");
1368
1369            assert_eq!(response.status, "error");
1370            assert_eq!(response.message.as_deref(), Some("adapter not found"));
1371        })
1372        .await;
1373    }
1374}
dynamo_runtime/system_status_server.rs

dynamo_runtime/
system_status_server.rs
