From e7ce4395e80b4796d8f31243ce31d63778fa0d95 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 16 May 2026 13:55:32 +0200 Subject: [PATCH] perf(precompute): parallel trades load + predecoded sidecar cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Flamegraph of precompute_features on 1Q ES showed 62% of CPU time in zstd decompression, 6% in DBN FSM parsing, and only 2% in the actual feature math — single-threaded zstd was the bottleneck, not compute. Two fixes: 1. Per-quarter parallelism on the volume-bar trades loop (was sequential `for file in &trade_files`); brings it in line with the OFI path that already used par_iter. 2. Predecoded sidecar cache in `crates/ml-features/src/predecoded.rs`: first call to a `.dbn.zst` writes a bincode'd Vec or Vec under `/predecoded/`. Subsequent calls deserialize the sidecar and skip zstd entirely. An mtime+size header self-invalidates the sidecar when the source changes — no manual flush needed when a quarter is re-downloaded. Local 1Q ES results: - cold (writes sidecar): 40.7s (was 39.3s; +1.4s for write) - warm (HIT): 4.7s (8.7× faster) - zstd in flat perf: 62% → 0% of CPU samples - sidecar disk per Q: ~150MB The sidecar layer also auto-dedupes within a single run: the OFI section re-loads trades, but the second call hits the sidecar that the volume-bar section wrote moments earlier. CLI: `--rebuild-predecoded` purges sidecars for cold-path testing or after a wire-format change to Mbp10Snapshot / DbnTrade. Sidecars also self-invalidate on format-version mismatch so old caches are skipped silently rather than mis-deserializing. Co-Authored-By: Claude Opus 4.7 --- Cargo.lock | 1 + crates/ml-features/Cargo.toml | 1 + crates/ml-features/src/lib.rs | 1 + crates/ml-features/src/predecoded.rs | 281 ++++++++++++++++++++++ crates/ml-features/src/trades_loader.rs | 3 +- crates/ml/examples/precompute_features.rs | 112 ++++++--- 6 files changed, 363 insertions(+), 36 deletions(-) create mode 100644 crates/ml-features/src/predecoded.rs diff --git a/Cargo.lock b/Cargo.lock index eb6ee110c..efc66e26b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6281,6 +6281,7 @@ dependencies = [ "rayon", "serde", "serde_json", + "tempfile", "tokio", "toml", "tracing", diff --git a/crates/ml-features/Cargo.toml b/crates/ml-features/Cargo.toml index e35044574..2e8d2ea87 100644 --- a/crates/ml-features/Cargo.toml +++ b/crates/ml-features/Cargo.toml @@ -49,6 +49,7 @@ rand.workspace = true tokio = { workspace = true, features = ["test-util", "macros"] } approx.workspace = true toml.workspace = true +tempfile = "3" [lints] workspace = true diff --git a/crates/ml-features/src/lib.rs b/crates/ml-features/src/lib.rs index e196dd5a2..bb987c531 100644 --- a/crates/ml-features/src/lib.rs +++ b/crates/ml-features/src/lib.rs @@ -43,6 +43,7 @@ pub mod ofi_calculator; pub mod order_events; pub mod pipeline; pub mod position_features; +pub mod predecoded; pub mod price_features; pub mod regime_adx; pub mod return_moments; diff --git a/crates/ml-features/src/predecoded.rs b/crates/ml-features/src/predecoded.rs new file mode 100644 index 000000000..eff1b6eba --- /dev/null +++ b/crates/ml-features/src/predecoded.rs @@ -0,0 +1,281 @@ +//! Predecoded sidecar cache for parsed DBN files. +//! +//! Wraps `load_trades_sync` and `parse_mbp10_streaming` with an on-disk cache. +//! First call decodes zstd-DBN and writes a sidecar; subsequent calls +//! deserialize the sidecar directly. zstd decompression was 62% of +//! precompute_features wall time on the 2026-05-16 1Q ES profile — moving it +//! to a one-shot cache eliminates that cost for all repeat runs. +//! +//! # Cache key +//! +//! The sidecar header records the source file's mtime and size. A sidecar is +//! considered valid only when both match the source on the current call; +//! re-downloading a quarter (which bumps mtime + usually size) silently +//! invalidates the sidecar. + +use std::fs::{self, File}; +use std::io::{BufReader, BufWriter, Read, Write}; +use std::path::{Path, PathBuf}; +use std::time::SystemTime; + +use bincode; +use data::providers::databento::dbn_parser::DbnParser; +use data::providers::databento::mbp10::Mbp10Snapshot; +use serde::{Deserialize, Serialize}; + +use crate::trades_loader::{load_trades_sync, DbnTrade}; +use crate::MLError; + +const MAGIC: u32 = u32::from_le_bytes(*b"PCD1"); +const VERSION: u32 = 1; + +fn sidecar_path(source: &Path, predecoded_dir: &Path, kind: &str) -> PathBuf { + let basename = source.file_name().unwrap_or_default().to_string_lossy(); + predecoded_dir.join(format!("{basename}.{kind}.predecoded.bin")) +} + +fn source_mtime_size(source: &Path) -> std::io::Result<(u128, u64)> { + let meta = fs::metadata(source)?; + let mtime = meta + .modified()? + .duration_since(SystemTime::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + Ok((mtime, meta.len())) +} + +fn write_header(w: &mut W, mtime_nanos: u128, size: u64) -> std::io::Result<()> { + w.write_all(&MAGIC.to_le_bytes())?; + w.write_all(&VERSION.to_le_bytes())?; + w.write_all(&mtime_nanos.to_le_bytes())?; + w.write_all(&size.to_le_bytes())?; + Ok(()) +} + +fn read_header_and_validate(r: &mut R, source: &Path) -> std::io::Result { + let mut magic = [0u8; 4]; + r.read_exact(&mut magic)?; + if u32::from_le_bytes(magic) != MAGIC { + return Ok(false); + } + let mut version = [0u8; 4]; + r.read_exact(&mut version)?; + if u32::from_le_bytes(version) != VERSION { + return Ok(false); + } + let mut mtime_bytes = [0u8; 16]; + r.read_exact(&mut mtime_bytes)?; + let mut size_bytes = [0u8; 8]; + r.read_exact(&mut size_bytes)?; + let stored_mtime = u128::from_le_bytes(mtime_bytes); + let stored_size = u64::from_le_bytes(size_bytes); + let (current_mtime, current_size) = source_mtime_size(source)?; + Ok(stored_mtime == current_mtime && stored_size == current_size) +} + +fn try_read_sidecar Deserialize<'de>>( + source: &Path, + sidecar: &Path, +) -> std::io::Result> { + if !sidecar.exists() { + return Ok(None); + } + let f = File::open(sidecar)?; + let mut r = BufReader::new(f); + if !read_header_and_validate(&mut r, source)? { + return Ok(None); + } + let payload: T = bincode::deserialize_from(r) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + Ok(Some(payload)) +} + +fn write_sidecar( + source: &Path, + sidecar: &Path, + payload: &T, +) -> std::io::Result<()> { + if let Some(parent) = sidecar.parent() { + fs::create_dir_all(parent)?; + } + let (mtime, size) = source_mtime_size(source)?; + let f = File::create(sidecar)?; + let mut w = BufWriter::new(f); + write_header(&mut w, mtime, size)?; + bincode::serialize_into(&mut w, payload) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + w.flush()?; + Ok(()) +} + +/// Load trades from a DBN file, using a predecoded sidecar when available. +/// +/// On miss (no sidecar / stale sidecar): decode `.dbn.zst` via +/// `load_trades_sync`, persist `Vec` as a sidecar, return trades. +/// On hit: deserialize sidecar (skips zstd). +/// +/// Sidecar write failures are logged and ignored — the caller still receives +/// the freshly-decoded `Vec`. +pub fn load_or_predecode_trades( + source: &Path, + predecoded_dir: &Path, +) -> Result, MLError> { + let sidecar = sidecar_path(source, predecoded_dir, "trades"); + match try_read_sidecar::>(source, &sidecar) { + Ok(Some(trades)) => { + tracing::info!( + "Predecoded HIT: {} ({} trades)", + sidecar.display(), + trades.len() + ); + return Ok(trades); + } + Ok(None) => {} + Err(e) => { + tracing::warn!( + "Predecoded trades sidecar read failed for {}: {} — falling back to zstd", + sidecar.display(), + e + ); + } + } + let trades = load_trades_sync(source)?; + if let Err(e) = write_sidecar(source, &sidecar, &trades) { + tracing::warn!( + "Failed to write predecoded trades sidecar {}: {}", + sidecar.display(), + e + ); + } else { + tracing::info!( + "Predecoded WRITE: {} ({} trades)", + sidecar.display(), + trades.len() + ); + } + Ok(trades) +} + +/// Load MBP-10 snapshots from a DBN file, using a predecoded sidecar when available. +/// +/// On miss: streams the file via `DbnParser::parse_mbp10_streaming`, accumulates +/// `Vec`, persists a sidecar, returns the snapshots. On hit: +/// deserializes the sidecar (skips zstd). +pub fn load_or_predecode_mbp10( + source: &Path, + predecoded_dir: &Path, +) -> Result, MLError> { + let sidecar = sidecar_path(source, predecoded_dir, "mbp10"); + match try_read_sidecar::>(source, &sidecar) { + Ok(Some(snapshots)) => { + tracing::info!( + "Predecoded HIT: {} ({} snapshots)", + sidecar.display(), + snapshots.len() + ); + return Ok(snapshots); + } + Ok(None) => {} + Err(e) => { + tracing::warn!( + "Predecoded mbp10 sidecar read failed for {}: {} — falling back to zstd", + sidecar.display(), + e + ); + } + } + let parser = DbnParser::new() + .map_err(|e| MLError::InsufficientData(format!("DbnParser init failed: {e}")))?; + let mut snapshots = Vec::new(); + parser + .parse_mbp10_streaming(source, 100, |snap| { + snapshots.push(snap.clone()); + }) + .map_err(|e| MLError::InsufficientData(format!("parse_mbp10_streaming failed: {e}")))?; + if let Err(e) = write_sidecar(source, &sidecar, &snapshots) { + tracing::warn!( + "Failed to write predecoded mbp10 sidecar {}: {}", + sidecar.display(), + e + ); + } else { + tracing::info!( + "Predecoded WRITE: {} ({} snapshots)", + sidecar.display(), + snapshots.len() + ); + } + Ok(snapshots) +} + +/// Remove all `*.predecoded.bin` files from `predecoded_dir`. Idempotent. +/// +/// Used by the `--rebuild-predecoded` CLI flag to force a fresh decode of every +/// source file. Returns the number of files removed. +pub fn purge_predecoded_dir(predecoded_dir: &Path) -> std::io::Result { + if !predecoded_dir.exists() { + return Ok(0); + } + let mut removed = 0usize; + for entry in fs::read_dir(predecoded_dir)? { + let entry = entry?; + let path = entry.path(); + if path + .file_name() + .and_then(|n| n.to_str()) + .is_some_and(|n| n.ends_with(".predecoded.bin")) + { + fs::remove_file(&path)?; + removed += 1; + } + } + Ok(removed) +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + #[test] + fn sidecar_path_uses_kind_suffix() { + let p = sidecar_path( + Path::new("/data/ES.FUT_2024-Q1.dbn.zst"), + Path::new("/tmp/predecoded"), + "trades", + ); + assert_eq!( + p, + PathBuf::from("/tmp/predecoded/ES.FUT_2024-Q1.dbn.zst.trades.predecoded.bin") + ); + } + + #[test] + fn header_round_trip_validates_source_metadata() { + let tmp = TempDir::new().unwrap(); + let source = tmp.path().join("source.bin"); + fs::write(&source, b"hello world").unwrap(); + let sidecar = tmp.path().join("source.bin.test.predecoded.bin"); + let payload: Vec = vec![1, 2, 3]; + write_sidecar(&source, &sidecar, &payload).unwrap(); + + let read: Option> = try_read_sidecar(&source, &sidecar).unwrap(); + assert_eq!(read, Some(vec![1, 2, 3])); + + // Mutate source → sidecar must be invalidated. + fs::write(&source, b"different content here").unwrap(); + let read: Option> = try_read_sidecar(&source, &sidecar).unwrap(); + assert_eq!(read, None); + } + + #[test] + fn purge_removes_only_predecoded_files() { + let tmp = TempDir::new().unwrap(); + fs::write(tmp.path().join("a.trades.predecoded.bin"), b"x").unwrap(); + fs::write(tmp.path().join("b.mbp10.predecoded.bin"), b"y").unwrap(); + fs::write(tmp.path().join("unrelated.txt"), b"z").unwrap(); + let removed = purge_predecoded_dir(tmp.path()).unwrap(); + assert_eq!(removed, 2); + assert!(tmp.path().join("unrelated.txt").exists()); + } +} diff --git a/crates/ml-features/src/trades_loader.rs b/crates/ml-features/src/trades_loader.rs index aa1934451..fad4eac42 100644 --- a/crates/ml-features/src/trades_loader.rs +++ b/crates/ml-features/src/trades_loader.rs @@ -10,11 +10,12 @@ use std::path::Path; use dbn::decode::{DbnDecoder, DecodeRecordRef}; use dbn::RecordRefEnum; +use serde::{Deserialize, Serialize}; use crate::MLError; /// A single parsed trade from a DBN file. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] pub struct DbnTrade { /// Event timestamp in nanoseconds since Unix epoch pub timestamp: u64, diff --git a/crates/ml/examples/precompute_features.rs b/crates/ml/examples/precompute_features.rs index 8bf55f24b..a80e18d23 100644 --- a/crates/ml/examples/precompute_features.rs +++ b/crates/ml/examples/precompute_features.rs @@ -184,6 +184,17 @@ struct Opts { /// Only used when `--row-unit snapshot`. #[arg(long, default_value_t = 100)] snapshot_warmup: usize, + + /// Purge the predecoded-sidecar cache before loading. + /// + /// Predecoded sidecars (`..predecoded.bin` under + /// `/predecoded/`) skip zstd-DBN decompression on subsequent + /// precompute runs. They self-invalidate when the source file's mtime or + /// size changes, so this flag is normally unnecessary — pass it only after + /// a format change to `Mbp10Snapshot` / `DbnTrade` (bumps fail to + /// deserialize) or when explicitly testing the cold path. + #[arg(long)] + rebuild_predecoded: bool, } // --------------------------------------------------------------------------- @@ -217,6 +228,20 @@ async fn main() -> Result<()> { let mbp10_dir = opts.mbp10_data_dir.as_ref().map(PathBuf::from); let trades_dir = opts.trades_data_dir.as_ref().map(PathBuf::from); + // Predecoded-sidecar cache lives under output_dir/predecoded/ so it sits + // next to the .fxcache outputs (same writability assumptions). First call + // for any source `.dbn.zst` pays zstd; subsequent calls deserialize the + // sidecar and skip decompression entirely (zstd was 62% of wall time on + // the 2026-05-16 1Q ES profile). + let predecoded_dir = output_dir.join("predecoded"); + std::fs::create_dir_all(&predecoded_dir) + .with_context(|| format!("Failed to create predecoded dir {}", predecoded_dir.display()))?; + if opts.rebuild_predecoded { + let removed = ml::features::predecoded::purge_predecoded_dir(&predecoded_dir) + .with_context(|| format!("Failed to purge {}", predecoded_dir.display()))?; + info!("--rebuild-predecoded: purged {} sidecar(s) from {}", removed, predecoded_dir.display()); + } + // ── Validate inputs ────────────────────────────────────────────────────── if !data_dir.exists() { anyhow::bail!("Data directory not found: {}", data_dir.display()); @@ -358,23 +383,40 @@ async fn main() -> Result<()> { // Load trades PER FILE and filter front-month within each file. // Each quarterly DBN file has a different front-month contract (e.g. ESH24, ESM24). // Filtering globally would select only ONE quarter's contract. - let mut all_front_month_trades: Vec = Vec::new(); - let mut total_raw = 0_usize; - for file in &trade_files { - match ml::features::load_trades_sync(file) { - Ok(trades) => { - let raw_count = trades.len(); - total_raw += raw_count; - let filtered = ml::features::filter_front_month(&trades); - info!(" {} -> {} trades, {} front-month", - file.file_name().unwrap_or_default().to_string_lossy(), - raw_count, filtered.len()); - all_front_month_trades.extend(filtered); - } - Err(e) => { - tracing::warn!(" Failed {:?}: {e}", file.file_name()); - } - } + // + // Parallel across quarters via rayon; each file goes through the + // predecoded-sidecar cache so a re-run of `precompute_features` skips + // zstd decompression entirely. + let per_file: Vec<(usize, Vec)> = { + use rayon::prelude::*; + trade_files + .par_iter() + .filter_map(|file| { + match ml::features::predecoded::load_or_predecode_trades(file, &predecoded_dir) { + Ok(trades) => { + let raw_count = trades.len(); + let filtered = ml::features::filter_front_month(&trades); + info!( + " {} -> {} trades, {} front-month", + file.file_name().unwrap_or_default().to_string_lossy(), + raw_count, + filtered.len() + ); + Some((raw_count, filtered)) + } + Err(e) => { + tracing::warn!(" Failed {:?}: {e}", file.file_name()); + None + } + } + }) + .collect() + }; + let total_raw: usize = per_file.iter().map(|(r, _)| r).sum(); + let mut all_front_month_trades: Vec = + Vec::with_capacity(per_file.iter().map(|(_, t)| t.len()).sum()); + for (_, t) in per_file { + all_front_month_trades.extend(t); } all_front_month_trades.sort_by_key(|t| t.timestamp); info!("Total trades: {} raw, {} front-month", total_raw, all_front_month_trades.len()); @@ -521,8 +563,7 @@ async fn main() -> Result<()> { const OFI_DIM: usize = ml_core::state_layout::OFI_DIM; let t2 = Instant::now(); let ofi: Vec<[f64; OFI_DIM]> = if let Some(ref mbp10_path) = mbp10_dir { - use ml::features::mbp10_loader::load_ofi_features_parallel; - use ml::features::trades_loader::load_trades_sync; + use ml::features::ofi_calculator::OFICalculator; info!("Loading MBP-10 snapshots from {}...", mbp10_path.display()); let mut mbp10_files = collect_dbn_files_recursive(mbp10_path); @@ -532,21 +573,17 @@ async fn main() -> Result<()> { info!("No MBP-10 files found, OFI will be zeros"); vec![[0.0; OFI_DIM]; n] } else { - // Load MBP-10 snapshots (parallel) - use data::providers::databento::dbn_parser::DbnParser; + // Load MBP-10 snapshots (parallel, with predecoded sidecar) let per_file: Vec> = { use rayon::prelude::*; mbp10_files.par_iter().filter_map(|file| { - let parser = match DbnParser::new() { - Ok(p) => p, - Err(e) => { tracing::warn!("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()); + match ml::features::predecoded::load_or_predecode_mbp10(file, &predecoded_dir) { + Ok(snapshots) => { + info!( + " {} -> {} snapshots", + file.file_name().unwrap_or_default().to_string_lossy(), + snapshots.len() + ); Some(snapshots) } Err(e) => { @@ -562,24 +599,29 @@ async fn main() -> Result<()> { all_snapshots.sort_by_key(|s| s.timestamp); info!("Loaded {} MBP-10 snapshots in {:.1}s", all_snapshots.len(), t2.elapsed().as_secs_f64()); - // Load trades (parallel) and filter front-month per-file. + // Load trades (parallel, with predecoded sidecar) and filter + // front-month per-file. // // Pre-2026-05-09: this branch loaded trades unfiltered, leaking // off-contract trade activity into the per-bar OFI/VPIN/Kyle's // Lambda computations during contract rollover windows. Volume - // bar formation already filters front-month at line 354 of this - // file; the OFI side did not. Audit confirmed the latent bug + // bar formation already filters front-month at the trades-for-bars + // load above; the OFI side did not. Audit confirmed the latent bug // affected every prior MBP-10+trades production run. // // Fix: mirror the per-file `filter_front_month` call from the // volume bar path so OFI windows only see in-month trades. + // + // Trade loads go through `load_or_predecode_trades` so the sidecar + // cache is shared with the bar-formation path above (one decode of + // each `.dbn.zst` per process — the OFI re-load is a sidecar HIT). let all_trades = if let Some(ref tdir) = trades_dir { let mut trade_files = collect_dbn_files_recursive(tdir); trade_files.sort(); let per_file_trades: Vec> = { use rayon::prelude::*; trade_files.par_iter().filter_map(|path| { - match load_trades_sync(path) { + match ml::features::predecoded::load_or_predecode_trades(path, &predecoded_dir) { Ok(t) => { let raw_count = t.len(); let filtered = ml::features::filter_front_month(&t);