From e2f6fc17f5e7198bb0812623cb2b31f4569c6530 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 20 Feb 2026 21:29:43 +0100 Subject: [PATCH] test(ml): add checkpoint-to-inference integration test Validates the complete DQN checkpoint lifecycle: train 5 epochs with DQNHyperparameters::conservative(), save via checkpoint callback, load into a fresh DQN with architecture auto-detected from checkpoint tensor metadata (noisy vs standard layers, state_dim), and run 100 inference passes asserting valid action indices, finite Q-values, and non-zero Q-values. Co-Authored-By: Claude Opus 4.6 --- ml/tests/dqn_inference_test.rs | 274 +++++++++++++++++++++++++++++++++ 1 file changed, 274 insertions(+) create mode 100644 ml/tests/dqn_inference_test.rs diff --git a/ml/tests/dqn_inference_test.rs b/ml/tests/dqn_inference_test.rs new file mode 100644 index 000000000..cea00d3ed --- /dev/null +++ b/ml/tests/dqn_inference_test.rs @@ -0,0 +1,274 @@ +//! DQN Checkpoint -> Inference Integration Test +//! +//! Proves the complete checkpoint -> load -> inference path: +//! 1. Train 5 epochs on real market data (fast, just to get a valid checkpoint) +//! 2. Save checkpoint to temp dir via the trainer's checkpoint callback +//! 3. Load checkpoint into a fresh DQN with matching architecture +//! 4. Run inference on 100 synthetic state vectors +//! 5. Assert: actions valid (0..45), Q-values finite, Q-values non-zero + +#![allow(unused_crate_dependencies)] + +use anyhow::{Context, Result}; +use candle_core::Tensor; +use ml::dqn::{DQNConfig, DQN}; +use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer}; +use std::path::PathBuf; + +/// Locate the small training data directory, returning an error if absent. +fn get_data_dir() -> Result { + let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .context("Failed to get workspace root")? + .to_path_buf(); + let data_dir = workspace_root.join("test_data/real/databento/ml_training_small"); + if !data_dir.exists() { + anyhow::bail!("Data not found: {}", data_dir.display()); + } + Ok(data_dir.to_string_lossy().to_string()) +} + +/// Build a `DQNConfig` that matches the trainer's internal config exactly. +/// +/// The trainer builds its DQNConfig from DQNHyperparameters::conservative() with +/// hardcoded architecture params: state_dim=54, num_actions=45, hidden_dims=[256,128,64]. +/// When `use_noisy_nets=true` (conservative default), the Sequential network uses +/// NoisyLinear layers with key prefix "noisy_hidden_*" / "noisy_output". +/// We detect the naming scheme from the checkpoint to guarantee a match. +fn build_matching_config(checkpoint_tensors: &std::collections::HashMap) -> DQNConfig { + // Detect whether the checkpoint was produced with noisy nets by inspecting + // the VarMap key prefixes. + let uses_noisy = checkpoint_tensors.keys().any(|k| k.starts_with("noisy_")); + + // Discover state_dim from the first layer weight. + // NoisyLinear: "noisy_hidden_0.mu_w" shape [hidden, input] + // Standard: "hidden_0.weight" shape [hidden, input] + let state_dim = if uses_noisy { + checkpoint_tensors + .iter() + .find(|(name, _)| name.contains("noisy_hidden_0") && name.contains("mu_w")) + .map(|(_, t)| { + let d = t.dims(); + if d.len() == 2 { d[1] } else { 54 } + }) + .unwrap_or(54) + } else { + checkpoint_tensors + .iter() + .find(|(name, _)| name.contains("hidden_0") && name.contains("weight")) + .map(|(_, t)| { + let d = t.dims(); + if d.len() == 2 { d[1] } else { 54 } + }) + .unwrap_or(54) + }; + + // Replicate the exact config the trainer constructs (see trainer.rs ~line 289). + let mut config = DQNConfig::conservative(); + config.state_dim = state_dim; + config.num_actions = 45; + config.hidden_dims = vec![256, 128, 64]; + config.use_noisy_nets = uses_noisy; + // Trainer hardcodes noisy_sigma_init = 0.5 via hyperparams + config.noisy_sigma_init = 0.5; + // Trainer hardcodes use_iqn = true, use_cql = true + config.use_iqn = true; + config.use_cql = true; + config +} + +#[tokio::test] +async fn test_checkpoint_to_inference() -> Result<()> { + // -- Skip gracefully when real data is absent (CI environments) -------- + let data_dir = match get_data_dir() { + Ok(dir) => dir, + Err(e) => { + eprintln!("Skipping test_checkpoint_to_inference: {e}"); + return Ok(()); + } + }; + + let checkpoint_dir = tempfile::tempdir()?; + + // ===================================================================== + // Phase 1: Train for 5 epochs to produce a valid checkpoint + // ===================================================================== + let mut hyperparams = DQNHyperparameters::conservative(); + hyperparams.epochs = 5; + hyperparams.batch_size = 32; + hyperparams.learning_rate = 0.0001; + hyperparams.early_stopping_enabled = false; + hyperparams.checkpoint_frequency = 5; // checkpoint on last epoch + + let mut trainer = DQNTrainer::new(hyperparams)?; + let mut best_checkpoint_path: Option = None; + + let _metrics = trainer + .train(&data_dir, |epoch, checkpoint_data, is_best| { + let name = if is_best { + "inference_best.safetensors".to_string() + } else { + format!("inference_epoch_{epoch}.safetensors") + }; + let path = checkpoint_dir.path().join(&name); + std::fs::write(&path, &checkpoint_data)?; + if is_best { + best_checkpoint_path = Some(path.clone()); + } + Ok(path.to_string_lossy().to_string()) + }) + .await?; + + // If no "best" was saved, fall back to any checkpoint in the directory. + let checkpoint_path = match best_checkpoint_path { + Some(p) => p, + None => { + // Find first .safetensors file in the temp dir + let mut entries: Vec<_> = std::fs::read_dir(checkpoint_dir.path())? + .filter_map(|e| e.ok()) + .filter(|e| { + e.path() + .extension() + .map_or(false, |ext| ext == "safetensors") + }) + .collect(); + anyhow::ensure!( + !entries.is_empty(), + "No checkpoint files were saved during training" + ); + entries.sort_by_key(|e| e.path()); + entries.pop().context("No checkpoint file")?.path() + } + }; + + assert!( + checkpoint_path.exists(), + "Checkpoint file does not exist: {}", + checkpoint_path.display() + ); + println!( + "Phase 1 complete: checkpoint at {} ({} bytes)", + checkpoint_path.display(), + std::fs::metadata(&checkpoint_path)?.len() + ); + + // ===================================================================== + // Phase 2: Load checkpoint into a fresh DQN + // ===================================================================== + // Load the checkpoint's tensor metadata to discover architecture params + // (state_dim, noisy vs standard layers) so we can construct a DQN whose + // VarMap key names and tensor shapes match exactly. + let checkpoint_path_str = checkpoint_path + .to_str() + .context("Non-UTF8 checkpoint path")?; + + let checkpoint_tensors = + candle_core::safetensors::load(checkpoint_path_str, &candle_core::Device::Cpu)?; + + println!( + "Checkpoint contains {} tensors: {:?}", + checkpoint_tensors.len(), + checkpoint_tensors.keys().collect::>() + ); + + let config = build_matching_config(&checkpoint_tensors); + let state_dim = config.state_dim; + println!( + "Built matching config: state_dim={}, num_actions={}, noisy={}", + config.state_dim, config.num_actions, config.use_noisy_nets + ); + + let mut fresh_dqn = DQN::new(config)?; + fresh_dqn.load_from_safetensors(checkpoint_path_str)?; + + println!("Phase 2 complete: fresh DQN loaded from checkpoint"); + + // ===================================================================== + // Phase 3: Inference on 100 synthetic state vectors + // ===================================================================== + let num_inference_samples: usize = 100; + let num_actions: usize = 45; + let device = fresh_dqn.device().clone(); + + let mut all_actions_valid = true; + let mut all_q_finite = true; + let mut any_q_nonzero = false; + let mut action_counts = vec![0usize; num_actions]; + + for i in 0..num_inference_samples { + // Deterministic synthetic state: slight variation per sample + let state_vec: Vec = (0..state_dim) + .map(|j| ((i * state_dim + j) as f32 * 0.01).sin()) + .collect(); + + let state_tensor = Tensor::from_vec(state_vec, (1, state_dim), &device)?; + + let q_values = fresh_dqn.forward(&state_tensor)?; + + // Shape check: [1, 45] + let dims = q_values.dims(); + assert_eq!( + dims, + &[1, num_actions], + "Q-value tensor shape mismatch: expected [1, {}], got {:?}", + num_actions, + dims + ); + + let q_vec: Vec = q_values.to_vec2::()?.into_iter().flatten().collect(); + + // Check finite + for &q in &q_vec { + if !q.is_finite() { + all_q_finite = false; + } + if q.abs() > 1e-9 { + any_q_nonzero = true; + } + } + + // Argmax to get action index + let best_action_idx = q_vec + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(idx, _)| idx) + .unwrap_or(0); + + if best_action_idx >= num_actions { + all_actions_valid = false; + } else if let Some(count) = action_counts.get_mut(best_action_idx) { + *count += 1; + } + } + + // ===================================================================== + // Phase 4: Assertions + // ===================================================================== + assert!(all_q_finite, "Some Q-values were NaN or Inf"); + assert!( + any_q_nonzero, + "All Q-values were zero -- model likely failed to load weights" + ); + assert!( + all_actions_valid, + "Some action indices were out of range (expected 0..45)" + ); + + // Report action distribution + let unique_actions = action_counts.iter().filter(|&&c| c > 0).count(); + println!( + "Phase 3 complete: {num_inference_samples} inference passes, \ + {unique_actions} unique actions selected" + ); + + // Sanity: at least 1 unique action (degenerate model is still valid for this test) + assert!( + unique_actions >= 1, + "No actions were selected -- inference loop failure" + ); + + println!("All assertions passed: checkpoint -> load -> inference pipeline verified"); + + Ok(()) +}