gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(data): parallelize MBP-10 and trade file loading with rayon

Browse Source 
Both walk-forward and full-training data loading paths now use rayon
par_iter to parse MBP-10 .dbn files concurrently. Previously 9 files
were parsed sequentially (~4 min on H100); parallel loading gives
~7-9x speedup proportional to file count. Trade file loading also
parallelized. Removed dead collect_dbn_files helper (superseded by
collect_dbn_files_recursive at module scope).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-12 21:52:45 +01:00
parent a0a85f2011
commit 12d85993cf
1 changed files with 123 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

215
crates/ml/src/trainers/dqn/data_loading.rs
View File

@@ -373,57 +373,57 @@ impl DQNTrainer {
None
};
// Recursive file discovery (data may be in symbol subdirectories)
fn collect_dbn_files(dir: &Path, out: &mut Vec<std::path::PathBuf>) {
if let Ok(entries) = std::fs::read_dir(dir) {
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
collect_dbn_files(&path, out);
} else {
let name = path.to_string_lossy();
if name.ends_with(".dbn") || name.ends_with(".dbn.zst") {
out.push(path);
}
}
}
}
}
// WAVE 2-A2: Load MBP-10 snapshots for OFI calculation
// WAVE 2-A2: Load MBP-10 snapshots for OFI calculation (parallel)
use data::providers::databento::dbn_parser::DbnParser;
let mbp10_dir_path = self.hyperparams.mbp10_data_dir.as_ref().map(Path::new);
let mbp10_snapshots = if let Some(mbp10_dir) = mbp10_dir_path {
if mbp10_dir.exists() {
info!("📊 Loading MBP-10 order book snapshots from {:?}...", mbp10_dir);
let mut all_snapshots = Vec::new();
let mut dbn_files = Vec::new();
collect_dbn_files(mbp10_dir, &mut dbn_files);
let mut dbn_files: Vec<_> = collect_dbn_files_recursive(mbp10_dir);
dbn_files.sort();
if !dbn_files.is_empty() {
let parser = DbnParser::new()
.context("Failed to create DBN parser for MBP-10 data")?;
for path in &dbn_files {
match parser.parse_mbp10_file(path).await {
Ok(mut snaps) => {
info!(" ✅ Loaded {} snapshots from {:?}", snaps.len(), path.file_name());
all_snapshots.append(&mut snaps);
}
Err(e) => {
warn!(" ⚠️ Failed to load {:?}: {}", path.file_name(), e);
}
}
}
}
let start_mbp10 = std::time::Instant::now();
let n_files = dbn_files.len();
info!("📊 Loading MBP-10 snapshots from {} files (parallel)...", n_files);
if all_snapshots.is_empty() {
warn!("⚠️ No MBP-10 snapshots loaded. OFI features will be zeros.");
None
} else {
let per_file: Vec<Vec<_>> = {
use rayon::prelude::*;
dbn_files.par_iter().filter_map(|file| {
let parser = match DbnParser::new() {
Ok(p) => p,
Err(e) => { warn!(" DBN parser init failed: {e}"); return None; }
};
let mut snapshots = Vec::new();
match parser.parse_mbp10_streaming(file, 100, |snap| {
snapshots.push(snap.clone());
}) {
Ok(_) => {
info!(" ✅ {} -> {} snapshots",
file.file_name().unwrap_or_default().to_string_lossy(),
snapshots.len());
Some(snapshots)
}
Err(e) => {
warn!(" ⚠️ Failed to load {:?}: {e}", file.file_name());
None
}
}
}).collect()
};
let total_snaps: usize = per_file.iter().map(|v| v.len()).sum();
let mut all_snapshots = Vec::with_capacity(total_snaps);
for snaps in per_file {
all_snapshots.extend(snaps);
}
all_snapshots.sort_by_key(|s| s.timestamp);
info!("✅ Total MBP-10 snapshots loaded: {} (sorted by timestamp)", all_snapshots.len());
Some(all_snapshots)
let mbp10_secs = start_mbp10.elapsed().as_secs_f64();
info!("✅ Loaded {} MBP-10 snapshots from {} files in {:.1}s", total_snaps, n_files, mbp10_secs);
if all_snapshots.is_empty() { None } else { Some(all_snapshots) }
} else {
warn!("⚠️ No MBP-10 .dbn files found. OFI features will be zeros.");
None
}
} else {
warn!("⚠️ MBP-10 directory not found at {:?}. OFI features disabled.", mbp10_dir);
@@ -443,36 +443,44 @@ impl DQNTrainer {
feature_vectors.len()
);
// Load trade data for real VPIN/Kyle's Lambda (Option B)
// Load trade data for real VPIN/Kyle's Lambda (parallel)
let trades = if let Some(ref trades_dir_str) = self.hyperparams.trades_data_dir {
use crate::features::trades_loader::load_trades_sync;
let trades_dir = Path::new(trades_dir_str);
if trades_dir.exists() {
info!("📊 Loading trade data from {:?} for VPIN/Kyle's Lambda...", trades_dir);
let mut all_trades = Vec::new();
let mut trade_files = Vec::new();
collect_dbn_files(trades_dir, &mut trade_files);
let mut trade_files: Vec<_> = collect_dbn_files_recursive(trades_dir);
trade_files.sort();
for path in &trade_files {
match load_trades_sync(path) {
Ok(mut t) => {
info!(" ✅ Loaded {} trades from {:?}", t.len(), path.file_name());
all_trades.append(&mut t);
}
Err(e) => {
warn!(" ⚠️ Failed to load trades from {:?}: {}", path.file_name(), e);
}
if !trade_files.is_empty() {
info!("📊 Loading trade data from {} files (parallel)...", trade_files.len());
let per_file_trades: Vec<Vec<_>> = {
use rayon::prelude::*;
trade_files.par_iter().filter_map(|path| {
match load_trades_sync(path) {
Ok(t) => {
info!(" ✅ {} trades from {:?}", t.len(), path.file_name());
Some(t)
}
Err(e) => {
warn!(" ⚠️ Failed to load {:?}: {e}", path.file_name());
None
}
}
}).collect()
};
let mut all_trades = Vec::new();
for t in per_file_trades {
all_trades.extend(t);
}
if all_trades.is_empty() {
warn!("⚠️ No trades loaded. VPIN/Kyle's Lambda will use tick-rule proxy.");
None
} else {
all_trades.sort_by_key(|t| t.timestamp);
info!("✅ Total trades loaded: {} (sorted by timestamp)", all_trades.len());
Some(all_trades)
}
}
if all_trades.is_empty() {
warn!("⚠️ No trades loaded. VPIN/Kyle's Lambda will use tick-rule proxy.");
None
} else {
all_trades.sort_by_key(|t| t.timestamp);
info!("✅ Total trades loaded: {} (sorted by timestamp)", all_trades.len());
Some(all_trades)
None
}
} else {
warn!("⚠️ Trades directory not found at {:?}. Using tick-rule proxy.", trades_dir);
@@ -671,6 +679,7 @@ impl DQNTrainer {
);
// Compute per-bar OFI features from MBP-10 snapshots (8 features per bar)
// Uses rayon par_iter to load files in parallel (~7-9x speedup over sequential)
if let Some(ref mbp10_dir_str) = self.hyperparams.mbp10_data_dir {
use crate::features::ofi_calculator::OFICalculator;
use crate::features::mbp10_loader::get_snapshots_for_timestamp;
@@ -679,46 +688,67 @@ impl DQNTrainer {
let mbp10_dir = Path::new(mbp10_dir_str);
if mbp10_dir.exists() {
info!("Loading MBP-10 snapshots for per-bar OFI...");
let mut dbn_files = Vec::new();
collect_dbn_files_recursive(mbp10_dir)
.iter()
.for_each(|f| dbn_files.push(f.clone()));
let mut dbn_files: Vec<_> = collect_dbn_files_recursive(mbp10_dir);
dbn_files.sort();
if !dbn_files.is_empty() {
let mut all_snapshots = Vec::new();
if let Ok(ref parser) = DbnParser::new() {
for file in &dbn_files {
match parser.parse_mbp10_file(file).await {
Ok(snaps) => {
info!(" Loaded {} MBP-10 snapshots from {:?}", snaps.len(), file.file_name());
all_snapshots.extend(snaps);
let start_mbp10 = std::time::Instant::now();
let n_files = dbn_files.len();
info!("Loading MBP-10 snapshots from {n_files} files (parallel)...");
// Parse all MBP-10 files in parallel with rayon.
// Each file produces a sorted Vec of snapshots; we merge after.
let per_file: Vec<Vec<_>> = {
use rayon::prelude::*;
dbn_files.par_iter().filter_map(|file| {
let parser = match DbnParser::new() {
Ok(p) => p,
Err(e) => { warn!(" DBN parser init failed: {e}"); return None; }
};
let mut snapshots = Vec::new();
match parser.parse_mbp10_streaming(file, 100, |snap| {
snapshots.push(snap.clone());
}) {
Ok(_) => {
info!(" {} -> {} snapshots",
file.file_name().unwrap_or_default().to_string_lossy(),
snapshots.len());
Some(snapshots)
}
Err(e) => {
warn!(" Failed to load {:?}: {e}", file.file_name());
None
}
Err(e) => warn!(" Failed to load MBP-10 from {:?}: {}", file.file_name(), e),
}
}
}).collect()
};
let total_snaps: usize = per_file.iter().map(|v| v.len()).sum();
let mut all_snapshots = Vec::with_capacity(total_snaps);
for snaps in per_file {
all_snapshots.extend(snaps);
}
all_snapshots.sort_by_key(|s| s.timestamp);
let mbp10_secs = start_mbp10.elapsed().as_secs_f64();
info!("Loaded {total_snaps} MBP-10 snapshots from {n_files} files in {mbp10_secs:.1}s");
if !all_snapshots.is_empty() {
all_snapshots.sort_by_key(|s| s.timestamp);
info!("Computing per-bar OFI from {} MBP-10 snapshots...", all_snapshots.len());
// Load trade data for VPIN/Kyle's Lambda enrichment
// Load trade data in parallel for VPIN/Kyle's Lambda enrichment
let trades = if let Some(ref trades_dir_str) = self.hyperparams.trades_data_dir {
use crate::features::trades_loader::load_trades_sync;
let trades_dir = Path::new(trades_dir_str);
if trades_dir.exists() {
let mut all_trades = Vec::new();
let mut trade_files = Vec::new();
collect_dbn_files_recursive(trades_dir)
.iter()
.for_each(|f| trade_files.push(f.clone()));
let mut trade_files: Vec<_> = collect_dbn_files_recursive(trades_dir);
trade_files.sort();
for path in &trade_files {
if let Ok(mut t) = load_trades_sync(path) {
all_trades.append(&mut t);
}
let per_file_trades: Vec<Vec<_>> = {
use rayon::prelude::*;
trade_files.par_iter().filter_map(|path| {
load_trades_sync(path).ok()
}).collect()
};
let mut all_trades = Vec::new();
for t in per_file_trades {
all_trades.extend(t);
}
(!all_trades.is_empty()).then(|| {
all_trades.sort_by_key(|t| t.timestamp);
@@ -731,6 +761,7 @@ impl DQNTrainer {
None
};
info!("Computing per-bar OFI from {} MBP-10 snapshots...", all_snapshots.len());
let mut ofi_calculator = OFICalculator::new();
let mut ofi_per_bar = Vec::with_capacity(feature_vectors.len());
const WARMUP: usize = 50;