foxhunt/crates/ml/tests/gpu_per_integration_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,
)]
#![cfg(feature = "cuda")]
//! Integration test: GPU-resident PER replay buffer
//!
//! Verifies that GpuReplayBuffer produces correct sampling dynamics:
//! 1. Priority-weighted sampling skews towards high-priority experiences
//! 2. Priority updates change sampling distribution
//! 3. Multi-epoch insert + sample + update cycle works end-to-end
//! 4. KS test: GPU PER distribution differs significantly from uniform

use candle_core::{Device, DType, Tensor};
use ml::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
use ml::dqn::replay_buffer_type::ReplayBufferType;

fn test_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfig {
    GpuReplayBufferConfig {
        capacity,
        state_dim,
        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,
    }
}

/// Insert N random experiences into the buffer.
fn fill_buffer(buf: &mut GpuReplayBuffer, n: usize, state_dim: usize) {
    let device = buf.device().clone();
    let states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], &device).unwrap();
    let next_states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], &device).unwrap();
    let actions = Tensor::zeros(&[n], DType::U32, &device).unwrap();
    let rewards = Tensor::randn(0.0_f32, 1.0, &[n], &device).unwrap();
    let dones = Tensor::zeros(&[n], DType::F32, &device).unwrap();
    buf.insert_batch(&states, &next_states, &actions, &rewards, &dones)
        .unwrap();
}

#[test]
fn test_gpu_per_training_loop_cycle() {
    // End-to-end: insert → sample → fake TD error → update priorities → repeat
    let device = Device::new_cuda(0).expect("CUDA required");
    let state_dim = 16;
    let capacity = 500;
    let batch_size = 32;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap();

    // Fill buffer
    fill_buffer(&mut buf, 200, state_dim);
    assert_eq!(buf.len(), 200);

    // Run 5 "training epochs"
    for epoch in 0..5 {
        let batch = buf.sample_proportional(batch_size).unwrap();

        // Verify batch shapes
        assert_eq!(batch.states.dims(), &[batch_size, state_dim]);
        assert_eq!(batch.actions.dims(), &[batch_size]);
        assert_eq!(batch.rewards.dims(), &[batch_size]);
        assert_eq!(batch.weights.dims(), &[batch_size]);
        assert_eq!(batch.indices.dims(), &[batch_size]);

        // Simulate TD errors: decreasing magnitude over epochs
        let scale = 1.0_f32 / (epoch as f32 + 1.0);
        let td_errors = Tensor::full(scale, &[batch_size], &device).unwrap();

        buf.update_priorities_gpu(&batch.indices, &td_errors).unwrap();
        buf.step();
    }

    // After updates, priorities should be non-uniform
    assert!(buf.len() == 200);
}

#[test]
fn test_gpu_per_priorities_affect_sampling() {
    // Verify that updating priorities actually changes what gets sampled.
    // Set one experience to very high priority, sample many times, check it appears often.
    let device = Device::new_cuda(0).expect("CUDA required");
    let state_dim = 4;
    let capacity = 100;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap();

    // Fill 100 experiences
    fill_buffer(&mut buf, 100, state_dim);

    // Set experience 0 to very high priority, everything else stays at 1.0
    let high_idx = Tensor::new(&[0_u32], &device).unwrap();
    let high_td = Tensor::new(&[100.0_f32], &device).unwrap();
    buf.update_priorities_gpu(&high_idx, &high_td).unwrap();

    // Sample 1000 times in batches of 10, count how often index 0 appears
    // GPU-only: compare each index to 0 via eq/sum, read back scalar count
    let mut count_idx_0_scalar = 0.0_f64;
    let total_samples = 100;
    let batch_size = 10;
    let zero_idx = Tensor::new(&[0_u32], &device).unwrap().broadcast_as(&[batch_size]).unwrap().contiguous().unwrap();
    for _ in 0..total_samples {
        let batch = buf.sample_proportional(batch_size).unwrap();
        // eq returns u8, cast to f32 for sum
        let matches = batch.indices.eq(&zero_idx).unwrap()
            .to_dtype(DType::F32).unwrap()
            .sum_all().unwrap()
            .to_scalar::<f32>().unwrap();
        count_idx_0_scalar += matches as f64;
    }
    let count_idx_0 = count_idx_0_scalar as usize;

    // With uniform sampling (alpha=0), expected = 1000 * 1/100 = 10
    // With priority 100^0.6 ≈ 15.85 vs rest at 1^0.6 = 1,
    // expected proportion = 15.85 / (15.85 + 99) ≈ 13.8%
    // So out of 1000 samples, index 0 should appear ~138 times
    // With alpha=0.6 and high priority, it should appear much more than uniform (10)
    assert!(
        count_idx_0 > 30,
        "Expected index 0 to be sampled frequently (got {}), priority weighting may not work",
        count_idx_0
    );
}

#[test]
fn test_gpu_per_sampling_distribution_differs_from_uniform() {
    // Two-sample KS-style test: compare GPU PER sampling with non-uniform priorities
    // against a uniform distribution. They should differ.
    let device = Device::new_cuda(0).expect("CUDA required");
    let state_dim = 4;
    let capacity = 50;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap();

    // Fill buffer
    fill_buffer(&mut buf, 50, state_dim);

    // Set skewed priorities: experience i gets priority (i+1)
    // This creates a strong left-to-right priority gradient
    for i in 0..50 {
        let idx = Tensor::new(&[i as u32], &device).unwrap();
        let td = Tensor::new(&[(i as f32 + 1.0)], &device).unwrap();
        buf.update_priorities_gpu(&idx, &td).unwrap();
    }

    // Sample 5000 indices via proportional sampling
    // GPU-only validation: verify mean index is above uniform mean (24.5),
    // proving the priority gradient skews sampling toward higher-priority (higher-index) items
    let n_samples = 500;
    let batch_size = 10;
    let mut sum_of_means = 0.0_f64;
    for _ in 0..n_samples {
        let batch = buf.sample_proportional(batch_size).unwrap();
        // Cast u32 indices to f32 for mean computation on GPU
        let mean_idx = batch.indices
            .to_dtype(DType::F32).unwrap()
            .mean_all().unwrap()
            .to_scalar::<f32>().unwrap();
        sum_of_means += mean_idx as f64;
    }
    let overall_mean = sum_of_means / n_samples as f64;

    // With priorities i+1 (1..50) and alpha=0.6, higher indices get sampled more.
    // Uniform mean = 24.5. Skewed mean should be significantly above that.
    assert!(
        overall_mean > 27.0,
        "Mean sampled index ({:.2}) not much above uniform 24.5 — PER should skew toward higher-priority items",
        overall_mean
    );
}

#[test]
fn test_gpu_per_is_weights_correct_range() {
    // IS weights should be in (0, 1] when beta < 1 and normalized by max weight
    let device = Device::new_cuda(0).expect("CUDA required");
    let state_dim = 8;
    let capacity = 200;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap();

    fill_buffer(&mut buf, 200, state_dim);

    // Set varied priorities
    for i in 0..200 {
        let idx = Tensor::new(&[i as u32], &device).unwrap();
        let td = Tensor::new(&[((i % 10) as f32 + 1.0)], &device).unwrap();
        buf.update_priorities_gpu(&idx, &td).unwrap();
    }

    let batch = buf.sample_proportional(64).unwrap();

    // GPU-only: check min > 0, max <= 1.0+eps, and max ≈ 1.0 (normalized)
    // weights is 1D [batch_size], so min/max over dim 0 produces a scalar directly
    let w_min = batch.weights
        .to_dtype(DType::F32).unwrap()
        .min(0).unwrap()
        .to_scalar::<f32>().unwrap();
    let w_max = batch.weights
        .to_dtype(DType::F32).unwrap()
        .max(0).unwrap()
        .to_scalar::<f32>().unwrap();

    assert!(w_min > 0.0, "Minimum weight is non-positive: {}", w_min);
    assert!(
        w_max <= 1.0 + 1e-6,
        "Maximum weight exceeds 1.0: {} (should be normalized)",
        w_max
    );
    // At least the max weight should be ~1.0 (the max weight normalizes to 1)
    assert!(
        (w_max - 1.0).abs() < 1e-4,
        "Max weight {} not close to 1.0 — normalization may be broken",
        w_max
    );
}

#[test]
fn test_gpu_per_ring_buffer_overwrites_correctly() {
    // Fill past capacity, verify size stays at capacity and sampling still works
    let device = Device::new_cuda(0).expect("CUDA required");
    let state_dim = 4;
    let capacity = 50;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &device).unwrap();

    // Insert 120 experiences in 3 batches of 40
    for _ in 0..3 {
        fill_buffer(&mut buf, 40, state_dim);
    }

    assert_eq!(buf.len(), capacity); // capped at 50
    assert!(buf.can_sample(50));

    // Sampling should still work
    let batch = buf.sample_proportional(20).unwrap();
    assert_eq!(batch.states.dims(), &[20, state_dim]);

    // GPU-only: verify all indices are in [0, capacity) via max < capacity
    let idx_max = batch.indices
        .to_dtype(DType::F32).unwrap()
        .max(0).unwrap()
        .to_scalar::<f32>().unwrap();
    let idx_min = batch.indices
        .to_dtype(DType::F32).unwrap()
        .min(0).unwrap()
        .to_scalar::<f32>().unwrap();
    assert!(
        idx_min >= 0.0,
        "Minimum index {} is negative",
        idx_min
    );
    assert!(
        (idx_max as usize) < capacity,
        "Maximum index {} out of range [0, {})",
        idx_max, capacity
    );
}

#[test]
fn test_gpu_per_oom_rejects_absurd_capacity() {
    // Absurd allocation that exceeds both GPU and CPU memory limits.
    // Must return Err — not panic or OOM-kill the system.
    let device = Device::new_cuda(0).expect("CUDA required");

    // 500M capacity × 4 state_dim: GPU pre-flight rejects (~24 GB),
    // CPU fallback pre-flight also rejects (~34 GB estimated).
    let result = ReplayBufferType::try_gpu_prioritized_with_fallback(
        500_000_000, // 500M capacity
        4,           // 4 state_dim
        0.6,
        0.4,
        1.0,
        1000,
        4 * 1024 * 1024 * 1024, // 4 GB limit for test
        &device,
    );

    assert!(
        result.is_err(),
        "500M capacity should be rejected by both GPU and CPU pre-flight checks"
    );
}