9aad6ff60e
Epoch duration self-balances: bigger GPU → bigger auto-scaled batch → fewer steps per epoch. The manual cap created 7 different values (0, 8, 64, 100, 200, 300, 2000) across configs/tests/examples, making behavior inconsistent between environments. Removed from: DQNHyperparameters, training profiles (smoketest, localdev, production), CLI args, Argo templates, hyperopt adapter, all test overrides, supervised example. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
919 lines
34 KiB
Rust
919 lines
34 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,
|
||
)]
|
||
//! Smoke tests with real Databento market data.
|
||
//!
|
||
//! These tests validate the full pipeline from DBN parsing through GPU kernel
|
||
//! forward pass using actual ES.FUT data downloaded from MinIO.
|
||
//!
|
||
//! Data requirements (in `data/smoke-test/` relative to workspace root):
|
||
//! - ohlcv/ES.FUT/ES.FUT_2024-Q1.dbn.zst (~700 KB)
|
||
//! - trades/ES.FUT/ES.FUT_2024-Q1.dbn.zst (~120 MB)
|
||
//! - trades/ES.FUT/ES.FUT_2025-Q2.dbn.zst (~505 MB, matches MBP-10 quarter)
|
||
//! - mbp10/ES.FUT/ES.FUT_2025-Q2.dbn.zst (~125 MB)
|
||
//!
|
||
//! Run: `cargo test -p ml --test smoke_test_real_data -- --nocapture`
|
||
|
||
use std::path::{Path, PathBuf};
|
||
use tracing::{info, warn};
|
||
|
||
/// Resolve workspace root from CARGO_MANIFEST_DIR (crates/ml -> ../..)
|
||
fn workspace_root() -> PathBuf {
|
||
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
||
.parent() // crates/
|
||
.and_then(|p| p.parent()) // workspace root
|
||
.expect("Failed to find workspace root")
|
||
.to_path_buf()
|
||
}
|
||
|
||
fn smoke_data_dir() -> PathBuf {
|
||
// CI: TEST_DATA_DIR points to test-data-pvc on H100
|
||
if let Ok(dir) = std::env::var("TEST_DATA_DIR") {
|
||
return PathBuf::from(dir);
|
||
}
|
||
// Local dev: relative path from workspace root
|
||
workspace_root().join("data").join("smoke-test")
|
||
}
|
||
|
||
/// Check if GPU has enough VRAM (MB). Returns false if no GPU or insufficient VRAM.
|
||
fn gpu_vram_mb() -> u64 {
|
||
let output = std::process::Command::new("nvidia-smi")
|
||
.args(["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
|
||
.output();
|
||
match output {
|
||
Ok(out) if out.status.success() => {
|
||
String::from_utf8_lossy(&out.stdout)
|
||
.trim()
|
||
.parse()
|
||
.unwrap_or(0)
|
||
}
|
||
_ => 0,
|
||
}
|
||
}
|
||
|
||
/// Return data path if it exists, `None` otherwise (test should skip).
|
||
fn try_data(relative: &str) -> Option<PathBuf> {
|
||
let path = smoke_data_dir().join(relative);
|
||
if !path.exists() {
|
||
warn!(path = %path.display(), "SKIP: smoke test data not found — download from MinIO first");
|
||
return None;
|
||
}
|
||
Some(path)
|
||
}
|
||
|
||
// ==========================================================================
|
||
// Test 1: Parse OHLCV data and extract 42-dim features
|
||
// ==========================================================================
|
||
|
||
#[test]
|
||
fn smoke_ohlcv_parse_and_extract_features() {
|
||
let Some(ohlcv_dir) = try_data("ohlcv") else { return };
|
||
|
||
// Load OHLCV bars from .dbn.zst
|
||
let bars = ml::hyperopt::adapters::dbn_loader::load_bars_from_dbn_dir(&ohlcv_dir)
|
||
.expect("Failed to load OHLCV bars from DBN");
|
||
|
||
info!(count = bars.len(), dir = %ohlcv_dir.display(), "Loaded OHLCV bars");
|
||
assert!(bars.len() > 100, "Expected >100 bars, got {}", bars.len());
|
||
|
||
// Verify bar structure
|
||
let first = &bars[0];
|
||
assert!(first.open > 0.0, "Open price should be positive");
|
||
assert!(first.high >= first.low, "High >= Low invariant");
|
||
assert!(first.volume >= 0.0, "Volume should be non-negative");
|
||
|
||
// Extract 42-dim features (requires 50-bar warmup)
|
||
let features = ml::features::extract_ml_features(&bars)
|
||
.expect("Feature extraction failed");
|
||
|
||
info!(count = features.len(), dim = 42, "Extracted feature vectors");
|
||
assert!(
|
||
features.len() >= bars.len() - 50,
|
||
"Expected at least {} features (bars - warmup), got {}",
|
||
bars.len().saturating_sub(50),
|
||
features.len()
|
||
);
|
||
|
||
// Verify features are finite and reasonable
|
||
for (i, feat) in features.iter().enumerate().take(10) {
|
||
for (j, &val) in feat.iter().enumerate() {
|
||
assert!(
|
||
val.is_finite(),
|
||
"Feature [{i}][{j}] = {val} is not finite"
|
||
);
|
||
}
|
||
}
|
||
|
||
// Spot-check: close price feature (index 3) should match OHLCV bars
|
||
// Features start at bar 50 (warmup period)
|
||
let close_feat = features[0][3]; // Normalized close
|
||
info!(close_feat, bar_50_close = bars[50].close, "First feature vector spot-check");
|
||
}
|
||
|
||
// ==========================================================================
|
||
// Test 2: Parse MBP-10 data and compute OFI features
|
||
// ==========================================================================
|
||
|
||
#[test]
|
||
fn smoke_mbp10_ofi_extraction() {
|
||
let Some(mbp10_dir) = try_data("mbp10/ES.FUT") else { return };
|
||
|
||
// Find the .dbn.zst file
|
||
let files: Vec<_> = std::fs::read_dir(&mbp10_dir)
|
||
.expect("Failed to read mbp10 dir")
|
||
.filter_map(|e| e.ok())
|
||
.map(|e| e.path())
|
||
.filter(|p| {
|
||
p.extension()
|
||
.map_or(false, |ext| ext == "zst" || ext == "dbn")
|
||
})
|
||
.collect();
|
||
|
||
assert!(!files.is_empty(), "No .dbn.zst files found in {}", mbp10_dir.display());
|
||
|
||
let file = &files[0];
|
||
info!(file = %file.display(), "Computing OFI");
|
||
|
||
let ofi_features = ml::features::mbp10_loader::compute_ofi_from_file(file)
|
||
.expect("OFI computation failed");
|
||
|
||
info!(count = ofi_features.len(), dim = 8, "Computed OFI feature vectors");
|
||
assert!(
|
||
ofi_features.len() > 1000,
|
||
"Expected >1000 OFI snapshots, got {}",
|
||
ofi_features.len()
|
||
);
|
||
|
||
// Verify OFI features are finite
|
||
let mut nan_count = 0;
|
||
let mut zero_count = 0;
|
||
for (i, feat) in ofi_features.iter().enumerate().take(100) {
|
||
for (j, &val) in feat.iter().enumerate() {
|
||
if !val.is_finite() {
|
||
nan_count += 1;
|
||
}
|
||
if val == 0.0 {
|
||
zero_count += 1;
|
||
}
|
||
// OFI values should not be astronomically large
|
||
if val.is_finite() {
|
||
assert!(
|
||
val.abs() < 1e12,
|
||
"OFI [{i}][{j}] = {val} unreasonably large"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
info!(nan_count, zero_count, sample = 800, "OFI quality (first 100 features)");
|
||
// Some NaN is expected in early warmup, but not all
|
||
assert!(
|
||
nan_count < 50,
|
||
"Too many NaN OFI features: {nan_count}/800"
|
||
);
|
||
}
|
||
|
||
// ==========================================================================
|
||
// Test 2b: OFI with trade enrichment — validates VPIN/Kyle's Lambda are non-zero
|
||
// ==========================================================================
|
||
|
||
#[test]
|
||
fn smoke_ofi_with_trade_enrichment() {
|
||
let Some(mbp10_file) = try_data("mbp10/ES.FUT/ES.FUT_2025-Q2.dbn.zst") else { return };
|
||
let Some(trades_file) = try_data("trades/ES.FUT/ES.FUT_2025-Q2.dbn.zst") else { return };
|
||
|
||
// Compute OFI WITHOUT trades (baseline — VPIN/Kyle's Lambda will be zero)
|
||
let ofi_without = ml::features::mbp10_loader::compute_ofi_from_file(&mbp10_file)
|
||
.expect("OFI without trades failed");
|
||
|
||
// Compute OFI WITH trades (enriched — VPIN/Kyle's Lambda should be non-zero)
|
||
let ofi_with = ml::features::mbp10_loader::compute_ofi_with_trades(&mbp10_file, &trades_file)
|
||
.expect("OFI with trades failed");
|
||
|
||
info!(without = ofi_without.len(), with = ofi_with.len(), "OFI feature counts");
|
||
|
||
// Both should produce the same number of features (same MBP-10 snapshots)
|
||
assert_eq!(
|
||
ofi_without.len(),
|
||
ofi_with.len(),
|
||
"Feature count should match: without={} vs with={}",
|
||
ofi_without.len(),
|
||
ofi_with.len()
|
||
);
|
||
|
||
// OFI feature indices:
|
||
// [0] ofi_level1 [1] ofi_level5 [2] depth_imbalance [3] vpin
|
||
// [4] kyle_lambda [5] bid_slope [6] ask_slope [7] trade_imbalance
|
||
//
|
||
// Without trades: indices 3, 4 should be ~all zero, 7 uses tick-rule fallback
|
||
// With trades: indices 3, 4 should have real values from trade data
|
||
|
||
// Count non-zero VPIN (index 3) in both
|
||
let skip = 100; // skip warmup
|
||
let sample_end = ofi_without.len().min(skip + 5000);
|
||
let sample = skip..sample_end;
|
||
|
||
let vpin_nonzero_without = ofi_without[sample.clone()]
|
||
.iter()
|
||
.filter(|f| f[3].abs() > 1e-10)
|
||
.count();
|
||
let vpin_nonzero_with = ofi_with[sample.clone()]
|
||
.iter()
|
||
.filter(|f| f[3].abs() > 1e-10)
|
||
.count();
|
||
|
||
let kyle_nonzero_without = ofi_without[sample.clone()]
|
||
.iter()
|
||
.filter(|f| f[4].abs() > 1e-10)
|
||
.count();
|
||
let kyle_nonzero_with = ofi_with[sample.clone()]
|
||
.iter()
|
||
.filter(|f| f[4].abs() > 1e-10)
|
||
.count();
|
||
|
||
let sample_size = sample.len();
|
||
info!(vpin_nonzero_without, vpin_nonzero_with, sample_size, "VPIN non-zero counts");
|
||
info!(kyle_nonzero_without, kyle_nonzero_with, sample_size, "Kyle's Lambda non-zero counts");
|
||
|
||
// The WITHOUT-trades path should have mostly zero VPIN/Kyle's Lambda
|
||
// (this is the bug we fixed)
|
||
assert!(
|
||
vpin_nonzero_without < sample_size / 10,
|
||
"Without trades, VPIN should be mostly zero but got {vpin_nonzero_without}/{sample_size}"
|
||
);
|
||
|
||
// The WITH-trades path should have significantly more non-zero VPIN
|
||
assert!(
|
||
vpin_nonzero_with > sample_size / 2,
|
||
"With trades, VPIN should be mostly non-zero but got {vpin_nonzero_with}/{sample_size}"
|
||
);
|
||
|
||
// Kyle's Lambda should also improve with trade data
|
||
assert!(
|
||
kyle_nonzero_with > kyle_nonzero_without,
|
||
"Kyle's Lambda should improve with trades: without={kyle_nonzero_without} vs with={kyle_nonzero_with}"
|
||
);
|
||
|
||
// Spot-check some enriched values
|
||
let sample_feat = &ofi_with[skip + 500];
|
||
info!(
|
||
ofi_l1 = sample_feat[0],
|
||
ofi_l5 = sample_feat[1],
|
||
depth_imb = sample_feat[2],
|
||
vpin = sample_feat[3],
|
||
kyle = sample_feat[4],
|
||
bid_slope = sample_feat[5],
|
||
ask_slope = sample_feat[6],
|
||
trade_imb = sample_feat[7],
|
||
"Sample enriched OFI"
|
||
);
|
||
}
|
||
|
||
// ==========================================================================
|
||
// Test 3: Parse trades data
|
||
// ==========================================================================
|
||
|
||
#[test]
|
||
fn smoke_trades_parse() {
|
||
let Some(trades_dir) = try_data("trades/ES.FUT") else { return };
|
||
|
||
let files: Vec<_> = std::fs::read_dir(&trades_dir)
|
||
.expect("Failed to read trades dir")
|
||
.filter_map(|e| e.ok())
|
||
.map(|e| e.path())
|
||
.filter(|p| {
|
||
p.extension()
|
||
.map_or(false, |ext| ext == "zst" || ext == "dbn")
|
||
})
|
||
.collect();
|
||
|
||
assert!(!files.is_empty(), "No .dbn.zst files in {}", trades_dir.display());
|
||
|
||
let file = &files[0];
|
||
info!(file = %file.display(), "Loading trades");
|
||
|
||
let trades = ml::features::trades_loader::load_trades_sync(file)
|
||
.expect("Trade loading failed");
|
||
|
||
info!(count = trades.len(), "Loaded trades");
|
||
assert!(
|
||
trades.len() > 10_000,
|
||
"Expected >10K trades for a full quarter, got {}",
|
||
trades.len()
|
||
);
|
||
|
||
// Verify trade structure
|
||
let first = &trades[0];
|
||
assert!(first.price > 0.0, "Trade price should be positive");
|
||
assert!(first.volume > 0, "Trade volume should be positive");
|
||
assert!(first.timestamp > 0, "Trade timestamp should be set");
|
||
}
|
||
|
||
// ==========================================================================
|
||
// Test 4: Full GPU pipeline — real data -> kernel forward pass -> Q-values
|
||
// ==========================================================================
|
||
|
||
mod gpu_smoke {
|
||
use super::*;
|
||
use std::sync::Arc;
|
||
|
||
// candle eliminated — use MlDevice for stream extraction
|
||
use ml_core::device::MlDevice;
|
||
use ml::cuda_pipeline::gpu_experience_collector::{
|
||
ExperienceCollectorConfig, GpuExperienceCollector,
|
||
};
|
||
use ml::dqn::branching::{BranchingConfig, BranchingDuelingQNetwork};
|
||
|
||
type CudaStream = cudarc::driver::CudaStream;
|
||
type CudaSlice<T> = cudarc::driver::CudaSlice<T>;
|
||
|
||
fn try_cuda() -> Option<(MlDevice, Arc<CudaStream>)> {
|
||
match MlDevice::cuda(0) {
|
||
Ok(dev) => {
|
||
let stream = dev.cuda_stream().expect("CUDA stream required").clone();
|
||
Some((dev, stream))
|
||
}
|
||
Err(_) => {
|
||
warn!("SKIP: CUDA not available");
|
||
None
|
||
}
|
||
}
|
||
}
|
||
|
||
fn has_enough_vram(min_mb: u64) -> bool {
|
||
let vram = gpu_vram_mb();
|
||
if vram < min_mb {
|
||
warn!(vram, min_mb, "SKIP: insufficient GPU VRAM");
|
||
false
|
||
} else {
|
||
true
|
||
}
|
||
}
|
||
|
||
/// Build flat GPU buffers from real OHLCV features + OFI.
|
||
///
|
||
/// Returns (market_buf [N*51], target_buf [N*4], num_bars).
|
||
fn real_market_data(
|
||
ohlcv_dir: &Path,
|
||
mbp10_dir: Option<&Path>,
|
||
stream: &Arc<CudaStream>,
|
||
) -> Result<(CudaSlice<half::bf16>, CudaSlice<half::bf16>, usize), anyhow::Error> {
|
||
// 1. Load OHLCV bars
|
||
let bars = ml::hyperopt::adapters::dbn_loader::load_bars_from_dbn_dir(ohlcv_dir)
|
||
.expect("Failed to load OHLCV bars");
|
||
info!(count = bars.len(), "Loaded OHLCV bars");
|
||
|
||
// 2. Extract 42-dim features
|
||
let features = ml::features::extract_ml_features(&bars)
|
||
.expect("Feature extraction failed");
|
||
let n = features.len();
|
||
info!(count = n, "Extracted feature vectors");
|
||
assert!(n > 200, "Need at least 200 bars for GPU test, got {n}");
|
||
|
||
// 3. Optionally load OFI features from MBP-10
|
||
let ofi_features: Option<Vec<[f64; 8]>> = mbp10_dir.and_then(|dir| {
|
||
let files: Vec<_> = std::fs::read_dir(dir)
|
||
.ok()?
|
||
.filter_map(|e| e.ok())
|
||
.map(|e| e.path())
|
||
.filter(|p| {
|
||
p.extension()
|
||
.map_or(false, |ext| ext == "zst" || ext == "dbn")
|
||
})
|
||
.collect();
|
||
let file = files.first()?;
|
||
info!(file = %file.display(), "Loading OFI");
|
||
ml::features::mbp10_loader::compute_ofi_from_file(file).ok()
|
||
});
|
||
|
||
let ofi_count = ofi_features.as_ref().map_or(0, |v| v.len());
|
||
info!(ofi_count, "OFI features loaded");
|
||
|
||
// 4. Build flat market buffer [n * MARKET_DIM(51)]
|
||
const MARKET_DIM: usize = 51;
|
||
let mut market_data = vec![0.0_f32; n * MARKET_DIM];
|
||
for (i, feat) in features.iter().enumerate() {
|
||
// Indices 0-39: market features
|
||
for (j, &val) in feat.iter().enumerate() {
|
||
market_data[i * MARKET_DIM + j] = val as f32;
|
||
}
|
||
// Indices 40-42: zeros (dead regime features)
|
||
// Indices 43-50: OFI features
|
||
if let Some(ref ofi) = ofi_features {
|
||
if let Some(ofi_row) = ofi.get(i) {
|
||
for (j, &val) in ofi_row.iter().enumerate() {
|
||
if val.is_finite() {
|
||
market_data[i * MARKET_DIM + 43 + j] = val as f32;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 5. Build flat target buffer [n * 4]
|
||
// Layout: [current_close, next_close, current_close_raw, next_close_raw]
|
||
// Features are offset by warmup (50 bars), so features[i] ~ bars[i + 50]
|
||
let warmup = bars.len() - n; // typically 50
|
||
let mut target_data = vec![0.0_f32; n * 4];
|
||
for i in 0..n {
|
||
let bar_idx = i + warmup;
|
||
let close = bars[bar_idx].close as f32;
|
||
let next_close = if bar_idx + 1 < bars.len() {
|
||
bars[bar_idx + 1].close as f32
|
||
} else {
|
||
close
|
||
};
|
||
target_data[i * 4] = close;
|
||
target_data[i * 4 + 1] = next_close;
|
||
target_data[i * 4 + 2] = close;
|
||
target_data[i * 4 + 3] = next_close;
|
||
}
|
||
|
||
// 6. Upload to GPU (convert f32 -> bf16 for kernel signature)
|
||
let market_bf16: Vec<half::bf16> = market_data.iter().map(|&x| half::bf16::from_f32(x)).collect();
|
||
let mut market_buf = stream.alloc_zeros::<half::bf16>(n * MARKET_DIM).unwrap();
|
||
stream.memcpy_htod(&market_bf16, &mut market_buf).unwrap();
|
||
|
||
let target_bf16: Vec<half::bf16> = target_data.iter().map(|&x| half::bf16::from_f32(x)).collect();
|
||
let mut target_buf = stream.alloc_zeros::<half::bf16>(n * 4).unwrap();
|
||
stream.memcpy_htod(&target_bf16, &mut target_buf).unwrap();
|
||
|
||
Ok((market_buf, target_buf, n))
|
||
}
|
||
|
||
// ------------------------------------------------------------------
|
||
// Test 4a: Real OHLCV data -> GPU kernel -> valid Q-values
|
||
// ------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn smoke_gpu_real_ohlcv_forward() -> Result<(), anyhow::Error> {
|
||
let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) };
|
||
let Some((_device, stream)) = try_cuda() else { return Ok(()) };
|
||
|
||
let (market_buf, target_buf, num_bars) =
|
||
real_market_data(&ohlcv_dir, None, &stream)?;
|
||
|
||
info!(num_bars, market_dim = 51, "GPU buffer allocated");
|
||
|
||
// Create dueling networks (state_dim=54: 51 market + 3 portfolio)
|
||
let config = BranchingConfig::trading_default(
|
||
54, vec![64, 64], 32, vec![9, 3, 3],
|
||
);
|
||
let online = BranchingDuelingQNetwork::new(config.clone(), stream.clone()).unwrap();
|
||
let target_net = BranchingDuelingQNetwork::new(config, stream.clone()).unwrap();
|
||
|
||
let mut collector = GpuExperienceCollector::new(
|
||
stream.clone(),
|
||
online.vars(),
|
||
target_net.vars(),
|
||
None,
|
||
100_000.0,
|
||
0.01, 0.05,
|
||
(64, 64, 32, 32),
|
||
(54, 51, 51),
|
||
4, 50,
|
||
).unwrap();
|
||
|
||
// Run 4 episodes of 50 timesteps each
|
||
let n_episodes = 4_usize;
|
||
let timesteps = 50;
|
||
let episode_starts: Vec<i32> = (0..n_episodes)
|
||
.map(|i| (i * (num_bars / n_episodes)) as i32)
|
||
.collect();
|
||
|
||
let exp_config = ExperienceCollectorConfig {
|
||
n_episodes: n_episodes as i32,
|
||
timesteps_per_episode: timesteps,
|
||
total_bars: num_bars as i32,
|
||
episode_length: timesteps,
|
||
epsilon: 0.1,
|
||
gamma: 0.99,
|
||
..Default::default()
|
||
};
|
||
|
||
let batch = collector
|
||
.collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &exp_config)
|
||
.expect("GPU experience collection failed on real data");
|
||
|
||
let expected_total = n_episodes * timesteps as usize;
|
||
assert_eq!(batch.n_episodes * batch.timesteps, expected_total);
|
||
|
||
// Download actions and rewards from GPU to host for validation
|
||
let mut actions_host = vec![0i32; expected_total];
|
||
stream.memcpy_dtoh(&batch.actions, &mut actions_host).unwrap();
|
||
let mut rewards_host = vec![0.0f32; expected_total];
|
||
stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();
|
||
|
||
// Branching DQN: 9×3×3 = 81 factored actions (0..81)
|
||
for (i, &a) in actions_host.iter().enumerate() {
|
||
assert!((0..81).contains(&a), "action[{i}] = {a} out of range (expected 0..81 for 9-action branching DQN)");
|
||
}
|
||
|
||
// Rewards should reflect real price movements (not all zero)
|
||
let nonzero_rewards = rewards_host.iter().filter(|&&r| r.abs() > 1e-8).count();
|
||
let r_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
|
||
let r_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
info!(
|
||
expected_total, nonzero_rewards,
|
||
r_min, r_max,
|
||
"Real OHLCV pipeline results"
|
||
);
|
||
Ok(())
|
||
}
|
||
|
||
// ------------------------------------------------------------------
|
||
// Test 4b: Real OHLCV + MBP-10 OFI data -> GPU kernel -> valid Q-values
|
||
// ------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn smoke_gpu_real_ohlcv_plus_ofi_forward() -> Result<(), anyhow::Error> {
|
||
if !has_enough_vram(2048) { return Ok(()) }
|
||
let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) };
|
||
let Some(mbp10_dir) = try_data("mbp10/ES.FUT") else { return Ok(()) };
|
||
let Some((_device, stream)) = try_cuda() else { return Ok(()) };
|
||
|
||
let (market_buf, target_buf, num_bars) =
|
||
real_market_data(&ohlcv_dir, Some(&mbp10_dir), &stream)?;
|
||
|
||
info!(num_bars, market_dim = 51, "GPU buffer allocated (with OFI)");
|
||
|
||
let config = BranchingConfig::trading_default(
|
||
54, vec![64, 64], 32, vec![9, 3, 3],
|
||
);
|
||
let online = BranchingDuelingQNetwork::new(config.clone(), stream.clone()).unwrap();
|
||
let target_net = BranchingDuelingQNetwork::new(config, stream.clone()).unwrap();
|
||
|
||
let mut collector = GpuExperienceCollector::new(
|
||
stream.clone(),
|
||
online.vars(),
|
||
target_net.vars(),
|
||
None,
|
||
100_000.0,
|
||
0.01, 0.05,
|
||
(64, 64, 32, 32),
|
||
(54, 51, 51),
|
||
4, 50,
|
||
).unwrap();
|
||
|
||
let n_episodes = 4_usize;
|
||
let timesteps = 50;
|
||
let episode_starts: Vec<i32> = (0..n_episodes)
|
||
.map(|i| (i * (num_bars / n_episodes)) as i32)
|
||
.collect();
|
||
|
||
let exp_config = ExperienceCollectorConfig {
|
||
n_episodes: n_episodes as i32,
|
||
timesteps_per_episode: timesteps,
|
||
total_bars: num_bars as i32,
|
||
episode_length: timesteps,
|
||
epsilon: 0.1,
|
||
gamma: 0.99,
|
||
..Default::default()
|
||
};
|
||
|
||
let batch = collector
|
||
.collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &exp_config)
|
||
.expect("GPU experience collection with OFI failed");
|
||
|
||
let expected_total = n_episodes * timesteps as usize;
|
||
assert_eq!(batch.n_episodes * batch.timesteps, expected_total);
|
||
|
||
// Download rewards from GPU to host for validation
|
||
let mut rewards_host = vec![0.0f32; expected_total];
|
||
stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();
|
||
|
||
for (i, &r) in rewards_host.iter().enumerate() {
|
||
assert!(r.is_finite(), "OFI reward[{i}] = {r} is not finite");
|
||
}
|
||
|
||
let r_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
|
||
let r_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
info!(expected_total, r_min, r_max, "OHLCV+OFI pipeline results");
|
||
Ok(())
|
||
}
|
||
|
||
// ------------------------------------------------------------------
|
||
// Test 4c: NoisyNet + C51 distributional on real data
|
||
// ------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn smoke_gpu_real_data_noisy_distributional() -> Result<(), anyhow::Error> {
|
||
use ml::dqn::branching::{BranchingConfig, BranchingDuelingQNetwork};
|
||
|
||
let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) };
|
||
let Some((_device, stream)) = try_cuda() else { return Ok(()) };
|
||
|
||
let (market_buf, target_buf, num_bars) =
|
||
real_market_data(&ohlcv_dir, None, &stream)?;
|
||
|
||
let config = BranchingConfig::trading_default(
|
||
54, vec![64, 64], 32, vec![9, 3, 3],
|
||
);
|
||
let online = BranchingDuelingQNetwork::new(config.clone(), stream.clone()).unwrap();
|
||
let target_net = BranchingDuelingQNetwork::new(config, stream.clone()).unwrap();
|
||
|
||
let mut collector = GpuExperienceCollector::new(
|
||
stream.clone(),
|
||
online.vars(),
|
||
target_net.vars(),
|
||
None,
|
||
100_000.0,
|
||
0.01, 0.05,
|
||
(64, 64, 32, 32),
|
||
(54, 51, 51),
|
||
4, 50,
|
||
).unwrap();
|
||
|
||
// Sync RMSNorm weights (C51 distributional)
|
||
collector.sync_online_rmsnorm(online.vars()).unwrap();
|
||
collector.sync_target_rmsnorm(target_net.vars()).unwrap();
|
||
|
||
let n_episodes = 4_usize;
|
||
let timesteps = 50;
|
||
let episode_starts: Vec<i32> = (0..n_episodes)
|
||
.map(|i| (i * (num_bars / n_episodes)) as i32)
|
||
.collect();
|
||
|
||
let exp_config = ExperienceCollectorConfig {
|
||
n_episodes: n_episodes as i32,
|
||
timesteps_per_episode: timesteps,
|
||
total_bars: num_bars as i32,
|
||
episode_length: timesteps,
|
||
epsilon: 0.1, // 10% random ensures ≥2 distinct actions in 200 samples
|
||
gamma: 0.99,
|
||
// sigma must be large enough that NoisyNet noise dominates Q-value gaps
|
||
// even on wide networks (256-dim). With random Xavier init on a 256-wide
|
||
// network, Q-values are O(1/sqrt(256)) ≈ 0.06 apart. sigma=2.0 ensures
|
||
// noise O(sigma/sqrt(fan_in)) ≈ 0.12 exceeds this gap.
|
||
noisy_sigma_init: 2.0,
|
||
num_atoms: 51,
|
||
v_min: -10.0,
|
||
v_max: 10.0,
|
||
..Default::default()
|
||
};
|
||
|
||
let batch = collector
|
||
.collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &exp_config)
|
||
.expect("Distributional+NoisyNet failed on real data");
|
||
|
||
let expected_total = n_episodes * timesteps as usize;
|
||
assert_eq!(batch.n_episodes * batch.timesteps, expected_total);
|
||
|
||
// Download actions and rewards from GPU to host for validation
|
||
let mut actions_host = vec![0i32; expected_total];
|
||
stream.memcpy_dtoh(&batch.actions, &mut actions_host).unwrap();
|
||
let mut rewards_host = vec![0.0f32; expected_total];
|
||
stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();
|
||
|
||
// Rewards should be finite
|
||
for (i, &r) in rewards_host.iter().enumerate() {
|
||
assert!(r.is_finite(), "C51 reward[{i}] = {r} is not finite");
|
||
}
|
||
|
||
// NoisyNet should produce diverse actions (not all the same)
|
||
// Branching DQN: 9×3×3 = 81 factored actions
|
||
let mut action_counts = [0_usize; 45];
|
||
for &a in &actions_host {
|
||
if let Some(count) = action_counts.get_mut(a as usize) {
|
||
*count += 1;
|
||
}
|
||
}
|
||
let unique_actions = action_counts.iter().filter(|&&c| c > 0).count();
|
||
info!(unique_actions, "C51+NoisyNet action distribution");
|
||
assert!(
|
||
unique_actions >= 2,
|
||
"NoisyNet should produce at least 2 distinct actions, got {unique_actions}"
|
||
);
|
||
|
||
let r_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
|
||
let r_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
info!(r_min, r_max, "C51 reward range");
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
// ==========================================================================
|
||
// Test 6: E2E DQN Training Loop — validates the full training pipeline
|
||
// using local smoke-test OHLCV data with branching DQN on CUDA GPU path.
|
||
//
|
||
// Validates:
|
||
// 1. All epochs complete (no crash, no premature early stopping)
|
||
// 2. Loss decreases (gradient flow works through branching network)
|
||
// 3. All losses finite (no NaN/Inf from gradient clipping)
|
||
// 4. Q-value divergence (model develops action preferences)
|
||
// 5. Action diversity uses 81-action factored space
|
||
// 6. Gradient norms are non-zero (clipping is functional)
|
||
// ==========================================================================
|
||
|
||
#[tokio::test]
|
||
async fn smoke_e2e_dqn_training_loop() {
|
||
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
||
|
||
if gpu_vram_mb() > 0 && gpu_vram_mb() < 2048 {
|
||
return;
|
||
}
|
||
let Some(ohlcv_dir) = try_data("ohlcv") else {
|
||
return;
|
||
};
|
||
let data_dir = ohlcv_dir.to_string_lossy().to_string();
|
||
|
||
// Configure for a fast smoke run: few epochs, small batch, low warmup
|
||
let mut hyperparams = DQNHyperparameters::conservative();
|
||
hyperparams.epochs = 2;
|
||
hyperparams.batch_size = 32;
|
||
hyperparams.learning_rate = 0.0003;
|
||
hyperparams.epsilon_start = 1.0;
|
||
hyperparams.epsilon_end = 0.1;
|
||
hyperparams.epsilon_decay = 0.7;
|
||
hyperparams.early_stopping_enabled = false;
|
||
hyperparams.checkpoint_frequency = 100;
|
||
hyperparams.warmup_steps = 0;
|
||
hyperparams.min_replay_size = 50;
|
||
// CI smoke: 16 episodes x 50 timesteps = 800 experiences/epoch.
|
||
// Exercises full fused CUDA kernel (branching+C51+NoisyNets+DSR).
|
||
hyperparams.gpu_timesteps_per_episode = 10;
|
||
// CI: cap training steps to avoid full 204K-bar dataset sweep (6375->64 steps).
|
||
// 64 steps x batch_size 32 = 2048 gradient updates — sufficient to validate
|
||
// finite loss, gradient flow, and action diversity without 400s/epoch overhead.
|
||
// Curiosity disabled: kernel crashes on RTX 3050 (needs investigation)
|
||
// C51 atom count from GPU profile (replaces hardcoded VRAM if/else)
|
||
let gpu_profile = ml_core::gpu::profile::GpuProfile::load();
|
||
hyperparams.num_atoms = gpu_profile.training.num_atoms;
|
||
|
||
let checkpoint_dir = tempfile::tempdir().expect("Failed to create temp dir");
|
||
let mut trainer = DQNTrainer::new(hyperparams).expect("Failed to create DQN trainer");
|
||
|
||
info!(data_dir, "Starting E2E DQN training smoke test");
|
||
|
||
let metrics = trainer
|
||
.train(&data_dir, |epoch, checkpoint_data, is_best| {
|
||
if is_best {
|
||
let path = checkpoint_dir.path().join("smoke_e2e_best.safetensors");
|
||
std::fs::write(&path, &checkpoint_data)?;
|
||
Ok(path.to_string_lossy().to_string())
|
||
} else {
|
||
Ok(format!("epoch_{epoch}"))
|
||
}
|
||
})
|
||
.await
|
||
.expect("DQN training failed");
|
||
|
||
// === ASSERTIONS ===
|
||
|
||
// 1. All epochs completed
|
||
assert_eq!(
|
||
metrics.epochs_trained, 2,
|
||
"Expected 2 epochs, got {}",
|
||
metrics.epochs_trained
|
||
);
|
||
|
||
// 2. Loss decreases (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 initial_loss = loss_history.first().copied().unwrap_or(f64::NAN);
|
||
let final_loss = loss_history.last().copied().unwrap_or(f64::NAN);
|
||
info!(initial_loss, final_loss, "Loss trajectory");
|
||
|
||
// 3. All losses must be finite (NaN would indicate broken gradient clipping)
|
||
for (i, &loss) in loss_history.iter().enumerate() {
|
||
assert!(
|
||
loss.is_finite(),
|
||
"Loss at epoch {i} is not finite: {loss}"
|
||
);
|
||
}
|
||
|
||
// 4. Q-value divergence (model has preferences)
|
||
let avg_q = metrics
|
||
.additional_metrics
|
||
.get("avg_q_value")
|
||
.copied()
|
||
.unwrap_or(0.0);
|
||
info!(avg_q, "Average Q-value");
|
||
|
||
// 5. Action diversity uses 81-action factored space
|
||
let action_space = metrics
|
||
.additional_metrics
|
||
.get("action_space_size")
|
||
.copied()
|
||
.unwrap_or(5.0);
|
||
let action_diversity = metrics
|
||
.additional_metrics
|
||
.get("action_diversity")
|
||
.copied()
|
||
.unwrap_or(0.0);
|
||
let total_actions = metrics
|
||
.additional_metrics
|
||
.get("total_actions")
|
||
.copied()
|
||
.unwrap_or(0.0);
|
||
info!(action_space, action_diversity, total_actions, "Action metrics");
|
||
// Branching DQN: action_space=81 (9×3×3) when factored counts populated,
|
||
// or 9 (exposure-only fallback) for very short runs.
|
||
assert!(
|
||
action_space >= 9.0 || total_actions == 0.0,
|
||
"Branching DQN should report >=9 action space, got {action_space}"
|
||
);
|
||
|
||
// 6. Gradient norms should be non-zero (clipping is functional)
|
||
let avg_grad_norm = metrics
|
||
.additional_metrics
|
||
.get("avg_gradient_norm")
|
||
.copied()
|
||
.unwrap_or(0.0);
|
||
info!(avg_grad_norm, "Average gradient norm");
|
||
// With the TensorId fallback fix, grad norms should be reported correctly
|
||
assert!(
|
||
avg_grad_norm.is_finite(),
|
||
"Gradient norm should be finite, got {avg_grad_norm}"
|
||
);
|
||
|
||
// === REPORT ===
|
||
let loss_change = if initial_loss > 0.0 && initial_loss.is_finite() {
|
||
(1.0 - final_loss / initial_loss) * 100.0
|
||
} else {
|
||
f64::NAN
|
||
};
|
||
info!(
|
||
epochs = metrics.epochs_trained,
|
||
initial_loss, final_loss, loss_change,
|
||
avg_q, action_space, action_diversity,
|
||
avg_grad_norm,
|
||
training_time_s = metrics.training_time_seconds,
|
||
"E2E DQN training smoke test report"
|
||
);
|
||
}
|