foxhunt/ml/examples/train_tlob.rs

//! TLOB Training Example
//!
//! Trains a TLOB transformer model on Level-2 order book data and saves checkpoints to disk.
//!
//! # Prerequisites
//!
//! - Level-2 order book data (MBP-10) from Agent 71
//! - Data directory: test_data/real/databento/ml_training_l2/
//! - GPU: RTX 3050 Ti (optional, will fall back to CPU)
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (500 epochs)
//! cargo run -p ml --example train_tlob --release --features cuda
//!
//! # Custom epochs and output path
//! cargo run -p ml --example train_tlob --release --features cuda -- \
//!   --epochs 1000 \
//!   --output ml/trained_models/tlob_model
//!
//! # Custom data directory and hyperparameters
//! cargo run -p ml --example train_tlob --release --features cuda -- \
//!   --data-dir test_data/real/databento/ml_training_l2 \
//!   --epochs 500 \
//!   --batch-size 16 \
//!   --learning-rate 0.0001 \
//!   --seq-len 128
//!
//! # CPU-only training (slower but works without GPU)
//! cargo run -p ml --example train_tlob --release -- \
//!   --no-gpu \
//!   --epochs 100
//! ```
//!
//! # Expected Output
//!
//! - Training checkpoints: ml/trained_models/tlob_epoch_*.safetensors
//! - Final model: ml/trained_models/tlob_final_epoch500.safetensors
//! - Training time: 5-8 hours (GPU), 20-30 hours (CPU) for 500 epochs

use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::trainers::tlob::{TLOBHyperparameters, TLOBTrainer, TLOBTrainingMetrics};

#[derive(Debug, Parser)]
#[command(
    name = "train_tlob",
    about = "Train TLOB transformer on Level-2 order book data"
)]
struct Opts {
    /// Number of training epochs
    #[arg(long, default_value = "500")]
    epochs: usize,

    /// Learning rate
    #[arg(long, default_value = "0.0001")]
    learning_rate: f64,

    /// Batch size (max 32 for RTX 3050 Ti 4GB)
    #[arg(long, default_value = "16")]
    batch_size: usize,

    /// Sequence length (number of order book snapshots)
    #[arg(long, default_value = "128")]
    seq_len: usize,

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

    /// Number of attention heads
    #[arg(long, default_value = "8")]
    num_heads: usize,

    /// Number of transformer layers
    #[arg(long, default_value = "4")]
    num_layers: usize,

    /// Dropout rate
    #[arg(long, default_value = "0.1")]
    dropout: f64,

    /// Gradient clipping threshold
    #[arg(long, default_value = "1.0")]
    grad_clip: f64,

    /// Weight decay for regularization
    #[arg(long, default_value = "0.0001")]
    weight_decay: f64,

    /// Checkpoint save frequency (epochs)
    #[arg(long, default_value = "10")]
    checkpoint_frequency: usize,

    /// Output directory for trained model
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// Data directory containing Level-2 order book files
    #[arg(long, default_value = "test_data/real/databento/ml_training_l2")]
    data_dir: String,

    /// Disable GPU acceleration (use CPU only)
    #[arg(long)]
    no_gpu: bool,

    /// Verbose logging
    #[arg(short, long)]
    verbose: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Parse CLI options
    let opts = Opts::parse();

    // Setup logging
    let level = if opts.verbose {
        tracing::Level::DEBUG
    } else {
        tracing::Level::INFO
    };

    let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    info!("🚀 Starting TLOB Transformer Training");
    info!("Configuration:");
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • Sequence length: {}", opts.seq_len);
    info!("  • Hidden dimension: {}", opts.d_model);
    info!("  • Attention heads: {}", opts.num_heads);
    info!("  • Transformer layers: {}", opts.num_layers);
    info!("  • Dropout: {}", opts.dropout);
    info!("  • Gradient clipping: {}", opts.grad_clip);
    info!("  • Weight decay: {}", opts.weight_decay);
    info!(
        "  • Checkpoint frequency: {} epochs",
        opts.checkpoint_frequency
    );
    info!("  • Output directory: {}", opts.output_dir);
    info!("  • Data directory: {}", opts.data_dir);
    info!("  • GPU enabled: {}", !opts.no_gpu);

    // Check if data directory exists
    let data_path = PathBuf::from(&opts.data_dir);
    if !data_path.exists() {
        warn!("⚠️  Data directory not found: {}", opts.data_dir);
        warn!("⚠️  This is expected if Agent 71 hasn't completed yet.");
        warn!("⚠️  Training will use dummy data for testing purposes.");
    }

    // Create output directory
    let output_path = PathBuf::from(&opts.output_dir);
    if !output_path.exists() {
        std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
        info!("✅ Created output directory: {}", opts.output_dir);
    }

    // Configure TLOB hyperparameters
    let hyperparams = TLOBHyperparameters {
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        seq_len: opts.seq_len,
        num_price_levels: 10, // MBP-10
        d_model: opts.d_model,
        num_heads: opts.num_heads,
        num_layers: opts.num_layers,
        dropout: opts.dropout,
        epochs: opts.epochs,
        checkpoint_frequency: opts.checkpoint_frequency,
        grad_clip: opts.grad_clip,
        weight_decay: opts.weight_decay,
    };

    // Create TLOB trainer
    let mut trainer = TLOBTrainer::new(hyperparams, &output_path, !opts.no_gpu)
        .context("Failed to create TLOB trainer")?;

    info!("✅ TLOB trainer initialized");

    // Track training progress
    let mut last_epoch = 0;
    let mut best_val_loss = f64::INFINITY;

    // Create progress callback
    let progress_callback = |metrics: TLOBTrainingMetrics| {
        if metrics.epoch != last_epoch {
            last_epoch = metrics.epoch;

            info!(
                "📊 Epoch {}/{}: train_loss={:.6}, val_loss={:.6}, mae={:.6}, grad_norm={:.6}",
                metrics.epoch,
                opts.epochs,
                metrics.train_loss,
                metrics.val_loss,
                metrics.avg_mae,
                metrics.gradient_norm
            );

            if metrics.val_loss < best_val_loss {
                best_val_loss = metrics.val_loss;
                info!("🌟 New best validation loss: {:.6}", best_val_loss);
            }
        }
    };

    // Train the model
    info!("\n🏋️  Starting training...\n");
    let start_time = std::time::Instant::now();

    let metrics = trainer
        .train(&opts.data_dir, progress_callback)
        .await
        .context("Training failed")?;

    let training_duration = start_time.elapsed();

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    info!("  • Final train loss: {:.6}", metrics.train_loss);
    info!("  • Final val loss: {:.6}", metrics.val_loss);
    info!("  • Best val loss: {:.6}", best_val_loss);
    info!("  • Final MAE: {:.6}", metrics.avg_mae);
    info!("  • Final gradient norm: {:.6}", metrics.gradient_norm);
    info!("  • Epochs trained: {}", metrics.epoch);
    info!(
        "  • Training time: {:.1}s ({:.1} min, {:.1} hours)",
        training_duration.as_secs_f64(),
        training_duration.as_secs_f64() / 60.0,
        training_duration.as_secs_f64() / 3600.0
    );

    // Calculate training speed
    let seconds_per_epoch = training_duration.as_secs_f64() / opts.epochs as f64;
    info!("  • Average time per epoch: {:.2}s", seconds_per_epoch);

    // Save final model
    let final_model_path = output_path.join(format!("tlob_final_epoch{}.safetensors", opts.epochs));
    info!("\n💾 Saving final model to: {}", final_model_path.display());

    // Get final model state
    let final_checkpoint_data = trainer
        .serialize_model()
        .await
        .context("Failed to serialize final model")?;

    std::fs::write(&final_model_path, &final_checkpoint_data)
        .context("Failed to save final model")?;

    info!(
        "✅ Final model saved: {} ({} bytes, {:.2} MB)",
        final_model_path.display(),
        final_checkpoint_data.len(),
        final_checkpoint_data.len() as f64 / 1_048_576.0
    );

    // Print summary
    info!("\n🎉 TLOB training complete!");
    info!("📁 Model files saved to: {}", opts.output_dir);
    info!("\n📈 Training Summary:");
    info!("  • Best validation loss: {:.6}", best_val_loss);
    info!(
        "  • Convergence: {}",
        if best_val_loss < 0.001 {
            "✅ Excellent"
        } else if best_val_loss < 0.01 {
            "✅ Good"
        } else {
            "⚠️  Needs more epochs"
        }
    );
    info!("  • Training efficiency: {:.2}s/epoch", seconds_per_epoch);

    // Estimate production inference latency
    let estimated_inference_us = seconds_per_epoch * 1_000_000.0 / 1000.0; // Rough estimate
    info!("\n🚀 Production Inference Estimate:");
    info!(
        "  • Expected latency: <{:.0}μs per prediction",
        estimated_inference_us.min(100.0)
    );
    info!("  • Target: <50μs (sub-50μs HFT requirement)");

    // Next steps
    info!("\n📋 Next Steps:");
    info!("  1. Validate model with test data");
    info!("  2. Convert to ONNX for production inference");
    info!("  3. Integrate with TLOB inference engine");
    info!("  4. Benchmark inference latency (<50μs target)");
    info!("  5. Deploy to ML Training Service");

    Ok(())
}