a29f109b67
Update all callers to match the new pure-cudarc APIs introduced by the hive agent CudaSlice migration. Key changes: - GpuTrainingGuard::new() now takes Arc<CudaStream>; callers use from_device() - check_and_accumulate/qvalue_stats/qvalue_divergence take &CudaSlice<f32> instead of &Tensor; callers convert via tensor_to_cuda_slice_f32() - accumulate_q_value takes f32 scalar, returns () (no Result) - GpuReplayBuffer::insert_batch gains batch_size arg, takes CudaSlice params - signal_adapter functions take &Arc<CudaStream> (cudarc 0.17 Arc requirement) - Add tensor_to_cuda_slice_u32() and cuda_f32_to_tensor() utility functions - Replace CudaView usage with owned CudaSlice via tensor_to_cuda_slice_f32() - Fix CudaStorage.device field access (was method call in older API) - Fix borrow-after-move in copy_actions_out via scoped DtoD copy Zero errors, zero warnings across lib + tests + examples + full workspace. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
354 lines
13 KiB
Rust
354 lines
13 KiB
Rust
#![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
|
||
|
||
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();
|
||
let sf = ml::cuda_pipeline::tensor_to_cuda_slice_f32(&states).unwrap();
|
||
let nsf = ml::cuda_pipeline::tensor_to_cuda_slice_f32(&next_states).unwrap();
|
||
let af = ml::cuda_pipeline::tensor_to_cuda_slice_u32(&actions).unwrap();
|
||
let rf = ml::cuda_pipeline::tensor_to_cuda_slice_f32(&rewards).unwrap();
|
||
let df = ml::cuda_pipeline::tensor_to_cuda_slice_f32(&dones).unwrap();
|
||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, n)
|
||
.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 GPU 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).
|
||
// No CPU fallback — GPU PER is mandatory.
|
||
let result = ReplayBufferType::try_gpu_with_halving(
|
||
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 GPU pre-flight checks"
|
||
);
|
||
}
|