foxhunt/crates/ml/src/registry/mod.rs

//! Model Registry for lifecycle management
//!
//! Tracks training runs, model versions, and promotions through
//! Candidate → Staging → Production → Archived lifecycle.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::SystemTime;

/// Model lifecycle stage
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelStage {
    /// Initial state after training
    Candidate,
    /// Passed validation, running canary
    Staging,
    /// Active in production ensemble
    Production,
    /// Replaced by newer version
    Archived,
}

impl std::fmt::Display for ModelStage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Candidate => write!(f, "candidate"),
            Self::Staging => write!(f, "staging"),
            Self::Production => write!(f, "production"),
            Self::Archived => write!(f, "archived"),
        }
    }
}

/// Record of a training run
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingRun {
    /// Unique run identifier
    pub run_id: String,
    /// Experiment name (e.g. "dqn-v3-sharpe-opt")
    pub experiment_name: String,
    /// Model type (DQN, PPO, TFT, etc.)
    pub model_type: String,
    /// Hyperparameters as JSON
    pub hyperparameters: serde_json::Value,
    /// Git commit hash at time of training
    pub git_commit: String,
    /// Hash of training data for reproducibility
    pub data_hash: String,
    /// When training started
    pub started_at: SystemTime,
    /// When training finished (None if still running)
    pub finished_at: Option<SystemTime>,
}

/// Metrics recorded for a model version
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetrics {
    /// Validation Sharpe ratio
    pub sharpe_ratio: f64,
    /// Validation accuracy
    pub accuracy: f64,
    /// Validation win rate
    pub win_rate: f64,
    /// Maximum drawdown on validation set
    pub max_drawdown: f64,
    /// Any additional metrics
    pub extra: HashMap<String, f64>,
}

/// A versioned model in the registry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelVersion {
    /// Unique version identifier
    pub version_id: String,
    /// Associated training run
    pub run_id: String,
    /// Model type
    pub model_type: String,
    /// Path to model artifact (safetensors file)
    pub artifact_path: String,
    /// Current lifecycle stage
    pub stage: ModelStage,
    /// Validation metrics
    pub metrics: Option<ModelMetrics>,
    /// When this version was registered
    pub registered_at: SystemTime,
    /// Who/what promoted this version
    pub promoted_by: Option<String>,
}

/// Side-by-side run comparison
#[derive(Debug, Clone)]
pub struct RunComparison {
    pub runs: Vec<(TrainingRun, Option<ModelMetrics>)>,
}

/// Model registry trait (async for database implementations)
#[async_trait::async_trait]
pub trait ModelRegistryTrait: Send + Sync {
    /// Log a new training run
    async fn log_run(&self, run: TrainingRun) -> Result<(), RegistryError>;

    /// Log metrics for a model version
    async fn log_metrics(
        &self,
        version_id: &str,
        metrics: ModelMetrics,
    ) -> Result<(), RegistryError>;

    /// Register a model version
    async fn register_version(&self, version: ModelVersion) -> Result<(), RegistryError>;

    /// Promote a model to a new stage
    async fn promote(
        &self,
        version_id: &str,
        to_stage: ModelStage,
        promoted_by: &str,
    ) -> Result<(), RegistryError>;

    /// Get the current production model for a given type
    async fn get_production_model(
        &self,
        model_type: &str,
    ) -> Result<Option<ModelVersion>, RegistryError>;

    /// Revert to previous production version
    async fn revert(&self, model_type: &str) -> Result<ModelVersion, RegistryError>;

    /// Compare metrics across runs
    async fn compare_runs(&self, run_ids: &[String]) -> Result<RunComparison, RegistryError>;
}

/// Registry errors
#[derive(Debug, thiserror::Error)]
pub enum RegistryError {
    #[error("Version not found: {0}")]
    VersionNotFound(String),
    #[error("Run not found: {0}")]
    RunNotFound(String),
    #[error("Invalid stage transition: {from} → {to}")]
    InvalidTransition { from: String, to: String },
    #[error("No previous version to revert to for model type: {0}")]
    NoPreviousVersion(String),
}

/// In-memory model registry for testing
#[derive(Debug, Default)]
pub struct InMemoryModelRegistry {
    runs: tokio::sync::RwLock<HashMap<String, TrainingRun>>,
    versions: tokio::sync::RwLock<HashMap<String, ModelVersion>>,
    metrics: tokio::sync::RwLock<HashMap<String, ModelMetrics>>,
}

impl InMemoryModelRegistry {
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait::async_trait]
impl ModelRegistryTrait for InMemoryModelRegistry {
    async fn log_run(&self, run: TrainingRun) -> Result<(), RegistryError> {
        self.runs.write().await.insert(run.run_id.clone(), run);
        Ok(())
    }

    async fn log_metrics(
        &self,
        version_id: &str,
        metrics: ModelMetrics,
    ) -> Result<(), RegistryError> {
        // Also update the version's metrics
        let mut versions = self.versions.write().await;
        if let Some(version) = versions.get_mut(version_id) {
            version.metrics = Some(metrics.clone());
        }
        self.metrics
            .write()
            .await
            .insert(version_id.to_string(), metrics);
        Ok(())
    }

    async fn register_version(&self, version: ModelVersion) -> Result<(), RegistryError> {
        self.versions
            .write()
            .await
            .insert(version.version_id.clone(), version);
        Ok(())
    }

    async fn promote(
        &self,
        version_id: &str,
        to_stage: ModelStage,
        promoted_by: &str,
    ) -> Result<(), RegistryError> {
        let mut versions = self.versions.write().await;
        let version = versions
            .get_mut(version_id)
            .ok_or_else(|| RegistryError::VersionNotFound(version_id.to_string()))?;

        // Validate transition
        let valid = matches!(
            (version.stage, to_stage),
            (ModelStage::Candidate, ModelStage::Staging)
                | (ModelStage::Staging, ModelStage::Production)
                | (ModelStage::Production, ModelStage::Archived)
                | (ModelStage::Staging, ModelStage::Archived)
                | (ModelStage::Candidate, ModelStage::Archived)
        );

        if !valid {
            return Err(RegistryError::InvalidTransition {
                from: version.stage.to_string(),
                to: to_stage.to_string(),
            });
        }

        // If promoting to Production, archive the current production model of same type
        if to_stage == ModelStage::Production {
            let model_type = version.model_type.clone();
            let current_prod: Vec<String> = versions
                .iter()
                .filter(|(id, v)| {
                    v.model_type == model_type
                        && v.stage == ModelStage::Production
                        && *id != version_id
                })
                .map(|(id, _)| id.clone())
                .collect();

            // Must drop the version borrow before modifying others.
            // version_id was confirmed to exist via the get_mut() + ok_or_else() above.
            let version = versions.get_mut(version_id).ok_or_else(|| {
                RegistryError::VersionNotFound(version_id.to_string())
            })?;
            version.stage = to_stage;
            version.promoted_by = Some(promoted_by.to_string());

            for old_id in current_prod {
                if let Some(old_version) = versions.get_mut(&old_id) {
                    old_version.stage = ModelStage::Archived;
                }
            }
        } else {
            version.stage = to_stage;
            version.promoted_by = Some(promoted_by.to_string());
        }

        Ok(())
    }

    async fn get_production_model(
        &self,
        model_type: &str,
    ) -> Result<Option<ModelVersion>, RegistryError> {
        let versions = self.versions.read().await;
        let prod = versions
            .values()
            .find(|v| v.model_type == model_type && v.stage == ModelStage::Production)
            .cloned();
        Ok(prod)
    }

    async fn revert(&self, model_type: &str) -> Result<ModelVersion, RegistryError> {
        let mut versions = self.versions.write().await;

        // Find the most recently archived version of this type
        let archived: Option<String> = versions
            .iter()
            .filter(|(_, v)| v.model_type == model_type && v.stage == ModelStage::Archived)
            .max_by_key(|(_, v)| v.registered_at)
            .map(|(id, _)| id.clone());

        let archived_id =
            archived.ok_or_else(|| RegistryError::NoPreviousVersion(model_type.to_string()))?;

        // Archive current production
        let current_prod: Vec<String> = versions
            .iter()
            .filter(|(_, v)| v.model_type == model_type && v.stage == ModelStage::Production)
            .map(|(id, _)| id.clone())
            .collect();

        for id in current_prod {
            if let Some(v) = versions.get_mut(&id) {
                v.stage = ModelStage::Archived;
            }
        }

        // Promote archived to production
        let version = versions
            .get_mut(&archived_id)
            .ok_or_else(|| RegistryError::VersionNotFound(archived_id.clone()))?;
        version.stage = ModelStage::Production;
        version.promoted_by = Some("revert".to_owned());

        Ok(version.clone())
    }

    async fn compare_runs(&self, run_ids: &[String]) -> Result<RunComparison, RegistryError> {
        let runs = self.runs.read().await;
        let metrics = self.metrics.read().await;

        let mut comparisons = Vec::new();
        for id in run_ids {
            let run = runs
                .get(id)
                .ok_or_else(|| RegistryError::RunNotFound(id.clone()))?
                .clone();
            let m = metrics.get(id).cloned();
            comparisons.push((run, m));
        }

        Ok(RunComparison {
            runs: comparisons,
        })
    }
}

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

    fn make_run(id: &str, model_type: &str) -> TrainingRun {
        TrainingRun {
            run_id: id.to_string(),
            experiment_name: format!("{}-exp", model_type),
            model_type: model_type.to_string(),
            hyperparameters: serde_json::json!({"lr": 0.001}),
            git_commit: "abc123".into(),
            data_hash: "sha256:deadbeef".into(),
            started_at: SystemTime::now(),
            finished_at: Some(SystemTime::now()),
        }
    }

    fn make_version(id: &str, run_id: &str, model_type: &str) -> ModelVersion {
        ModelVersion {
            version_id: id.to_string(),
            run_id: run_id.to_string(),
            model_type: model_type.to_string(),
            artifact_path: format!("models/{}/{}.safetensors", model_type, id),
            stage: ModelStage::Candidate,
            metrics: None,
            registered_at: SystemTime::now(),
            promoted_by: None,
        }
    }

    #[test]
    fn test_model_stage_display() {
        assert_eq!(ModelStage::Candidate.to_string(), "candidate");
        assert_eq!(ModelStage::Production.to_string(), "production");
    }

    #[tokio::test]
    async fn test_log_and_register() {
        let registry = InMemoryModelRegistry::new();
        let run = make_run("run-1", "DQN");
        registry.log_run(run).await.unwrap();

        let version = make_version("v1", "run-1", "DQN");
        registry.register_version(version).await.unwrap();

        let prod = registry.get_production_model("DQN").await.unwrap();
        assert!(prod.is_none()); // Not promoted yet
    }

    #[tokio::test]
    async fn test_promote_lifecycle() {
        let registry = InMemoryModelRegistry::new();
        let run = make_run("run-1", "DQN");
        registry.log_run(run).await.unwrap();

        let version = make_version("v1", "run-1", "DQN");
        registry.register_version(version).await.unwrap();

        // Candidate → Staging
        registry
            .promote("v1", ModelStage::Staging, "ci")
            .await
            .unwrap();

        // Staging → Production
        registry
            .promote("v1", ModelStage::Production, "ci")
            .await
            .unwrap();

        let prod = registry.get_production_model("DQN").await.unwrap();
        assert!(prod.is_some());
        assert_eq!(prod.unwrap().version_id, "v1");
    }

    #[tokio::test]
    async fn test_invalid_transition() {
        let registry = InMemoryModelRegistry::new();
        let version = make_version("v1", "run-1", "DQN");
        registry.register_version(version).await.unwrap();

        // Candidate → Production is invalid (must go through Staging)
        let result = registry
            .promote("v1", ModelStage::Production, "ci")
            .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_promotion_archives_old() {
        let registry = InMemoryModelRegistry::new();

        // Register and promote v1 to production
        registry
            .register_version(make_version("v1", "run-1", "DQN"))
            .await
            .unwrap();
        registry
            .promote("v1", ModelStage::Staging, "ci")
            .await
            .unwrap();
        registry
            .promote("v1", ModelStage::Production, "ci")
            .await
            .unwrap();

        // Register and promote v2 to production
        registry
            .register_version(make_version("v2", "run-2", "DQN"))
            .await
            .unwrap();
        registry
            .promote("v2", ModelStage::Staging, "ci")
            .await
            .unwrap();
        registry
            .promote("v2", ModelStage::Production, "ci")
            .await
            .unwrap();

        // v1 should be archived, v2 should be production
        let prod = registry.get_production_model("DQN").await.unwrap();
        assert_eq!(prod.unwrap().version_id, "v2");
    }

    #[tokio::test]
    async fn test_revert() {
        let registry = InMemoryModelRegistry::new();

        // v1 → production → archived (when v2 promoted)
        registry
            .register_version(make_version("v1", "run-1", "DQN"))
            .await
            .unwrap();
        registry
            .promote("v1", ModelStage::Staging, "ci")
            .await
            .unwrap();
        registry
            .promote("v1", ModelStage::Production, "ci")
            .await
            .unwrap();

        registry
            .register_version(make_version("v2", "run-2", "DQN"))
            .await
            .unwrap();
        registry
            .promote("v2", ModelStage::Staging, "ci")
            .await
            .unwrap();
        registry
            .promote("v2", ModelStage::Production, "ci")
            .await
            .unwrap();

        // Revert should bring v1 back
        let reverted = registry.revert("DQN").await.unwrap();
        assert_eq!(reverted.version_id, "v1");
        assert_eq!(reverted.stage, ModelStage::Production);
    }

    #[tokio::test]
    async fn test_log_metrics() {
        let registry = InMemoryModelRegistry::new();
        registry
            .register_version(make_version("v1", "run-1", "DQN"))
            .await
            .unwrap();

        let metrics = ModelMetrics {
            sharpe_ratio: 1.5,
            accuracy: 0.62,
            win_rate: 0.58,
            max_drawdown: -0.05,
            extra: HashMap::new(),
        };
        registry.log_metrics("v1", metrics).await.unwrap();

        // Metrics should be attached to version
        let versions = registry.versions.read().await;
        let v = versions.get("v1").unwrap();
        assert!(v.metrics.is_some());
        assert!((v.metrics.as_ref().unwrap().sharpe_ratio - 1.5).abs() < 1e-10);
    }

    #[tokio::test]
    async fn test_compare_runs() {
        let registry = InMemoryModelRegistry::new();
        registry
            .log_run(make_run("run-1", "DQN"))
            .await
            .unwrap();
        registry
            .log_run(make_run("run-2", "DQN"))
            .await
            .unwrap();

        let comparison = registry
            .compare_runs(&["run-1".into(), "run-2".into()])
            .await
            .unwrap();
        assert_eq!(comparison.runs.len(), 2);
    }
}