chronix_analytics/forecast/
storage.rs

1//! Model storage: serialization, metadata, and in-memory catalog.
2//!
3//! `ModelStore` enables round-trip serialization of fitted models via
4//! `serde` + `postcard`. `ModelCatalog` provides an in-memory registry
5//! keyed by `(measurement, model_name)`.
6
7use std::collections::HashMap;
8use std::time::{SystemTime, UNIX_EPOCH};
9
10use serde::{Deserialize, Serialize};
11
12use crate::forecast::result::ModelType;
13
14// ─── ModelStore trait ───────────────────────────────────────────────
15
16/// Trait for serializing and deserializing fitted models.
17///
18/// Implementors must be `Serialize + Deserialize`. The default
19/// implementation uses `postcard` for compact binary encoding.
20pub trait ModelStore: Serialize + for<'de> Deserialize<'de> {
21    /// Serializes this model to a byte vector.
22    fn to_bytes(&self) -> Result<Vec<u8>, StorageError> {
23        postcard::to_stdvec(self).map_err(|e| StorageError::Serialize(e.to_string()))
24    }
25
26    /// Deserializes a model from bytes.
27    fn from_bytes(bytes: &[u8]) -> Result<Self, StorageError>
28    where
29        Self: Sized,
30    {
31        postcard::from_bytes(bytes).map_err(|e| StorageError::Deserialize(e.to_string()))
32    }
33}
34
35// ─── StorageError ───────────────────────────────────────────────────
36
37/// Errors from model storage operations.
38#[derive(Debug, thiserror::Error)]
39pub enum StorageError {
40    /// The model could not be serialized to bytes.
41    #[error("serialization failed: {0}")]
42    Serialize(String),
43    /// The byte payload could not be deserialized into a model.
44    #[error("deserialization failed: {0}")]
45    Deserialize(String),
46    /// No model with the given key exists in the store.
47    #[error("model not found: {0}")]
48    NotFound(String),
49}
50
51// ─── ModelMetadata ──────────────────────────────────────────────────
52
53/// Metadata describing a stored model.
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct ModelMetadata {
56    /// Type of the forecast model (SES, Holt, ARIMA, …).
57    pub model_type: ModelType,
58    /// User-supplied name for this model instance.
59    pub model_name: String,
60    /// Measurement (metric) on which the model was trained.
61    pub measurement: String,
62    /// Monotonically increasing version counter.
63    pub version: u32,
64    /// UNIX timestamp (seconds) when the model was stored.
65    pub created_at: i64,
66    /// Start of the training data range (nanoseconds).
67    pub training_data_start: i64,
68    /// End of the training data range (nanoseconds).
69    pub training_data_end: i64,
70    /// Number of data points used for training.
71    pub training_points: usize,
72    /// Mean Squared Error on training data.
73    pub mse: f64,
74    /// Mean Absolute Error on training data.
75    pub mae: f64,
76    /// Coefficient of determination R² on training data.
77    pub r_squared: f64,
78}
79
80// ─── CatalogEntry ───────────────────────────────────────────────────
81
82/// A single entry in the catalog: metadata + serialized model bytes.
83#[derive(Debug, Clone)]
84struct CatalogEntry {
85    metadata: ModelMetadata,
86    bytes: Vec<u8>,
87}
88
89// ─── ModelCatalog ───────────────────────────────────────────────────
90
91/// In-memory model catalog keyed by `(measurement, model_name)`.
92///
93/// # Thread Safety
94///
95/// `ModelCatalog` is not internally synchronised. The server wraps it
96/// in `Arc<RwLock<ModelCatalog>>` (see `AppState::model_catalog`),
97/// which provides correct concurrent access. This is intentional:
98/// embedding a lock would impose overhead on single-threaded forecast
99/// pipelines and prevents read-side parallelism that `RwLock` provides.
100#[derive(Debug, Default)]
101pub struct ModelCatalog {
102    /// `(measurement, model_name)` → `CatalogEntry`
103    entries: HashMap<(String, String), CatalogEntry>,
104}
105
106impl ModelCatalog {
107    /// Creates an empty catalog.
108    pub fn new() -> Self {
109        Self::default()
110    }
111
112    /// Saves a model to the catalog.
113    ///
114    /// If a model with the same `(measurement, name)` already exists
115    /// the version is auto-incremented.
116    #[allow(clippy::too_many_arguments)]
117    pub fn save_model<M: ModelStore>(
118        &mut self,
119        measurement: &str,
120        name: &str,
121        model: &M,
122        model_type: ModelType,
123        training_range: (i64, i64),
124        training_points: usize,
125        mse: f64,
126        mae: f64,
127        r_squared: f64,
128    ) -> Result<ModelMetadata, StorageError> {
129        let bytes = model.to_bytes()?;
130        let key = (measurement.to_string(), name.to_string());
131
132        let version = self.entries.get(&key).map_or(1, |e| e.metadata.version + 1);
133
134        let now_ns = i64::try_from(
135            SystemTime::now()
136                .duration_since(UNIX_EPOCH)
137                .unwrap_or_default()
138                .as_nanos(),
139        )
140        .unwrap_or(i64::MAX);
141
142        let metadata = ModelMetadata {
143            model_type,
144            model_name: name.to_string(),
145            measurement: measurement.to_string(),
146            version,
147            created_at: now_ns,
148            training_data_start: training_range.0,
149            training_data_end: training_range.1,
150            training_points,
151            mse,
152            mae,
153            r_squared,
154        };
155
156        self.entries.insert(
157            key,
158            CatalogEntry {
159                metadata: metadata.clone(),
160                bytes,
161            },
162        );
163
164        Ok(metadata)
165    }
166
167    /// Loads a model from the catalog, deserializing it.
168    pub fn load_model<M: ModelStore>(
169        &self,
170        measurement: &str,
171        name: &str,
172    ) -> Result<M, StorageError> {
173        let key = (measurement.to_string(), name.to_string());
174        let entry = self
175            .entries
176            .get(&key)
177            .ok_or_else(|| StorageError::NotFound(format!("{measurement}/{name}")))?;
178        M::from_bytes(&entry.bytes)
179    }
180
181    /// Returns metadata for a stored model (if it exists).
182    pub fn get_metadata(&self, measurement: &str, name: &str) -> Option<&ModelMetadata> {
183        let key = (measurement.to_string(), name.to_string());
184        self.entries.get(&key).map(|e| &e.metadata)
185    }
186
187    /// Lists metadata for all models under a given measurement.
188    pub fn list_models(&self, measurement: &str) -> Vec<&ModelMetadata> {
189        self.entries
190            .values()
191            .filter(|e| e.metadata.measurement == measurement)
192            .map(|e| &e.metadata)
193            .collect()
194    }
195
196    /// Lists metadata for all models in the catalog.
197    pub fn list_all(&self) -> Vec<&ModelMetadata> {
198        self.entries.values().map(|e| &e.metadata).collect()
199    }
200
201    /// Deletes a model from the catalog.
202    pub fn delete_model(&mut self, measurement: &str, name: &str) -> bool {
203        let key = (measurement.to_string(), name.to_string());
204        self.entries.remove(&key).is_some()
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use crate::forecast::{ForecastModel, SesModel};
212
213    #[test]
214    fn ses_round_trip() {
215        let ts: Vec<i64> = (0..100).map(|i| i * 1_000_000_000).collect();
216        let vals: Vec<f64> = (0..100).map(|i| 50.0 + i as f64 * 0.1).collect();
217        let mut model = SesModel::new(Some(0.3));
218        model.fit(&ts, &vals).unwrap();
219
220        let orig_pred = model.predict(10).unwrap();
221        let bytes = model.to_bytes().unwrap();
222        let restored = SesModel::from_bytes(&bytes).unwrap();
223        let rest_pred = restored.predict(10).unwrap();
224
225        assert_eq!(orig_pred.values.len(), rest_pred.values.len());
226        for (a, b) in orig_pred.values.iter().zip(rest_pred.values.iter()) {
227            assert!((a - b).abs() < 1e-10, "values diverged: {a} vs {b}");
228        }
229    }
230
231    #[test]
232    fn catalog_save_load_delete() {
233        let ts: Vec<i64> = (0..50).map(|i| i * 1_000_000_000).collect();
234        let vals: Vec<f64> = vec![42.0; 50];
235        let mut model = SesModel::new(Some(0.5));
236        model.fit(&ts, &vals).unwrap();
237
238        let mut catalog = ModelCatalog::new();
239        let meta = catalog
240            .save_model(
241                "cpu",
242                "ses_v1",
243                &model,
244                ModelType::Ses,
245                (ts[0], *ts.last().unwrap()),
246                vals.len(),
247                0.0,
248                0.0,
249                1.0,
250            )
251            .unwrap();
252        assert_eq!(meta.version, 1);
253
254        // List
255        let listed = catalog.list_models("cpu");
256        assert_eq!(listed.len(), 1);
257
258        // Load
259        let loaded: SesModel = catalog.load_model("cpu", "ses_v1").unwrap();
260        let pred = loaded.predict(5).unwrap();
261        assert_eq!(pred.values.len(), 5);
262
263        // Version bump
264        let meta2 = catalog
265            .save_model(
266                "cpu",
267                "ses_v1",
268                &model,
269                ModelType::Ses,
270                (0, 0),
271                0,
272                0.0,
273                0.0,
274                0.0,
275            )
276            .unwrap();
277        assert_eq!(meta2.version, 2);
278
279        // Delete
280        assert!(catalog.delete_model("cpu", "ses_v1"));
281        assert!(catalog.load_model::<SesModel>("cpu", "ses_v1").is_err());
282    }
283
284    #[test]
285    fn not_found() {
286        let catalog = ModelCatalog::new();
287        let result = catalog.load_model::<SesModel>("cpu", "missing");
288        assert!(result.is_err());
289    }
290
291    #[test]
292    fn list_all_models() {
293        let ts: Vec<i64> = (0..50).map(|i| i * 1_000_000_000).collect();
294        let vals: Vec<f64> = vec![42.0; 50];
295        let mut model = SesModel::new(Some(0.5));
296        model.fit(&ts, &vals).unwrap();
297
298        let mut catalog = ModelCatalog::new();
299        catalog
300            .save_model(
301                "cpu",
302                "m1",
303                &model,
304                ModelType::Ses,
305                (0, 49),
306                50,
307                0.0,
308                0.0,
309                1.0,
310            )
311            .unwrap();
312        catalog
313            .save_model(
314                "mem",
315                "m2",
316                &model,
317                ModelType::Ses,
318                (0, 49),
319                50,
320                0.0,
321                0.0,
322                1.0,
323            )
324            .unwrap();
325        catalog
326            .save_model(
327                "cpu",
328                "m3",
329                &model,
330                ModelType::Ses,
331                (0, 49),
332                50,
333                0.0,
334                0.0,
335                1.0,
336            )
337            .unwrap();
338
339        // list_all returns entries from all measurements.
340        let all = catalog.list_all();
341        assert_eq!(all.len(), 3);
342
343        // list_models filters by measurement.
344        assert_eq!(catalog.list_models("cpu").len(), 2);
345        assert_eq!(catalog.list_models("mem").len(), 1);
346        assert_eq!(catalog.list_models("disk").len(), 0);
347    }
348}
chronix_analytics/forecast/storage.rs

chronix_analytics/forecast/
storage.rs
