foxhunt/ml/examples/benchmark_streaming_vs_batch.rs

//! Benchmark: Streaming vs Batch Data Loading
//!
//! Compares memory usage and performance between batch and streaming loaders.
//!
//! ## Metrics Compared
//!
//! - Peak memory usage (RSS)
//! - Loading time
//! - Sequences per second
//! - Memory efficiency ratio
//!
//! ## Usage
//!
//! ```bash
//! # Small dataset (4 files, ~400KB)
//! cargo run -p ml --example benchmark_streaming_vs_batch --release -- \
//!     --data-dir test_data/real/databento/ml_training_small
//!
//! # Large dataset (360 files, ~15MB)
//! cargo run -p ml --example benchmark_streaming_vs_batch --release -- \
//!     --data-dir test_data/real/databento/ml_training
//! ```

use anyhow::Result;
use clap::Parser;
use ml::data_loaders::{DbnSequenceLoader, StreamingDbnLoader};
use std::path::PathBuf;
use std::time::Instant;

#[derive(Parser, Debug)]
#[command(author, version, about)]
struct Args {
    /// Directory containing DBN files
    #[arg(long, default_value = "test_data/real/databento/ml_training_small")]
    data_dir: PathBuf,

    /// Sequence length
    #[arg(long, default_value = "60")]
    seq_len: usize,

    /// Model dimension
    #[arg(long, default_value = "256")]
    d_model: usize,

    /// Train/validation split
    #[arg(long, default_value = "0.9")]
    train_split: f64,

    /// Streaming batch size (bars)
    #[arg(long, default_value = "10000")]
    batch_size: usize,

    /// Stride for sliding window
    #[arg(long, default_value = "100")]
    stride: usize,
}

/// Memory statistics
#[derive(Debug, Clone)]
struct MemoryStats {
    rss_kb: usize,
    vms_kb: usize,
}

impl MemoryStats {
    /// Get current memory usage
    fn current() -> Result<Self> {
        let status = std::fs::read_to_string("/proc/self/status")?;

        let mut rss_kb = 0;
        let mut vms_kb = 0;

        for line in status.lines() {
            if line.starts_with("VmRSS:") {
                rss_kb = line
                    .split_whitespace()
                    .nth(1)
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(0);
            } else if line.starts_with("VmSize:") {
                vms_kb = line
                    .split_whitespace()
                    .nth(1)
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(0);
            }
        }

        Ok(Self { rss_kb, vms_kb })
    }

    fn rss_mb(&self) -> f64 {
        self.rss_kb as f64 / 1024.0
    }

    fn vms_mb(&self) -> f64 {
        self.vms_kb as f64 / 1024.0
    }
}

/// Benchmark results
#[derive(Debug)]
struct BenchmarkResult {
    name: String,
    total_sequences: usize,
    duration_secs: f64,
    sequences_per_sec: f64,
    peak_memory_mb: f64,
    memory_efficiency_ratio: f64,
}

impl BenchmarkResult {
    fn print_report(&self) {
        println!("\n{:=<60}", "");
        println!("  {} BENCHMARK RESULTS", self.name.to_uppercase());
        println!("{:=<60}", "");
        println!("  Total Sequences:     {}", self.total_sequences);
        println!("  Duration:            {:.2}s", self.duration_secs);
        println!(
            "  Throughput:          {:.0} sequences/sec",
            self.sequences_per_sec
        );
        println!("  Peak Memory (RSS):   {:.1} MB", self.peak_memory_mb);
        println!(
            "  Memory Efficiency:   {:.2}x",
            self.memory_efficiency_ratio
        );
        println!("{:=<60}\n", "");
    }
}

/// Benchmark batch loading
async fn benchmark_batch(args: &Args) -> Result<BenchmarkResult> {
    println!("\n🔄 BATCH LOADING BENCHMARK");
    println!("   Loading all data into memory...\n");

    // Measure baseline memory
    let baseline_memory = MemoryStats::current()?;
    println!("   Baseline memory: {:.1} MB RSS", baseline_memory.rss_mb());

    let start = Instant::now();
    let mut peak_memory = baseline_memory.clone();

    // Create batch loader
    let mut loader = DbnSequenceLoader::with_limits(
        args.seq_len,
        args.d_model,
        Some(1_000), // Limit sequences to prevent OOM
        args.stride,
    )
    .await?;

    // Load all sequences at once
    let (train_data, val_data) = loader
        .load_sequences(&args.data_dir, args.train_split)
        .await?;

    // Measure peak memory
    let current_memory = MemoryStats::current()?;
    if current_memory.rss_kb > peak_memory.rss_kb {
        peak_memory = current_memory;
    }

    let duration = start.elapsed();
    let total_sequences = train_data.len() + val_data.len();

    println!("   ✅ Loaded {} sequences", total_sequences);
    println!("   Peak memory: {:.1} MB RSS", peak_memory.rss_mb());
    println!("   Duration: {:.2}s", duration.as_secs_f64());

    Ok(BenchmarkResult {
        name: "Batch Loading".to_string(),
        total_sequences,
        duration_secs: duration.as_secs_f64(),
        sequences_per_sec: total_sequences as f64 / duration.as_secs_f64(),
        peak_memory_mb: peak_memory.rss_mb() - baseline_memory.rss_mb(),
        memory_efficiency_ratio: 1.0, // Baseline
    })
}

/// Benchmark streaming loading
async fn benchmark_streaming(
    args: &Args,
    batch_baseline: &BenchmarkResult,
) -> Result<BenchmarkResult> {
    println!("\n🌊 STREAMING LOADING BENCHMARK");
    println!(
        "   Loading data in batches of {} bars...\n",
        args.batch_size
    );

    // Measure baseline memory
    let baseline_memory = MemoryStats::current()?;
    println!("   Baseline memory: {:.1} MB RSS", baseline_memory.rss_mb());

    let start = Instant::now();
    let mut peak_memory = baseline_memory.clone();

    // Create streaming loader
    let loader =
        StreamingDbnLoader::with_config(args.seq_len, args.d_model, args.batch_size, args.stride)
            .await?;

    // Stream sequences
    let mut stream = loader
        .stream_sequences(&args.data_dir, args.train_split)
        .await?;
    let mut total_sequences = 0;
    let mut batch_count = 0;

    // Process training data
    loop {
        // Measure memory before batch
        let current_memory = MemoryStats::current()?;
        if current_memory.rss_kb > peak_memory.rss_kb {
            peak_memory = current_memory;
        }

        match stream.next_batch().await? {
            Some(batch) => {
                batch_count += 1;
                total_sequences += batch.len();

                if batch_count % 10 == 0 {
                    let current_mem = MemoryStats::current()?;
                    println!(
                        "   Batch {}: {} sequences (memory: {:.1} MB)",
                        batch_count,
                        batch.len(),
                        current_mem.rss_mb() - baseline_memory.rss_mb()
                    );
                }
            },
            None => break,
        }
    }

    let duration = start.elapsed();

    println!(
        "   ✅ Processed {} sequences in {} batches",
        total_sequences, batch_count
    );
    println!("   Peak memory: {:.1} MB RSS", peak_memory.rss_mb());
    println!("   Duration: {:.2}s", duration.as_secs_f64());

    // Calculate memory efficiency ratio
    let memory_used = peak_memory.rss_mb() - baseline_memory.rss_mb();
    let memory_efficiency = batch_baseline.peak_memory_mb / memory_used.max(1.0);

    Ok(BenchmarkResult {
        name: "Streaming Loading".to_string(),
        total_sequences,
        duration_secs: duration.as_secs_f64(),
        sequences_per_sec: total_sequences as f64 / duration.as_secs_f64(),
        peak_memory_mb: memory_used,
        memory_efficiency_ratio: memory_efficiency,
    })
}

/// Print comparison table
fn print_comparison(batch: &BenchmarkResult, streaming: &BenchmarkResult) {
    println!("\n{:=<80}", "");
    println!("  COMPREHENSIVE COMPARISON");
    println!("{:=<80}", "");
    println!();
    println!("  {:<30} {:>20} {:>20}", "Metric", "Batch", "Streaming");
    println!("  {:-<30} {:-<20} {:-<20}", "", "", "");

    println!(
        "  {:<30} {:>20} {:>20}",
        "Sequences Loaded", batch.total_sequences, streaming.total_sequences
    );

    println!(
        "  {:<30} {:>18.2}s {:>18.2}s",
        "Duration", batch.duration_secs, streaming.duration_secs
    );

    let speed_ratio = streaming.duration_secs / batch.duration_secs;
    let speed_pct = (speed_ratio - 1.0) * 100.0;
    println!(
        "  {:<30} {:>20.0} {:>20.0} ({:+.1}%)",
        "Throughput (seq/s)", batch.sequences_per_sec, streaming.sequences_per_sec, -speed_pct
    );

    println!(
        "  {:<30} {:>18.1} MB {:>18.1} MB",
        "Peak Memory", batch.peak_memory_mb, streaming.peak_memory_mb
    );

    let memory_reduction = (1.0 - streaming.peak_memory_mb / batch.peak_memory_mb) * 100.0;
    println!(
        "  {:<30} {:>20} {:>18.2}x ({:.0}% reduction)",
        "Memory Efficiency", "1.0x", streaming.memory_efficiency_ratio, memory_reduction
    );

    println!("\n{:=<80}", "");

    // Success criteria check
    println!("\n  SUCCESS CRITERIA:");
    println!("  {:-<80}", "");

    let memory_ok = streaming.peak_memory_mb < 512.0;
    let speed_ok = speed_pct.abs() < 10.0;

    println!(
        "  ✓ Memory < 512MB:       {} ({:.1} MB)",
        if memory_ok { "✅ PASS" } else { "❌ FAIL" },
        streaming.peak_memory_mb
    );

    println!(
        "  ✓ Speed penalty < 10%:  {} ({:+.1}%)",
        if speed_ok { "✅ PASS" } else { "❌ FAIL" },
        speed_pct
    );

    if memory_ok && speed_ok {
        println!("\n  🎉 ALL SUCCESS CRITERIA MET!");
    } else {
        println!("\n  ⚠️  Some criteria not met - may need tuning");
    }

    println!("{:=<80}\n", "");
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    let args = Args::parse();

    println!("\n{:=<80}", "");
    println!("  STREAMING VS BATCH DATA LOADING BENCHMARK");
    println!("{:=<80}", "");
    println!("  Data Directory:  {:?}", args.data_dir);
    println!("  Sequence Length: {}", args.seq_len);
    println!("  Model Dimension: {}", args.d_model);
    println!("  Batch Size:      {} bars", args.batch_size);
    println!("  Stride:          {}", args.stride);
    println!("  Train Split:     {:.0}%", args.train_split * 100.0);
    println!("{:=<80}\n", "");

    // Verify data directory exists
    if !args.data_dir.exists() {
        eprintln!("❌ Error: Data directory not found: {:?}", args.data_dir);
        eprintln!("\nAvailable test directories:");
        eprintln!("  - test_data/real/databento/ml_training_small (4 files, ~400KB)");
        eprintln!("  - test_data/real/databento/ml_training (360 files, ~15MB)");
        std::process::exit(1);
    }

    // Run benchmarks
    println!("🚀 Starting benchmarks...\n");

    let batch_result = benchmark_batch(&args).await?;
    batch_result.print_report();

    // Force garbage collection between benchmarks
    println!("🧹 Cleaning up memory...");
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;

    let streaming_result = benchmark_streaming(&args, &batch_result).await?;
    streaming_result.print_report();

    // Print comparison
    print_comparison(&batch_result, &streaming_result);

    Ok(())
}