Files
foxhunt/docs/rust-logging-best-practices-ml-training.md
jgrusewski 4080f73ba4 feat(ml): WAVE 30 - Implement optimized DQN logging module
Add comprehensive logging utilities for DQN training with best practices:

- LoggingConfig: Configurable log levels, intervals, and sampling rates
- MetricsAggregator: Windowed statistics (mean, std_dev) for training metrics
- SampledLogger: Rate-limited logging for high-frequency events

Key features:
- Structured logging with tracing crate (info/debug/trace hierarchy)
- 23 unit tests for full coverage
- Integration with existing DQN training pipeline

Bug fixes:
- Fix u8 overflow in prioritized_replay.rs test (500 > u8::MAX)
- Fix GradStore assertion in residual.rs (no is_empty method)
- Fix Tensor::get() Option/Result handling in quantile_regression.rs
- Fix Device PartialEq comparison in ensemble_network.rs

Documentation:
- Add Rust logging best practices guide for ML training
- Add DQN logging analysis and design summary

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
2025-11-28 00:15:23 +01:00

30 KiB
Raw Blame History

Rust Logging Best Practices for ML Training Systems

Comprehensive Guide for Foxhunt ML Training Infrastructure

Table of Contents

  1. Log Level Guidelines
  2. Tracing vs Log Crate
  3. Structured Logging with Spans
  4. Performance Optimization
  5. ML Training Specific Patterns
  6. Production Deployment

1. Log Level Guidelines

Overview

The tracing crate provides five log levels: error!, warn!, info!, debug!, and trace!. For ML training systems, use these strategically to balance observability and performance.

Level Definitions for ML Training

error! - Unrecoverable Errors (Training Failures)

Use for failures that require immediate attention and halt training:

use tracing::error;

// Training termination
if gradient_norm.is_nan() {
    error!(
        epoch = current_epoch,
        gradient_norm = %gradient_norm,
        "🛑 TERMINATING: Gradient explosion detected - NaN values in gradient"
    );
    return Err(MLError::TrainingFailure {
        reason: "Gradient NaN".to_string(),
    });
}

// Checkpoint save failure
if let Err(e) = checkpoint_manager.save(model_state).await {
    error!(
        error = %e,
        checkpoint_path = %path,
        "Failed to save checkpoint - training progress may be lost"
    );
    return Err(e.into());
}

// Critical validation failure
if accuracy < 0.5 && epoch > 100 {
    error!(
        epoch = epoch,
        accuracy = %accuracy,
        "Model performing worse than random - critical training bug"
    );
    return Err(MLError::ValidationFailure);
}

When to use:

  • Training crashes (OOM, GPU errors, NaN/Inf values)
  • Checkpoint save/load failures
  • Critical validation failures (accuracy worse than baseline)
  • Data pipeline corruption
  • Unrecoverable circuit breaker trips

warn! - Recoverable Issues (Degraded Performance)

Use for issues that affect training quality but don't halt execution:

use tracing::warn;

// Early stopping triggered
if let Err(e) = early_stopping.check(val_loss) {
    warn!(
        epoch = current_epoch,
        val_loss = %val_loss,
        patience_remaining = early_stopping.patience_counter,
        "⚠️  Early stopping triggered - validation loss plateau"
    );
    break; // Exit training loop gracefully
}

// Action diversity warning
if action_percentage < 0.5 {
    warn!(
        epoch = epoch,
        action = %action_name,
        percentage = %action_percentage,
        "⚠️  LOW ACTION DIVERSITY: {} only {:.1}% of actions",
        action_name, action_percentage
    );
}

// Reward signal quality warning
if reward_std < 0.01 {
    warn!(
        epoch = epoch,
        reward_mean = %reward_mean,
        reward_std = %reward_std,
        consecutive_epochs = consecutive_constant_epochs,
        "⚠️  CONSTANT REWARDS DETECTED - possible reward calculation bug"
    );
}

// Q-value range warning (DQN-specific)
let (q_min, q_max, q_mean) = monitor.get_q_value_stats();
if q_max - q_min < 0.1 {
    warn!(
        epoch = epoch,
        q_min = %q_min,
        q_max = %q_max,
        q_range = %(q_max - q_min),
        "⚠️  Q-VALUES COLLAPSED - very narrow range"
    );
}

When to use:

  • Early stopping triggers
  • Gradient clipping activations
  • Action diversity issues
  • Q-value collapse warnings
  • Memory pressure (approaching limits)
  • Degraded throughput (below target)
  • Circuit breaker warnings (pre-trip)

info! - High-Level Progress (Epoch Start/End, Checkpoints)

Use for milestone events and key metrics:

use tracing::info;

// Epoch start
info!(
    epoch = current_epoch,
    total_epochs = total_epochs,
    learning_rate = %lr,
    batch_size = batch_size,
    "🚀 Starting epoch {}/{}", current_epoch, total_epochs
);

// Epoch completion with aggregated metrics
info!(
    epoch = current_epoch,
    train_loss = %train_loss,
    val_loss = %val_loss,
    val_accuracy = %val_accuracy,
    q_value_mean = %q_value_mean,
    epsilon = %epsilon,
    duration_secs = epoch_duration.as_secs(),
    "✅ Epoch complete - Loss: {:.4}, Acc: {:.2}%, ε={:.3}",
    train_loss, val_accuracy * 100.0, epsilon
);

// Checkpoint save
info!(
    epoch = current_epoch,
    checkpoint_path = %path,
    model_size_mb = size_mb,
    "💾 Checkpoint saved (best val_loss: {:.4})", best_val_loss
);

// Training initialization
info!(
    model_type = "DQN",
    hidden_dims = ?[256, 256, 256],
    device = ?device,
    total_params = total_params,
    "Initializing DQN with {} parameters on {:?}", total_params, device
);

// Final training summary
info!(
    total_epochs = final_epoch,
    best_epoch = best_epoch,
    best_val_loss = %best_val_loss,
    total_duration_mins = total_duration.as_secs() / 60,
    "🎉 Training complete - Best loss: {:.4} at epoch {}",
    best_val_loss, best_epoch
);

When to use:

  • Epoch start/end events
  • Checkpoint saves
  • Model initialization
  • Training completion summary
  • Configuration changes (learning rate updates)
  • Major phase transitions (calibration → training)
  • System resource status (GPU memory, etc.)

debug! - Detailed Diagnostics (Batch Metrics, Gradient Stats)

Use for troubleshooting and performance analysis (disabled in production):

use tracing::debug;

// Batch-level metrics
debug!(
    epoch = epoch,
    batch = batch_idx,
    batch_loss = %batch_loss,
    gradient_norm = %grad_norm,
    "Batch {}/{}: loss={:.4}, grad_norm={:.4}",
    batch_idx, total_batches, batch_loss, grad_norm
);

// Action selection details
debug!(
    step = step,
    action = ?selected_action,
    q_values = ?q_values,
    epsilon = %epsilon,
    exploration = is_exploration,
    "Action selected: {:?}, Q-values: {:?}", selected_action, q_values
);

// Feature statistics
debug!(
    feature_name = "volatility",
    mean = %feature_mean,
    std = %feature_std,
    min = %feature_min,
    max = %feature_max,
    "Feature stats: μ={:.4}, σ={:.4}, range=[{:.4}, {:.4}]",
    feature_mean, feature_std, feature_min, feature_max
);

// Replay buffer sampling
debug!(
    buffer_size = buffer.len(),
    batch_size = batch_size,
    sampling_strategy = "prioritized",
    "Sampled {} transitions from buffer (size={})",
    batch_size, buffer.len()
);

When to use:

  • Per-batch loss and gradient statistics
  • Action selection details (Q-values, exploration)
  • Feature normalization statistics
  • Replay buffer operations
  • Model layer activations
  • Optimizer state changes
  • Data augmentation results

trace! - Very Verbose (Per-Step Q-Values, Individual Actions)

Use sparingly for deep debugging (highest performance cost):

use tracing::trace;

// Per-step Q-value tracking
trace!(
    step = step,
    state_hash = state_hash,
    q_values = ?q_values,
    max_q = %max_q_value,
    "Q-values for state {}: {:?}", state_hash, q_values
);

// Individual experience transitions
trace!(
    step = step,
    state = ?state,
    action = ?action,
    reward = %reward,
    next_state = ?next_state,
    done = done,
    "Experience: s={:?}, a={:?}, r={:.4}, s'={:?}, done={}",
    state, action, reward, next_state, done
);

// Tensor shape tracking
trace!(
    operation = "forward_pass",
    input_shape = ?input.shape(),
    output_shape = ?output.shape(),
    layer = "dense_1",
    "Layer forward: input={:?} → output={:?}",
    input.shape(), output.shape()
);

When to use:

  • Per-step Q-value logging (for debugging Q-value collapse)
  • Individual experience transitions (for reward debugging)
  • Tensor shape verification during development
  • Network layer-by-layer outputs
  • AVOID IN PRODUCTION - extremely high overhead

Log Level Selection Matrix

Event Type Frequency Production Level Development Level
Training start/end Once per run info! info!
Epoch summary Once per epoch info! info!
Checkpoint save Every N epochs info! info!
Batch metrics Every batch disabled debug!
Gradient stats Every 10 batches disabled debug!
Action selection Every step disabled trace!
Q-value tracking Every step disabled trace!
Early stopping When triggered warn! warn!
NaN/Inf detection When detected error! error!
Reward warnings When triggered warn! debug!

2. Tracing vs Log Crate

Why tracing Over log?

The tracing crate is superior for ML training systems due to:

  1. Async-friendly: Fully compatible with tokio runtime
  2. Structured logging: First-class support for key-value fields
  3. Span-based tracing: Track execution through training loops
  4. Zero-cost when disabled: Compiler eliminates disabled logs
  5. Rich ecosystem: Integration with OpenTelemetry, Jaeger, etc.

Dependency Setup

[dependencies]
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }

[dev-dependencies]
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }

Basic Initialization

use tracing_subscriber::FmtSubscriber;

fn main() {
    // Initialize tracing subscriber (once per process)
    let subscriber = FmtSubscriber::builder()
        .with_max_level(tracing::Level::INFO)  // Filter level
        .with_target(true)                      // Show module path
        .with_thread_ids(true)                  // Show thread ID
        .with_line_number(true)                 // Show line numbers
        .finish();

    tracing::subscriber::set_global_default(subscriber)
        .expect("setting default subscriber failed");

    // Your training code here
}

Environment-based Filtering

use tracing_subscriber::EnvFilter;

fn main() {
    // RUST_LOG=ml::trainers::dqn=debug,ml=info cargo run
    let filter = EnvFilter::from_default_env()
        .add_directive("ml::trainers::dqn=debug".parse().unwrap())
        .add_directive("ml=info".parse().unwrap());

    let subscriber = FmtSubscriber::builder()
        .with_env_filter(filter)
        .finish();

    tracing::subscriber::set_global_default(subscriber)
        .expect("setting default subscriber failed");
}

Comparison Table

Feature log tracing
Async support Limited Native
Structured fields No Yes
Span tracking No Yes
Zero-cost disabled No Yes
OpenTelemetry Via adapter Native
Performance Good Excellent
Ecosystem Mature Growing fast

3. Structured Logging with Spans

Span Basics

Spans represent periods of time in your code (e.g., an epoch, a batch, a forward pass). They automatically track:

  • Start and end times
  • Nested relationships (parent-child spans)
  • Structured context (key-value pairs)

Epoch-Level Spans

use tracing::{info, info_span};

pub fn train(&mut self, num_epochs: usize) -> Result<()> {
    for epoch in 0..num_epochs {
        // Create epoch span with context
        let _epoch_span = info_span!(
            "train_epoch",
            epoch = epoch,
            total_epochs = num_epochs
        ).entered();

        info!("Starting epoch {}/{}", epoch, num_epochs);

        // All logs within this scope inherit epoch context
        let train_loss = self.train_one_epoch()?;
        let val_loss = self.validate()?;

        info!(
            train_loss = %train_loss,
            val_loss = %val_loss,
            "Epoch complete"
        );
    }

    Ok(())
}

Output:

INFO train_epoch{epoch=0 total_epochs=100}: Starting epoch 0/100
INFO train_epoch{epoch=0 total_epochs=100}: Epoch complete train_loss=0.234 val_loss=0.189

Batch-Level Spans (Nested)

use tracing::debug_span;

fn train_one_epoch(&mut self) -> Result<f32> {
    let _epoch_span = info_span!("train_one_epoch").entered();

    let mut total_loss = 0.0;
    for (batch_idx, batch) in self.data_loader.iter().enumerate() {
        // Nested span for batch processing
        let _batch_span = debug_span!(
            "process_batch",
            batch = batch_idx,
            batch_size = batch.len()
        ).entered();

        let loss = self.forward_backward(&batch)?;
        total_loss += loss;

        debug!(loss = %loss, "Batch processed");
    }

    Ok(total_loss / self.data_loader.len() as f32)
}

Output with spans:

DEBUG train_one_epoch{}: process_batch{batch=0 batch_size=32}: Batch processed loss=0.234
DEBUG train_one_epoch{}: process_batch{batch=1 batch_size=32}: Batch processed loss=0.189

Function-Level Instrumentation

The #[instrument] macro automatically creates spans for functions:

use tracing::instrument;

// Automatic span creation with function arguments as fields
#[instrument(skip(self, batch))]  // Skip large arguments
pub fn forward_backward(&mut self, batch: &Batch) -> Result<f32> {
    let output = self.model.forward(&batch.input)?;
    let loss = self.criterion.compute(&output, &batch.target)?;

    self.optimizer.zero_grad();
    loss.backward()?;
    self.optimizer.step();

    Ok(loss.to_scalar())
}

// With custom field names
#[instrument(
    skip(self, state),
    fields(
        state_dim = state.len(),
        epsilon = %self.epsilon
    )
)]
pub fn select_action(&mut self, state: &Tensor) -> Result<Action> {
    if rand::random::<f32>() < self.epsilon {
        Ok(Action::random())
    } else {
        let q_values = self.q_network.forward(state)?;
        Ok(Action::from_q_values(&q_values))
    }
}

Key advantages:

  • Automatic entry/exit logging
  • Exception/error tracking
  • Duration measurement
  • Low boilerplate

Structured Fields

Always use structured fields instead of string formatting:

// ❌ BAD: String formatting (loses structure)
info!("Epoch {} complete: loss={:.4}", epoch, loss);

// ✅ GOOD: Structured fields (queryable, filterable)
info!(
    epoch = epoch,
    loss = %loss,  // Display format
    "Epoch complete"
);

// ✅ EVEN BETTER: Multiple fields with semantic names
info!(
    epoch = epoch,
    train_loss = %train_loss,
    val_loss = %val_loss,
    val_accuracy = %val_accuracy,
    learning_rate = %lr,
    duration_secs = duration.as_secs(),
    "Epoch complete"
);

Span Context Propagation

Spans automatically propagate context to nested function calls:

pub fn train(&mut self) -> Result<()> {
    let _train_span = info_span!("dqn_training", model="DQN-Rainbow").entered();

    for epoch in 0..self.num_epochs {
        let _epoch_span = info_span!("epoch", epoch=epoch).entered();

        // All nested function calls inherit epoch context
        self.collect_experience()?;  // Logs: dqn_training{model="DQN-Rainbow"}:epoch{epoch=0}
        self.update_q_network()?;    // Logs: dqn_training{model="DQN-Rainbow"}:epoch{epoch=0}
        self.update_target_network()?; // Logs: dqn_training{model="DQN-Rainbow"}:epoch{epoch=0}
    }

    Ok(())
}

4. Performance Optimization

Lazy Evaluation with Format Arguments

The tracing macros use lazy evaluation - arguments are only evaluated if the log level is enabled:

// ✅ GOOD: Expensive computation only if debug enabled
debug!(
    q_values = ?compute_q_values(&state),  // Only computed if debug! is enabled
    "Q-values computed"
);

// ❌ BAD: Always computes, even if debug! is disabled
let q_vals = compute_q_values(&state);  // Always runs
debug!(q_values = ?q_vals, "Q-values computed");

Conditional Compilation with cfg!

For extremely hot paths, use conditional compilation:

// High-frequency logging (per-step) - only compiled in debug builds
#[cfg(debug_assertions)]
trace!(
    step = step,
    q_values = ?q_values,
    "Step Q-values"
);

// Or use feature flags
#[cfg(feature = "verbose-logging")]
trace!(
    state = ?state,
    action = ?action,
    "State-action pair"
);

Add to Cargo.toml:

[features]
verbose-logging = []

Then enable with: cargo run --features verbose-logging

Log Sampling for High-Frequency Events

Sample logs to reduce overhead:

// Log every 100 steps instead of every step
if step % 100 == 0 {
    debug!(
        step = step,
        avg_loss = recent_losses.iter().sum::<f32>() / 100.0,
        "Averaged loss over last 100 steps"
    );
}

// Log every 10 batches
if batch_idx % 10 == 0 {
    debug!(
        batch = batch_idx,
        gradient_norm = %grad_norm,
        "Gradient norm (sampled)"
    );
}

Bounded History to Prevent Memory Leaks

Limit history size in training monitors:

pub struct TrainingMonitor {
    reward_history: Vec<f32>,
    q_value_history: Vec<f64>,
}

impl TrainingMonitor {
    fn track_reward(&mut self, reward: f32) {
        self.reward_history.push(reward);

        // MEMORY LEAK FIX: Limit to last 1000 entries
        if self.reward_history.len() > 1000 {
            self.reward_history.drain(0..500);  // Remove oldest 500
        }
    }

    fn track_q_value(&mut self, q_value: f64) {
        self.q_value_history.push(q_value);

        // Keep only recent values for statistics
        if self.q_value_history.len() > 1000 {
            self.q_value_history.drain(0..500);
        }
    }
}

Async Logging for I/O-Bound Writes

Use async logging to avoid blocking training:

use tracing_subscriber::fmt::time::UtcTime;

// Async subscriber (non-blocking writes)
let (non_blocking, _guard) = tracing_appender::non_blocking(std::io::stdout());

let subscriber = tracing_subscriber::fmt()
    .with_writer(non_blocking)
    .with_timer(UtcTime::rfc_3339())
    .finish();

tracing::subscriber::set_global_default(subscriber)
    .expect("setting default subscriber failed");

Performance Benchmarks

Log Level Overhead per Call Use Case
error! ~50-100ns Always enabled
warn! ~50-100ns Production
info! ~100-200ns Production (epoch-level)
debug! ~200-500ns Development only
trace! ~500-1000ns Avoid in tight loops
Disabled 0ns (compiled out) Production

Key insight: debug! and trace! should never be in per-step loops in production.


5. ML Training Specific Patterns

Epoch-Level Logging (Standard Pattern)

use tracing::{info, info_span};

pub fn train(&mut self, num_epochs: usize) -> Result<()> {
    for epoch in 0..num_epochs {
        let epoch_start = Instant::now();
        let _epoch_span = info_span!("epoch", epoch=epoch).entered();

        info!("🚀 Starting epoch {}/{}", epoch, num_epochs);

        // Training phase
        let train_loss = self.train_epoch()?;

        // Validation phase
        let (val_loss, val_metrics) = self.validate()?;

        // Epoch summary with aggregated metrics
        info!(
            epoch = epoch,
            train_loss = %train_loss,
            val_loss = %val_loss,
            val_accuracy = %val_metrics.accuracy,
            q_value_mean = %val_metrics.q_value_mean,
            epsilon = %self.epsilon,
            duration_secs = epoch_start.elapsed().as_secs(),
            "✅ Epoch complete"
        );

        // Checkpoint saving
        if val_loss < self.best_val_loss {
            self.best_val_loss = val_loss;
            info!(
                checkpoint_path = %self.checkpoint_path,
                "💾 New best model saved (val_loss: {:.4})", val_loss
            );
            self.save_checkpoint()?;
        }
    }

    Ok(())
}

Batch-Level Logging (Debug Only)

use tracing::debug;

fn train_epoch(&mut self) -> Result<f32> {
    let mut total_loss = 0.0;

    for (batch_idx, batch) in self.data_loader.iter().enumerate() {
        let batch_start = Instant::now();

        // Forward + backward pass
        let loss = self.forward_backward(&batch)?;
        total_loss += loss;

        // Sample: Log every 10 batches in debug mode
        if batch_idx % 10 == 0 {
            debug!(
                batch = batch_idx,
                batch_loss = %loss,
                gradient_norm = %self.last_gradient_norm,
                batch_time_ms = batch_start.elapsed().as_millis(),
                "Batch processed (sampled)"
            );
        }
    }

    Ok(total_loss / self.data_loader.len() as f32)
}

Training vs Evaluation Mode Logging

#[instrument(skip(self))]
pub fn set_mode(&mut self, mode: TrainingMode) {
    match mode {
        TrainingMode::Train => {
            info!("Switching to TRAINING mode (exploration enabled)");
            self.model.train();
            self.epsilon = self.epsilon_train;
        }
        TrainingMode::Eval => {
            info!("Switching to EVALUATION mode (exploitation only)");
            self.model.eval();
            self.epsilon = 0.0;  // No exploration
        }
    }
}

Metric Aggregation Pattern

Instead of logging every step, aggregate and log periodically:

pub struct MetricAggregator {
    losses: Vec<f32>,
    q_values: Vec<f64>,
    rewards: Vec<f32>,
    window_size: usize,
}

impl MetricAggregator {
    pub fn add_step_metrics(&mut self, loss: f32, q_value: f64, reward: f32) {
        self.losses.push(loss);
        self.q_values.push(q_value);
        self.rewards.push(reward);

        // Log aggregated metrics every window_size steps
        if self.losses.len() >= self.window_size {
            self.log_and_reset();
        }
    }

    fn log_and_reset(&mut self) {
        let avg_loss = self.losses.iter().sum::<f32>() / self.losses.len() as f32;
        let avg_q = self.q_values.iter().sum::<f64>() / self.q_values.len() as f64;
        let avg_reward = self.rewards.iter().sum::<f32>() / self.rewards.len() as f32;

        debug!(
            window_size = self.window_size,
            avg_loss = %avg_loss,
            avg_q_value = %avg_q,
            avg_reward = %avg_reward,
            "Aggregated metrics over {} steps", self.window_size
        );

        self.losses.clear();
        self.q_values.clear();
        self.rewards.clear();
    }
}

Progress Reporting for Long Operations

use tracing::info;

pub fn load_training_data(&self, path: &Path) -> Result<Dataset> {
    info!(
        data_path = %path.display(),
        "Loading training data..."
    );

    let data = load_parquet(path)?;

    info!(
        num_samples = data.len(),
        num_features = data.num_features(),
        size_mb = data.size_in_bytes() / 1_000_000,
        "✅ Data loaded successfully"
    );

    Ok(data)
}

Error Context with Tracing

use tracing::error;
use anyhow::{Context, Result};

pub fn save_checkpoint(&self, path: &Path) -> Result<()> {
    self.checkpoint_manager
        .save(self.model.state_dict())
        .await
        .with_context(|| {
            error!(
                checkpoint_path = %path.display(),
                epoch = self.current_epoch,
                "Failed to save checkpoint"
            );
            format!("Failed to save checkpoint to {}", path.display())
        })
}

6. Production Deployment

Production Log Level Configuration

# Production: INFO level only (minimal overhead)
RUST_LOG=info cargo run --release

# Troubleshooting: DEBUG for specific modules
RUST_LOG=ml::trainers::dqn=debug,ml=info cargo run --release

# Development: Full verbosity
RUST_LOG=trace cargo run

Production Subscriber Setup

use tracing_subscriber::{fmt, EnvFilter};

pub fn init_production_logging() {
    let filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new("info"));

    fmt()
        .with_env_filter(filter)
        .with_target(true)           // Include module path
        .with_thread_ids(false)      // Disable for cleaner logs
        .with_line_number(false)     // Disable for cleaner logs
        .with_level(true)
        .compact()                   // Compact format for production
        .init();
}

JSON Logging for Production Monitoring

use tracing_subscriber::fmt::format::JsonFields;

pub fn init_json_logging() {
    fmt()
        .json()  // JSON output for log aggregation
        .with_current_span(true)
        .with_span_list(true)
        .init();
}

Output:

{
  "timestamp": "2025-01-27T10:30:45.123Z",
  "level": "INFO",
  "target": "ml::trainers::dqn",
  "fields": {
    "epoch": 42,
    "train_loss": 0.234,
    "val_loss": 0.189,
    "message": "Epoch complete"
  },
  "span": {
    "name": "train_epoch",
    "epoch": 42
  }
}

File-Based Logging with Rotation

use tracing_appender::rolling::{RollingFileAppender, Rotation};

pub fn init_file_logging() {
    let file_appender = RollingFileAppender::new(
        Rotation::DAILY,
        "/var/log/ml-training",
        "training.log"
    );

    let (non_blocking, _guard) = tracing_appender::non_blocking(file_appender);

    fmt()
        .with_writer(non_blocking)
        .with_ansi(false)  // No ANSI colors in files
        .init();
}

Performance Monitoring Integration

use tracing::{info, instrument};

#[instrument(skip(self))]
pub fn train_with_metrics(&mut self) -> Result<()> {
    let start = Instant::now();

    // Training loop
    for epoch in 0..self.num_epochs {
        self.train_epoch()?;
    }

    let duration = start.elapsed();

    info!(
        total_epochs = self.num_epochs,
        total_duration_secs = duration.as_secs(),
        avg_epoch_time_secs = duration.as_secs() / self.num_epochs as u64,
        throughput_samples_per_sec = self.total_samples / duration.as_secs(),
        "Training complete - Performance summary"
    );

    Ok(())
}

Summary: Quick Reference

Development (Local Training)

RUST_LOG=debug cargo run
  • Use info! for epoch summaries
  • Use debug! for batch metrics
  • Use warn! for all warnings

Production (GPU Cluster)

RUST_LOG=info cargo run --release
  • Use info! for epoch summaries only
  • Use warn! for degraded performance
  • Use error! for training failures
  • DISABLE debug! and trace! completely

Troubleshooting (Debugging Specific Issues)

RUST_LOG=ml::trainers::dqn=trace,ml=info cargo run
  • Enable trace! for single problematic module
  • Keep rest of system at info!

Real-World Example: DQN Training Loop

use tracing::{debug, error, info, info_span, warn};

pub struct DQNTrainer {
    agent: DQN,
    hyperparams: DQNHyperparameters,
    metrics: TrainingMetrics,
}

impl DQNTrainer {
    #[instrument(skip(self, train_data, val_data))]
    pub fn train(&mut self, train_data: &Dataset, val_data: &Dataset) -> Result<()> {
        info!(
            num_epochs = self.hyperparams.num_epochs,
            batch_size = self.hyperparams.batch_size,
            learning_rate = %self.hyperparams.learning_rate,
            "🚀 Starting DQN training"
        );

        for epoch in 0..self.hyperparams.num_epochs {
            let epoch_start = Instant::now();
            let _epoch_span = info_span!("epoch", epoch=epoch).entered();

            // Training phase
            let train_loss = self.train_epoch(train_data)?;

            // Validation phase
            let (val_loss, q_stats) = self.validate(val_data)?;

            // Check for Q-value collapse
            if q_stats.range() < 0.1 {
                warn!(
                    epoch = epoch,
                    q_min = %q_stats.min,
                    q_max = %q_stats.max,
                    q_range = %q_stats.range(),
                    "⚠️  Q-VALUE COLLAPSE detected"
                );
            }

            // Epoch summary
            info!(
                epoch = epoch,
                train_loss = %train_loss,
                val_loss = %val_loss,
                q_value_mean = %q_stats.mean,
                epsilon = %self.agent.epsilon,
                duration_secs = epoch_start.elapsed().as_secs(),
                "✅ Epoch {}/{} complete", epoch, self.hyperparams.num_epochs
            );

            // Checkpoint best model
            if val_loss < self.best_val_loss {
                self.best_val_loss = val_loss;
                self.save_checkpoint(epoch)?;
            }

            // Early stopping check
            if let Err(e) = self.check_early_stopping(val_loss) {
                warn!(
                    epoch = epoch,
                    patience_remaining = self.patience_counter,
                    "⚠️  Early stopping: {}", e
                );
                break;
            }
        }

        info!(
            total_epochs = self.hyperparams.num_epochs,
            best_val_loss = %self.best_val_loss,
            "🎉 Training complete"
        );

        Ok(())
    }

    fn train_epoch(&mut self, data: &Dataset) -> Result<f32> {
        let mut total_loss = 0.0;

        for (batch_idx, batch) in data.iter_batches(self.hyperparams.batch_size).enumerate() {
            let loss = self.train_batch(&batch)?;
            total_loss += loss;

            // Sample debug logs (every 100 batches)
            if batch_idx % 100 == 0 {
                debug!(
                    batch = batch_idx,
                    batch_loss = %loss,
                    gradient_norm = %self.last_gradient_norm,
                    "Batch processed (sampled)"
                );
            }
        }

        Ok(total_loss / data.num_batches() as f32)
    }
}

References


Document Version: 1.0 Last Updated: 2025-01-27 Author: Research Agent (Claude Code) Project: Foxhunt ML Training Infrastructure