foxhunt/crates/ml/src/trainers/dqn/smoke_tests/performance.rs

//! Performance smoke tests.
//!
//! Speed benchmarks for the DQN training pipeline. Requires GPU and
//! real DBN test data in `test_data/ES.FUT/`.
//!
//! ```sh
//! cargo test -p ml --lib smoke_tests::performance -- --nocapture
//! ```
use super::helpers::{assert_finite, smoke_params, test_data_dir};
use crate::trainers::dqn::DQNTrainer;
use ml_core::device::MlDevice;
use std::sync::Arc;
use std::time::Instant;
use tracing::info;

/// Resolve test data dir — hard error if missing.
fn data_dir() -> String {
    test_data_dir()
        .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback")
}

/// Measure end-to-end training throughput (3 epochs, real data).
///
/// Purpose: verify the training pipeline runs at non-trivial speed AND produces
/// plausible metrics. Previously only checked `loss.is_finite()` which passes on
/// trivial zeros or NaN-less garbage — giving no signal that anything healthy
/// happened. We now assert a conservative steps/second floor and bound the
/// Q-value magnitude.
#[test]
#[ignore] // Loads real training data — run via nightly CI or manual trigger
fn test_training_throughput_measurement() -> anyhow::Result<()> {
    let mut trainer = DQNTrainer::new(smoke_params())?; // auto-detect GPU
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;

    let start = Instant::now();
    let metrics = rt.block_on(trainer.train(&data_dir(), "ES.FUT", |_epoch, _bytes, _best| {
        Ok("skip".to_owned())
    }))?;
    let elapsed = start.elapsed();

    // ── loss: finite and non-negative ──
    assert_finite(metrics.loss, "loss");
    assert!(
        metrics.loss >= 0.0,
        "loss should be non-negative (sign-bug indicator), got {}",
        metrics.loss,
    );

    // ── epochs must have run ──
    assert!(
        metrics.epochs_trained >= 1,
        "no epoch completed — throughput cannot be measured. metrics: {:?}",
        metrics,
    );

    // ── throughput floor ──
    // RTX 3050 Ti smoketest target: > 10 epochs/s is aspirational, > 0.05 epochs/s
    // (i.e. each epoch < 20s) is a safe local floor.  Raise in CI if desired — the
    // point of this floor is to catch catastrophic regressions (e.g. accidentally
    // falling back to CPU, or a kernel that ran but pinned a CPU loop).
    let epochs_per_sec = metrics.epochs_trained as f64 / elapsed.as_secs_f64().max(1e-6);
    assert!(
        epochs_per_sec > 0.05,
        "Training throughput {:.3} epochs/s below floor 0.05 (laptop target). \
         {} epochs in {:.2}s — CI may set a higher floor.",
        epochs_per_sec,
        metrics.epochs_trained,
        elapsed.as_secs_f64(),
    );

    // ── avg_q_value must be present, finite, and bounded ──
    // Rules out NaN-like explosions that happen to be finite-but-huge.
    let avg_q = metrics
        .additional_metrics
        .get("avg_q_value")
        .copied()
        .unwrap_or(f64::NAN);
    assert!(
        avg_q.is_finite(),
        "avg_q_value missing or NaN — diagnostic regressed: {:?}",
        metrics.additional_metrics.get("avg_q_value"),
    );
    assert!(
        avg_q.abs() < 1e6,
        "avg_q_value {avg_q} is implausibly large (C51 atoms exhausted or \
         NaN-like explosion). v_range is finite by construction.",
    );

    info!(
        elapsed_secs = elapsed.as_secs_f64(),
        epochs_trained = metrics.epochs_trained,
        epochs_per_sec = epochs_per_sec,
        "Training throughput"
    );
    info!(loss = metrics.loss, "Avg loss");
    info!(avg_q_value = avg_q, "Avg Q-value");
    if let Some(&final_eps) = metrics.additional_metrics.get("final_epsilon") {
        info!(final_epsilon = final_eps, "Final epsilon");
    }
    drop(trainer);
    drop(rt);
    Ok(())
}

/// Measure per-sample latency of the GPU PER replay buffer.
#[test]
fn test_per_sample_latency() -> anyhow::Result<()> {
    use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};

    let dev = MlDevice::new_cuda(0)?;
    let stream = Arc::clone(dev.cuda_stream()?);
    let state_dim: usize = 48;
    let capacity: usize = 10_000;
    let config = GpuReplayBufferConfig {
        capacity,
        alpha: 0.6,
        beta_start: 0.4,
        beta_max: 1.0,
        beta_annealing_steps: 1000,
        epsilon: 1e-6,
        max_memory_bytes: 4 * 1024 * 1024 * 1024,
        max_batch_size: 1024,
    };
    let mut buf = GpuReplayBuffer::new(config, &stream)?;

    // Fill with 5000 experiences
    let fill_count: usize = 5000;
    let batch_insert: usize = 100;
    for _ in 0..(fill_count / batch_insert) {
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let host_states: Vec<f32> = (0..batch_insert * state_dim).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();
        let host_rewards: Vec<f32> = (0..batch_insert).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();

        let states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?;
        let next_states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?;
        let actions = stream.alloc_zeros::<u32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
        let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?;
        let dones = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
        let aux_sign = stream.alloc_zeros::<i32>(batch_insert).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
        let aux_conf = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?;
        buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, batch_insert)?;
    }
    assert_eq!(buf.len(), fill_count);

    // Warmup
    for _ in 0..5 {
        let _ = buf.sample_proportional(64)?;
    }

    // Timed: 100 proportional samples of batch_size=64
    let sample_rounds: usize = 100;
    let start = Instant::now();
    for _ in 0..sample_rounds {
        let _ = buf.sample_proportional(64)?;
    }
    let elapsed = start.elapsed();
    let us_per_sample = elapsed.as_micros() as f64 / sample_rounds as f64;

    info!(
        us_per_sample = us_per_sample,
        rounds = sample_rounds,
        batch_size = 64,
        "PER proportional sample latency"
    );
    assert!(
        us_per_sample < 50_000.0,
        "sample latency {us_per_sample:.1} us is unreasonably high"
    );

    Ok(())
}

/// Single-epoch training on real data.
/// Loads full 163K-bar dataset — too slow for CI (run via nightly or manual trigger).
///
/// Purpose: one epoch on real ES.FUT data produces valid training metrics.
/// Previously asserted only `loss.is_finite()`, which passes on `loss == 0.0`
/// (a sign-bug / accumulation-bug tell) or any bounded NaN-less value. A
/// healthy epoch must complete at least one step, produce bounded non-negative
/// loss, and a bounded finite avg_q_value.
#[test]
#[ignore]
fn test_real_data_single_epoch() -> anyhow::Result<()> {
    let mut params = smoke_params();
    params.epochs = 1;
    let mut trainer = DQNTrainer::new(params)?;
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;

    let start = Instant::now();
    let metrics = rt.block_on(trainer.train(&data_dir(), "ES.FUT", |_epoch, _bytes, _best| {
        Ok("skip".to_owned())
    }))?;
    let elapsed = start.elapsed();

    // ── loss: finite, bounded, non-negative ──
    // Negative loss would signal a sign bug (loss is ≥0 by construction — MSE + C51
    // KL + CQL are all non-negative). Zero loss on a real DQN training run is also
    // a tell: either the network is predicting exactly the targets (impossible at
    // step 1) or the backward pass is returning nothing (accumulation bug). The
    // upper bound 1e8 rules out NaN-propagation disguised as "finite".
    assert_finite(metrics.loss, "real_data_loss");
    assert!(
        metrics.loss >= 0.0,
        "real_data loss is negative ({}) — sign bug in MSE / C51 / CQL accumulator",
        metrics.loss,
    );
    assert!(
        metrics.loss < 1e8,
        "real_data loss {} exceeds bound 1e8 — finite but implausibly huge, check \
         NaN propagation or reward scaling",
        metrics.loss,
    );

    // ── at least one epoch completed ──
    assert!(
        metrics.epochs_trained >= 1,
        "real_data single-epoch test produced 0 epochs. metrics: {:?}",
        metrics,
    );

    // ── avg_q_value: finite + bounded ──
    let avg_q = metrics
        .additional_metrics
        .get("avg_q_value")
        .copied()
        .unwrap_or(f64::NAN);
    assert!(
        avg_q.is_finite(),
        "real_data avg_q_value missing/NaN: {:?}",
        metrics.additional_metrics.get("avg_q_value"),
    );
    assert!(
        avg_q.abs() < 1e6,
        "real_data avg_q_value {avg_q} is implausibly large — C51 atoms exhausted \
         or NaN-like explosion",
    );

    info!(
        elapsed_secs = elapsed.as_secs_f64(),
        loss = metrics.loss,
        epochs_trained = metrics.epochs_trained,
        avg_q_value = avg_q,
        "Real-data single epoch"
    );
    drop(trainer);
    drop(rt);
    Ok(())
}