8b0122a5ac
Adds `has_ofi: bool` to `FxCacheData` and the fxcache binary header (stored in `reserved[0]`). Propagates through load_fxcache, write_fxcache, all callers (precompute_features, train_baseline_rl, hyperopt dqn adapter), existing roundtrip tests, and adds two new has_ofi-specific roundtrip tests. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
469 lines
19 KiB
Rust
469 lines
19 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,
|
||
)]
|
||
//! Precompute Features — DBN to .fxcache pipeline
|
||
//!
|
||
//! Reads OHLCV + MBP-10 + trades DBN data, runs the full feature extraction
|
||
//! pipeline (standalone, no GPU required), and writes a `.fxcache` binary file for
|
||
//! zero-overhead GPU loading during training.
|
||
//!
|
||
//! Usage:
|
||
//! cargo run -p ml --example precompute_features --release -- \
|
||
//! --data-dir test_data/futures-baseline \
|
||
//! --mbp10-data-dir test_data/futures-baseline-mbp10 \
|
||
//! --trades-data-dir test_data/futures-baseline-trades \
|
||
//! --output-dir test_data/feature-cache \
|
||
//! --symbol ES.FUT \
|
||
//! --bf16 \
|
||
//! --yes
|
||
|
||
use std::path::{Path, PathBuf};
|
||
use std::time::Instant;
|
||
|
||
use anyhow::{Context, Result};
|
||
use clap::Parser;
|
||
use tracing::info;
|
||
|
||
use ml::trainers::dqn::{
|
||
collect_dbn_files_filtered, collect_dbn_files_recursive, extract_features_from_bars,
|
||
extract_ohlcv_bars_from_dbn,
|
||
};
|
||
use ml::features::extraction::OHLCVBar;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// CLI
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(
|
||
name = "precompute_features",
|
||
about = "Precompute features from DBN data and write .fxcache binary"
|
||
)]
|
||
struct Opts {
|
||
/// Directory containing OHLCV .dbn/.dbn.zst files (e.g. test_data/futures-baseline)
|
||
#[arg(long)]
|
||
data_dir: String,
|
||
|
||
/// Directory containing MBP-10 .dbn/.dbn.zst files (optional)
|
||
#[arg(long)]
|
||
mbp10_data_dir: Option<String>,
|
||
|
||
/// Directory containing trades .dbn/.dbn.zst files (optional)
|
||
#[arg(long)]
|
||
trades_data_dir: Option<String>,
|
||
|
||
/// Output directory for .fxcache files (default: sibling of --data-dir)
|
||
#[arg(long)]
|
||
output_dir: Option<String>,
|
||
|
||
/// Symbol name (used in output filename)
|
||
#[arg(long, default_value = "ES.FUT")]
|
||
symbol: String,
|
||
|
||
/// Write version 2 (bf16) instead of version 1 (f64)
|
||
#[arg(long)]
|
||
bf16: bool,
|
||
|
||
/// Skip confirmation prompt
|
||
#[arg(long)]
|
||
yes: bool,
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Cache key computation
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Compute a SHA256 cache key from all DBN files across the given directories.
|
||
///
|
||
// ---------------------------------------------------------------------------
|
||
// Main
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[tokio::main]
|
||
async fn main() -> Result<()> {
|
||
// Initialize tracing
|
||
tracing_subscriber::fmt()
|
||
.with_env_filter(
|
||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||
)
|
||
.init();
|
||
|
||
let opts = Opts::parse();
|
||
|
||
let data_dir = PathBuf::from(&opts.data_dir);
|
||
// Default output_dir: sibling of data_dir (e.g. test_data/futures-baseline → test_data/feature-cache)
|
||
let output_dir = match opts.output_dir {
|
||
Some(ref d) => PathBuf::from(d),
|
||
None => data_dir.parent().unwrap_or(&data_dir).join("feature-cache"),
|
||
};
|
||
let mbp10_dir = opts.mbp10_data_dir.as_ref().map(PathBuf::from);
|
||
let trades_dir = opts.trades_data_dir.as_ref().map(PathBuf::from);
|
||
|
||
// ── Validate inputs ──────────────────────────────────────────────────────
|
||
if !data_dir.exists() {
|
||
anyhow::bail!("Data directory not found: {}", data_dir.display());
|
||
}
|
||
|
||
// ── Print config summary ─────────────────────────────────────────────────
|
||
println!("================================================================================");
|
||
println!("Precompute Features — DBN to .fxcache Pipeline");
|
||
println!("================================================================================");
|
||
println!();
|
||
println!("Data dir: {}", data_dir.display());
|
||
if let Some(ref d) = mbp10_dir {
|
||
println!("MBP-10 dir: {}", d.display());
|
||
}
|
||
if let Some(ref d) = trades_dir {
|
||
println!("Trades dir: {}", d.display());
|
||
}
|
||
println!("Output dir: {}", output_dir.display());
|
||
println!("Symbol: {}", opts.symbol);
|
||
println!("Format: {}", if opts.bf16 { "bf16 (v2)" } else { "f64 (v1)" });
|
||
println!();
|
||
|
||
// ── Confirmation ─────────────────────────────────────────────────────────
|
||
if !opts.yes {
|
||
print!("Proceed with feature extraction? (yes/no): ");
|
||
std::io::Write::flush(&mut std::io::stdout())?;
|
||
|
||
let mut input = String::new();
|
||
std::io::stdin().read_line(&mut input)?;
|
||
let trimmed = input.trim();
|
||
if !trimmed.eq_ignore_ascii_case("yes") && !trimmed.eq_ignore_ascii_case("y") {
|
||
println!("Cancelled.");
|
||
return Ok(());
|
||
}
|
||
println!();
|
||
}
|
||
|
||
let t0 = Instant::now();
|
||
|
||
// ── Early exit if cache already exists for this data ─────────────────────
|
||
let hex_key_early = ml::feature_cache::calculate_dbn_cache_key_full(
|
||
&data_dir,
|
||
mbp10_dir.as_deref(),
|
||
trades_dir.as_deref(),
|
||
&opts.symbol,
|
||
"ohlcv",
|
||
).context("Failed to compute cache key")?;
|
||
let early_check_path = output_dir.join(format!("{hex_key_early}.fxcache"));
|
||
if early_check_path.exists() {
|
||
let size_mb = std::fs::metadata(&early_check_path)
|
||
.map(|m| m.len() as f64 / 1_048_576.0)
|
||
.unwrap_or(0.0);
|
||
println!("Cache already exists: {} ({:.1} MB) — skipping extraction",
|
||
early_check_path.display(), size_mb);
|
||
return Ok(());
|
||
}
|
||
|
||
// ── Step 1: Load OHLCV bars from DBN files ──────────────────────────────
|
||
info!("Loading OHLCV bars from {} (symbol={})...", data_dir.display(), opts.symbol);
|
||
let mut dbn_files = collect_dbn_files_filtered(&data_dir, Some(&opts.symbol));
|
||
if dbn_files.is_empty() {
|
||
anyhow::bail!(
|
||
"No DBN files found for symbol '{}' in {}. Available: {:?}",
|
||
opts.symbol,
|
||
data_dir.display(),
|
||
std::fs::read_dir(&data_dir)
|
||
.ok()
|
||
.map(|d| d.filter_map(|e| e.ok())
|
||
.filter(|e| e.path().is_dir())
|
||
.map(|e| e.file_name().to_string_lossy().into_owned())
|
||
.collect::<Vec<_>>())
|
||
.unwrap_or_default()
|
||
);
|
||
}
|
||
dbn_files.sort();
|
||
info!("Found {} DBN files for symbol '{}'", dbn_files.len(), opts.symbol);
|
||
|
||
let mut all_bars: Vec<OHLCVBar> = Vec::new();
|
||
for file in &dbn_files {
|
||
match extract_ohlcv_bars_from_dbn(file) {
|
||
Ok(bars) => all_bars.extend(bars),
|
||
Err(e) => tracing::warn!("Failed to load {:?}: {e}", file.file_name()),
|
||
}
|
||
}
|
||
all_bars.sort_by_key(|b| b.timestamp);
|
||
info!("Loaded {} OHLCV bars in {:.1}s", all_bars.len(), t0.elapsed().as_secs_f64());
|
||
|
||
// ── Step 2: Extract 42-dim feature vectors ──────────────────────────────
|
||
let t1 = Instant::now();
|
||
info!("Extracting features...");
|
||
let feature_vectors = extract_features_from_bars(&all_bars)
|
||
.context("Feature extraction failed")?;
|
||
info!("Extracted {} feature vectors in {:.1}s", feature_vectors.len(), t1.elapsed().as_secs_f64());
|
||
|
||
// ── Step 3: Build targets [preproc_close, preproc_next, raw_close, raw_next]
|
||
const WARMUP: usize = 50;
|
||
let n = feature_vectors.len().saturating_sub(1); // need next bar for target
|
||
let features: Vec<[f64; 42]> = feature_vectors[..n].to_vec();
|
||
let mut targets: Vec<[f64; 4]> = Vec::with_capacity(n);
|
||
for i in 0..n {
|
||
let raw_curr = all_bars[i + WARMUP].close;
|
||
let raw_next = all_bars[i + WARMUP + 1].close;
|
||
targets.push([raw_curr, raw_next, raw_curr, raw_next]);
|
||
}
|
||
|
||
// ── Step 3b: Build per-bar timestamps ─────────────────────────────────
|
||
let timestamps: Vec<i64> = (0..n)
|
||
.map(|i| all_bars[i + WARMUP].timestamp.timestamp_nanos_opt().unwrap_or(0))
|
||
.collect();
|
||
|
||
// ── Step 4: Compute OFI from MBP-10 + trades ────────────────────────────
|
||
let t2 = Instant::now();
|
||
let ofi: Vec<[f64; 8]> = 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);
|
||
mbp10_files.sort();
|
||
|
||
if mbp10_files.is_empty() {
|
||
info!("No MBP-10 files found, OFI will be zeros");
|
||
vec![[0.0; 8]; n]
|
||
} else {
|
||
// Load MBP-10 snapshots (parallel)
|
||
use data::providers::databento::dbn_parser::DbnParser;
|
||
let per_file: Vec<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());
|
||
Some(snapshots)
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(" Failed {:?}: {e}", file.file_name());
|
||
None
|
||
}
|
||
}
|
||
}).collect()
|
||
};
|
||
|
||
let mut all_snapshots = Vec::new();
|
||
for snaps in per_file { all_snapshots.extend(snaps); }
|
||
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)
|
||
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<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) => { tracing::warn!(" Failed {:?}: {e}", path.file_name()); None }
|
||
}
|
||
}).collect()
|
||
};
|
||
let mut trades = Vec::new();
|
||
for t in per_file_trades { trades.extend(t); }
|
||
trades.sort_by_key(|t| t.timestamp);
|
||
if trades.is_empty() { None } else { Some(trades) }
|
||
} else {
|
||
None
|
||
};
|
||
|
||
// Compute per-bar OFI
|
||
use ml::features::mbp10_loader::get_snapshots_for_timestamp;
|
||
use ml::features::trades_loader::get_trades_for_bar;
|
||
let mut calculator = OFICalculator::new();
|
||
let mut ofi_per_bar = Vec::with_capacity(n);
|
||
|
||
for i in 0..n {
|
||
let bar = &all_bars[i + WARMUP];
|
||
let bar_ts = bar.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64;
|
||
|
||
// Feed trades for this bar window
|
||
if let Some(ref trades) = all_trades {
|
||
let bar_end_ts = all_bars.get(i + WARMUP + 1)
|
||
.map(|b| b.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64)
|
||
.unwrap_or(bar_ts + 60_000_000_000);
|
||
for trade in get_trades_for_bar(trades, bar_ts, bar_end_ts) {
|
||
calculator.feed_trade(trade.price, trade.volume, trade.is_buy);
|
||
}
|
||
}
|
||
|
||
let window = get_snapshots_for_timestamp(&all_snapshots, bar_ts, 1);
|
||
if let Some(snap) = window.first() {
|
||
match calculator.calculate(snap) {
|
||
Ok(f) if f.is_valid() => ofi_per_bar.push(f.to_array()),
|
||
_ => ofi_per_bar.push([0.0; 8]),
|
||
}
|
||
} else {
|
||
ofi_per_bar.push([0.0; 8]);
|
||
}
|
||
}
|
||
|
||
let non_zero = ofi_per_bar.iter().filter(|f| f.iter().any(|&v| v != 0.0)).count();
|
||
info!("OFI computed: {} bars, {} non-zero ({:.1}%) in {:.1}s",
|
||
ofi_per_bar.len(), non_zero,
|
||
if ofi_per_bar.is_empty() { 0.0 } else { non_zero as f64 / ofi_per_bar.len() as f64 * 100.0 },
|
||
t2.elapsed().as_secs_f64());
|
||
ofi_per_bar
|
||
}
|
||
} else {
|
||
info!("No MBP-10 directory, OFI will be zeros");
|
||
vec![[0.0; 8]; n]
|
||
};
|
||
|
||
let total_len = n;
|
||
|
||
// ── Compute cache key ────────────────────────────────────────────────────
|
||
let hex_key = ml::feature_cache::calculate_dbn_cache_key_full(
|
||
&data_dir,
|
||
mbp10_dir.as_deref(),
|
||
trades_dir.as_deref(),
|
||
&opts.symbol,
|
||
"ohlcv",
|
||
).context("Failed to compute cache key")?;
|
||
let cache_key: [u8; 32] = hex::decode(&hex_key)
|
||
.context("Invalid hex key")?
|
||
.try_into()
|
||
.map_err(|_| anyhow::anyhow!("Cache key wrong length"))?;
|
||
|
||
// ── Normalize features (z-score) ──────────────────────────────────────────
|
||
// Raw features contain OHLCV prices (up to 25,000) that overflow bf16 in
|
||
// cuBLAS SGEMM. Normalize once at precompute time so every consumer
|
||
// (training, hyperopt, inference) gets consistent normalized features.
|
||
let norm_stats = ml::walk_forward::NormStats::from_features(&features);
|
||
let features = norm_stats.normalize_batch(&features);
|
||
info!("Features z-score normalized ({} bars × 42 dims)", features.len());
|
||
|
||
// ── Write .fxcache ───────────────────────────────────────────────────────
|
||
std::fs::create_dir_all(&output_dir)
|
||
.with_context(|| format!("Failed to create output directory: {}", output_dir.display()))?;
|
||
|
||
let output_path = output_dir.join(format!("{hex_key}.fxcache"));
|
||
|
||
// Save NormStats alongside cache for inference denormalization
|
||
let norm_path = output_dir.join(format!("{hex_key}.norm_stats.json"));
|
||
let norm_json = serde_json::to_string_pretty(&norm_stats)
|
||
.context("Failed to serialize NormStats")?;
|
||
std::fs::write(&norm_path, &norm_json)
|
||
.with_context(|| format!("Failed to write NormStats: {}", norm_path.display()))?;
|
||
info!("NormStats saved to {}", norm_path.display());
|
||
|
||
let t2 = Instant::now();
|
||
info!("Writing .fxcache to {}...", output_path.display());
|
||
|
||
let has_ofi = mbp10_dir.is_some();
|
||
let bytes_written = ml::fxcache::write_fxcache(
|
||
&output_path,
|
||
&features,
|
||
&targets,
|
||
&ofi,
|
||
×tamps,
|
||
cache_key,
|
||
opts.bf16,
|
||
has_ofi,
|
||
)
|
||
.context("Failed to write .fxcache file")?;
|
||
|
||
let write_secs = t2.elapsed().as_secs_f64();
|
||
let total_secs = t0.elapsed().as_secs_f64();
|
||
|
||
// ── Summary ──────────────────────────────────────────────────────────────
|
||
println!();
|
||
println!("================================================================================");
|
||
println!("PRECOMPUTE SUMMARY");
|
||
println!("================================================================================");
|
||
println!();
|
||
println!("Bars: {}", total_len);
|
||
println!("Features: 42-dim");
|
||
println!("Targets: 4-dim");
|
||
println!("OFI: 8-dim ({})", if has_ofi { "from MBP-10" } else { "zero-padded" });
|
||
println!("Format: {} (v{})", if opts.bf16 { "bf16" } else { "f64" }, if opts.bf16 { 2 } else { 1 });
|
||
println!("Cache key: {}", hex_key);
|
||
println!("Output: {}", output_path.display());
|
||
println!("Size: {:.2} MB", bytes_written as f64 / 1_048_576.0);
|
||
println!("Write time: {:.1}s", write_secs);
|
||
println!("Total time: {:.1}s", total_secs);
|
||
println!();
|
||
|
||
Ok(())
|
||
}
|