diff --git a/crates/ml-core/src/device.rs b/crates/ml-core/src/device.rs index 66ae20cb6..eebb84c36 100644 --- a/crates/ml-core/src/device.rs +++ b/crates/ml-core/src/device.rs @@ -47,6 +47,10 @@ impl MlDevice { let context = CudaContext::new(ordinal).map_err(|e| { MLError::DeviceError(format!("Failed to open CUDA device {ordinal}: {e}")) })?; + // Drain any stale CUDA errors from a previous context user in the same process. + // CudaContext::new reuses the primary context (cuDevicePrimaryCtxRetain), + // so deferred errors from a previous test's Drop persist in error_state. + let _ = context.check_err(); let stream = context.new_stream().map_err(|e| { MLError::DeviceError(format!("Failed to create CUDA stream on device {ordinal}: {e}")) })?; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index b919fa4e2..e42d840a4 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -506,6 +506,11 @@ impl GpuDqnTrainer { stream: Arc, config: GpuDqnTrainConfig, ) -> Result { + // Drain any stale CUDA errors from a previous trainer instance + // (e.g., deferred errors from CudaSlice/CudaEvent Drop in the same process). + // Without this, bind_to_thread() fails with the stale error on the next API call. + let _ = stream.context().check_err(); + let b = config.batch_size; let total_params = compute_total_params(&config); diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index c7e9af590..edb1e3784 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -424,6 +424,9 @@ impl GpuExperienceCollector { n_episodes: usize, timesteps_per_episode: usize, ) -> Result { + // Drain stale CUDA errors from previous test/trainer instances. + let _ = stream.context().check_err(); + let alloc_episodes = n_episodes.min(MAX_EPISODES_LIMIT); let alloc_timesteps = timesteps_per_episode.min(MAX_TIMESTEPS_LIMIT); let (shared_h1, shared_h2, value_h, adv_h) = network_dims; diff --git a/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs b/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs index 8c44c8262..bd8b89d4c 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs @@ -5,22 +5,25 @@ //! circuit breaker + n-step + noisy nets + PER + dueling. //! Real ES.FUT data, GPU-only. -use super::helpers::{assert_finite, smoke_trainer, test_data_dir}; +use super::helpers::{assert_finite, log_gpu_memory, smoke_trainer, test_data_dir}; /// Full production DQN config on real ES.FUT data — 3 epochs. /// /// All features are always on in production. There is no code path /// without curiosity, IQN, branching, etc. This single test validates /// the entire fused CUDA training pipeline end-to-end. -#[tokio::test] +#[test] #[ignore] // Loads real training data — run via nightly CI or manual trigger -async fn test_production_config_trains() -> anyhow::Result<()> { +fn test_production_config_trains() -> anyhow::Result<()> { let data_dir = test_data_dir() .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback"); let mut trainer = smoke_trainer()?; - let metrics = trainer.train(&data_dir, |_epoch, _bytes, _best| { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _best| { Ok("skip".to_owned()) - }).await?; + }))?; assert_finite(metrics.loss, "production_loss"); assert!(metrics.loss >= 0.0, "loss should be non-negative"); assert!( @@ -28,5 +31,8 @@ async fn test_production_config_trains() -> anyhow::Result<()> { "expected at least 1 epoch, got {}", metrics.epochs_trained, ); + drop(trainer); + drop(rt); + log_gpu_memory("after test_production_config_trains"); Ok(()) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index 675db29dc..855eb58b4 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -84,8 +84,8 @@ fn dtoh_u32(slice: &cudarc::driver::CudaSlice, stream: &Arc anyhow::Result<()> { +#[test] +fn test_gpu_replay_buffer_insert_and_len() -> anyhow::Result<()> { let stream = smoke_stream(); let config = test_buffer_config(100, 48); let mut buf = GpuReplayBuffer::new(config, &stream)?; @@ -93,11 +93,12 @@ async fn test_gpu_replay_buffer_insert_and_len() -> anyhow::Result<()> { insert_random_batch(&mut buf, 20, 48)?; assert_eq!(buf.len(), 20); + log_gpu_memory("after test_gpu_replay_buffer_insert_and_len"); Ok(()) } -#[tokio::test] -async fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> { +#[test] +fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> { let stream = smoke_stream(); let config = test_buffer_config(100, 48); let mut buf = GpuReplayBuffer::new(config, &stream)?; @@ -129,11 +130,12 @@ async fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<() // Shape sanity assert_eq!(indices_host.len(), 16); + log_gpu_memory("after test_gpu_replay_buffer_proportional_sample_valid"); Ok(()) } -#[tokio::test] -async fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> { +#[test] +fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> { let stream = smoke_stream(); let config = test_buffer_config(100, 48); let mut buf = GpuReplayBuffer::new(config, &stream)?; @@ -165,11 +167,12 @@ async fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> // Shape sanity assert_eq!(indices_host.len(), 16); + log_gpu_memory("after test_gpu_replay_buffer_rank_based_sample_valid"); Ok(()) } -#[tokio::test] -async fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> { +#[test] +fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> { let stream = smoke_stream(); let config = test_buffer_config(100, 48); let epsilon = config.epsilon; @@ -196,11 +199,12 @@ async fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> { // Epsilon is the priority floor -- verify it's reasonable assert!(epsilon > 0.0, "epsilon should be positive"); + log_gpu_memory("after test_gpu_replay_buffer_priority_update_valid"); Ok(()) } -#[tokio::test] -async fn test_searchsorted_gpu_correctness() -> anyhow::Result<()> { +#[test] +fn test_searchsorted_gpu_correctness() -> anyhow::Result<()> { let stream = smoke_stream(); let config = test_buffer_config(50, 4); let mut buf = GpuReplayBuffer::new(config, &stream)?; @@ -222,36 +226,40 @@ async fn test_searchsorted_gpu_correctness() -> anyhow::Result<()> { assert_eq!(indices_host.len(), 10); } + log_gpu_memory("after test_searchsorted_gpu_correctness"); Ok(()) } -#[tokio::test] +#[test] #[ignore] // Loads real training data — run via nightly CI or manual trigger -async fn test_train_step_produces_finite_metrics() -> anyhow::Result<()> { +fn test_train_step_produces_finite_metrics() -> anyhow::Result<()> { let data_dir = test_data_dir() .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback"); let mut trainer = smoke_trainer()?; - let metrics = trainer.train(&data_dir, |_epoch, _bytes, _is_best| { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _is_best| { Ok(String::new()) - }).await?; - - // Final loss must be finite (not NaN, not Inf) + }))?; assert_finite(metrics.loss, "final_loss"); - + drop(trainer); + drop(rt); + log_gpu_memory("after test_train_step_produces_finite_metrics"); Ok(()) } /// Verify training dtype is BF16 on CUDA. -#[tokio::test] -async fn test_gpu_training_dtype_bf16() -> anyhow::Result<()> { +#[test] +fn test_gpu_training_dtype_bf16() -> anyhow::Result<()> { let dtype = ml_core::native_types::NativeDType::BF16; assert_eq!(dtype, ml_core::native_types::NativeDType::BF16, "CUDA should use BF16 training dtype"); Ok(()) } /// Diagnose: GPU training dtype is BF16. -#[tokio::test] -async fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> { +#[test] +fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> { let dev = cuda_device(); let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); info!(device = ?dev, "GPU training dtype"); @@ -263,6 +271,7 @@ async fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> { let cublas = cudarc::cublas::CudaBlas::new(Arc::clone(&stream)).map_err(|e| anyhow::anyhow!("cuBLAS: {e}"))?; let (out, _activations) = layer.forward(&input, &vs, &cublas, &stream)?; info!(dims = ?out.dims(), "Linear forward OK"); + log_gpu_memory("after test_gpu_training_dtype_diagnosis"); Ok(()) } @@ -271,8 +280,8 @@ async fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> { // Testing "what if we disable a mandatory component" is testing dead code. /// Diagnose: can we create + step an AdamW optimizer on GPU? -#[tokio::test] -async fn test_gpu_adamw_creation() -> anyhow::Result<()> { +#[test] +fn test_gpu_adamw_creation() -> anyhow::Result<()> { let dev = cuda_device(); let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); info!(device = ?dev, "GPU device"); @@ -308,5 +317,6 @@ async fn test_gpu_adamw_creation() -> anyhow::Result<()> { } let _grad_norm = opt.step(&mut vs, &grads)?; info!("AdamW step OK"); + log_gpu_memory("after test_gpu_adamw_creation"); Ok(()) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs index c4f538e7d..9092cf468 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs @@ -57,25 +57,15 @@ pub(super) fn smoke_params() -> DQNHyperparameters { p } -/// Shared CUDA device for all smoke tests — reuses one cuBLAS handle -/// instead of creating a new one per test, preventing handle pool exhaustion. -static SMOKE_CUDA: std::sync::OnceLock = std::sync::OnceLock::new(); - -/// Get CUDA device. Panics if no GPU is available. -/// Tests a small matmul to detect broken CUDA runtime before trainer creation. +/// Create a fresh CUDA device for each test. +/// +/// NOT cached in a static — each test gets its own CudaContext + CudaStream. +/// A static OnceLock caused sequential test failures: CudaSlice Drop +/// records errors on the shared context's error_state, poisoning it for the next test. pub(super) fn cuda_device() -> MlDevice { init_tracing(); - SMOKE_CUDA - .get_or_init(|| { - let dev = - MlDevice::new_cuda(0).expect("CUDA device required — no CPU fallback in smoke tests"); - // Probe: allocate a small tensor to verify GPU runtime works - let stream = dev.cuda_stream().expect("cuda stream"); - let probe = ml_core::cuda_autograd::GpuTensor::zeros(&[2, 2], stream); - probe.expect("CUDA alloc probe failed — GPU runtime is broken"); - dev - }) - .clone() + MlDevice::new_cuda(0) + .expect("CUDA device required — no CPU fallback in smoke tests") } /// Create trainer on CUDA device (no CPU fallback). @@ -88,6 +78,27 @@ pub(super) fn smoke_trainer_with(params: DQNHyperparameters) -> anyhow::Result String { } /// Measure end-to-end training throughput (3 epochs, real data). -#[tokio::test] +#[test] #[ignore] // Loads real training data — run via nightly CI or manual trigger -async fn test_training_throughput_measurement() -> anyhow::Result<()> { +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 = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + let metrics = rt.block_on(trainer.train(&data_dir(), |_epoch, _bytes, _best| { Ok("skip".to_owned()) - }).await?; + }))?; let elapsed = start.elapsed(); assert_finite(metrics.loss, "loss"); @@ -45,12 +48,15 @@ async fn test_training_throughput_measurement() -> anyhow::Result<()> { if let Some(&final_eps) = metrics.additional_metrics.get("final_epsilon") { info!(final_epsilon = final_eps, "Final epsilon"); } + drop(trainer); + drop(rt); + log_gpu_memory("after test_training_throughput_measurement"); Ok(()) } /// Measure per-sample latency of the GPU PER replay buffer. -#[tokio::test] -async fn test_per_sample_latency() -> anyhow::Result<()> { +#[test] +fn test_per_sample_latency() -> anyhow::Result<()> { use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig}; let dev = MlDevice::new_cuda(0)?; @@ -113,22 +119,26 @@ async fn test_per_sample_latency() -> anyhow::Result<()> { "sample latency {us_per_sample:.1} us is unreasonably high" ); + log_gpu_memory("after test_per_sample_latency"); Ok(()) } /// Single-epoch training on real data. /// Loads full 163K-bar dataset — too slow for CI (run via nightly or manual trigger). -#[tokio::test] +#[test] #[ignore] -async fn test_real_data_single_epoch() -> anyhow::Result<()> { +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 = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + let metrics = rt.block_on(trainer.train(&data_dir(), |_epoch, _bytes, _best| { Ok("skip".to_owned()) - }).await?; + }))?; let elapsed = start.elapsed(); assert_finite(metrics.loss, "real_data_loss"); @@ -140,6 +150,8 @@ async fn test_real_data_single_epoch() -> anyhow::Result<()> { if let Some(&avg_q) = metrics.additional_metrics.get("avg_q_value") { info!(avg_q_value = avg_q, "Avg Q-value"); } - + drop(trainer); + drop(rt); + log_gpu_memory("after test_real_data_single_epoch"); Ok(()) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index 9f139ae03..232873408 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -3,18 +3,21 @@ //! Verifies that production-config DQN training produces sane, finite metrics //! on real ES.FUT data. GPU-only, all features on. -use super::helpers::{assert_finite, cuda_device, smoke_trainer, smoke_params, smoke_trainer_with, test_data_dir}; +use super::helpers::{assert_finite, cuda_device, log_gpu_memory, smoke_trainer, smoke_params, smoke_trainer_with, test_data_dir}; /// Full production training: loss finite, gradients flowing, Q-values bounded, epsilon decays. -#[tokio::test] +#[test] #[ignore] // Loads real training data — run via nightly CI or manual trigger -async fn test_production_training_stability() -> anyhow::Result<()> { +fn test_production_training_stability() -> anyhow::Result<()> { let data_dir = test_data_dir() .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback"); let mut trainer = smoke_trainer()?; - let metrics = trainer.train(&data_dir, |_epoch, _bytes, _best| { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _best| { Ok("skip".to_owned()) - }).await?; + }))?; // Loss must be finite and non-negative assert_finite(metrics.loss, "loss"); @@ -37,12 +40,15 @@ async fn test_production_training_stability() -> anyhow::Result<()> { // Must have completed all epochs assert!(metrics.epochs_trained > 0, "epochs_trained should be > 0"); + drop(trainer); + drop(rt); + log_gpu_memory("after test_production_training_stability"); Ok(()) } /// PER importance-sampling weights must be finite and positive. -#[tokio::test] -async fn test_per_weights_valid() -> anyhow::Result<()> { +#[test] +fn test_per_weights_valid() -> anyhow::Result<()> { use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig}; use std::sync::Arc; @@ -82,12 +88,13 @@ async fn test_per_weights_valid() -> anyhow::Result<()> { } let w_sum: f32 = weights_host.iter().sum(); assert!(w_sum.is_finite(), "weight sum not finite: {w_sum}"); + log_gpu_memory("after test_per_weights_valid"); Ok(()) } /// PER indices must be within buffer bounds. -#[tokio::test] -async fn test_per_indices_valid() -> anyhow::Result<()> { +#[test] +fn test_per_indices_valid() -> anyhow::Result<()> { use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig}; use std::sync::Arc; @@ -125,23 +132,30 @@ async fn test_per_indices_valid() -> anyhow::Result<()> { assert!(idx < 50, "index[{i}] out of bounds: {idx}"); } assert_eq!(indices_host.len(), 16); + log_gpu_memory("after test_per_indices_valid"); Ok(()) } /// CUDA training must reject missing GPU experience collector. -#[tokio::test] +#[test] #[ignore] // Loads real training data — run via nightly CI or manual trigger -async fn test_rejects_missing_gpu_collector() -> anyhow::Result<()> { +fn test_rejects_missing_gpu_collector() -> anyhow::Result<()> { let data_dir = test_data_dir() .expect("FOXHUNT_TEST_DATA or test_data/ must exist"); let mut params = smoke_params(); params.enable_gpu_experience_collector = false; let mut trainer = smoke_trainer_with(params)?; - let result = trainer.train(&data_dir, |_epoch, _bytes, _is_best| { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let result = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _is_best| { Ok(String::new()) - }).await; + })); assert!(result.is_err(), "Must reject missing GPU experience collector"); let err = result.unwrap_err().to_string(); assert!(err.contains("GPU experience collector MUST be active"), "Got: {err}"); + drop(trainer); + drop(rt); + log_gpu_memory("after test_rejects_missing_gpu_collector"); Ok(()) } diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 03b7b8546..dce63e1d0 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -120,8 +120,6 @@ impl DQNTrainer { let device = if let Some(dev) = override_device { dev } else { - // Synchronize device before creating new trainer. - unsafe { cudarc::driver::sys::cuCtxSynchronize(); } match MlDevice::cuda(0) { Ok(dev) => { dev @@ -142,6 +140,10 @@ impl DQNTrainer { #[cfg(feature = "cuda")] if let MlDevice::Cuda { ref context, .. } = device { use cudarc::driver::sys::{cuCtxSetLimit, CUlimit, CUresult}; + // Drain any stale error_state set by CudaSlice/CudaEvent Drop + // from a previous trainer in the same process. check_err() atomically + // swaps error_state to 0, preventing bind_to_thread() from failing. + let _ = context.check_err(); context.bind_to_thread().map_err(|e| anyhow::anyhow!("bind: {e}"))?; let gpu_profile = ml_core::gpu::profile::GpuProfile::load(); let stack_bytes: usize = gpu_profile.cuda.cuda_stack_bytes;