From 586d1e78207276fb43000a49e5e06d42cf415534 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 00:19:06 +0200 Subject: [PATCH] perf(ml-alpha): cache loaded file across next_sequence calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MultiHorizonLoader was calling load_or_predecode_mbp10 on every next_sequence() call, deserializing millions of MBP-10 snapshots per sequence. With 8000 sequences and 9 files this gave ~50+ hour training time on what should be IO-trivial work. Now keep one file cached (LoadedFile { snapshots, labels_full }), yield ceil(n_max_sequences / n_files) sequences from it before advancing. Per-horizon labels are computed once per file load and sliced cheaply per anchor. With 8000/9: ~890 loads → 9 loads. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/data/loader.rs | 142 ++++++++++++++++++----------- 1 file changed, 89 insertions(+), 53 deletions(-) diff --git a/crates/ml-alpha/src/data/loader.rs b/crates/ml-alpha/src/data/loader.rs index d06343754..76aa849d3 100644 --- a/crates/ml-alpha/src/data/loader.rs +++ b/crates/ml-alpha/src/data/loader.rs @@ -48,12 +48,32 @@ pub struct LabeledSequence { pub labels: [Vec; 5], } +/// Loaded-file cache. Holds the deserialized snapshots + pre-computed +/// per-horizon labels so subsequent `next_sequence()` calls on the +/// same file avoid the multi-second bincode deserialization. +struct LoadedFile { + snapshots: Vec, + /// `labels_full[h]` holds the full-stream label vector for horizon + /// `cfg.horizons[h]`, indexed the same as `snapshots`. Sliced per + /// anchor on each `next_sequence()` call. + labels_full: [Vec; 5], +} + pub struct MultiHorizonLoader { cfg: MultiHorizonLoaderConfig, files: Vec, rng: ChaCha8Rng, file_idx: usize, yielded: usize, + + /// Currently-loaded file. Cleared and replaced when we advance. + current: Option, + /// Sequences yielded from `current` since it was loaded. + yielded_from_current: usize, + /// How many sequences to yield from each file before advancing. Set + /// from `n_max_sequences / n_files` (rounded up) at construction, so + /// each file is loaded exactly once per "pass" through the loader. + sequences_per_file: usize, } impl MultiHorizonLoader { @@ -75,12 +95,20 @@ impl MultiHorizonLoader { ); let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed); files.shuffle(&mut rng); + let n_files = files.len(); + // Yield (n_max_sequences / n_files) sequences per file, rounded + // up so we always honor n_max_sequences. With 8000 max and 9 + // files: 8000/9 = 889/file, total 9 loads instead of 8000. + let sequences_per_file = cfg.n_max_sequences.div_ceil(n_files).max(1); Ok(Self { cfg: cfg.clone(), files, rng, file_idx: 0, yielded: 0, + current: None, + yielded_from_current: 0, + sequences_per_file, }) } @@ -96,64 +124,72 @@ impl MultiHorizonLoader { if self.yielded >= self.cfg.n_max_sequences { return Ok(None); } - // Loop until we either yield or exhaust the file list (re-shuffle and - // try again if files are exhausted but we haven't hit n_max_sequences). - for _ in 0..self.files.len() * 2 { - if self.file_idx >= self.files.len() { - self.files.shuffle(&mut self.rng); - self.file_idx = 0; - } - let path = self.files[self.file_idx].clone(); - self.file_idx += 1; + let max_horizon = *self.cfg.horizons.iter().max().expect("non-empty horizons"); - let snapshots = load_or_predecode_mbp10(&path, &self.cfg.predecoded_dir) - .with_context(|| format!("load mbp10 {}", path.display()))?; - let max_horizon = *self.cfg.horizons.iter().max().expect("non-empty horizons"); - if snapshots.len() < self.cfg.seq_len + max_horizon + 1 { - continue; - } - - let max_anchor = snapshots.len() - self.cfg.seq_len - max_horizon; - let anchor: usize = self.rng.gen_range(0..max_anchor); - - // Compute mid prices over the entire file (label generation needs - // the trailing window past the anchor + seq_len). - let prices: Vec = snapshots - .iter() - .map(|s| mid_price_f32(s)) - .collect(); - - let mut labels: [Vec; 5] = Default::default(); - for (h_idx, &h) in self.cfg.horizons.iter().enumerate() { - let mut out = vec![f32::NAN; self.cfg.seq_len]; - let raw = generate_labels(&prices, h); - for (i, &t) in raw.valid_indices.iter().enumerate() { - if t >= anchor && t < anchor + self.cfg.seq_len { - out[t - anchor] = raw.labels[i]; - } + // Advance to a new file if we have none cached or have yielded our + // budget from the current one. + let need_new = self.current.is_none() + || self.yielded_from_current >= self.sequences_per_file; + if need_new { + // Try files until one is big enough; advance file_idx each time. + let mut found: Option = None; + for _ in 0..self.files.len() * 2 { + if self.file_idx >= self.files.len() { + self.files.shuffle(&mut self.rng); + self.file_idx = 0; + } + let path = &self.files[self.file_idx]; + let snapshots = load_or_predecode_mbp10(path, &self.cfg.predecoded_dir) + .with_context(|| format!("load mbp10 {}", path.display()))?; + self.file_idx += 1; + if snapshots.len() < self.cfg.seq_len + max_horizon + 1 { + continue; } - labels[h_idx] = out; - } - // Build sequence of Mbp10RawInput. Each input carries its own - // prev_mid + prev_ts_ns + trade_signed_vol derived from the - // PRIOR snapshot in the source stream (not from `anchor`). - let mut sequence = Vec::with_capacity(self.cfg.seq_len); - for k in 0..self.cfg.seq_len { - let idx = anchor + k; - let cur = &snapshots[idx]; - let prev_idx = if idx == 0 { 0 } else { idx - 1 }; - let prev = &snapshots[prev_idx]; - sequence.push(convert(cur, prev)); + // Precompute mid prices + per-horizon labels for the WHOLE file. + // We pay this once per file load, then slice cheaply per + // anchor on subsequent next_sequence calls. + let prices: Vec = snapshots.iter().map(mid_price_f32).collect(); + let mut labels_full: [Vec; 5] = Default::default(); + for (h_idx, &h) in self.cfg.horizons.iter().enumerate() { + let mut full = vec![f32::NAN; snapshots.len()]; + let raw = generate_labels(&prices, h); + for (i, &t) in raw.valid_indices.iter().enumerate() { + full[t] = raw.labels[i]; + } + labels_full[h_idx] = full; + } + found = Some(LoadedFile { snapshots, labels_full }); + break; + } + self.current = found; + self.yielded_from_current = 0; + if self.current.is_none() { + return Ok(None); // no file in the inventory has enough snapshots } - - self.yielded += 1; - return Ok(Some(LabeledSequence { - snapshots: sequence, - labels, - })); } - Ok(None) // exhausted files and none had enough snapshots + + // Sample an anchor inside the current file and build the sequence. + let lf = self.current.as_ref().expect("set above"); + let max_anchor = lf.snapshots.len() - self.cfg.seq_len - max_horizon; + let anchor: usize = self.rng.gen_range(0..max_anchor); + + let mut labels: [Vec; 5] = Default::default(); + for h in 0..5 { + labels[h] = lf.labels_full[h][anchor..anchor + self.cfg.seq_len].to_vec(); + } + let mut sequence = Vec::with_capacity(self.cfg.seq_len); + for k in 0..self.cfg.seq_len { + let idx = anchor + k; + let cur = &lf.snapshots[idx]; + let prev_idx = if idx == 0 { 0 } else { idx - 1 }; + let prev = &lf.snapshots[prev_idx]; + sequence.push(convert(cur, prev)); + } + + self.yielded += 1; + self.yielded_from_current += 1; + Ok(Some(LabeledSequence { snapshots: sequence, labels })) } }