axonml-server 0.4.2

REST API server for AxonML Machine Learning Framework
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//! Metrics API endpoints for AxonML
//!
//! # File
//! `crates/axonml-server/src/api/metrics.rs`
//!
//! # Author
//! Andrew Jewell Sr - AutomataNexus
//!
//! # Updated
//! March 8, 2026
//!
//! # Disclaimer
//! Use at own risk. This software is provided "as is", without warranty of any
//! kind, express or implied. The author and AutomataNexus shall not be held
//! liable for any damages arising from the use of this software.

use crate::api::AppState;
use crate::auth::{AuthError, AuthUser};
use crate::db::models::ModelRepository;
use crate::db::runs::RunRepository;
use crate::db::users::UserRepository;
use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
use serde::Serialize;

// ============================================================================
// Response Types
// ============================================================================

#[derive(Debug, Serialize)]
pub struct SystemStats {
    pub users_total: u64,
    pub runs_total: u64,
    pub runs_running: u64,
    pub models_total: u64,
    pub endpoints_total: u64,
    pub endpoints_running: u64,
}

// ============================================================================
// Handlers
// ============================================================================

/// Get Prometheus-style metrics
pub async fn get_metrics(
    State(state): State<AppState>,
    _user: AuthUser,
) -> Result<impl IntoResponse, AuthError> {
    let user_repo = UserRepository::new(&state.db);
    let run_repo = RunRepository::new(&state.db);
    let model_repo = ModelRepository::new(&state.db);

    // Collect metrics
    let users_total = user_repo.count().await.unwrap_or(0);
    let runs_running = run_repo.count_running().await.unwrap_or(0);
    let endpoints_running = model_repo.count_running_endpoints().await.unwrap_or(0);

    // Build Prometheus format
    let metrics = format!(
        r#"# HELP axonml_users_total Total number of registered users
# TYPE axonml_users_total gauge
axonml_users_total {}

# HELP axonml_runs_running Number of currently running training runs
# TYPE axonml_runs_running gauge
axonml_runs_running {}

# HELP axonml_endpoints_running Number of currently running inference endpoints
# TYPE axonml_endpoints_running gauge
axonml_endpoints_running {}

# HELP axonml_server_info Server information
# TYPE axonml_server_info gauge
axonml_server_info{{version="0.1.0"}} 1
"#,
        users_total, runs_running, endpoints_running
    );

    Ok((
        StatusCode::OK,
        [("Content-Type", "text/plain; charset=utf-8")],
        metrics,
    ))
}

/// Get system statistics (admin only)
pub async fn get_stats(State(state): State<AppState>) -> Result<Json<SystemStats>, AuthError> {
    let user_repo = UserRepository::new(&state.db);
    let run_repo = RunRepository::new(&state.db);
    let model_repo = ModelRepository::new(&state.db);

    // Collect stats
    let users_total = user_repo.count().await.unwrap_or(0);

    let runs = run_repo
        .list_all(None, Some(10000), Some(0))
        .await
        .unwrap_or_default();
    let runs_total = runs.len() as u64;
    let runs_running = run_repo.count_running().await.unwrap_or(0);

    let models = model_repo
        .list_all(Some(10000), Some(0))
        .await
        .unwrap_or_default();
    let models_total = models.len() as u64;

    let endpoints = model_repo.list_endpoints().await.unwrap_or_default();
    let endpoints_total = endpoints.len() as u64;
    let endpoints_running = model_repo.count_running_endpoints().await.unwrap_or(0);

    Ok(Json(SystemStats {
        users_total,
        runs_total,
        runs_running,
        models_total,
        endpoints_total,
        endpoints_running,
    }))
}