foxhunt/crates/ml/tests/supervised_gpu_smoke_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Supervised Model GPU Smoke Tests
//!
//! Validates all 8 supervised models through the UnifiedTrainable pipeline
//! on a CUDA device: construct -> train 10 epochs -> checkpoint roundtrip.
//!
//! Models: TFT, Mamba2, TGGN, TLOB, Liquid, KAN, xLSTM, Diffusion
//!
//! Uses synthetic data with small configs to keep VRAM usage minimal.
//! Skips gracefully if no CUDA device is available.

#![allow(unused_crate_dependencies)]

// candle eliminated — test uses native cudarc + UnifiedTrainable APIs
use ml::training::unified_trainer::UnifiedTrainable;
use std::sync::Arc;
use cudarc::driver::{CudaContext, CudaStream};
use tracing::{info, warn};

fn require_cuda() -> Arc<CudaStream> {
    match CudaContext::new(0) {
        Ok(ctx) => match ctx.new_stream() {
            Ok(stream) => {
                info!("CUDA device available");
                stream
            }
            Err(e) => {
                warn!("CUDA stream creation failed: {}, skipping GPU smoke tests", e);
                std::process::exit(0);
            }
        },
        Err(e) => {
            warn!("CUDA not available ({}), skipping GPU smoke tests", e);
            std::process::exit(0);
        }
    }
}

/// Generate random f32 data on host for synthetic inputs/targets.
fn random_f32_data(n: usize) -> Vec<f32> {
    use rand::Rng;
    let mut rng = rand::thread_rng();
    (0..n).map(|_| rng.gen::<f32>() - 0.5).collect()
}

/// Generic train-checkpoint pipeline for any UnifiedTrainable model.
/// Uses forward_loss (f32 slices) -> backward (f64 loss) -> optimizer_step.
fn smoke_pipeline(
    adapter: &mut dyn UnifiedTrainable,
    input: &[f32],
    target: &[f32],
    model_name: &str,
) {
    // 1. Train 10 epochs
    let mut first_loss = None;
    let mut last_loss = 0.0_f64;
    for epoch in 0..10 {
        let loss_val = adapter.forward_loss(input, target).unwrap();
        assert!(
            loss_val.is_finite(),
            "{} epoch {}: loss is NaN/Inf ({})",
            model_name,
            epoch,
            loss_val,
        );
        if first_loss.is_none() {
            first_loss = Some(loss_val);
        }
        last_loss = loss_val;

        // Some models (TFT, xLSTM, Diffusion) don't support backward via
        // UnifiedTrainable — they use their own train() methods. Skip if Err.
        let grad_result = adapter.backward(loss_val);
        if grad_result.is_err() {
            adapter.zero_grad().ok();
            continue;
        }
        let grad_norm = grad_result.unwrap();
        assert!(
            grad_norm.is_finite(),
            "{} epoch {}: grad_norm is NaN/Inf ({})",
            model_name,
            epoch,
            grad_norm,
        );

        adapter.optimizer_step().unwrap();
        adapter.zero_grad().unwrap();
    }
    let first = first_loss.unwrap();
    info!(
        model = model_name,
        first_loss = first,
        last_loss,
        reduction_pct = (1.0 - last_loss / first) * 100.0,
        "Train complete"
    );
    // Models with backward support increment step in backward/optimizer_step.
    // Models without (TFT, xLSTM, Diffusion) stay at step 0 — that's expected.
    let step = adapter.get_step();
    assert!(step == 0 || step == 10, "{} expected 0 or 10 steps, got {}", model_name, step);

    // 2. Checkpoint roundtrip
    let tmp_dir = std::env::temp_dir().join(format!("gpu_smoke_{}", model_name.to_lowercase()));
    std::fs::create_dir_all(&tmp_dir).unwrap();
    let ckpt_path = tmp_dir.join("ckpt");
    let save_result = adapter.save_checkpoint(ckpt_path.to_str().unwrap());
    assert!(
        save_result.is_ok(),
        "{} checkpoint save failed: {:?}",
        model_name,
        save_result.err(),
    );

    // 3. Metrics
    let metrics = adapter.collect_metrics();
    assert!(
        metrics.learning_rate > 0.0,
        "{} learning rate should be > 0",
        model_name,
    );

    let _ = std::fs::remove_dir_all(&tmp_dir);
}

// ────────────────────────────────────────────
// TFT
// ────────────────────────────────────────────
#[test]
fn test_tft_gpu_smoke() {
    let _stream = require_cuda();
    info!("=== TFT GPU Smoke ===");

    use ml::tft::trainable_adapter::TrainableTFT;
    use ml::tft::TFTConfig;

    let feature_dim = 10;
    let mut config = TFTConfig::default();
    config.input_dim = feature_dim;
    config.hidden_dim = 32;
    config.num_heads = 2;
    config.num_layers = 1;
    config.num_quantiles = 3;
    config.num_static_features = 0;
    config.num_known_features = 0;
    config.num_unknown_features = feature_dim;
    config.sequence_length = 1;
    config.prediction_horizon = 1;
    config.learning_rate = 1e-3;
    config.dropout_rate = 0.0;

    // TFT auto-selects device (cuda_if_available) in TemporalFusionTransformer::new()
    let mut adapter = TrainableTFT::new(config).unwrap();
    assert!(
        adapter.device_name().contains("cuda"),
        "TFT should be on CUDA, got: {}",
        adapter.device_name()
    );

    // TFT forward_loss processes one sample at a time via UnifiedTrainable.
    // Input: [feature_dim] (single sample), target: [quantiles * horizon]
    let input = random_f32_data(feature_dim);
    let target = random_f32_data(1 * 3); // horizon=1, quantiles=3
    smoke_pipeline(&mut adapter, &input, &target, "TFT");
}

// ────────────────────────────────────────────
// Mamba2
// ────────────────────────────────────────────
#[test]
fn test_mamba2_gpu_smoke() {
    let stream = require_cuda();
    info!("=== Mamba2 GPU Smoke ===");

    use ml::mamba::trainable_adapter::Mamba2TrainableAdapter;
    use ml::mamba::Mamba2Config;

    let mut config = Mamba2Config::default();
    config.d_model = 32;
    config.num_layers = 1;
    config.d_state = 8;
    config.max_seq_len = 8;

    let mut adapter = Mamba2TrainableAdapter::new(config, &stream).unwrap();
    assert!(
        adapter.device_name().contains("uda") || adapter.device_name().contains("CUDA"),
        "Mamba2 should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, seq_len=8, d_model=32] flattened
    let input = random_f32_data(16 * 8 * 32);
    // Mamba2 compute_loss narrows to last step -> target is [batch=16, 1] flattened
    let target = random_f32_data(16 * 1);
    smoke_pipeline(&mut adapter, &input, &target, "Mamba2");
}

// ────────────────────────────────────────────
// TGGN
// ────────────────────────────────────────────
#[test]
fn test_tggn_gpu_smoke() {
    let _stream = require_cuda();
    info!("=== TGGN GPU Smoke ===");

    use ml::tgnn::trainable_adapter::TGGNTrainableAdapter;
    use ml::tgnn::TGGNConfig;

    let config = TGGNConfig {
        node_dim: 10,
        hidden_dim: 16,
        num_layers: 2,
        max_nodes: 8,
        max_edges: 16,
        edge_dim: 4,
        temporal_decay: 0.99,
        update_frequency_ns: 1_000_000,
        use_simd: false,
    };

    let mut adapter = TGGNTrainableAdapter::new(config).unwrap();
    assert!(
        adapter.device_name().contains("cuda"),
        "TGGN should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, node_dim=10] flattened, Output: [batch=16, 1] flattened
    let input = random_f32_data(16 * 10);
    let target = random_f32_data(16 * 1);
    smoke_pipeline(&mut adapter, &input, &target, "TGGN");
}

// ────────────────────────────────────────────
// TLOB
// ────────────────────────────────────────────
#[test]
fn test_tlob_gpu_smoke() {
    let _stream = require_cuda();
    info!("=== TLOB GPU Smoke ===");

    use ml::tlob::trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter};

    let config = TLOBAdapterConfig {
        d_model: 16,
        num_heads: 2,
        num_layers: 1,
        seq_len: 1,
        feature_dim: 10,
    };

    let mut adapter = TLOBTrainableAdapter::new(config).unwrap();
    assert!(
        adapter.device_name().contains("cuda"),
        "TLOB should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, seq_len*feature_dim=10] flattened, Output: [batch=16, 1] flattened
    let input = random_f32_data(16 * 10);
    let target = random_f32_data(16 * 1);
    smoke_pipeline(&mut adapter, &input, &target, "TLOB");
}

// ────────────────────────────────────────────
// Liquid (CfC)
// ────────────────────────────────────────────
#[test]
fn test_liquid_gpu_smoke() {
    let _stream = require_cuda();
    info!("=== Liquid GPU Smoke ===");

    use ml::liquid::adapter::LiquidTrainableAdapter;
    use ml::liquid::CfCTrainConfig;
    use ml::gpu::DeviceConfig;

    let config = CfCTrainConfig {
        input_size: 10,
        hidden_size: 16,
        output_size: 1,
        backbone_hidden_sizes: vec![16, 8],
        learning_rate: 1e-3,
        device: DeviceConfig::Cuda(0),
        ..CfCTrainConfig::default()
    };

    let mut adapter = LiquidTrainableAdapter::new(config).unwrap();
    assert!(
        adapter.device_name().contains("cuda"),
        "Liquid should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, seq_len=4, input_size=10] flattened (3D)
    let input = random_f32_data(16 * 4 * 10);
    // Output: [batch=16, output_size=1] flattened
    let target = random_f32_data(16 * 1);
    smoke_pipeline(&mut adapter, &input, &target, "Liquid");
}

// ────────────────────────────────────────────
// KAN
// ────────────────────────────────────────────
#[test]
fn test_kan_gpu_smoke() {
    let _stream = require_cuda();
    info!("=== KAN GPU Smoke ===");

    use ml::kan::config::KANConfig;
    use ml::kan::trainable::KANTrainableAdapter;

    let config = KANConfig {
        layer_widths: vec![10, 8, 4, 1],
        grid_size: 3,
        spline_order: 3,
        learning_rate: 1e-3,
        weight_decay: 1e-5,
        grad_clip: 1.0,
    };

    let mut adapter = KANTrainableAdapter::new(config).unwrap();
    assert!(
        adapter.device_name().contains("cuda"),
        "KAN should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, 10] flattened, Output: [batch=16, 1] flattened
    let input = random_f32_data(16 * 10);
    let target = random_f32_data(16 * 1);
    smoke_pipeline(&mut adapter, &input, &target, "KAN");
}

// ────────────────────────────────────────────
// xLSTM
// ────────────────────────────────────────────
#[test]
fn test_xlstm_gpu_smoke() {
    let _stream = require_cuda();
    info!("=== xLSTM GPU Smoke ===");

    use ml::xlstm::trainable::XLSTMTrainableAdapter;
    use ml::xlstm::config::XLSTMConfig;

    let config = XLSTMConfig {
        input_dim: 10,
        hidden_dim: 16,
        num_blocks: 2,
        num_heads: 2,
        slstm_ratio: 0.5,
        output_dim: 1,
        dropout: 0.0,
        learning_rate: 1e-3,
        weight_decay: 1e-5,
        grad_clip: 1.0,
    };

    let mut adapter = XLSTMTrainableAdapter::new(config).unwrap();
    assert!(
        adapter.device_name().contains("cuda"),
        "xLSTM should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, seq_len=4, input_dim=10] flattened (3D)
    let input = random_f32_data(16 * 4 * 10);
    // Output: [batch=16, 1] flattened
    let target = random_f32_data(16 * 1);
    smoke_pipeline(&mut adapter, &input, &target, "xLSTM");
}

// ────────────────────────────────────────────
// Diffusion
// ────────────────────────────────────────────
#[test]
fn test_diffusion_gpu_smoke() {
    let stream = require_cuda();
    info!("=== Diffusion GPU Smoke ===");

    use ml::diffusion::config::DiffusionConfig;
    use ml::diffusion::trainable::DiffusionTrainableAdapter;

    let config = DiffusionConfig {
        num_timesteps: 50,
        sampling_steps: 5,
        seq_len: 8,
        feature_dim: 1,
        hidden_dim: 16,
        num_layers: 1,
        time_embed_dim: 8,
        learning_rate: 1e-3,
        weight_decay: 1e-5,
        grad_clip: 1.0,
        ..Default::default()
    };
    let data_dim = config.data_dim(); // 8

    let mut adapter = DiffusionTrainableAdapter::new(config, &stream).unwrap();
    assert!(
        adapter.device_name().contains("cuda") || adapter.device_name().contains("Cuda"),
        "Diffusion should be on CUDA, got: {}",
        adapter.device_name()
    );

    // Input: [batch=16, data_dim=8] flattened
    let input = random_f32_data(16 * data_dim);
    // Diffusion: use same-sized pseudo-target (noise prediction, loss won't be meaningful)
    let target = random_f32_data(16 * data_dim);
    smoke_pipeline(&mut adapter, &input, &target, "Diffusion");
}