mold-ai-server 0.15.0

HTTP inference server for mold
Documentation 
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

#[cfg(all(test, feature = "metrics"))]
mod tests {
    use axum::{
        body::Body,
        http::{Request, StatusCode},
        middleware,
        routing::get,
        Router,
    };
    use tower::ServiceExt;

    use crate::{
        metrics::{self, MetricsState},
        routes::create_router,
        state::AppState,
    };

    /// Build a test app with the /metrics route and HTTP metrics middleware.
    fn app_with_metrics() -> Router {
        let handle = metrics::install_recorder_for_test();

        let (tx, _rx) = tokio::sync::mpsc::channel(16);
        let queue = crate::state::QueueHandle::new(tx);
        let gpu_pool = std::sync::Arc::new(crate::gpu_pool::GpuPool {
            workers: Vec::new(),
        });
        let state = AppState::empty(mold_core::Config::default(), queue, gpu_pool, 200);

        let start_time = state.start_time;
        let metrics_state = MetricsState { handle, start_time };

        create_router(state)
            .layer(middleware::from_fn(metrics::http_metrics_middleware))
            .route(
                "/metrics",
                get(metrics::metrics_endpoint).with_state(metrics_state),
            )
    }

    #[tokio::test]
    async fn metrics_endpoint_returns_prometheus_format() {
        let app = app_with_metrics();

        let resp = app
            .oneshot(
                Request::builder()
                    .uri("/metrics")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(resp.status(), StatusCode::OK);

        let content_type = resp
            .headers()
            .get("content-type")
            .unwrap()
            .to_str()
            .unwrap();
        assert!(
            content_type.contains("text/plain"),
            "expected text/plain content type, got: {content_type}"
        );

        let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
            .await
            .unwrap();
        let text = String::from_utf8(body.to_vec()).unwrap();

        // Should contain uptime gauge (recorded on each scrape).
        assert!(
            text.contains("mold_uptime_seconds"),
            "expected mold_uptime_seconds in metrics output:\n{text}"
        );
    }

    #[tokio::test]
    async fn metrics_records_http_requests() {
        let app = app_with_metrics();

        // Hit /health to generate an HTTP metric.
        let resp = app
            .clone()
            .oneshot(
                Request::builder()
                    .uri("/health")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        // Now scrape /metrics and check for the recorded request.
        let resp = app
            .oneshot(
                Request::builder()
                    .uri("/metrics")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
            .await
            .unwrap();
        let text = String::from_utf8(body.to_vec()).unwrap();

        assert!(
            text.contains("mold_http_requests_total"),
            "expected mold_http_requests_total in output:\n{text}"
        );
        assert!(
            text.contains("mold_http_request_duration_seconds"),
            "expected mold_http_request_duration_seconds in output:\n{text}"
        );
    }

    #[tokio::test]
    async fn metrics_records_queue_submit() {
        // Ensure the recorder is installed regardless of test execution order.
        metrics::install_recorder_for_test();
        metrics::record_queue_submit();
        metrics::record_queue_depth(3);
        metrics::record_generation("test-model", 1.5);
        metrics::record_model_load("test-model", 2.0);
        metrics::set_model_loaded("test-model");
        metrics::record_gpu_memory(1_000_000_000);
        metrics::record_generation_error("test-model");
        metrics::record_cache_eviction("capacity");
        metrics::record_cache_eviction("idle-ttl");
        metrics::set_cache_size(2);
        metrics::record_queue_reorder();

        // If we got here without panicking, the metric calls are valid.
    }
}