Skip to main content

entrenar/storage/registry/
traits.rs

1//! Model registry trait definition
2
3use std::collections::HashMap;
4
5use super::comparison::VersionComparison;
6use super::error::Result;
7use super::policy::{PolicyCheckResult, PromotionPolicy};
8use super::stage::ModelStage;
9use super::transition::StageTransition;
10use super::version::ModelVersion;
11
12/// Model registry trait
13pub trait ModelRegistry: Send + Sync {
14    /// Register a new model version
15    fn register_model(&mut self, name: &str, artifact_uri: &str) -> Result<ModelVersion>;
16
17    /// Get a model version
18    fn get_model(&self, name: &str, version: u32) -> Result<ModelVersion>;
19
20    /// Get latest version of a model
21    fn get_latest(&self, name: &str) -> Result<ModelVersion>;
22
23    /// Get latest version at a specific stage
24    fn get_latest_by_stage(&self, name: &str, stage: ModelStage) -> Option<ModelVersion>;
25
26    /// List all versions of a model
27    fn list_versions(&self, name: &str) -> Result<Vec<ModelVersion>>;
28
29    /// Transition model to new stage
30    fn transition_stage(
31        &mut self,
32        name: &str,
33        version: u32,
34        target_stage: ModelStage,
35        user: Option<&str>,
36    ) -> Result<()>;
37
38    /// Compare two versions
39    fn compare_versions(&self, name: &str, v1: u32, v2: u32) -> Result<VersionComparison>;
40
41    /// Log metrics for a model version
42    fn log_metrics(
43        &mut self,
44        name: &str,
45        version: u32,
46        metrics: HashMap<String, f64>,
47    ) -> Result<()>;
48
49    /// Get transition history for a model
50    fn get_transition_history(&self, name: &str) -> Result<Vec<StageTransition>>;
51
52    /// Set promotion policy for a stage
53    fn set_policy(&mut self, policy: PromotionPolicy);
54
55    /// Get promotion policy for a stage
56    fn get_policy(&self, stage: ModelStage) -> Option<&PromotionPolicy>;
57
58    /// Check if model can be promoted (with policy check)
59    fn can_promote(
60        &self,
61        name: &str,
62        version: u32,
63        target_stage: ModelStage,
64        approvals: u32,
65    ) -> Result<PolicyCheckResult>;
66}