From 0e1c8be186ba371702ec7af72ff5ecb17a38b26d Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sat, 21 Feb 2026 19:23:40 +0100
Subject: [PATCH] feat(ml): wire TFT safetensors checkpoint save/load

Replace empty placeholder safetensors file with actual model weight
serialization using the VarMap, matching the DQN trainable adapter pattern.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ml/src/tft/trainable_adapter.rs | 76 ++++++++++++++++++++++++++++-----
 1 file changed, 66 insertions(+), 10 deletions(-)

diff --git a/ml/src/tft/trainable_adapter.rs b/ml/src/tft/trainable_adapter.rs
index a6958f035..bfe5cd45c 100644
--- a/ml/src/tft/trainable_adapter.rs
+++ b/ml/src/tft/trainable_adapter.rs
@@ -438,7 +438,7 @@ impl UnifiedTrainable for TrainableTFT {
 
     /// Save model checkpoint in standardized format
     ///
-    /// Format: JSON for metadata (model weights require TFT refactoring)
+    /// Saves model weights via safetensors and metadata via JSON.
     ///
     /// # Arguments
     /// * `checkpoint_path` - Path to save checkpoint (without extension)
@@ -446,9 +446,6 @@ impl UnifiedTrainable for TrainableTFT {
     /// # Returns
     /// Path to saved checkpoint
     fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
-        // TODO: Implement safetensors checkpoint save when TFT exposes VarMap
-        // For now, save only metadata
-
         // Create and save checkpoint metadata
         let metadata = CheckpointMetadata {
             model_type: "TFT".to_string(),
@@ -464,34 +461,93 @@ impl UnifiedTrainable for TrainableTFT {
         // Save metadata to JSON
         crate::training::unified_trainer::checkpoint::save_metadata(&metadata, checkpoint_path)?;
 
-        // Create placeholder safetensors file for compatibility
+        // Save model weights to safetensors format
         let safetensors_path = format!("{}.safetensors", checkpoint_path);
-        std::fs::write(&safetensors_path, b"")
-            .map_err(|e| MLError::ModelError(format!("Failed to create checkpoint file: {}", e)))?;
+
+        // Extract tensors from VarMap
+        let vars_data = self
+            .model
+            .varmap
+            .data()
+            .lock()
+            .map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap for checkpoint save: {}", e)))?;
+
+        let mut tensors: HashMap<String, Tensor> = HashMap::new();
+        for (name, var) in vars_data.iter() {
+            tensors.insert(name.clone(), var.as_tensor().clone());
+        }
+
+        // Save using safetensors
+        candle_core::safetensors::save(&tensors, &safetensors_path)
+            .map_err(|e| MLError::ModelError(format!("Failed to save safetensors: {}", e)))?;
+
+        tracing::info!(
+            "Saved TFT checkpoint to {} (step {}, {} tensors)",
+            checkpoint_path,
+            self.step_count,
+            tensors.len()
+        );
 
         Ok(safetensors_path)
     }
 
     /// Load model checkpoint from standardized format
     ///
+    /// Loads model weights from safetensors and metadata from JSON.
+    ///
     /// # Arguments
     /// * `checkpoint_path` - Path to checkpoint (without extension)
     ///
     /// # Returns
     /// Loaded checkpoint metadata
     fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
-        // TODO: Implement safetensors checkpoint load when TFT exposes VarMap
-        // For now, load only metadata
-
         // Load metadata from JSON
         let metadata =
             crate::training::unified_trainer::checkpoint::load_metadata(checkpoint_path)?;
 
+        // Validate model type
+        if metadata.model_type != "TFT" {
+            return Err(MLError::ModelError(format!(
+                "Invalid model type in checkpoint: expected TFT, got {}",
+                metadata.model_type
+            )));
+        }
+
+        // Load model weights from safetensors
+        let safetensors_path = format!("{}.safetensors", checkpoint_path);
+        let tensors = candle_core::safetensors::load(&safetensors_path, &self.model.device)
+            .map_err(|e| MLError::ModelError(format!("Failed to load safetensors: {}", e)))?;
+
+        // Load tensors into VarMap
+        let vars_data = self
+            .model
+            .varmap
+            .data()
+            .lock()
+            .map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap for checkpoint load: {}", e)))?;
+
+        for (name, tensor) in &tensors {
+            if let Some(var) = vars_data.get(name) {
+                var.set(tensor).map_err(|e| {
+                    MLError::ModelError(format!("Failed to set var {}: {}", name, e))
+                })?;
+            } else {
+                tracing::warn!("Checkpoint contains unknown variable: {}", name);
+            }
+        }
+
         // Update model state from metadata
         self.step_count = metadata.step;
         self.model.is_trained = true;
         self.learning_rate = metadata.metrics.learning_rate;
 
+        tracing::info!(
+            "Loaded TFT checkpoint from {} (step {}, {} tensors)",
+            checkpoint_path,
+            metadata.step,
+            tensors.len()
+        );
+
         Ok(metadata)
     }