coglet 0.19.0

High-performance prediction server for Cog ML models
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

//! Predictor traits and prediction lifecycle types.

use std::collections::HashMap;
use std::time::{Duration, Instant};

pub use crate::prediction::{CancellationToken, PredictionOutput};

/// Result of a completed prediction.
#[derive(Debug, Clone)]
pub struct PredictionResult {
    pub output: PredictionOutput,
    pub predict_time: Option<Duration>,
    pub logs: String,
    /// User-emitted metrics from the prediction.
    pub metrics: HashMap<String, serde_json::Value>,
}

/// Metrics collected during prediction.
#[derive(Debug, Clone, Default)]
pub struct PredictionMetrics {
    pub predict_time: Option<Duration>,
}

/// RAII guard for prediction lifecycle timing.
pub struct PredictionGuard {
    start_time: Instant,
    metrics: PredictionMetrics,
    cancel_token: CancellationToken,
}

impl PredictionGuard {
    pub fn new() -> Self {
        Self {
            start_time: Instant::now(),
            metrics: PredictionMetrics::default(),
            cancel_token: CancellationToken::new(),
        }
    }

    pub fn cancel_token(&self) -> CancellationToken {
        self.cancel_token.clone()
    }

    pub fn is_cancelled(&self) -> bool {
        self.cancel_token.is_cancelled()
    }

    pub fn cancel(&self) {
        self.cancel_token.cancel();
    }

    pub fn finish(mut self) -> PredictionMetrics {
        self.metrics.predict_time = Some(self.start_time.elapsed());
        self.metrics
    }
}

impl Default for PredictionGuard {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, thiserror::Error)]
pub enum PredictionError {
    #[error("Prediction failed: {0}")]
    Failed(String),

    #[error("Input validation error: {0}")]
    InvalidInput(String),

    #[error(
        "Setup has not finished yet. Wait until it has finished, or GET /health-check for status."
    )]
    NotReady,

    #[error("Prediction was cancelled")]
    Cancelled,
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn prediction_output_single_is_not_stream() {
        let output = PredictionOutput::Single(json!("hello"));
        assert!(!output.is_stream());
    }

    #[test]
    fn prediction_output_stream_is_stream() {
        let output = PredictionOutput::Stream(vec![json!("a"), json!("b")]);
        assert!(output.is_stream());
    }

    #[test]
    fn prediction_output_serializes_untagged() {
        let single = PredictionOutput::Single(json!("hello"));
        insta::assert_json_snapshot!("output_single", single);

        let stream = PredictionOutput::Stream(vec![json!(1), json!(2)]);
        insta::assert_json_snapshot!("output_stream", stream);
    }

    #[test]
    fn prediction_guard_tracks_time() {
        let guard = PredictionGuard::new();
        std::thread::sleep(std::time::Duration::from_millis(10));
        let metrics = guard.finish();

        assert!(metrics.predict_time.is_some());
        let time = metrics.predict_time.unwrap();
        assert!(time.as_millis() >= 10);
        assert!(time.as_secs() < 1);
    }

    #[test]
    fn prediction_error_display() {
        let err = PredictionError::Failed("something broke".to_string());
        assert_eq!(format!("{}", err), "Prediction failed: something broke");

        let err = PredictionError::InvalidInput("bad json".to_string());
        assert_eq!(format!("{}", err), "Input validation error: bad json");

        let err = PredictionError::NotReady;
        assert_eq!(
            format!("{}", err),
            "Setup has not finished yet. Wait until it has finished, or GET /health-check for status."
        );
    }
}