9aad6ff60e
Epoch duration self-balances: bigger GPU → bigger auto-scaled batch → fewer steps per epoch. The manual cap created 7 different values (0, 8, 64, 100, 200, 300, 2000) across configs/tests/examples, making behavior inconsistent between environments. Removed from: DQNHyperparameters, training profiles (smoketest, localdev, production), CLI args, Argo templates, hyperopt adapter, all test overrides, supervised example. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
376 lines
14 KiB
Rust
376 lines
14 KiB
Rust
//! Hyperopt campaign configuration and runner.
|
||
//!
|
||
//! Defines campaign parameters for systematic hyperparameter search,
|
||
//! and provides a `run_campaign()` orchestrator that drives the optimizer
|
||
//! against the configured model adapter, persisting results to disk.
|
||
|
||
use std::path::PathBuf;
|
||
|
||
use serde::{Deserialize, Serialize};
|
||
|
||
use crate::gpu::capabilities::cached_capabilities;
|
||
use crate::gpu::memory_profile;
|
||
use crate::batch_size_resolver::resolve_batch_size;
|
||
|
||
// Use canonical ModelType from crate root (re-exported from common)
|
||
pub use crate::ModelType;
|
||
|
||
/// Campaign configuration for multi-trial hyperparameter optimization.
|
||
/// Campaign execution mode.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum CampaignMode {
|
||
/// Quick: fewer trials, fewer epochs. For local validation.
|
||
Quick,
|
||
/// Standard: default trials and epochs. For single-phase hyperopt.
|
||
Standard,
|
||
/// Full: all three training phases (BC → Online RL → Refinement).
|
||
/// Runs max_epochs=100 per trial to cover all phases.
|
||
/// Phase 1 (BC): epochs 0..c51_warmup_epochs (MSE + expert demos + DT)
|
||
/// Phase 2 (RL): epochs warmup..80% (all features, C51 ramp, HER)
|
||
/// Phase 3 (Refine): last 20% (pure C51, shrink-and-perturb)
|
||
Full,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct CampaignConfig {
|
||
/// Which model to optimize.
|
||
pub model_type: ModelType,
|
||
/// Total number of trials to run.
|
||
pub num_trials: usize,
|
||
/// Path to training data directory.
|
||
pub data_dir: PathBuf,
|
||
/// Maximum batch size (bounded by GPU VRAM).
|
||
pub max_batch_size: usize,
|
||
/// Reduction factor η for SHA/Hyperband.
|
||
pub early_stopping_eta: usize,
|
||
/// Maximum epochs per trial.
|
||
pub max_epochs_per_trial: usize,
|
||
/// Base directory for results output.
|
||
pub results_base_dir: PathBuf,
|
||
/// Campaign execution mode.
|
||
pub mode: CampaignMode,
|
||
/// Optional MBP-10 order book data directory for OFI features.
|
||
/// When set, enables 8 additional OFI features (VPIN, Kyle's Lambda, etc.)
|
||
/// producing state_dim=80 instead of 72.
|
||
pub mbp10_data_dir: Option<PathBuf>,
|
||
/// Optional trades data directory for VPIN/Kyle's Lambda computation.
|
||
/// Required alongside mbp10_data_dir for full OFI feature extraction.
|
||
pub trades_data_dir: Option<PathBuf>,
|
||
}
|
||
|
||
impl CampaignConfig {
|
||
/// DQN campaign defaults (50 trials, SHA with η=3, GPU-adaptive batch size).
|
||
pub fn dqn_default() -> Self {
|
||
let caps = cached_capabilities();
|
||
let max_batch = resolve_batch_size(caps, &memory_profile::estimates::DQN, 512);
|
||
Self {
|
||
model_type: ModelType::DQN,
|
||
num_trials: 50,
|
||
data_dir: PathBuf::from("test_data/real/databento/ml_training"),
|
||
max_batch_size: max_batch,
|
||
early_stopping_eta: 3,
|
||
max_epochs_per_trial: 81,
|
||
results_base_dir: PathBuf::from("ml/hyperopt_results"),
|
||
mode: CampaignMode::Standard,
|
||
mbp10_data_dir: None,
|
||
trades_data_dir: None,
|
||
}
|
||
}
|
||
|
||
/// DQN full pipeline: all 3 training phases (BC → RL → Refinement).
|
||
/// 20 trials × 50 epochs — covers Phase 1 (warmup), Phase 2 (full RL),
|
||
/// and Phase 3 (refinement with shrink-and-perturb).
|
||
pub fn dqn_full() -> Self {
|
||
let caps = cached_capabilities();
|
||
let max_batch = resolve_batch_size(caps, &memory_profile::estimates::DQN, 512);
|
||
Self {
|
||
model_type: ModelType::DQN,
|
||
num_trials: 20,
|
||
data_dir: PathBuf::from("test_data/real/databento/ml_training"),
|
||
max_batch_size: max_batch,
|
||
early_stopping_eta: 3,
|
||
max_epochs_per_trial: 50,
|
||
results_base_dir: PathBuf::from("ml/hyperopt_results"),
|
||
mode: CampaignMode::Full,
|
||
mbp10_data_dir: None,
|
||
trades_data_dir: None,
|
||
}
|
||
}
|
||
|
||
/// DQN local dev campaign (10 trials × 20 epochs — RTX 3050 friendly).
|
||
pub fn dqn_localdev() -> Self {
|
||
Self {
|
||
model_type: ModelType::DQN,
|
||
num_trials: 10,
|
||
data_dir: PathBuf::from("test_data/futures-baseline"),
|
||
max_batch_size: 64,
|
||
early_stopping_eta: 3,
|
||
max_epochs_per_trial: 20,
|
||
results_base_dir: PathBuf::from("ml/hyperopt_results"),
|
||
mode: CampaignMode::Standard,
|
||
mbp10_data_dir: None,
|
||
trades_data_dir: None,
|
||
}
|
||
}
|
||
|
||
/// PPO campaign defaults (30 trials, Hyperband, GPU-adaptive batch size).
|
||
pub fn ppo_default() -> Self {
|
||
let caps = cached_capabilities();
|
||
let max_batch = resolve_batch_size(caps, &memory_profile::estimates::PPO, 512);
|
||
Self {
|
||
model_type: ModelType::PPO,
|
||
num_trials: 30,
|
||
data_dir: PathBuf::from("test_data/real/databento/ml_training"),
|
||
max_batch_size: max_batch,
|
||
early_stopping_eta: 3,
|
||
max_epochs_per_trial: 81,
|
||
results_base_dir: PathBuf::from("ml/hyperopt_results"),
|
||
mode: CampaignMode::Standard,
|
||
mbp10_data_dir: None,
|
||
trades_data_dir: None,
|
||
}
|
||
}
|
||
|
||
/// Generate timestamped results directory path.
|
||
pub fn results_dir(&self) -> String {
|
||
let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
|
||
format!(
|
||
"{}/{}/{}",
|
||
self.results_base_dir.display(),
|
||
self.model_type,
|
||
timestamp
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Results from a completed campaign.
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct CampaignResults {
|
||
/// Number of trials that completed.
|
||
pub trials_completed: usize,
|
||
/// Best objective value found (lower is better).
|
||
pub best_loss: f64,
|
||
/// Best parameters serialized as JSON.
|
||
pub best_params_json: String,
|
||
/// Path to best checkpoint file, if saved.
|
||
pub best_checkpoint_path: Option<String>,
|
||
/// Total wall-clock time in seconds.
|
||
pub total_time_seconds: f64,
|
||
}
|
||
|
||
/// Run a hyperopt campaign with the given configuration.
|
||
///
|
||
/// Creates an `ArgminOptimizer`, runs it against the configured model adapter,
|
||
/// and saves results to the campaign results directory.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns an error if:
|
||
/// - The results directory cannot be created
|
||
/// - The model adapter fails to initialize
|
||
/// - The optimization run fails
|
||
/// - Result serialization or file writes fail
|
||
pub fn run_campaign(config: &CampaignConfig) -> anyhow::Result<CampaignResults> {
|
||
let start = std::time::Instant::now();
|
||
|
||
// Create results directory
|
||
let results_dir = config.results_dir();
|
||
std::fs::create_dir_all(&results_dir)?;
|
||
|
||
match config.model_type {
|
||
ModelType::DQN => run_dqn_campaign(config, &results_dir, start),
|
||
ModelType::PPO => {
|
||
anyhow::bail!("PPO campaign not yet implemented")
|
||
}
|
||
other @ ModelType::CompactDQN
|
||
| other @ ModelType::DistilledMicroNet
|
||
| other @ ModelType::RainbowDQN
|
||
| other @ ModelType::MAMBA
|
||
| other @ ModelType::TFT
|
||
| other @ ModelType::TGGN
|
||
| other @ ModelType::LNN
|
||
| other @ ModelType::TLOB
|
||
| other @ ModelType::Transformer
|
||
| other @ ModelType::Mamba
|
||
| other @ ModelType::LiquidNet
|
||
| other @ ModelType::TGNN
|
||
| other @ ModelType::Ensemble
|
||
| other @ ModelType::KAN
|
||
| other @ ModelType::XLSTM
|
||
| other @ ModelType::Diffusion => {
|
||
anyhow::bail!("Hyperopt campaign not supported for model type: {}", other)
|
||
}
|
||
}
|
||
}
|
||
|
||
fn run_dqn_campaign(
|
||
config: &CampaignConfig,
|
||
results_dir: &str,
|
||
start: std::time::Instant,
|
||
) -> anyhow::Result<CampaignResults> {
|
||
use crate::hyperopt::adapters::dqn::DQNTrainer;
|
||
use crate::hyperopt::optimizer::ArgminOptimizer;
|
||
|
||
// Build optimizer
|
||
let optimizer = ArgminOptimizer::builder()
|
||
.max_trials(config.num_trials)
|
||
.n_initial((config.num_trials / 2).max(1).min(5))
|
||
.seed(42)
|
||
.build();
|
||
|
||
// Create DQN adapter with optional MBP-10 + trades data for OFI features
|
||
let adapter = DQNTrainer::new(&config.data_dir, config.max_epochs_per_trial)?
|
||
.with_ofi_data_dirs(
|
||
config.mbp10_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
|
||
config.trades_data_dir.as_ref().map(|p| p.to_string_lossy().into_owned()),
|
||
);
|
||
|
||
// Run optimization
|
||
let result = optimizer.optimize(adapter)?;
|
||
|
||
// Serialize best params
|
||
let best_params_json = serde_json::to_string_pretty(&result.best_params)?;
|
||
|
||
// Save best parameters
|
||
let params_path = format!("{}/best_params.json", results_dir);
|
||
std::fs::write(¶ms_path, &best_params_json)?;
|
||
|
||
let campaign_results = CampaignResults {
|
||
trials_completed: result.all_trials.len(),
|
||
best_loss: result.best_objective,
|
||
best_params_json,
|
||
best_checkpoint_path: None, // Checkpoints saved by trainer internally
|
||
total_time_seconds: start.elapsed().as_secs_f64(),
|
||
};
|
||
|
||
// Save campaign summary
|
||
let summary_path = format!("{}/campaign_summary.json", results_dir);
|
||
std::fs::write(
|
||
&summary_path,
|
||
serde_json::to_string_pretty(&campaign_results)?,
|
||
)?;
|
||
|
||
Ok(campaign_results)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_campaign_config_dqn_defaults() {
|
||
let config = CampaignConfig::dqn_default();
|
||
assert_eq!(config.model_type, ModelType::DQN);
|
||
assert_eq!(config.num_trials, 50);
|
||
assert!(config.max_batch_size > 0, "max_batch_size should be positive");
|
||
// No longer asserts <= 230 -- batch size is dynamic based on GPU
|
||
}
|
||
|
||
#[test]
|
||
fn test_campaign_config_ppo_defaults() {
|
||
let config = CampaignConfig::ppo_default();
|
||
assert_eq!(config.model_type, ModelType::PPO);
|
||
assert_eq!(config.num_trials, 30);
|
||
}
|
||
|
||
#[test]
|
||
fn test_results_dir_creation() {
|
||
let config = CampaignConfig::dqn_default();
|
||
let dir = config.results_dir();
|
||
assert!(dir.starts_with("ml/hyperopt_results/dqn/"));
|
||
// Should contain timestamp
|
||
assert!(dir.len() > "ml/hyperopt_results/dqn/".len());
|
||
}
|
||
|
||
#[test]
|
||
fn test_model_type_display() {
|
||
assert_eq!(format!("{}", ModelType::DQN), "dqn");
|
||
assert_eq!(format!("{}", ModelType::PPO), "ppo");
|
||
}
|
||
|
||
#[test]
|
||
fn test_campaign_results_serialization() {
|
||
let results = CampaignResults {
|
||
trials_completed: 10,
|
||
best_loss: 0.42,
|
||
best_params_json: r#"{"learning_rate": 0.001}"#.to_string(),
|
||
best_checkpoint_path: Some("path/to/best.safetensors".into()),
|
||
total_time_seconds: 120.0,
|
||
};
|
||
let json = serde_json::to_string(&results).expect("serialize should succeed");
|
||
let back: CampaignResults =
|
||
serde_json::from_str(&json).expect("deserialize should succeed");
|
||
assert_eq!(back.trials_completed, 10);
|
||
assert!((back.best_loss - 0.42).abs() < 1e-10);
|
||
assert_eq!(
|
||
back.best_checkpoint_path,
|
||
Some("path/to/best.safetensors".to_owned())
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_campaign_results_without_checkpoint() {
|
||
let results = CampaignResults {
|
||
trials_completed: 5,
|
||
best_loss: 0.55,
|
||
best_params_json: "{}".into(),
|
||
best_checkpoint_path: None,
|
||
total_time_seconds: 60.0,
|
||
};
|
||
let json = serde_json::to_string(&results).expect("serialize");
|
||
assert!(json.contains("\"best_checkpoint_path\":null"));
|
||
}
|
||
|
||
/// Local hyperopt: 5 PSO trials × 50 epochs on real ES data.
|
||
/// Run with: FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib --release -- test_local_hyperopt --ignored --nocapture
|
||
#[test]
|
||
#[ignore] // GPU + real data, ~30 min on RTX 3050
|
||
fn test_local_hyperopt() {
|
||
// Initialize tracing so info!/warn!/error! output is visible in test runs
|
||
let _ = tracing_subscriber::fmt()
|
||
.with_env_filter("info")
|
||
.with_test_writer()
|
||
.try_init();
|
||
|
||
let data_dir = std::env::var("FOXHUNT_TEST_DATA")
|
||
.unwrap_or_else(|_| "test_data/futures-baseline".to_owned());
|
||
let data_path = std::path::Path::new(&data_dir);
|
||
|
||
// Auto-detect: if data_dir has ES.FUT subdirectory, use it directly
|
||
let effective_dir = if data_path.join("ES.FUT").exists() {
|
||
data_path.join("ES.FUT")
|
||
} else {
|
||
data_path.to_path_buf()
|
||
};
|
||
|
||
if !effective_dir.exists() {
|
||
eprintln!("Skipping: data dir not found: {}", effective_dir.display());
|
||
return;
|
||
}
|
||
|
||
let mut config = CampaignConfig::dqn_localdev();
|
||
config.data_dir = effective_dir;
|
||
// Env var overrides for CI (quick) vs manual (full) runs
|
||
if let Ok(t) = std::env::var("FOXHUNT_HYPEROPT_TRIALS") {
|
||
config.num_trials = t.parse().unwrap_or(config.num_trials);
|
||
}
|
||
if let Ok(e) = std::env::var("FOXHUNT_HYPEROPT_EPOCHS") {
|
||
config.max_epochs_per_trial = e.parse().unwrap_or(config.max_epochs_per_trial);
|
||
}
|
||
|
||
println!("Starting local DQN hyperopt: {} trials × {} epochs ({:?} mode)",
|
||
config.num_trials, config.max_epochs_per_trial, config.mode);
|
||
println!("Data: {}", config.data_dir.display());
|
||
|
||
let result = run_campaign(&config).expect("campaign should complete");
|
||
|
||
println!("\n=== HYPEROPT RESULTS ===");
|
||
println!("Trials completed: {}", result.trials_completed);
|
||
println!("Best loss (negative Sharpe): {:.4}", result.best_loss);
|
||
println!("Total time: {:.1}s", result.total_time_seconds);
|
||
println!("Best params:\n{}", result.best_params_json);
|
||
|
||
assert!(result.trials_completed >= 1, "Should complete at least 1 trial");
|
||
}
|
||
}
|