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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
274
ml/tests/dqn_inference_test.rs
Normal file
274
ml/tests/dqn_inference_test.rs
Normal file
@@ -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<String> {
|
||||
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<String, Tensor>) -> 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<PathBuf> = 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::<Vec<_>>()
|
||||
);
|
||||
|
||||
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<f32> = (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<f32> = q_values.to_vec2::<f32>()?.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(())
|
||||
}
|
||||
Reference in New Issue
Block a user