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,
)]
//! 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

// candle eliminated — test uses native cudarc APIs
use ml::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
use std::sync::Arc;
use cudarc::driver::{CudaContext, CudaSlice, CudaStream};

fn require_cuda_stream() -> Arc<CudaStream> {
    let ctx = CudaContext::new(0).expect("CUDA required");
    ctx.new_stream().expect("CUDA stream required")
}

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

/// Insert N random experiences into the buffer via CudaSlice uploads.
fn fill_buffer(buf: &mut GpuReplayBuffer, n: usize, state_dim: usize) {
    use rand::Rng;
    let stream = buf.stream();
    let mut rng = rand::thread_rng();

    // Generate random f32 data on host, upload to GPU
    let states_host: Vec<f32> = (0..n * state_dim).map(|_| rng.gen::<f32>()).collect();
    let next_states_host: Vec<f32> = (0..n * state_dim).map(|_| rng.gen::<f32>()).collect();
    let actions_host: Vec<u32> = vec![0_u32; n];
    let rewards_host: Vec<f32> = (0..n).map(|_| rng.gen::<f32>()).collect();
    let dones_host: Vec<f32> = vec![0.0_f32; n];

    let sf = stream.clone_htod(&states_host).unwrap();
    let nsf = stream.clone_htod(&next_states_host).unwrap();
    let af: CudaSlice<u32> = stream.clone_htod(&actions_host).unwrap();
    let rf = stream.clone_htod(&rewards_host).unwrap();
    let df = stream.clone_htod(&dones_host).unwrap();
    let aux_sign: CudaSlice<i32> = stream.alloc_zeros::<i32>(n).unwrap();
    let aux_conf: CudaSlice<f32> = stream.alloc_zeros::<f32>(n).unwrap();

    buf.insert_batch(&sf, &nsf, &af, &rf, &df, &aux_sign, &aux_conf, n).unwrap();
}

#[test]
fn test_gpu_per_training_loop_cycle() {
    // End-to-end: insert -> sample -> fake TD error -> update priorities -> repeat
    let stream = require_cuda_stream();
    let state_dim = 16;
    let capacity = 500;
    let batch_size = 32;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).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 via CudaSlice lengths
        assert_eq!(batch.batch_size * batch.state_dim, batch_size * state_dim); // f32 states [bs*sd]
        assert_eq!(batch.batch_size, batch_size);
        assert_eq!(batch.batch_size, batch_size);
        assert_eq!(batch.batch_size, batch_size);
        assert_eq!(batch.batch_size, batch_size);

        // Simulate TD errors: decreasing magnitude over epochs
        let scale = 1.0_f32 / (epoch as f32 + 1.0);
        let td_host: Vec<f32> = vec![scale; batch_size];
        let td_errors: CudaSlice<f32> = stream.clone_htod(&td_host).unwrap();

        let idx_ptr = batch.indices_ptr;
        buf.update_priorities_gpu_raw(idx_ptr, &td_errors, batch_size, None).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 stream = require_cuda_stream();
    let state_dim = 4;
    let capacity = 100;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).unwrap();

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

    // Set experience 0 to very high priority
    let high_idx: CudaSlice<u32> = stream.clone_htod(&[0_u32]).unwrap();
    let high_td: CudaSlice<f32> = stream.clone_htod(&[100.0]).unwrap();
    buf.update_priorities_gpu(&high_idx, &high_td, 1, None).unwrap();

    // Sample 1000 times in batches of 10, count how often index 0 appears
    // Download indices to host, count matches
    let mut count_idx_0 = 0_usize;
    let total_samples = 100;
    let batch_size = 10;
    for _ in 0..total_samples {
        let batch = buf.sample_proportional(batch_size).unwrap();
        // Download u32 indices to host
        let mut indices_host = vec![0_u32; batch_size];
        stream.memcpy_dtoh(&buf.sample_indices_ref(), &mut indices_host).unwrap();
        count_idx_0 += indices_host.iter().filter(|&&idx| idx == 0).count();
    }

    // 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 stream = require_cuda_stream();
    let state_dim = 4;
    let capacity = 50;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).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: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
        let td: CudaSlice<f32> = stream.clone_htod(&[i as f32 + 1.0]).unwrap();
        buf.update_priorities_gpu(&idx, &td, 1, None).unwrap();
    }

    // Sample 5000 indices via proportional sampling
    // Validate 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();
        // Download u32 indices to host, compute mean
        let mut indices_host = vec![0_u32; batch_size];
        stream.memcpy_dtoh(&buf.sample_indices_ref(), &mut indices_host).unwrap();
        let mean_idx: f64 = indices_host.iter().map(|&i| i as f64).sum::<f64>() / batch_size as f64;
        sum_of_means += mean_idx;
    }
    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 stream = require_cuda_stream();
    let state_dim = 8;
    let capacity = 200;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).unwrap();

    fill_buffer(&mut buf, 200, state_dim);

    // Set varied priorities
    for i in 0..200 {
        let idx: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
        let td: CudaSlice<f32> = stream.clone_htod(&[((i % 10) as f32 + 1.0)]).unwrap();
        buf.update_priorities_gpu(&idx, &td, 1, None).unwrap();
    }

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

    // Download f32 weights to host for validation
    let mut weights_host = vec![0.0_f32; 64];
    stream.memcpy_dtoh(&buf.sample_weights_ref(), &mut weights_host).unwrap();

    let w_min = weights_host.iter().copied().fold(f32::INFINITY, f32::min);
    let w_max = weights_host.iter().copied().fold(f32::NEG_INFINITY, f32::max);

    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 stream = require_cuda_stream();
    let state_dim = 4;
    let capacity = 50;
    let mut buf = GpuReplayBuffer::new(test_config(capacity, state_dim), &stream).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.batch_size * batch.state_dim, 20 * state_dim); // f32 states [bs*sd]

    // Download u32 indices to host for range validation
    let mut indices_host = vec![0_u32; 20];
    stream.memcpy_dtoh(&buf.sample_indices_ref(), &mut indices_host).unwrap();

    let idx_min = *indices_host.iter().min().unwrap();
    let idx_max = *indices_host.iter().max().unwrap();
    assert!(
        idx_min < capacity as u32,
        "Minimum index {} is out of range",
        idx_min
    );
    assert!(
        idx_max < capacity as u32,
        "Maximum index {} out of range [0, {})",
        idx_max, capacity
    );
}