foxhunt/ml/tests/tft_int8_integration_test.rs

//! Integration test for TFT INT8 quantization workflow
//!
//! Tests the complete flow:
//! 1. Train FP32 model
//! 2. Automatic INT8 quantization
//! 3. Checkpoint saving with metadata
//! 4. Verify memory savings

use ml::checkpoint::FileSystemStorage;
use ml::tft::training::{TFTBatch, TFTDataLoader};
use ml::trainers::tft::{TFTTrainer, TFTTrainerConfig};
use ndarray::Array2;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::TempDir;

#[tokio::test]
async fn test_tft_int8_quantization_integration() {
    // Create temporary directory for checkpoints
    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let checkpoint_dir = temp_dir.path().to_str().unwrap().to_string();

    // Create trainer config with INT8 quantization enabled
    let config = TFTTrainerConfig {
        epochs: 2, // Small number for testing
        batch_size: 2,
        hidden_dim: 32,
        num_attention_heads: 2,
        checkpoint_dir: checkpoint_dir.clone(),
        use_int8_quantization: true, // Enable INT8
        ..Default::default()
    };

    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(&checkpoint_dir)));
    let mut trainer = TFTTrainer::new(config, storage).expect("Failed to create trainer");

    // Create minimal training data
    let train_loader = create_minimal_dataloader(2);
    let val_loader = create_minimal_dataloader(1);

    // Train model (should automatically quantize to INT8 after FP32 training)
    let result = trainer.train(train_loader, val_loader).await;
    assert!(result.is_ok(), "Training failed: {:?}", result.err());

    // Verify trainer switched to INT8 model
    assert!(
        trainer.is_int8(),
        "Trainer should be using INT8 model after training"
    );

    // Verify checkpoint file exists with INT8 suffix
    let checkpoint_path = PathBuf::from(&checkpoint_dir).join("tft_225_int8_epoch_1.safetensors");
    assert!(
        checkpoint_path.exists(),
        "INT8 checkpoint file does not exist: {:?}",
        checkpoint_path
    );

    // Verify metadata indicates INT8
    let metadata_path = checkpoint_path.with_extension("json");
    assert!(metadata_path.exists(), "Metadata file does not exist");

    let metadata_content =
        std::fs::read_to_string(&metadata_path).expect("Failed to read metadata");
    let metadata: serde_json::Value =
        serde_json::from_str(&metadata_content).expect("Failed to parse metadata JSON");

    assert_eq!(metadata["model_name"], "TFT-INT8");
    assert_eq!(metadata["hyperparameters"]["quantization"], "int8");
    assert_eq!(metadata["custom_metadata"]["model_type"], "int8");

    println!("✅ INT8 quantization integration test passed");
    println!("✅ Checkpoint saved: {}", checkpoint_path.display());
    println!("✅ Metadata verified: INT8 model type");
}

#[tokio::test]
async fn test_tft_fp32_no_quantization() {
    // Create temporary directory for checkpoints
    let temp_dir = TempDir::new().expect("Failed to create temp dir");
    let checkpoint_dir = temp_dir.path().to_str().unwrap().to_string();

    // Create trainer config WITHOUT INT8 quantization
    let config = TFTTrainerConfig {
        epochs: 2,
        batch_size: 2,
        hidden_dim: 32,
        num_attention_heads: 2,
        checkpoint_dir: checkpoint_dir.clone(),
        use_int8_quantization: false, // Disable INT8
        ..Default::default()
    };

    let storage = Arc::new(FileSystemStorage::new(PathBuf::from(&checkpoint_dir)));
    let mut trainer = TFTTrainer::new(config, storage).expect("Failed to create trainer");

    // Create minimal training data
    let train_loader = create_minimal_dataloader(2);
    let val_loader = create_minimal_dataloader(1);

    // Train model (should remain FP32)
    let result = trainer.train(train_loader, val_loader).await;
    assert!(result.is_ok(), "Training failed: {:?}", result.err());

    // Verify trainer is still using FP32 model
    assert!(
        !trainer.is_int8(),
        "Trainer should be using FP32 model when quantization disabled"
    );

    // Verify checkpoint file exists with FP32 suffix
    let checkpoint_path = PathBuf::from(&checkpoint_dir).join("tft_225_fp32_epoch_1.safetensors");
    assert!(
        checkpoint_path.exists(),
        "FP32 checkpoint file does not exist: {:?}",
        checkpoint_path
    );

    // Verify metadata indicates FP32
    let metadata_path = checkpoint_path.with_extension("json");
    let metadata_content =
        std::fs::read_to_string(&metadata_path).expect("Failed to read metadata");
    let metadata: serde_json::Value =
        serde_json::from_str(&metadata_content).expect("Failed to parse metadata JSON");

    assert_eq!(metadata["model_name"], "TFT");
    assert_eq!(metadata["hyperparameters"]["quantization"], "fp32");

    println!("✅ FP32 no-quantization test passed");
}

/// Helper: Create minimal TFTDataLoader for testing
fn create_minimal_dataloader(num_batches: usize) -> TFTDataLoader {
    let mut batches = Vec::new();

    for _ in 0..num_batches {
        let batch = TFTBatch {
            static_features: Array2::zeros((2, 5)), // [batch=2, static=5]
            historical_features: Array2::zeros((2, 39)), // [batch=2, unknown=39]
            future_features: Array2::zeros((2, 10)), // [batch=2, known=10]
            targets: Array2::zeros((2, 10)),        // [batch=2, horizon=10]
        };
        batches.push(batch);
    }

    TFTDataLoader::new(batches)
}