foxhunt/crates/ml/tests/dqn_early_stopping_termination_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,
)]
//! WAVE 23 P0: Early Stopping Termination Integration Test
//!
//! Validates that early stopping properly terminates training with error code
//! when gradient collapse or Q-value divergence is detected.
//!
//! This test addresses the bug where early stopping was detected and logged
//! but training continued to completion with exit code 0.
//!
//! Expected behavior:
//! - Gradient collapse detected for 5+ consecutive epochs → training terminates immediately
//! - Per-epoch early stopping detected → training terminates immediately
//! - Exit code is non-zero (error returned)
//! - Checkpoint is saved before termination

#![allow(unused_crate_dependencies)]

use tracing::{info, warn};

fn init_test_tracing() {
    use tracing_subscriber::EnvFilter;
    let _ = tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .with_test_writer()
        .try_init();
}

/// Resolve workspace root and return path to futures-baseline DBN data.
fn get_dbn_data_dir() -> Option<String> {
    // CI: TEST_DATA_DIR points to test-data-pvc on H100
    if let Ok(dir) = std::env::var("TEST_DATA_DIR") {
        let ohlcv = std::path::PathBuf::from(&dir).join("ohlcv");
        if ohlcv.exists() {
            return Some(ohlcv.to_string_lossy().to_string());
        }
        return Some(dir);
    }
    let workspace_root = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .ancestors()
        .find(|p| p.join("test_data").exists())?
        .to_path_buf();
    let data_dir = workspace_root.join("test_data/futures-baseline");
    if !data_dir.exists() {
        warn!(path = %data_dir.display(), "SKIP: DBN data not found");
        return None;
    }
    Some(data_dir.to_string_lossy().to_string())
}

#[tokio::test]
async fn test_early_stopping_terminates_with_error() {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    init_test_tracing();
    let Some(data_dir) = get_dbn_data_dir() else { return };

    // Configure DQN with aggressive early stopping (will trigger quickly)
    let mut hyperparams = DQNHyperparameters::conservative();
    // Apply smoketest profile for consistent hidden_dim across RTX 3050 and H100.
    ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);

    // Force gradient collapse detection by setting threshold above typical grad norm.
    // Threshold = lr × multiplier. With v_range ±240 (C51), gradient norms can reach
    // 100-10000 depending on epoch/data. Set threshold = 1e-8 × 1e12 = 10000.0 to
    // guarantee the norm stays below threshold → collapse triggers reliably.
    hyperparams.learning_rate = 1e-8;
    hyperparams.epochs = 10;
    hyperparams.gradient_collapse_multiplier = 1e12;
    hyperparams.gradient_collapse_patience = 3;
    hyperparams.min_epochs_before_stopping = 1;
    hyperparams.early_stopping_enabled = true;
    hyperparams.checkpoint_frequency = 1;
    hyperparams.batch_size = 32;
    hyperparams.buffer_size = 1024;
    hyperparams.min_replay_size = 32;
    hyperparams.warmup_steps = 0;
    hyperparams.replay_buffer_vram_fraction = 0.0;

    hyperparams.gpu_timesteps_per_episode = 50;
    let mut trainer = DQNTrainer::new(hyperparams.clone())
        .expect("Failed to create DQN trainer");

    // Track checkpoints saved (atomic — no async race)
    let checkpoint_counter = Arc::new(AtomicUsize::new(0));
    let counter_clone = checkpoint_counter.clone();

    let checkpoint_callback = move |epoch: usize, data: Vec<u8>, _is_best: bool| {
        let checkpoint_path = format!("/tmp/early_stop_test_epoch_{}.safetensors", epoch);
        std::fs::write(&checkpoint_path, data).expect("Failed to write checkpoint");
        counter_clone.fetch_add(1, Ordering::SeqCst);
        Ok(checkpoint_path)
    };

    // Run training - should terminate with error
    let result = trainer.train(&data_dir, "ES.FUT", checkpoint_callback).await;

    // Validate early stopping behavior
    assert!(
        result.is_err(),
        "Training should terminate with error when early stopping triggers, got Ok(metrics) instead"
    );

    let error = result.unwrap_err();
    let error_msg = error.to_string();

    // Verify error message indicates gradient collapse
    assert!(
        error_msg.contains("collapse") || error_msg.contains("Collapse")
            || error_msg.contains("early stopping") || error_msg.contains("Early stopping"),
        "Error message should mention collapse or early stopping, got: {}",
        error_msg
    );

    let checkpoint_count = checkpoint_counter.load(Ordering::SeqCst);
    info!(
        error = %error_msg,
        checkpoints_saved = checkpoint_count,
        "Early stopping correctly terminated training with error"
    );
}

#[tokio::test]
async fn test_gradient_collapse_propagates_error() {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

    init_test_tracing();
    let Some(data_dir) = get_dbn_data_dir() else { return };

    // Configure DQN to trigger per-step gradient collapse detection
    let mut hyperparams = DQNHyperparameters::conservative();
    ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);

    // Force gradient collapse detection: threshold = lr × multiplier = 5e-9 × 1e12 = 5000.0
    // With v_range ±240 (C51), grad norms can reach 100-10000. Threshold must be above
    // the maximum expected norm to guarantee collapse detection fires.
    // GPU PER requires buffer_size >= 1024. Warmup = 1024 × 0.2 = 204 steps.
    hyperparams.learning_rate = 5e-9;
    hyperparams.epochs = 10;
    hyperparams.gradient_collapse_multiplier = 1e12;
    hyperparams.gradient_collapse_patience = 3;
    hyperparams.min_epochs_before_stopping = 1;
    hyperparams.early_stopping_enabled = true;
    hyperparams.batch_size = 32;
    hyperparams.buffer_size = 1024;
    hyperparams.min_replay_size = 32;
    hyperparams.warmup_steps = 0;
    hyperparams.replay_buffer_vram_fraction = 0.0;

    hyperparams.gpu_timesteps_per_episode = 50;
    let mut trainer = DQNTrainer::new(hyperparams.clone())
        .expect("Failed to create DQN trainer");

    let checkpoint_callback = |epoch: usize, data: Vec<u8>, _is_best: bool| {
        let checkpoint_path = format!("/tmp/gradient_collapse_test_epoch_{}.safetensors", epoch);
        std::fs::write(&checkpoint_path, data).expect("Failed to write checkpoint");
        Ok(checkpoint_path)
    };

    // Run training
    let result = trainer.train(&data_dir, "ES.FUT", checkpoint_callback).await;

    // Should terminate with error
    assert!(
        result.is_err(),
        "Training should terminate with error on gradient collapse"
    );

    let error = result.unwrap_err();
    let error_msg = error.to_string();

    // Verify error indicates training stopped (gradient collapse or patience-based)
    assert!(
        error_msg.contains("collapse") || error_msg.contains("Collapse")
            || error_msg.contains("early stopping") || error_msg.contains("Early stopping"),
        "Error should mention collapse or early stopping, got: {}",
        error_msg
    );

    info!(error = %error_msg, "Gradient collapse correctly propagated as error");
}

#[tokio::test]
async fn test_healthy_training_completes_successfully() {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

    init_test_tracing();
    let Some(data_dir) = get_dbn_data_dir() else { return };

    // Configure with healthy parameters (should NOT early stop)
    let mut hyperparams = DQNHyperparameters::conservative();
    ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);

    // Normal hyperopt-validated parameters — should NOT early stop
    // Threshold = 1e-5 × 100 = 0.001 — well below typical grad_norm (~1-4), no false collapse.
    hyperparams.learning_rate = 1e-5;
    hyperparams.epochs = 2;
    hyperparams.gradient_collapse_multiplier = 100.0;
    hyperparams.gradient_collapse_patience = 5;
    hyperparams.min_epochs_before_stopping = 1;
    hyperparams.early_stopping_enabled = false; // Explicitly disable — this test validates healthy training
    hyperparams.batch_size = 32;
    hyperparams.gamma = 0.961042;
    hyperparams.buffer_size = 1024;  // Minimum for GPU PER (MIN_GPU_CAPACITY)
    hyperparams.min_replay_size = 32;
    hyperparams.warmup_steps = 0;  // Skip train_step() warmup so training_steps increments immediately
    hyperparams.replay_buffer_vram_fraction = 0.0;  // Disable AutoReplaySizer for test determinism

    hyperparams.gpu_timesteps_per_episode = 50;
    // Disable CQL for stability — CQL penalty on a tiny [64,64] network causes NaN.
    // CQL correctness is tested in dqn_action_collapse_fix_test.
    hyperparams.cql_alpha = 0.0;
    let mut trainer = DQNTrainer::new(hyperparams.clone())
        .expect("Failed to create DQN trainer");

    let checkpoint_callback = |epoch: usize, data: Vec<u8>, _is_best: bool| {
        let checkpoint_path = format!("/tmp/healthy_training_test_epoch_{}.safetensors", epoch);
        std::fs::write(&checkpoint_path, data).expect("Failed to write checkpoint");
        Ok(checkpoint_path)
    };

    // Run training
    let result = trainer.train(&data_dir, "ES.FUT", checkpoint_callback).await;

    // Should complete successfully (no early stopping)
    assert!(
        result.is_ok(),
        "Healthy training should complete successfully, got error: {:?}",
        result.err()
    );

    let metrics = result.unwrap();

    // Verify all epochs completed
    assert_eq!(
        metrics.epochs_trained, 2,
        "Should complete all 2 epochs, got: {}",
        metrics.epochs_trained
    );

    // Verify early_stopped flag is not set
    let early_stopped = metrics.additional_metrics.get("early_stopped").copied().unwrap_or(0.0);
    assert_eq!(
        early_stopped, 0.0,
        "early_stopped flag should not be set for healthy training"
    );

    info!(
        epochs_trained = metrics.epochs_trained,
        final_loss = metrics.loss,
        "Healthy training completed successfully"
    );
}

#[test]
fn test_early_stopping_config_fields_exist() {
    use ml::trainers::dqn::DQNHyperparameters;

    // Verify WAVE 23 P0 fields exist and have correct defaults
    let hyperparams = DQNHyperparameters::conservative();

    assert_eq!(
        hyperparams.gradient_collapse_multiplier, 100.0,
        "gradient_collapse_multiplier should default to 100.0"
    );

    assert_eq!(
        hyperparams.gradient_collapse_patience, 5,
        "gradient_collapse_patience should default to 5"
    );

    // Verify adaptive threshold calculation.
    // gradient_collapse_multiplier is f32 (kernel-facing); learning_rate stays f64 for precision.
    let threshold = hyperparams.learning_rate * f64::from(hyperparams.gradient_collapse_multiplier);
    assert!(
        threshold > 0.0,
        "Adaptive threshold should be positive: {}",
        threshold
    );

    info!(
        multiplier = hyperparams.gradient_collapse_multiplier,
        patience = hyperparams.gradient_collapse_patience,
        learning_rate = hyperparams.learning_rate,
        adaptive_threshold = threshold,
        "Early stopping configuration fields validated"
    );
}