foxhunt/crates/ml-supervised/src/checkpoint.rs

//! Checkpoint implementations for supervised models (Mamba2, TGGN).

use std::collections::HashMap;

use async_trait::async_trait;
use ml_core::{Checkpointable, MLError, ModelType};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tracing::debug;

use crate::mamba::{Mamba2Config, Mamba2SSM};
use crate::tgnn::{TGGNConfig, TGGN};

// ========== MAMBA2 ==========

/// Serializable state for Mamba2 model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MambaCheckpointState {
    pub config: Mamba2Config,
    pub epoch: Option<u64>,
    pub step: Option<u64>,
    pub training_loss: f64,
    pub validation_loss: f64,
    pub ssd_layer_weights: Vec<Vec<f32>>,
    pub input_projection_weights: Vec<f32>,
    pub output_projection_weights: Vec<f32>,
    pub layer_norm_weights: Vec<Vec<f32>>,
    pub ssm_a_matrices: Vec<Vec<f32>>,
    pub ssm_b_matrices: Vec<Vec<f32>>,
    pub ssm_c_matrices: Vec<Vec<f32>>,
    pub ssm_delta_params: Vec<f32>,
    pub total_inferences: u64,
    pub avg_latency_us: f64,
    pub throughput_pps: f64,
}

#[async_trait]
impl Checkpointable for Mamba2SSM {
    fn model_type(&self) -> ModelType {
        ModelType::MAMBA
    }

    fn model_name(&self) -> &str {
        &self.metadata.model_id
    }

    fn model_version(&self) -> &str {
        &self.metadata.version
    }

    async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
        let (current_epoch, current_step) = self.get_current_training_state();
        let training_metrics = self.get_training_metrics();
        let performance_stats = self.get_inference_stats();

        let state = MambaCheckpointState {
            config: self.config.clone(),
            epoch: current_epoch,
            step: current_step,
            training_loss: training_metrics
                .get("training_loss")
                .copied()
                .unwrap_or(0.0),
            validation_loss: training_metrics
                .get("validation_loss")
                .copied()
                .unwrap_or(0.0),
            ssd_layer_weights: self.extract_ssd_weights(),
            input_projection_weights: self.extract_input_projection_weights(),
            output_projection_weights: self.extract_output_projection_weights(),
            layer_norm_weights: self.extract_layer_norm_weights(),
            ssm_a_matrices: self.extract_ssm_matrices("A"),
            ssm_b_matrices: self.extract_ssm_matrices("B"),
            ssm_c_matrices: self.extract_ssm_matrices("C"),
            ssm_delta_params: self.extract_delta_params(),
            total_inferences: self
                .total_inferences
                .load(std::sync::atomic::Ordering::Relaxed),
            avg_latency_us: performance_stats
                .get("avg_latency_us")
                .copied()
                .unwrap_or(0.0),
            throughput_pps: performance_stats
                .get("throughput_pps")
                .copied()
                .unwrap_or(0.0),
        };

        let serialized = serde_json::to_vec(&state)
            .map_err(|e| MLError::ModelError(format!("MAMBA serialization failed: {}", e)))?;

        debug!("Serialized MAMBA state: {} bytes", serialized.len());
        Ok(serialized)
    }

    async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
        let state: MambaCheckpointState = serde_json::from_slice(data)
            .map_err(|e| MLError::ModelError(format!("MAMBA deserialization failed: {}", e)))?;

        self.config = state.config;
        self.restore_ssd_weights(&state.ssd_layer_weights);
        self.restore_input_projection_weights(&state.input_projection_weights);
        self.restore_output_projection_weights(&state.output_projection_weights);
        self.restore_layer_norm_weights(&state.layer_norm_weights);
        self.restore_ssm_matrices("A", &state.ssm_a_matrices);
        self.restore_ssm_matrices("B", &state.ssm_b_matrices);
        self.restore_ssm_matrices("C", &state.ssm_c_matrices);
        self.restore_delta_params(&state.ssm_delta_params);

        debug!("Deserialized MAMBA state from {} bytes", data.len());
        Ok(())
    }

    fn get_training_state(&self) -> (Option<u64>, Option<u64>, Option<f64>, Option<f64>) {
        if let Some(last_epoch) = self.metadata.training_history.back() {
            let total_steps = self
                .metadata
                .training_history
                .iter()
                .map(|epoch| epoch.epoch as u64 * 1000)
                .max()
                .unwrap_or(0);

            (
                Some(last_epoch.epoch as u64),
                Some(total_steps),
                Some(last_epoch.loss),
                Some(last_epoch.accuracy),
            )
        } else {
            (Some(0), Some(0), Some(f64::INFINITY), Some(0.0))
        }
    }

    fn get_hyperparameters(&self) -> HashMap<String, Value> {
        let mut params = HashMap::new();
        params.insert("d_model".to_owned(), Value::from(self.config.d_model));
        params.insert("d_state".to_owned(), Value::from(self.config.d_state));
        params.insert("d_head".to_owned(), Value::from(self.config.d_head));
        params.insert("num_heads".to_owned(), Value::from(self.config.num_heads));
        params.insert("expand".to_owned(), Value::from(self.config.expand));
        params.insert(
            "num_layers".to_owned(),
            Value::from(self.config.num_layers),
        );
        params.insert("dropout".to_owned(), Value::from(self.config.dropout));
        params.insert(
            "learning_rate".to_owned(),
            Value::from(self.config.learning_rate),
        );
        params.insert(
            "target_latency_us".to_owned(),
            Value::from(self.config.target_latency_us),
        );
        params
    }

    fn get_metrics(&self) -> HashMap<String, f64> {
        self.get_performance_metrics()
    }

    fn get_architecture_info(&self) -> HashMap<String, Value> {
        let mut info = HashMap::new();
        info.insert("model_type".to_owned(), Value::from("MAMBA-2"));
        info.insert("input_dim".to_owned(), Value::from(self.config.d_model));
        info.insert(
            "hidden_dim".to_owned(),
            Value::from(self.config.d_model * self.config.expand),
        );
        info.insert("state_dim".to_owned(), Value::from(self.config.d_state));
        info.insert(
            "num_layers".to_owned(),
            Value::from(self.config.num_layers),
        );
        info.insert("use_ssd".to_owned(), Value::from(self.config.use_ssd));
        info.insert(
            "use_selective_state".to_owned(),
            Value::from(self.config.use_selective_state),
        );
        info.insert(
            "hardware_aware".to_owned(),
            Value::from(self.config.hardware_aware),
        );
        info
    }
}

// ========== TGGN ==========

/// Serializable state for TGGN model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TGGNCheckpointState {
    pub config: TGGNConfig,
    pub epoch: Option<u64>,
    pub step: Option<u64>,
    pub training_loss: f64,
    pub validation_loss: f64,
    pub node_embeddings: HashMap<String, Vec<f32>>,
    pub edge_embeddings: HashMap<String, Vec<f32>>,
    pub graph_statistics: HashMap<String, f64>,
    pub message_passing_weights: Vec<Vec<f32>>,
    pub gating_weights: Vec<f32>,
    pub total_inferences: u64,
    pub graph_updates: u64,
    pub avg_inference_latency_ns: u64,
    pub avg_graph_update_latency_ns: u64,
}

#[async_trait]
impl Checkpointable for TGGN {
    fn model_type(&self) -> ModelType {
        ModelType::TGGN
    }

    fn model_name(&self) -> &str {
        &self.metadata.version
    }

    fn model_version(&self) -> &str {
        &self.metadata.version
    }

    async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
        let node_embeddings: HashMap<String, Vec<f32>> = self
            .node_embeddings()
            .iter()
            .map(|entry| {
                let key = format!("{:?}", entry.key());
                let value = entry.value().iter().map(|&x| x as f32).collect();
                (key, value)
            })
            .collect();

        let edge_embeddings: HashMap<String, Vec<f32>> = self
            .edge_embeddings()
            .iter()
            .map(|entry| {
                let key = format!("{:?}", entry.key());
                let value = entry.value().iter().map(|&x| x as f32).collect();
                (key, value)
            })
            .collect();

        let state = TGGNCheckpointState {
            config: self.config().clone(),
            epoch: None,
            step: None,
            training_loss: 0.0,
            validation_loss: 0.0,
            node_embeddings,
            edge_embeddings,
            graph_statistics: HashMap::new(),
            message_passing_weights: Vec::new(),
            gating_weights: Vec::new(),
            total_inferences: self
                .inference_count()
                .load(std::sync::atomic::Ordering::Relaxed),
            graph_updates: self
                .graph_updates()
                .load(std::sync::atomic::Ordering::Relaxed),
            avg_inference_latency_ns: self.calculate_avg_inference_latency(),
            avg_graph_update_latency_ns: self.calculate_avg_graph_update_latency(),
        };

        let serialized = serde_json::to_vec(&state)
            .map_err(|e| MLError::ModelError(format!("TGGN serialization failed: {}", e)))?;

        debug!("Serialized TGGN state: {} bytes", serialized.len());
        Ok(serialized)
    }

    async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
        let state: TGGNCheckpointState = serde_json::from_slice(data)
            .map_err(|e| MLError::ModelError(format!("TGGN deserialization failed: {}", e)))?;

        self.restore_node_embeddings(&state.node_embeddings)?;
        self.restore_edge_embeddings(&state.edge_embeddings)?;
        self.restore_graph_statistics(&state.graph_statistics)?;
        self.restore_message_passing_weights(&state.message_passing_weights)?;

        self.inference_count()
            .store(state.total_inferences, std::sync::atomic::Ordering::Relaxed);
        self.graph_updates()
            .store(state.graph_updates, std::sync::atomic::Ordering::Relaxed);

        debug!("Deserialized TGGN state from {} bytes", data.len());
        Ok(())
    }

    fn get_hyperparameters(&self) -> HashMap<String, Value> {
        let mut params = HashMap::new();
        let config = self.config();
        params.insert("max_nodes".to_owned(), Value::from(config.max_nodes));
        params.insert("max_edges".to_owned(), Value::from(config.max_edges));
        params.insert("node_dim".to_owned(), Value::from(config.node_dim));
        params.insert("edge_dim".to_owned(), Value::from(config.edge_dim));
        params.insert("hidden_dim".to_owned(), Value::from(config.hidden_dim));
        params.insert("num_layers".to_owned(), Value::from(config.num_layers));
        params.insert(
            "temporal_decay".to_owned(),
            Value::from(config.temporal_decay),
        );
        params.insert("use_simd".to_owned(), Value::from(config.use_simd));
        params
    }

    fn get_architecture_info(&self) -> HashMap<String, Value> {
        let mut info = HashMap::new();
        let config = self.config();
        info.insert("model_type".to_owned(), Value::from("TGGN"));
        info.insert("max_nodes".to_owned(), Value::from(config.max_nodes));
        info.insert("max_edges".to_owned(), Value::from(config.max_edges));
        info.insert("node_feature_dim".to_owned(), Value::from(config.node_dim));
        info.insert("edge_feature_dim".to_owned(), Value::from(config.edge_dim));
        info.insert("gnn_layers".to_owned(), Value::from(config.num_layers));
        info.insert("supports_temporal_decay".to_owned(), Value::from(true));
        info
    }
}