Files
foxhunt/crates/ml/tests/dqn_training_pipeline_test.rs
jgrusewski 673b04a8d4 perf(checkpoint): async best-ckpt serialize via spawn_blocking + mapped-pinned param snapshot
Best-checkpoint save (val Sharpe improvement, ~30% of epochs in
convergent runs) blocked the epoch loop for 20-40s on each improvement:
serialize_model() chained N (~26) per-tensor DtoH downloads via
GpuTensor::to_host → memcpy_dtoh, each forcing an implicit stream sync
on a busy training stream. The DtoH chain also violates
feedback_no_htod_htoh_only_mapped_pinned.md (only cuMemHostAlloc
DEVICEMAP allowed for CPU↔GPU paths).

Plan B:

- Introduce snapshot_model_to_pinned (mod.rs): allocate one
  MappedF32Buffer sized for all named weight slices concatenated,
  cuMemcpyDtoDAsync each slice into the buffer's device pointer
  (which aliases the host page), single stream sync, copy bytes out
  to a Send + 'static Vec<u8>. One sync per snapshot, replaces N.

- serialize_snapshot_bytes (mod.rs): pure-CPU safetensors construction
  from CheckpointSnapshot. Static — callable without &self, so the
  worker can move the snapshot across thread boundary.

- handle_epoch_checkpoints_and_early_stopping on val-Sharpe
  improvement: save_best_gpu_params (DtoD, fast) + snapshot to pinned
  + tokio::task::spawn_blocking the safetensors construction +
  checkpoint_callback invocation. JoinHandle parked on
  pending_checkpoint_handles. Training loop continues immediately.

- await_pending_checkpoint_handles drains in-flight workers at
  training end (success branch + early-stop branches) and before
  any synchronous cold-path checkpoint write to keep disk ordering
  deterministic.

- F bound on train / train_walk_forward / train_fold_from_slices
  gains + 'static so the callback can be moved into the worker.
  All public callers already use 'static-compatible move closures
  (test fixtures with shared mutable state migrate to Arc<Mutex<T>>).
  Internal pipeline uses CheckpointCallbackHandle =
  Arc<std::sync::Mutex<Box<dyn FnMut + Send + 'static>>> so the
  same callback flows through multi-fold walk-forward into every
  fold's worker.

- serialize_model itself rewritten via the snapshot path: the
  no-DtoH rule now holds across ALL checkpoint paths (best, periodic,
  early-stop, plateau-exhausted). The pre-existing GpuTensor::to_host
  path is no longer reachable from the DQN trainer.

The audit's spec called for an mpsc channel(1) drop-old worker, but
the multi-fold + &mut F pre-existing API made the simpler
fire-and-forget spawn_blocking pattern a cleaner fit (Mutex
serialises any concurrent invocations; Vec<JoinHandle> drain at end
guarantees disk writes complete before the trainer returns). Same
overlap benefit (training rolls while serialize+disk run on a
blocking thread); upper bound on in-flight work is one-per-improved-
epoch which approximates the spec's depth=1 in realistic training
runs.

Per feedback_no_partial_refactor: every site that constructs a
checkpoint payload migrated in lockstep — best-improvement uses
the worker; periodic / plateau-exhausted / early-stop call the
shared Arc<Mutex<F>> handle inline. All paths read params via
snapshot_model_to_pinned, so the no-DtoH rule applies uniformly.
Test fixtures (8 .rs files) updated for the + 'static bound (move
closures + cloned PathBufs / Arc<Mutex<T>> for shared mutable state).

Verified: SQLX_OFFLINE=true cargo check --workspace --tests clean
(warnings unchanged from baseline). cargo test -p ml --lib --no-run
clean. No fingerprint change.

Wire-up audit entry extended with Plan B file:line edit sites
(rides under the same Async-validation overlap section started by
the companion Plan A commit).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-28 18:54:00 +02:00

653 lines
23 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,
)]
//! **DQN Training Pipeline Test Suite**
//!
//! TDD implementation for DQN training on real ES.FUT market data.
//!
//! **Test Strategy**:
//! 1. Load real market data from DBN files
//! 2. Train DQN model for multiple epochs
//! 3. Verify loss decreases (>30% improvement)
//! 4. Save and load checkpoints
//! 5. Validate inference pipeline
//!
//! **Expected Outcomes**:
//! - All tests pass (6/6)
//! - Loss reduction >30% over training
//! - Checkpoint save/load functional
//! - Inference latency <1ms
#![allow(unused_crate_dependencies)]
use anyhow::{Context, Result};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use std::path::PathBuf;
use std::time::Instant;
use tracing::info;
use tracing::warn;
fn init_test_tracing() {
let _ = tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.with_test_writer()
.try_init();
}
/// Helper: Get path to ES.FUT test data (DBN format)
fn get_es_fut_data_dir() -> Result<String> {
// CI: TEST_DATA_DIR points to test-data-pvc on H100
if let Ok(dir) = std::env::var("TEST_DATA_DIR") {
let ohlcv = std::path::PathBuf::from(&dir).join("ohlcv");
if ohlcv.exists() {
return Ok(ohlcv.to_string_lossy().to_string());
}
return Ok(dir);
}
let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.ancestors()
.find(|p| p.join("test_data").exists())
.context("Failed to find workspace root with test_data/")?
.to_path_buf();
let data_dir = workspace_root.join("test_data/futures-baseline");
if !data_dir.exists() {
anyhow::bail!(
"ES.FUT data directory not found: {}. Run data acquisition first.",
data_dir.display()
);
}
Ok(data_dir.to_string_lossy().to_string())
}
/// Helper: Create checkpoint directory
fn create_checkpoint_dir() -> Result<PathBuf> {
let checkpoint_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("checkpoints");
std::fs::create_dir_all(&checkpoint_dir)?;
Ok(checkpoint_dir)
}
// ============================================================================
// TEST 1: Core Training Pipeline (RED → GREEN)
// ============================================================================
/// **TEST 1 (PRIMARY)**: Train DQN on ES.FUT data and verify loss decreases
///
/// **Expected**: This test should FAIL initially (RED phase) until we implement
/// the training pipeline. Once implemented, loss should decrease >30%.
#[tokio::test]
async fn test_dqn_trains_on_es_fut() -> Result<()> {
init_test_tracing();
info!("TEST 1: DQN Training Pipeline on ES.FUT");
let start_time = Instant::now();
// ========================================================================
// ARRANGE: Setup training configuration
// ========================================================================
info!("ARRANGE: Setting up DQN training configuration...");
let data_dir = match get_es_fut_data_dir() {
Ok(dir) => dir,
Err(e) => {
warn!(reason = %e, "Skipping test: data not available");
return Ok(());
},
};
let checkpoint_dir = create_checkpoint_dir()?;
// Configure hyperparameters from smoketest profile (all sizing + lr from TOML)
let mut hyperparams = DQNHyperparameters::conservative();
ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);
hyperparams.epochs = 2;
hyperparams.early_stopping_enabled = false;
hyperparams.checkpoint_frequency = 1;
hyperparams.cql_alpha = 0.0;
info!(data_dir = %data_dir, checkpoint_dir = %checkpoint_dir.display(), epochs = hyperparams.epochs, "Configuration ready");
// ========================================================================
// ACT: Create trainer and run training
// ========================================================================
info!("ACT: Running DQN training...");
let mut trainer = DQNTrainer::new(hyperparams.clone())?;
// Shared state across the async-checkpoint worker boundary — wrap in
// Arc<Mutex<>> to satisfy the `+ 'static` bound on the callback (worker
// can't borrow the test fn's stack frame).
let checkpoint_saved = std::sync::Arc::new(std::sync::Mutex::new(false));
let final_checkpoint_path = std::sync::Arc::new(
std::sync::Mutex::new(PathBuf::new()),
);
let saved_clone = std::sync::Arc::clone(&checkpoint_saved);
let final_clone = std::sync::Arc::clone(&final_checkpoint_path);
let ckpt_dir_clone = checkpoint_dir.clone();
let metrics = trainer
.train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, _is_best| {
let path = ckpt_dir_clone.join(format!("dqn_test_epoch_{}.safetensors", epoch));
std::fs::write(&path, checkpoint_data)?;
*saved_clone.lock().unwrap() = true;
*final_clone.lock().unwrap() = path.clone();
info!(epoch, "Checkpoint saved");
Ok(path.to_string_lossy().to_string())
})
.await?;
let checkpoint_saved = *checkpoint_saved.lock().unwrap();
let final_checkpoint_path = final_checkpoint_path.lock().unwrap().clone();
let training_time = start_time.elapsed();
info!(training_secs = training_time.as_secs_f64(), "Training completed");
// ========================================================================
// ASSERT: Verify training results
// ========================================================================
info!("ASSERT: Validating training results...");
// 1. Check that training completed all epochs
info!(
epochs = metrics.epochs_trained,
loss = metrics.loss,
training_time_seconds = metrics.training_time_seconds,
convergence = metrics.convergence_achieved,
"Training Metrics"
);
assert_eq!(
metrics.epochs_trained, hyperparams.epochs as u32,
"Should complete all {} epochs",
hyperparams.epochs
);
// 2. Check that loss is reasonable (not NaN, not infinite)
assert!(
metrics.loss.is_finite(),
"Loss should be finite, got: {}",
metrics.loss
);
assert!(
metrics.loss > 0.0,
"Loss should be positive, got: {}",
metrics.loss
);
// 3. Check Q-value metrics exist
if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
info!(avg_q_value, "Q-value metric");
assert!(
avg_q_value.is_finite(),
"Q-value should be finite, got: {}",
avg_q_value
);
} else {
panic!("Missing avg_q_value metric");
}
// 4. Check that checkpoint was saved
assert!(checkpoint_saved, "Checkpoint should have been saved");
assert!(
final_checkpoint_path.exists(),
"Checkpoint file should exist: {}",
final_checkpoint_path.display()
);
let checkpoint_size = std::fs::metadata(&final_checkpoint_path)?.len();
info!(checkpoint_kb = checkpoint_size / 1024, "Checkpoint size");
assert!(
checkpoint_size > 1024,
"Checkpoint should be >1KB, got: {} bytes",
checkpoint_size
);
info!("TEST 1 PASSED: DQN Training Pipeline Functional");
Ok(())
}
// ============================================================================
// TEST 2: Loss Convergence Validation
// ============================================================================
/// **TEST 2**: Verify DQN loss decreases during training (>5% improvement via loss_history)
#[tokio::test]
async fn test_dqn_loss_decreases() -> Result<()> {
info!("TEST 2: DQN Loss Convergence Test");
let data_dir = match get_es_fut_data_dir() {
Ok(dir) => dir,
Err(e) => {
warn!(reason = %e, "Skipping test: data not available");
return Ok(());
},
};
let checkpoint_dir = create_checkpoint_dir()?;
// Train for 3 epochs — CI validates gradient flow, not full convergence
let mut hyperparams = DQNHyperparameters::conservative();
ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);
hyperparams.epochs = 3;
hyperparams.early_stopping_enabled = false;
hyperparams.cql_alpha = 0.0;
let mut trainer = DQNTrainer::new(hyperparams)?;
// Track losses per epoch (would need to modify trainer to expose this)
let ckpt_dir_clone = checkpoint_dir.clone();
let metrics = trainer
.train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, _is_best| {
let path = ckpt_dir_clone.join(format!("dqn_loss_test_epoch_{}.safetensors", epoch));
std::fs::write(&path, checkpoint_data)?;
Ok(path.to_string_lossy().to_string())
})
.await?;
info!(loss = metrics.loss, convergence = metrics.convergence_achieved, "Training result");
// Use loss_history to verify loss decreased (gradient flow works)
let loss_history = trainer.loss_history();
assert!(
loss_history.len() >= 2,
"Need at least 2 epochs of loss history, got {}",
loss_history.len()
);
let first_loss = loss_history.first().copied().unwrap_or(f64::MAX);
let last_loss = loss_history.last().copied().unwrap_or(f64::MAX);
let min_loss = loss_history.iter().copied().fold(f64::MAX, f64::min);
// With noisy nets (epsilon=0), loss may not decrease monotonically in 20 epochs
// because Q-value-driven actions change the experience distribution.
// What matters: (1) all losses are finite, (2) min loss < first loss (model CAN learn)
for (i, loss) in loss_history.iter().enumerate() {
assert!(loss.is_finite(), "Loss at epoch {} is not finite: {}", i, loss);
}
info!(first_loss, last_loss, min_loss, "Loss history summary");
// Final loss should be finite and reasonable (not NaN/Inf/extreme)
assert!(
metrics.loss.is_finite() && metrics.loss < 100.0,
"Final loss should be finite and <100, got: {}",
metrics.loss
);
info!(decrease_pct = (1.0 - last_loss / first_loss) * 100.0, "Loss convergence validated");
Ok(())
}
// ============================================================================
// TEST 3: Checkpoint Save/Load Cycle
// ============================================================================
/// **TEST 3**: Save DQN checkpoint and reload it successfully
#[tokio::test]
async fn test_dqn_checkpoint_save_load() -> Result<()> {
info!("TEST 3: DQN Checkpoint Save/Load Test");
let data_dir = match get_es_fut_data_dir() {
Ok(dir) => dir,
Err(e) => {
warn!(reason = %e, "Skipping test: data not available");
return Ok(());
},
};
let checkpoint_dir = create_checkpoint_dir()?;
// Train for 2 epochs and save checkpoint
let mut hyperparams = DQNHyperparameters::conservative();
ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);
hyperparams.epochs = 2;
hyperparams.early_stopping_enabled = false;
hyperparams.checkpoint_frequency = 2;
hyperparams.cql_alpha = 0.0;
let mut trainer = DQNTrainer::new(hyperparams)?;
let saved_checkpoint_path = std::sync::Arc::new(std::sync::Mutex::new(PathBuf::new()));
let saved_clone = std::sync::Arc::clone(&saved_checkpoint_path);
let ckpt_dir_clone = checkpoint_dir.clone();
let _metrics = trainer
.train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, _is_best| {
let path =
ckpt_dir_clone.join(format!("dqn_checkpoint_test_epoch_{}.safetensors", epoch));
std::fs::write(&path, checkpoint_data)?;
*saved_clone.lock().unwrap() = path.clone();
info!(path = %path.display(), "Checkpoint saved");
Ok(path.to_string_lossy().to_string())
})
.await?;
let saved_checkpoint_path = saved_checkpoint_path.lock().unwrap().clone();
// Verify checkpoint exists
assert!(
saved_checkpoint_path.exists(),
"Checkpoint should exist: {}",
saved_checkpoint_path.display()
);
// Verify checkpoint size
let checkpoint_size = std::fs::metadata(&saved_checkpoint_path)?.len();
info!(checkpoint_kb = checkpoint_size / 1024, "Checkpoint size");
assert!(checkpoint_size > 1024, "Checkpoint should be >1KB");
// Verify the checkpoint bytes round-trip back from disk. A dedicated
// `load_checkpoint` entry point is not part of this end-to-end test;
// roundtrip coverage for the loader lives in `dqn_checkpoint_tests`.
let checkpoint_data = std::fs::read(&saved_checkpoint_path)?;
assert!(
checkpoint_data.len() == checkpoint_size as usize,
"Checkpoint data should match file size"
);
info!("Checkpoint save/load validated");
Ok(())
}
// ============================================================================
// TEST 4: Q-Value Predictions
// ============================================================================
/// **TEST 4**: Verify DQN produces valid Q-values for given states
#[tokio::test]
async fn test_dqn_q_value_predictions() -> Result<()> {
info!("TEST 4: DQN Q-Value Prediction Test");
let data_dir = match get_es_fut_data_dir() {
Ok(dir) => dir,
Err(e) => {
warn!(reason = %e, "Skipping test: data not available");
return Ok(());
},
};
let checkpoint_dir = create_checkpoint_dir()?;
// Train minimal model
let mut hyperparams = DQNHyperparameters::conservative();
ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);
hyperparams.epochs = 2;
hyperparams.early_stopping_enabled = false;
hyperparams.checkpoint_frequency = 1;
hyperparams.cql_alpha = 0.0;
let mut trainer = DQNTrainer::new(hyperparams)?;
let ckpt_dir_clone = checkpoint_dir.clone();
let metrics = trainer
.train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, _is_best| {
let path = ckpt_dir_clone.join(format!("dqn_qvalue_test_epoch_{}.safetensors", epoch));
std::fs::write(&path, checkpoint_data)?;
Ok(path.to_string_lossy().to_string())
})
.await?;
// Check Q-value metrics
if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
info!(avg_q_value, "Q-value metric");
// Q-values should be finite and within reasonable range
assert!(avg_q_value.is_finite(), "Q-value should be finite");
assert!(
*avg_q_value > -100.0 && *avg_q_value < 100.0,
"Q-value should be in reasonable range [-100, 100], got: {}",
avg_q_value
);
info!("Q-value predictions validated");
} else {
panic!("Missing avg_q_value metric");
}
Ok(())
}
// ============================================================================
// TEST 5: Epsilon-Greedy Exploration
// ============================================================================
/// **TEST 5**: Verify epsilon-greedy exploration behavior
#[tokio::test]
async fn test_dqn_epsilon_greedy() -> Result<()> {
info!("TEST 5: DQN Epsilon-Greedy Exploration Test");
let data_dir = match get_es_fut_data_dir() {
Ok(dir) => dir,
Err(e) => {
warn!(reason = %e, "Skipping test: data not available");
return Ok(());
},
};
let checkpoint_dir = create_checkpoint_dir()?;
// Configure with high epsilon decay
let mut hyperparams = DQNHyperparameters::conservative();
ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);
hyperparams.epochs = 2;
hyperparams.early_stopping_enabled = false;
hyperparams.cql_alpha = 0.0;
hyperparams.epsilon_start = 1.0;
hyperparams.epsilon_end = 0.01;
hyperparams.epsilon_decay = 0.9; // Fast decay
let mut trainer = DQNTrainer::new(hyperparams)?;
let ckpt_dir_clone = checkpoint_dir.clone();
let _metrics = trainer
.train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, _is_best| {
let path = ckpt_dir_clone.join(format!("dqn_epsilon_test_epoch_{}.safetensors", epoch));
std::fs::write(&path, checkpoint_data)?;
Ok(path.to_string_lossy().to_string())
})
.await?;
// BUG #40 VERIFIED: With noisy nets enabled (conservative() default),
// epsilon is fixed at noisy_epsilon_floor (0.05) — not decayed, not 1.0.
// Before the fix, epsilon stayed at 1.0 (100% random actions throughout training).
// The noisy_epsilon_floor provides a minimum exploration rate to prevent action collapse
// while NoisyNets provide the primary learned exploration signal.
let final_epsilon = trainer.get_agent_epsilon().await;
info!(final_epsilon, "Final epsilon");
assert!(
final_epsilon < 0.10,
"BUG #40: Epsilon should be at noisy_epsilon_floor (~0.05) with noisy nets (got {:.4}). \
If this fails, noisy net epsilon override is broken.",
final_epsilon
);
// Verify training completed and model learned (Q-values non-zero)
let avg_q = _metrics.additional_metrics.get("avg_q_value").copied().unwrap_or(0.0);
assert!(
avg_q.abs() > 0.001,
"Model should develop Q-value preferences, got avg_q={:.6}",
avg_q
);
info!(avg_q, "Noisy nets exploration validated (epsilon=floor, Q-values+noise drive actions)");
Ok(())
}
// ============================================================================
// TEST 6: Production Training (50 epochs)
// ============================================================================
/// **TEST 6**: Full production training run (50 epochs)
///
/// **Note**: This test takes ~5-10 minutes. Run separately for production validation.
#[tokio::test]
#[ignore = "Ignore by default due to long runtime"]
async fn test_dqn_full_production_training() -> Result<()> {
info!("TEST 6: DQN Full Production Training (50 epochs) — expected runtime: 5-10 minutes");
let start_time = Instant::now();
let data_dir = match get_es_fut_data_dir() {
Ok(dir) => dir,
Err(e) => {
warn!(reason = %e, "Skipping test: data not available");
return Ok(());
},
};
let checkpoint_dir = create_checkpoint_dir()?;
let production_checkpoint_path = checkpoint_dir.join("dqn_es_fut_v1.safetensors");
// Production hyperparameters (smoketest profile for GPU sizing)
let mut hyperparams = DQNHyperparameters::conservative();
ml::training_profile::DqnTrainingProfile::load("dqn-smoketest").apply_to(&mut hyperparams);
hyperparams.replay_buffer_vram_fraction = 0.0;
// GPU PER is mandatory for fused CUDA training — do not disable
hyperparams.epochs = 50;
hyperparams.batch_size = 128;
hyperparams.learning_rate = 0.0001;
hyperparams.gamma = 0.99;
hyperparams.epsilon_start = 1.0;
hyperparams.epsilon_end = 0.01;
hyperparams.epsilon_decay = 0.995;
hyperparams.checkpoint_frequency = 10;
hyperparams.early_stopping_enabled = true;
let mut trainer = DQNTrainer::new(hyperparams.clone())?;
let final_epoch = hyperparams.epochs;
let prod_path_clone = production_checkpoint_path.clone();
let ckpt_dir_clone = checkpoint_dir.clone();
let metrics = trainer
.train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, _is_best| {
let path = if epoch == final_epoch {
prod_path_clone.clone()
} else {
ckpt_dir_clone.join(format!("dqn_production_epoch_{}.safetensors", epoch))
};
std::fs::write(&path, checkpoint_data)?;
info!(epoch, "Checkpoint saved");
Ok(path.to_string_lossy().to_string())
})
.await?;
let training_time = start_time.elapsed();
let avg_q_value = metrics.additional_metrics.get("avg_q_value").copied();
let final_epsilon = metrics.additional_metrics.get("final_epsilon").copied();
info!(
epochs = metrics.epochs_trained,
loss = metrics.loss,
training_secs = training_time.as_secs_f64(),
convergence = metrics.convergence_achieved,
avg_q_value,
final_epsilon,
"Production training results"
);
// Verify production checkpoint exists
assert!(
production_checkpoint_path.exists(),
"Production checkpoint should exist: {}",
production_checkpoint_path.display()
);
let checkpoint_size = std::fs::metadata(&production_checkpoint_path)?.len();
info!(checkpoint_kb = checkpoint_size / 1024, "Production checkpoint size");
// Production assertions
assert!(
metrics.loss < 2.0,
"Production loss should be <2.0, got: {}",
metrics.loss
);
assert!(
checkpoint_size > 10_000,
"Production checkpoint should be >10KB"
);
info!("Production training validation passed");
Ok(())
}