a5ea8713b6
Branching DQN (Tavakoli et al., 2018):
- 3 independent advantage heads (exposure=5, order=3, urgency=3) in CUDA kernel
- `use_branching` as hyperopt-tunable parameter (index 11 in 29D space)
- Branch weight pointers packed into KernelWeightPack struct
KernelWeightPack refactor:
- 87→38 CUDA kernel params by packing 48 weight pointers into 384-byte #[repr(C)] struct
- UNPACK_WEIGHT_PTRS macro in CUDA header for clean kernel-side access
- Single .arg(&weight_pack) replaces 48 individual .arg() calls
Hyperopt metrics in JSON output:
- Added `metrics: Option<serde_json::Value>` to TrialResult<P>
- `extract_metrics()` trait method with default None (DQN overrides)
- JSON output now includes per-trial backtest metrics (Sharpe, Sortino,
Calmar, Omega, drawdown, win_rate, trades) + top-level best_metrics
Prometheus backtest gauges (Rust-side, no CUDA):
- 8 new gauges: foxhunt_hyperopt_best_{sharpe,sortino,calmar,omega,
max_drawdown_pct,win_rate,total_trades,total_return_pct}
Dynamic episode scaling:
- Removed hardcoded 256 episode cap on small GPUs
- VRAM budget calculation in optimal_n_episodes() handles scaling naturally
- Removed stale .min(256) in PPO trainer
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1683 lines
64 KiB
Rust
1683 lines
64 KiB
Rust
//! Generic Bayesian Optimizer using Argmin
|
|
//!
|
|
//! This module provides a production-ready optimization implementation that works
|
|
//! with any model implementing `HyperparameterOptimizable`. It uses the argmin
|
|
//! library with Particle Swarm Optimization (PSO) for efficient parameter search.
|
|
//!
|
|
//! ## Key Features
|
|
//!
|
|
//! - **Particle Swarm Optimization**: Derivative-free, population-based optimization
|
|
//! - **Latin Hypercube Sampling**: Smart initialization for first N trials
|
|
//! - **Type-safe**: Works with any `ParameterSpace` implementation
|
|
//! - **Production-ready**: Comprehensive error handling and logging
|
|
//!
|
|
//! ## Algorithm Overview
|
|
//!
|
|
//! 1. **Initialization**: Generate `n_initial` diverse samples via LHS
|
|
//! 2. **Evaluation**: Train model with each parameter set, record objective
|
|
//! 3. **Optimization**: Use Particle Swarm to find global optima
|
|
//! 4. **Multi-restart**: Swarm automatically explores multiple regions
|
|
//! 5. **Repeat**: Continue until `max_trials` reached
|
|
//!
|
|
//! ## Usage Example
|
|
//!
|
|
//! ```rust,no_run
|
|
//! use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
|
|
//! use ml::hyperopt::adapters::mamba2::{Mamba2Trainer, Mamba2Params};
|
|
//!
|
|
//! # fn example() -> anyhow::Result<()> {
|
|
//! let trainer = Mamba2Trainer::new("data.parquet", 50)?;
|
|
//!
|
|
//! let optimizer = ArgminOptimizer::builder()
|
|
//! .max_trials(30)
|
|
//! .n_initial(5)
|
|
//! .seed(42)
|
|
//! .build();
|
|
//!
|
|
//! let result = optimizer.optimize(trainer)?;
|
|
//! println!("Best loss: {:.6}", result.best_objective);
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
use anyhow::{Context, Result};
|
|
use argmin::core::{CostFunction, Executor, State};
|
|
use argmin::solver::particleswarm::ParticleSwarm;
|
|
use ndarray::Array2;
|
|
use rand::prelude::*;
|
|
use rand::SeedableRng;
|
|
use rand_distr::Uniform;
|
|
use std::sync::{Arc, Mutex};
|
|
use std::time::Instant;
|
|
use tracing::{info, warn};
|
|
|
|
use super::traits::{
|
|
HyperparameterOptimizable, ObjectiveMode, OptimizationResult, ParameterSpace, TrialResult,
|
|
};
|
|
use crate::MLError;
|
|
|
|
/// Bayesian optimizer using argmin library
|
|
///
|
|
/// This optimizer uses Particle Swarm Optimization (PSO) with Latin Hypercube Sampling
|
|
/// for initialization. It's designed to work with any model implementing
|
|
/// `HyperparameterOptimizable`.
|
|
///
|
|
/// ## Configuration
|
|
///
|
|
/// - `max_trials`: Total evaluations (including initial samples)
|
|
/// - `n_initial`: Number of Latin Hypercube samples for initialization
|
|
/// - `n_particles`: Number of particles in the swarm (default: 20)
|
|
/// - `seed`: Random seed for reproducibility (optional)
|
|
///
|
|
/// ## Performance Characteristics
|
|
///
|
|
/// - Initial samples: ~O(n_initial * training_time)
|
|
/// - Per-iteration overhead: ~1-10ms for swarm updates
|
|
/// - Memory usage: O(max_trials) for trial history
|
|
///
|
|
/// ## Advantages over Nelder-Mead
|
|
///
|
|
/// - Supports multi-dimensional vector parameters (Vec<f64>)
|
|
/// - Better at avoiding local minima through swarm intelligence
|
|
/// - Scales well to higher dimensions (20-50 parameters)
|
|
/// - No gradient information required
|
|
#[derive(Debug, Clone)]
|
|
pub struct ArgminOptimizer {
|
|
/// Maximum number of trials (evaluations)
|
|
pub max_trials: usize,
|
|
/// Number of initial samples via Latin Hypercube Sampling
|
|
pub n_initial: usize,
|
|
/// Number of particles in the swarm
|
|
pub n_particles: usize,
|
|
/// Random seed for reproducibility
|
|
pub seed: Option<u64>,
|
|
/// Maximum iterations per restart
|
|
pub max_iters_per_restart: usize,
|
|
/// Model name for Prometheus metrics progress reporting (e.g. "dqn", "ppo")
|
|
pub model_name: Option<String>,
|
|
}
|
|
|
|
impl Default for ArgminOptimizer {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_trials: 30,
|
|
n_initial: 5,
|
|
n_particles: 20,
|
|
seed: None,
|
|
max_iters_per_restart: 50,
|
|
model_name: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ArgminOptimizer {
|
|
/// Create a new optimizer with default settings
|
|
///
|
|
/// Uses 30 max trials and 5 initial samples.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Create optimizer with specified settings
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `max_trials` - Total number of trials (must be > n_initial)
|
|
/// * `n_initial` - Number of initial LHS samples (typically 3-10)
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Panics if `max_trials <= n_initial` or if either is zero.
|
|
pub fn with_trials(max_trials: usize, n_initial: usize) -> Self {
|
|
assert!(max_trials > n_initial, "max_trials must be > n_initial");
|
|
assert!(n_initial > 0, "n_initial must be > 0");
|
|
|
|
Self {
|
|
max_trials,
|
|
n_initial,
|
|
n_particles: 20,
|
|
seed: None,
|
|
max_iters_per_restart: 50,
|
|
model_name: None,
|
|
}
|
|
}
|
|
|
|
/// Set random seed for reproducibility
|
|
pub fn with_seed(mut self, seed: u64) -> Self {
|
|
self.seed = Some(seed);
|
|
self
|
|
}
|
|
|
|
/// Start building an optimizer with custom configuration
|
|
pub fn builder() -> ArgminOptimizerBuilder {
|
|
ArgminOptimizerBuilder::new()
|
|
}
|
|
|
|
/// Generate Latin Hypercube samples for initialization
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_samples` - Number of samples to generate
|
|
/// * `bounds` - Parameter bounds [(min, max), ...]
|
|
/// * `rng` - Random number generator
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Array of shape (n_samples, n_params) with LHS samples
|
|
pub fn latin_hypercube_sampling<R: Rng>(
|
|
n_samples: usize,
|
|
bounds: &[(f64, f64)],
|
|
rng: &mut R,
|
|
) -> Array2<f64> {
|
|
let n_params = bounds.len();
|
|
let mut samples = Array2::zeros((n_samples, n_params));
|
|
|
|
// For each parameter dimension
|
|
for dim in 0..n_params {
|
|
// Create permutation of [0, 1, 2, ..., n_samples-1]
|
|
let mut indices: Vec<usize> = (0..n_samples).collect();
|
|
indices.shuffle(rng);
|
|
|
|
// Generate stratified samples
|
|
let uniform = Uniform::new(0.0, 1.0);
|
|
for (sample_idx, &perm_idx) in indices.iter().enumerate() {
|
|
let segment_start = perm_idx as f64 / n_samples as f64;
|
|
let segment_end = (perm_idx + 1) as f64 / n_samples as f64;
|
|
let random_offset: f64 = rng.sample(uniform);
|
|
let normalized = segment_start + random_offset * (segment_end - segment_start);
|
|
|
|
// Denormalize to actual bounds
|
|
let (min, max) = bounds[dim];
|
|
samples[[sample_idx, dim]] = min + normalized * (max - min);
|
|
}
|
|
}
|
|
|
|
samples
|
|
}
|
|
|
|
/// Run optimization on a model
|
|
///
|
|
/// This method:
|
|
/// 1. Extracts parameter bounds from `P::continuous_bounds()`
|
|
/// 2. Generates initial samples via Latin Hypercube
|
|
/// 3. Evaluates each sample by calling `model.train_with_params()`
|
|
/// 4. Uses Nelder-Mead to optimize from best initial point
|
|
/// 5. Performs multi-restart from other good points
|
|
/// 6. Repeats until `max_trials` reached
|
|
///
|
|
/// # Type Parameters
|
|
///
|
|
/// - `M`: Model type implementing `HyperparameterOptimizable`
|
|
/// - `P`: Parameter type implementing `ParameterSpace`
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `model` - Mutable reference to model (needed for training)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `OptimizationResult<P>` containing best parameters and full trial history
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns error if:
|
|
/// - Parameter space has zero dimensions
|
|
/// - Initial sampling fails
|
|
/// - All trials produce errors (no valid observations)
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
|
|
/// # use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
|
|
/// # fn example() -> anyhow::Result<()> {
|
|
/// let mut trainer = Mamba2Trainer::new("data.parquet", 50)?;
|
|
/// let optimizer = ArgminOptimizer::new();
|
|
/// let result = optimizer.optimize(trainer)?;
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
#[allow(clippy::unwrap_in_result)]
|
|
#[allow(clippy::cognitive_complexity)]
|
|
pub fn optimize<M>(&self, mut model: M) -> Result<OptimizationResult<M::Params>>
|
|
where
|
|
M: HyperparameterOptimizable + Send,
|
|
M::Params: ParameterSpace + Send,
|
|
{
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Bayesian Hyperparameter Optimization (Argmin) ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
// Detect hardware budget and extract VRAM-aware parameter space bounds
|
|
let budget = crate::traits::HardwareBudget::detect();
|
|
info!("Hardware budget: {}MB GPU memory", budget.gpu_memory_mb);
|
|
let bounds = M::Params::continuous_bounds_for(&budget);
|
|
let n_params = bounds.len();
|
|
|
|
if n_params == 0 {
|
|
return Err(MLError::ConfigError("Parameter space has zero dimensions".to_owned())
|
|
.into());
|
|
}
|
|
|
|
info!("Configuration:");
|
|
info!(" Max Trials: {}", self.max_trials);
|
|
info!(" Initial Samples: {}", self.n_initial);
|
|
info!(" Swarm Particles: {}", self.n_particles);
|
|
info!(" Parameters: {}", n_params);
|
|
info!(" Max Iters/Restart: {}", self.max_iters_per_restart);
|
|
|
|
let param_names = M::Params::param_names();
|
|
for (i, name) in param_names.iter().enumerate() {
|
|
info!(" {} - [{:.6}, {:.6}]", name, bounds[i].0, bounds[i].1);
|
|
}
|
|
|
|
// Initialize RNG
|
|
let mut rng = match self.seed {
|
|
Some(seed) => StdRng::seed_from_u64(seed),
|
|
None => StdRng::from_entropy(),
|
|
};
|
|
|
|
// Generate initial samples using Latin Hypercube
|
|
info!(
|
|
"Generating {} initial samples with Latin Hypercube Sampling...",
|
|
self.n_initial
|
|
);
|
|
let initial_samples = Self::latin_hypercube_sampling(self.n_initial, &bounds, &mut rng);
|
|
info!("✓ Generated {} initial samples", initial_samples.nrows());
|
|
|
|
// Shared state for collecting trial results
|
|
let trial_results = Arc::new(Mutex::new(Vec::new()));
|
|
let trial_counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
|
|
|
// Evaluate initial samples
|
|
info!("Evaluating initial samples...");
|
|
for i in 0..self.n_initial {
|
|
let continuous_vec: Vec<f64> = initial_samples.row(i).to_vec();
|
|
Self::evaluate_point(
|
|
&continuous_vec,
|
|
&mut model,
|
|
&trial_results,
|
|
&trial_counter,
|
|
¶m_names,
|
|
)?;
|
|
if let Some(ref name) = self.model_name {
|
|
common::metrics::training_metrics::set_hyperopt_trial(
|
|
name,
|
|
(i + 1) as f64,
|
|
self.max_trials as f64,
|
|
);
|
|
}
|
|
}
|
|
|
|
// Create trial budget observer BEFORE creating cost function
|
|
// This observer enforces strict trial limits (fixes PSO infinite iteration bug)
|
|
let observer = crate::TrialBudgetObserver::new(self.max_trials);
|
|
|
|
// Create cost function wrapper
|
|
let optimization_start = Instant::now();
|
|
let cost_fn = ObjectiveFunction {
|
|
model: Arc::new(Mutex::new(model)),
|
|
trial_results: Arc::clone(&trial_results),
|
|
trial_counter: Arc::clone(&trial_counter),
|
|
param_names: param_names.clone(),
|
|
bounds: bounds.clone(),
|
|
observer: observer.clone(),
|
|
model_name: self.model_name.clone(),
|
|
max_trials: self.max_trials,
|
|
optimization_start,
|
|
};
|
|
|
|
// Find best initial point to start optimization
|
|
let trials = trial_results.lock().map_err(|e| {
|
|
anyhow::anyhow!("Failed to lock trial results: {}", e)
|
|
})?.clone();
|
|
let best_initial = trials
|
|
.iter()
|
|
.min_by(|a, b| {
|
|
a.objective
|
|
.partial_cmp(&b.objective)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
})
|
|
.ok_or_else(|| {
|
|
anyhow::anyhow!("No initial trials completed successfully -- cannot start PSO")
|
|
})?;
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Starting Particle Swarm Optimization ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
info!("Best initial objective: {:.6}", best_initial.objective);
|
|
info!("Execution mode: Sequential trials (model locked by Mutex, rayon for swarm only)");
|
|
|
|
// Determine how many iterations we can afford
|
|
let trials_used = trials.len();
|
|
let remaining_trials = self.max_trials.saturating_sub(trials_used);
|
|
|
|
// FIX: Each PSO iteration evaluates n_particles candidates (sequentially via mutex)
|
|
// PSO evaluates ALL particles in swarm per iteration, so divide remaining budget
|
|
// by swarm size to prevent trial count overflow (fixes 962 trial bug)
|
|
// CRITICAL FIX (2025-11-07): Use CEILING division to ensure all trials complete
|
|
// Example: 8 remaining ÷ 20 particles = 0.4 → ceil to 1 iteration (not 0)
|
|
let max_iters_by_budget =
|
|
((remaining_trials as f64) / (self.n_particles as f64)).ceil() as usize;
|
|
|
|
let max_iters = max_iters_by_budget.min(self.max_iters_per_restart);
|
|
|
|
info!(
|
|
"PSO Budget: {} iterations ({} remaining trials ÷ {} particles = {} max iters)",
|
|
max_iters, remaining_trials, self.n_particles, max_iters_by_budget
|
|
);
|
|
|
|
if max_iters > 0 {
|
|
// Create bounds vectors for ParticleSwarm
|
|
let lower_bounds: Vec<f64> = bounds.iter().map(|(min, _)| *min).collect();
|
|
let upper_bounds: Vec<f64> = bounds.iter().map(|(_, max)| *max).collect();
|
|
|
|
// Create Particle Swarm solver
|
|
let solver = ParticleSwarm::new((lower_bounds, upper_bounds), self.n_particles);
|
|
|
|
// Run optimization (parallel execution enabled via rayon feature)
|
|
// CRITICAL FIX (2025-11-03): Removed .target_cost(0.0) to prevent early termination
|
|
// PSO must run for exactly max_iters iterations to complete all requested trials
|
|
// CRITICAL FIX (2025-11-07): Added TrialBudgetObserver to enforce strict trial limits
|
|
let res = Executor::new(cost_fn, solver)
|
|
.configure(|state| state.max_iters(max_iters as u64))
|
|
.add_observer(
|
|
observer,
|
|
argmin::core::observers::ObserverMode::Always,
|
|
)
|
|
.run();
|
|
|
|
// Handle budget exhaustion gracefully (not an error condition)
|
|
match res {
|
|
Ok(result) => {
|
|
info!("Optimization complete:");
|
|
info!(" Final cost: {:.6}", result.state().get_best_cost());
|
|
info!(" Iterations: {}", result.state().get_iter());
|
|
info!(" Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed));
|
|
},
|
|
Err(e) if e.to_string().contains("Trial budget exhausted") => {
|
|
info!(
|
|
"PSO terminated: trial budget reached ({} trials)",
|
|
self.max_trials
|
|
);
|
|
info!(" Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed));
|
|
},
|
|
Err(e) => return Err(e),
|
|
}
|
|
} else {
|
|
info!("No remaining budget for Particle Swarm optimization");
|
|
}
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Optimization Complete ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
// Extract final results
|
|
let trials = match Arc::try_unwrap(trial_results) {
|
|
Ok(mutex) => mutex.into_inner().map_err(|e| {
|
|
anyhow::anyhow!("Failed to unwrap trial results mutex: {}", e)
|
|
})?,
|
|
Err(arc) => arc.lock().map_err(|e| {
|
|
anyhow::anyhow!("Failed to lock trial results: {}", e)
|
|
})?.clone(),
|
|
};
|
|
|
|
if trials.is_empty() {
|
|
return Err(MLError::TrainingError(
|
|
"All trials failed - no valid observations".to_owned(),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
// Create optimization result
|
|
let result = OptimizationResult::from_trials(trials);
|
|
|
|
info!("Best Parameters Found:");
|
|
let best_continuous = result.best_params.to_continuous();
|
|
for (i, name) in param_names.iter().enumerate() {
|
|
info!(" {}: {:.6}", name, best_continuous[i]);
|
|
}
|
|
info!("Best Objective: {:.6}", result.best_objective);
|
|
info!("Total Improvement: {:.6}", result.total_improvement());
|
|
info!("Improvement: {:.2}%", result.improvement_percentage());
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Evaluate a single point in the parameter space.
|
|
///
|
|
/// Converts `continuous_vec` to typed parameters, trains the model,
|
|
/// records the trial result, and returns the objective value.
|
|
#[allow(clippy::unwrap_in_result)]
|
|
#[allow(clippy::cognitive_complexity)]
|
|
pub(crate) fn evaluate_point<M>(
|
|
continuous_vec: &[f64],
|
|
model: &mut M,
|
|
trial_results: &Arc<Mutex<Vec<TrialResult<M::Params>>>>,
|
|
trial_counter: &Arc<std::sync::atomic::AtomicUsize>,
|
|
_param_names: &[&'static str],
|
|
) -> Result<f64>
|
|
where
|
|
M: HyperparameterOptimizable,
|
|
M::Params: ParameterSpace,
|
|
{
|
|
let start_time = Instant::now();
|
|
|
|
// Increment trial counter
|
|
let trial_num = trial_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!(
|
|
"║ Trial {}: Evaluating Parameters ║",
|
|
trial_num
|
|
);
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
// Convert continuous vector to parameters BEFORE logging
|
|
let params =
|
|
M::Params::from_continuous(continuous_vec).context("Failed to convert parameters")?;
|
|
|
|
// Log CONVERTED parameters (shows actual values: learning_rate ~1e-4, not -11)
|
|
info!(" Parameters (converted): {:?}", params);
|
|
|
|
// Train model with parameters
|
|
let metrics = model
|
|
.train_with_params(params.clone())
|
|
.context(format!("Training failed for trial {}", trial_num))?;
|
|
|
|
// Extract objective (replace NaN/Inf with large penalty)
|
|
let raw_objective = M::extract_objective(&metrics);
|
|
let objective = if raw_objective.is_finite() {
|
|
raw_objective
|
|
} else {
|
|
warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
|
|
1e6
|
|
};
|
|
let trial_metrics = M::extract_metrics(&metrics);
|
|
let duration_secs = start_time.elapsed().as_secs_f64();
|
|
|
|
info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);
|
|
info!(" Objective: {:.6}", objective);
|
|
|
|
// Record trial result
|
|
{
|
|
let mut results = trial_results.lock().map_err(|e| {
|
|
anyhow::anyhow!("Failed to lock trial results: {}", e)
|
|
})?;
|
|
results.push(TrialResult {
|
|
trial_num,
|
|
params,
|
|
objective,
|
|
duration_secs,
|
|
metrics: trial_metrics,
|
|
});
|
|
}
|
|
|
|
Ok(objective)
|
|
}
|
|
|
|
/// Run two-phase optimization: episode reward then Sharpe ratio.
|
|
///
|
|
/// Phase A: `max_trials/2` trials optimizing for fast convergence (episode reward).
|
|
/// Phase B: remaining trials optimizing for financial quality (Sharpe ratio),
|
|
/// seeded from Phase A's best parameter configuration.
|
|
///
|
|
/// This reuses the existing `optimize()` method for each phase. The objective
|
|
/// mode is set on the model before each phase via `TwoPhaseObjective`.
|
|
///
|
|
/// # Current Limitations
|
|
///
|
|
/// Full two-phase with Phase B seeding requires `M: Clone`. Currently runs
|
|
/// Phase A only and returns its result. Phase B will be enabled once the
|
|
/// model types implement `Clone`.
|
|
#[allow(clippy::cognitive_complexity, clippy::similar_names)]
|
|
pub fn optimize_two_phase<M>(&self, mut model: M) -> Result<OptimizationResult<M::Params>>
|
|
where
|
|
M: HyperparameterOptimizable + TwoPhaseObjective + Send,
|
|
M::Params: ParameterSpace + Send,
|
|
{
|
|
let phase_a_trials = self.max_trials / 2;
|
|
let phase_b_trials = self.max_trials - phase_a_trials;
|
|
|
|
info!("======= Two-Phase Optimization =======");
|
|
info!("Phase A: {} trials optimizing episode reward", phase_a_trials);
|
|
info!("Phase B: {} trials optimizing Sharpe ratio (pending Clone support)", phase_b_trials);
|
|
|
|
// Phase A: Episode Reward
|
|
model.set_objective_mode(ObjectiveMode::EpisodeReward);
|
|
let phase_a_optimizer = ArgminOptimizer::with_trials(
|
|
phase_a_trials,
|
|
self.n_initial.min(phase_a_trials - 1).max(1),
|
|
);
|
|
let phase_a_result = phase_a_optimizer.optimize(model)?;
|
|
|
|
info!(
|
|
"Phase A complete: best_objective={:.6}, evaluated {} trials",
|
|
phase_a_result.best_objective,
|
|
phase_a_result.all_trials.len()
|
|
);
|
|
|
|
// TODO: Phase B requires M: Clone to reconstruct model from Phase A result.
|
|
// Once DQNTrainer implements Clone, Phase B will:
|
|
// 1. Extract the model back (currently consumed by optimize())
|
|
// 2. Set objective mode to Sharpe
|
|
// 3. Seed Phase B with top-3 parameter configs from Phase A
|
|
// 4. Run remaining trials with Sharpe-based objective
|
|
|
|
info!(
|
|
"======= Two-Phase Complete: best_objective={:.6} =======",
|
|
phase_a_result.best_objective
|
|
);
|
|
|
|
Ok(phase_a_result)
|
|
}
|
|
|
|
/// Run optimization with parallel trial evaluation.
|
|
///
|
|
/// Requires `M: Clone` so each evaluation thread gets its own model instance.
|
|
/// Parallelism is controlled by `RAYON_NUM_THREADS` (for PSO phase) and
|
|
/// internal thread scoping (for LHS phase, capped at `n_initial` threads).
|
|
///
|
|
/// DQN/PPO trials are CPU-bound (~1 core each, 278MB VRAM). On an 8-vCPU L4
|
|
/// with 24GB VRAM, 4-6 concurrent trials safely fit. This gives ~4-5x speedup
|
|
/// over sequential execution.
|
|
#[allow(clippy::unwrap_in_result)]
|
|
#[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
|
|
pub fn optimize_parallel<M>(&self, model: M) -> Result<OptimizationResult<M::Params>>
|
|
where
|
|
M: HyperparameterOptimizable + Clone + Send,
|
|
M::Params: ParameterSpace + Send,
|
|
{
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Parallel Hyperparameter Optimization (Argmin + Rayon) ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
let rayon_threads = rayon::current_num_threads();
|
|
info!("Parallel execution: {} rayon threads", rayon_threads);
|
|
|
|
// Detect hardware budget and extract VRAM-aware parameter space bounds
|
|
let budget = crate::traits::HardwareBudget::detect();
|
|
info!("Hardware budget: {}MB GPU memory", budget.gpu_memory_mb);
|
|
|
|
// Compute VRAM-aware concurrency strategy
|
|
let strategy = budget.plan_hyperopt(
|
|
M::Params::model_overhead_mb(),
|
|
M::Params::per_sample_mb(),
|
|
64.0,
|
|
4096.0,
|
|
);
|
|
info!("HyperoptStrategy: {} concurrent trials, batch bounds [{}, {}], {:.0}MB/trial, {:.0}MB total",
|
|
strategy.max_concurrent_trials,
|
|
strategy.batch_size_bounds.0, strategy.batch_size_bounds.1,
|
|
strategy.per_trial_memory_mb, strategy.total_reserved_mb);
|
|
|
|
// Auto-scale PSO particles to match GPU concurrency for maximum
|
|
// hardware utilization. Each PSO iteration evaluates all particles,
|
|
// so matching particle count to VRAM slots fills every GPU slot per
|
|
// iteration. Cap at max_trials — no benefit from more particles than
|
|
// the total trial budget. On CPU (max_concurrent = 1), fall back to
|
|
// configured n_particles for sequential PSO exploration.
|
|
let effective_particles = if strategy.max_concurrent_trials > 1 {
|
|
let scaled = strategy.max_concurrent_trials.min(self.max_trials);
|
|
if scaled != self.n_particles {
|
|
info!("Auto-scaled PSO particles: {} -> {} (GPU fits {} concurrent, budget {})",
|
|
self.n_particles, scaled, strategy.max_concurrent_trials, self.max_trials);
|
|
}
|
|
scaled
|
|
} else {
|
|
self.n_particles.min(self.max_trials)
|
|
};
|
|
|
|
// Trial budget is user-configured — never auto-inflate.
|
|
let effective_max_trials = self.max_trials;
|
|
|
|
let bounds = M::Params::continuous_bounds_for(&budget);
|
|
let n_params = bounds.len();
|
|
|
|
if n_params == 0 {
|
|
return Err(MLError::ConfigError("Parameter space has zero dimensions".to_owned())
|
|
.into());
|
|
}
|
|
|
|
info!("Configuration:");
|
|
info!(" Max Trials: {} (configured: {})", effective_max_trials, self.max_trials);
|
|
info!(" Initial Samples: {}", self.n_initial);
|
|
info!(" Swarm Particles: {} (configured: {})", effective_particles, self.n_particles);
|
|
info!(" Parameters: {}", n_params);
|
|
info!(" Max Iters/Restart: {}", self.max_iters_per_restart);
|
|
info!(" Parallel Threads: {}", rayon_threads);
|
|
|
|
let param_names = M::Params::param_names();
|
|
for (i, name) in param_names.iter().enumerate() {
|
|
info!(" {} - [{:.6}, {:.6}]", name, bounds[i].0, bounds[i].1);
|
|
}
|
|
|
|
// Initialize RNG
|
|
let mut rng = match self.seed {
|
|
Some(seed) => StdRng::seed_from_u64(seed),
|
|
None => StdRng::from_entropy(),
|
|
};
|
|
|
|
// Generate initial samples using Latin Hypercube
|
|
info!(
|
|
"Generating {} initial samples with Latin Hypercube Sampling...",
|
|
self.n_initial
|
|
);
|
|
let initial_samples = Self::latin_hypercube_sampling(self.n_initial, &bounds, &mut rng);
|
|
info!("Generated {} initial samples", initial_samples.nrows());
|
|
|
|
// Shared state for collecting trial results
|
|
let trial_results = Arc::new(Mutex::new(Vec::new()));
|
|
let trial_counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
|
|
|
// OOM Layer 2: Start at ceil(max/2), ramp up after first success
|
|
let initial_concurrency = strategy.max_concurrent_trials.div_ceil(2).max(1);
|
|
let concurrency = std::sync::atomic::AtomicUsize::new(initial_concurrency);
|
|
|
|
info!(
|
|
"Evaluating {} initial samples with ramp-up concurrency (start: {}, max: {})...",
|
|
self.n_initial, initial_concurrency, strategy.max_concurrent_trials
|
|
);
|
|
|
|
let initial_vecs: Vec<Vec<f64>> = (0..self.n_initial)
|
|
.map(|i| initial_samples.row(i).to_vec())
|
|
.collect();
|
|
|
|
for chunk in initial_vecs.chunks(initial_concurrency) {
|
|
// OOM Layer 4: VRAM watchdog — check free memory before spawning
|
|
if let Ok((_total_mb, free_mb, _name)) = ml_core::memory_optimization::auto_batch_size::detect_gpu_memory() {
|
|
let total_mb = budget.gpu_memory_mb as f64;
|
|
if total_mb > 0.0 && free_mb / total_mb < 0.10 {
|
|
warn!("VRAM watchdog: only {:.0}MB free ({:.0}%), reducing concurrency",
|
|
free_mb, (free_mb / total_mb) * 100.0);
|
|
let current = concurrency.load(std::sync::atomic::Ordering::Relaxed);
|
|
concurrency.store((current / 2).max(1), std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
}
|
|
|
|
let lhs_errors: Vec<Option<anyhow::Error>> = std::thread::scope(|scope| {
|
|
let handles: Vec<_> = chunk.iter().map(|continuous_vec| {
|
|
let mut model_clone = model.clone();
|
|
let results = Arc::clone(&trial_results);
|
|
let counter = Arc::clone(&trial_counter);
|
|
let names = param_names.clone();
|
|
let vec = continuous_vec.clone();
|
|
scope.spawn(move || {
|
|
Self::evaluate_point(&vec, &mut model_clone, &results, &counter, &names).err()
|
|
})
|
|
}).collect();
|
|
handles.into_iter().map(|h| h.join().unwrap_or(None)).collect()
|
|
});
|
|
|
|
for err in lhs_errors.into_iter().flatten() {
|
|
warn!("LHS evaluation error (continuing): {}", err);
|
|
}
|
|
|
|
// OOM Layer 2: Ramp up after first chunk succeeds
|
|
let current = concurrency.load(std::sync::atomic::Ordering::Relaxed);
|
|
if current < strategy.max_concurrent_trials {
|
|
info!("Ramp-up: {} -> {} concurrent trials (first chunk succeeded)",
|
|
current, strategy.max_concurrent_trials);
|
|
concurrency.store(strategy.max_concurrent_trials, std::sync::atomic::Ordering::Relaxed);
|
|
}
|
|
|
|
if let Some(ref name) = self.model_name {
|
|
let done = trial_counter.load(std::sync::atomic::Ordering::Relaxed);
|
|
common::metrics::training_metrics::set_hyperopt_trial(
|
|
name,
|
|
done as f64,
|
|
self.max_trials as f64,
|
|
);
|
|
}
|
|
}
|
|
|
|
// Create trial budget observer BEFORE creating cost function
|
|
let observer = crate::TrialBudgetObserver::new(effective_max_trials);
|
|
|
|
// Create parallel cost function — clones model per evaluation instead of holding mutex
|
|
let cost_fn = ParallelObjectiveFunction {
|
|
model: Arc::new(Mutex::new(model)),
|
|
trial_results: Arc::clone(&trial_results),
|
|
trial_counter: Arc::clone(&trial_counter),
|
|
param_names: param_names.clone(),
|
|
bounds: bounds.clone(),
|
|
observer: observer.clone(),
|
|
model_name: self.model_name.clone(),
|
|
max_trials: self.max_trials,
|
|
optimization_start: Instant::now(),
|
|
};
|
|
|
|
// Find best initial point
|
|
let trials = trial_results.lock().map_err(|e| {
|
|
anyhow::anyhow!("Failed to lock trial results: {}", e)
|
|
})?.clone();
|
|
let best_initial = trials
|
|
.iter()
|
|
.min_by(|a, b| {
|
|
a.objective
|
|
.partial_cmp(&b.objective)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
})
|
|
.ok_or_else(|| {
|
|
anyhow::anyhow!("No initial trials completed successfully -- cannot start PSO")
|
|
})?;
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Starting Parallel Particle Swarm Optimization ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
info!("Best initial objective: {:.6}", best_initial.objective);
|
|
let final_concurrency = concurrency.load(std::sync::atomic::Ordering::Relaxed);
|
|
info!("Execution mode: Parallel trials (clone-per-particle, {} VRAM-concurrent, {} rayon threads)",
|
|
final_concurrency, rayon_threads);
|
|
if rayon_threads > strategy.max_concurrent_trials {
|
|
warn!("Rayon threads ({}) exceed VRAM concurrency ({}). Consider capping rayon pool at startup.",
|
|
rayon_threads, strategy.max_concurrent_trials);
|
|
}
|
|
|
|
// Determine PSO iterations
|
|
let trials_used = trials.len();
|
|
let remaining_trials = effective_max_trials.saturating_sub(trials_used);
|
|
let max_iters_by_budget =
|
|
((remaining_trials as f64) / (effective_particles as f64)).ceil() as usize;
|
|
let max_iters = max_iters_by_budget.min(self.max_iters_per_restart);
|
|
|
|
info!(
|
|
"PSO Budget: {} iterations ({} remaining trials / {} particles = {} max iters)",
|
|
max_iters, remaining_trials, effective_particles, max_iters_by_budget
|
|
);
|
|
|
|
if max_iters > 0 {
|
|
let lower_bounds: Vec<f64> = bounds.iter().map(|(min, _)| *min).collect();
|
|
let upper_bounds: Vec<f64> = bounds.iter().map(|(_, max)| *max).collect();
|
|
|
|
let solver = ParticleSwarm::new((lower_bounds, upper_bounds), effective_particles);
|
|
|
|
let res = Executor::new(cost_fn, solver)
|
|
.configure(|state| state.max_iters(max_iters as u64))
|
|
.add_observer(
|
|
observer,
|
|
argmin::core::observers::ObserverMode::Always,
|
|
)
|
|
.run();
|
|
|
|
match res {
|
|
Ok(result) => {
|
|
info!("Optimization complete:");
|
|
info!(" Final cost: {:.6}", result.state().get_best_cost());
|
|
info!(" Iterations: {}", result.state().get_iter());
|
|
info!(" Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed));
|
|
},
|
|
Err(e) if e.to_string().contains("Trial budget exhausted") => {
|
|
info!(
|
|
"PSO terminated: trial budget reached ({} trials, configured {})",
|
|
effective_max_trials, self.max_trials
|
|
);
|
|
info!(" Evaluations: {}", trial_counter.load(std::sync::atomic::Ordering::Relaxed));
|
|
},
|
|
Err(e) => return Err(e),
|
|
}
|
|
} else {
|
|
info!("No remaining budget for Particle Swarm optimization");
|
|
}
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Optimization Complete ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
// Extract final results
|
|
let trials = match Arc::try_unwrap(trial_results) {
|
|
Ok(mutex) => mutex.into_inner().map_err(|e| {
|
|
anyhow::anyhow!("Failed to unwrap trial results mutex: {}", e)
|
|
})?,
|
|
Err(arc) => arc.lock().map_err(|e| {
|
|
anyhow::anyhow!("Failed to lock trial results: {}", e)
|
|
})?.clone(),
|
|
};
|
|
|
|
if trials.is_empty() {
|
|
return Err(MLError::TrainingError(
|
|
"All trials failed - no valid observations".to_owned(),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
let result = OptimizationResult::from_trials(trials);
|
|
|
|
info!("Best Parameters Found:");
|
|
let best_continuous = result.best_params.to_continuous();
|
|
for (i, name) in param_names.iter().enumerate() {
|
|
info!(" {}: {:.6}", name, best_continuous[i]);
|
|
}
|
|
info!("Best Objective: {:.6}", result.best_objective);
|
|
info!("Total Improvement: {:.6}", result.total_improvement());
|
|
info!("Improvement: {:.2}%", result.improvement_percentage());
|
|
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
/// Run optimization using Tree-Parzen Estimator (TPE) instead of PSO.
|
|
///
|
|
/// TPE builds separate density models for "good" and "bad" trials,
|
|
/// then suggests new points by maximizing Expected Improvement.
|
|
/// Better sample efficiency than PSO in 20-50D spaces.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `model` - Model implementing `HyperparameterOptimizable` (consumed)
|
|
/// * `max_trials` - Total number of trials to evaluate
|
|
/// * `n_initial` - Number of initial LHS samples before TPE kicks in
|
|
/// * `seed` - Optional random seed for reproducibility
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `OptimizationResult<M::Params>` containing best parameters and full trial history
|
|
#[allow(clippy::cognitive_complexity)]
|
|
pub fn optimize_with_tpe<M>(
|
|
mut model: M,
|
|
max_trials: usize,
|
|
n_initial: usize,
|
|
seed: Option<u64>,
|
|
model_name: Option<&str>,
|
|
) -> Result<OptimizationResult<M::Params>>
|
|
where
|
|
M: HyperparameterOptimizable + Send,
|
|
M::Params: ParameterSpace + Send,
|
|
{
|
|
use crate::tpe::{TpeConfig, TpeOptimizer};
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!("║ Bayesian Hyperparameter Optimization (TPE) ║");
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
let budget = crate::traits::HardwareBudget::detect();
|
|
let bounds = M::Params::continuous_bounds_for(&budget);
|
|
let n_params = bounds.len();
|
|
|
|
if n_params == 0 {
|
|
return Err(MLError::ConfigError("Parameter space has zero dimensions".to_owned())
|
|
.into());
|
|
}
|
|
|
|
// Clamp n_initial to valid range: at least 1, at most max_trials - 1
|
|
let n_initial = n_initial.max(1).min(max_trials.saturating_sub(1).max(1));
|
|
|
|
info!("Configuration:");
|
|
info!(" Optimizer: TPE (Tree-Parzen Estimator)");
|
|
info!(" Max Trials: {}", max_trials);
|
|
info!(" Initial LHS Samples: {}", n_initial);
|
|
info!(" Parameters: {}", n_params);
|
|
let n_candidates = 256_usize.max(8 * n_params);
|
|
let gamma = if max_trials < 50 { 0.15 } else { 0.25 };
|
|
info!(" Gamma (good quantile): {:.2}", gamma);
|
|
info!(" EI Candidates: {} (scaled for {}D)", n_candidates, n_params);
|
|
|
|
let param_names = M::Params::param_names();
|
|
for (i, name) in param_names.iter().enumerate() {
|
|
if let Some(&(lo, hi)) = bounds.get(i) {
|
|
info!(" {} - [{:.6}, {:.6}]", name, lo, hi);
|
|
}
|
|
}
|
|
|
|
let tpe_config = TpeConfig {
|
|
n_dims: n_params,
|
|
max_trials,
|
|
n_initial,
|
|
gamma,
|
|
n_candidates,
|
|
seed,
|
|
};
|
|
let mut tpe = TpeOptimizer::new(tpe_config);
|
|
|
|
let trial_results = Arc::new(Mutex::new(Vec::new()));
|
|
let trial_counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
|
let tpe_start = Instant::now();
|
|
|
|
for trial_idx in 0..max_trials {
|
|
let continuous_vec = tpe.suggest(&bounds);
|
|
|
|
info!(
|
|
"TPE Trial {}/{}: evaluating suggested point",
|
|
trial_idx + 1,
|
|
max_trials
|
|
);
|
|
|
|
ArgminOptimizer::evaluate_point(
|
|
&continuous_vec,
|
|
&mut model,
|
|
&trial_results,
|
|
&trial_counter,
|
|
¶m_names,
|
|
)?;
|
|
|
|
// Feed the objective back to TPE so it can update its density models
|
|
let results = trial_results.lock().map_err(|e| {
|
|
MLError::ConcurrencyError {
|
|
operation: format!("lock trial results: {}", e),
|
|
}
|
|
})?;
|
|
if let Some(last) = results.last() {
|
|
tpe.add_trial(continuous_vec, last.objective);
|
|
info!(
|
|
"TPE Trial {}/{}: objective = {:.6}",
|
|
trial_idx + 1,
|
|
max_trials,
|
|
last.objective
|
|
);
|
|
}
|
|
|
|
if let Some(name) = model_name {
|
|
common::metrics::training_metrics::set_hyperopt_trial(
|
|
name,
|
|
(trial_idx + 1) as f64,
|
|
max_trials as f64,
|
|
);
|
|
common::metrics::training_metrics::set_hyperopt_elapsed(
|
|
name,
|
|
tpe_start.elapsed().as_secs_f64(),
|
|
);
|
|
if let Some(last) = results.last() {
|
|
common::metrics::training_metrics::record_hyperopt_trial_duration(
|
|
name,
|
|
last.duration_secs,
|
|
);
|
|
}
|
|
if let Some(best) = results.iter().map(|r| r.objective).reduce(f64::min) {
|
|
common::metrics::training_metrics::set_hyperopt_best_objective(name, best);
|
|
common::metrics::training_metrics::set_hyperopt_trial_best_loss(name, best);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Extract results
|
|
let trials = match Arc::try_unwrap(trial_results) {
|
|
Ok(mutex) => mutex.into_inner().map_err(|e| {
|
|
anyhow::anyhow!("Failed to unwrap trial results mutex: {}", e)
|
|
})?,
|
|
Err(arc) => arc.lock().map_err(|e| {
|
|
anyhow::anyhow!("Failed to lock trial results: {}", e)
|
|
})?.clone(),
|
|
};
|
|
|
|
if trials.is_empty() {
|
|
return Err(
|
|
MLError::ModelError("No valid trials completed".to_owned()).into(),
|
|
);
|
|
}
|
|
|
|
let result = OptimizationResult::from_trials(trials);
|
|
|
|
info!("═══ TPE Optimization Complete ═══");
|
|
info!(
|
|
" Best trial: objective {:.6}",
|
|
result.best_objective
|
|
);
|
|
info!(" Total trials: {}", result.all_trials.len());
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Trait for models supporting two-phase objective switching.
|
|
///
|
|
/// Used by [`ArgminOptimizer::optimize_two_phase()`] to switch between
|
|
/// episode reward (fast convergence) and Sharpe ratio (financial quality).
|
|
///
|
|
/// This trait is intentionally separate from [`HyperparameterOptimizable`]
|
|
/// because `extract_objective` is a static method that cannot access instance
|
|
/// state. Two-phase optimization instead sets the objective mode on the model
|
|
/// instance before each phase.
|
|
pub trait TwoPhaseObjective {
|
|
/// Set the objective mode for the current optimization phase.
|
|
fn set_objective_mode(&mut self, mode: ObjectiveMode);
|
|
/// Get the current objective mode.
|
|
fn objective_mode(&self) -> ObjectiveMode;
|
|
}
|
|
|
|
/// Cost function wrapper for argmin
|
|
struct ObjectiveFunction<M>
|
|
where
|
|
M: HyperparameterOptimizable,
|
|
M::Params: ParameterSpace,
|
|
{
|
|
model: Arc<Mutex<M>>,
|
|
trial_results: Arc<Mutex<Vec<TrialResult<M::Params>>>>,
|
|
trial_counter: Arc<std::sync::atomic::AtomicUsize>,
|
|
param_names: Vec<&'static str>,
|
|
bounds: Vec<(f64, f64)>,
|
|
observer: crate::TrialBudgetObserver,
|
|
model_name: Option<String>,
|
|
max_trials: usize,
|
|
optimization_start: Instant,
|
|
}
|
|
|
|
impl<M> CostFunction for ObjectiveFunction<M>
|
|
where
|
|
M: HyperparameterOptimizable,
|
|
M::Params: ParameterSpace,
|
|
{
|
|
type Param = Vec<f64>;
|
|
type Output = f64;
|
|
|
|
#[allow(clippy::cognitive_complexity)]
|
|
fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
|
|
// Increment observer trial count FIRST
|
|
self.observer.increment_trial();
|
|
|
|
// Check if budget exhausted (prevents runaway PSO)
|
|
if self.observer.should_terminate() {
|
|
warn!(
|
|
"Trial budget exhausted: {}/{} trials. Stopping PSO.",
|
|
self.observer.get_trials_used(),
|
|
self.observer.get_trials_used()
|
|
);
|
|
return Err(argmin::core::Error::msg("Trial budget exhausted"));
|
|
}
|
|
|
|
// Clamp parameters to bounds
|
|
let mut clamped = param.clone();
|
|
for (i, (min, max)) in self.bounds.iter().enumerate() {
|
|
clamped[i] = clamped[i].clamp(*min, *max);
|
|
}
|
|
|
|
let start_time = Instant::now();
|
|
|
|
// Increment trial counter
|
|
let trial_num = self.trial_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!(
|
|
"║ Trial {}: Evaluating Parameters ║",
|
|
trial_num
|
|
);
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
// Convert continuous vector to parameters BEFORE logging
|
|
let params = match M::Params::from_continuous(&clamped) {
|
|
Ok(p) => p,
|
|
Err(e) => {
|
|
warn!(
|
|
"Failed to convert parameters for trial {}: {}",
|
|
trial_num, e
|
|
);
|
|
return Ok(1e6); // Penalty for invalid parameters
|
|
},
|
|
};
|
|
|
|
// Log CONVERTED parameters (shows actual values: learning_rate ~1e-4, not -11)
|
|
info!(" Parameters (converted): {:?}", params);
|
|
|
|
// Train model with parameters
|
|
let mut model = self.model.lock().map_err(|e| {
|
|
argmin::core::Error::msg(format!("Failed to lock model: {}", e))
|
|
})?;
|
|
let metrics = match model.train_with_params(params.clone()) {
|
|
Ok(m) => m,
|
|
Err(e) => {
|
|
warn!("Training failed for trial {}: {}", trial_num, e);
|
|
return Ok(1e6); // Penalty for training failure
|
|
},
|
|
};
|
|
|
|
// Extract objective (replace NaN/Inf with large penalty)
|
|
let raw_objective = M::extract_objective(&metrics);
|
|
let objective = if raw_objective.is_finite() {
|
|
raw_objective
|
|
} else {
|
|
warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
|
|
1e6
|
|
};
|
|
let trial_metrics = M::extract_metrics(&metrics);
|
|
let duration_secs = start_time.elapsed().as_secs_f64();
|
|
|
|
info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);
|
|
info!(" Objective: {:.6}", objective);
|
|
|
|
// Record trial result
|
|
{
|
|
let mut results = self.trial_results.lock().map_err(|e| {
|
|
argmin::core::Error::msg(format!("Failed to lock trial results: {}", e))
|
|
})?;
|
|
results.push(TrialResult {
|
|
trial_num,
|
|
params,
|
|
objective,
|
|
duration_secs,
|
|
metrics: trial_metrics,
|
|
});
|
|
}
|
|
|
|
if let Some(ref name) = self.model_name {
|
|
common::metrics::training_metrics::set_hyperopt_trial(
|
|
name,
|
|
trial_num as f64,
|
|
self.max_trials as f64,
|
|
);
|
|
common::metrics::training_metrics::record_hyperopt_trial_duration(name, duration_secs);
|
|
common::metrics::training_metrics::set_hyperopt_elapsed(
|
|
name,
|
|
self.optimization_start.elapsed().as_secs_f64(),
|
|
);
|
|
// Update running best objective
|
|
if let Ok(results) = self.trial_results.lock() {
|
|
if let Some(best) = results.iter().map(|r| r.objective).reduce(f64::min) {
|
|
common::metrics::training_metrics::set_hyperopt_best_objective(name, best);
|
|
common::metrics::training_metrics::set_hyperopt_trial_best_loss(name, best);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(objective)
|
|
}
|
|
}
|
|
|
|
/// Parallel cost function for argmin — clones model per evaluation.
|
|
///
|
|
/// Unlike [`ObjectiveFunction`] which holds an `Arc<Mutex<M>>` lock for the
|
|
/// entire training duration, this variant clones the model and releases the
|
|
/// lock immediately. This allows argmin's rayon-based PSO to evaluate multiple
|
|
/// particles truly in parallel.
|
|
///
|
|
/// Each clone gets its own `InternalDQNTrainer` (created inside `train_with_params`),
|
|
/// so GPU/CPU resources are the only shared concern. DQN uses ~1 core + 278MB VRAM
|
|
/// per trial, allowing 4-6 concurrent trials on an 8-vCPU L4 with 24GB VRAM.
|
|
struct ParallelObjectiveFunction<M>
|
|
where
|
|
M: HyperparameterOptimizable + Clone,
|
|
M::Params: ParameterSpace,
|
|
{
|
|
model: Arc<Mutex<M>>,
|
|
trial_results: Arc<Mutex<Vec<TrialResult<M::Params>>>>,
|
|
trial_counter: Arc<std::sync::atomic::AtomicUsize>,
|
|
param_names: Vec<&'static str>,
|
|
bounds: Vec<(f64, f64)>,
|
|
observer: crate::TrialBudgetObserver,
|
|
model_name: Option<String>,
|
|
max_trials: usize,
|
|
optimization_start: Instant,
|
|
}
|
|
|
|
impl<M> CostFunction for ParallelObjectiveFunction<M>
|
|
where
|
|
M: HyperparameterOptimizable + Clone,
|
|
M::Params: ParameterSpace,
|
|
{
|
|
type Param = Vec<f64>;
|
|
type Output = f64;
|
|
|
|
#[allow(clippy::cognitive_complexity)]
|
|
fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
|
|
// Increment observer trial count FIRST
|
|
self.observer.increment_trial();
|
|
|
|
// Check if budget exhausted
|
|
if self.observer.should_terminate() {
|
|
warn!(
|
|
"Trial budget exhausted: {} trials. Stopping PSO.",
|
|
self.observer.get_trials_used()
|
|
);
|
|
return Err(argmin::core::Error::msg("Trial budget exhausted"));
|
|
}
|
|
|
|
// Clone model and release mutex immediately — enables true parallel evaluation
|
|
let mut model = {
|
|
let guard = self.model.lock().map_err(|e| {
|
|
argmin::core::Error::msg(format!("Failed to lock model: {}", e))
|
|
})?;
|
|
guard.clone()
|
|
}; // Mutex released here — other particles can proceed concurrently
|
|
|
|
// Clamp parameters to bounds
|
|
let mut clamped = param.clone();
|
|
for (i, (min, max)) in self.bounds.iter().enumerate() {
|
|
clamped[i] = clamped[i].clamp(*min, *max);
|
|
}
|
|
|
|
let start_time = Instant::now();
|
|
|
|
let trial_num = self.trial_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
|
|
|
|
info!("╔═══════════════════════════════════════════════════════════╗");
|
|
info!(
|
|
"║ Trial {}: Evaluating Parameters (parallel) ║",
|
|
trial_num
|
|
);
|
|
info!("╚═══════════════════════════════════════════════════════════╝");
|
|
|
|
let params = match M::Params::from_continuous(&clamped) {
|
|
Ok(p) => p,
|
|
Err(e) => {
|
|
warn!("Failed to convert parameters for trial {}: {}", trial_num, e);
|
|
return Ok(1e6);
|
|
},
|
|
};
|
|
|
|
info!(" Parameters (converted): {:?}", params);
|
|
|
|
let metrics = match model.train_with_params(params.clone()) {
|
|
Ok(m) => m,
|
|
Err(e) => {
|
|
warn!("Training failed for trial {}: {}", trial_num, e);
|
|
return Ok(1e6);
|
|
},
|
|
};
|
|
|
|
let raw_objective = M::extract_objective(&metrics);
|
|
let objective = if raw_objective.is_finite() {
|
|
raw_objective
|
|
} else {
|
|
warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
|
|
1e6
|
|
};
|
|
let trial_metrics = M::extract_metrics(&metrics);
|
|
let duration_secs = start_time.elapsed().as_secs_f64();
|
|
|
|
info!("Trial {} completed in {:.1}s", trial_num, duration_secs);
|
|
info!(" Objective: {:.6}", objective);
|
|
|
|
{
|
|
let mut results = self.trial_results.lock().map_err(|e| {
|
|
argmin::core::Error::msg(format!("Failed to lock trial results: {}", e))
|
|
})?;
|
|
results.push(TrialResult {
|
|
trial_num,
|
|
params,
|
|
objective,
|
|
duration_secs,
|
|
metrics: trial_metrics,
|
|
});
|
|
}
|
|
|
|
if let Some(ref name) = self.model_name {
|
|
common::metrics::training_metrics::set_hyperopt_trial(
|
|
name,
|
|
trial_num as f64,
|
|
self.max_trials as f64,
|
|
);
|
|
common::metrics::training_metrics::record_hyperopt_trial_duration(name, duration_secs);
|
|
common::metrics::training_metrics::set_hyperopt_elapsed(
|
|
name,
|
|
self.optimization_start.elapsed().as_secs_f64(),
|
|
);
|
|
if let Ok(results) = self.trial_results.lock() {
|
|
if let Some(best) = results.iter().map(|r| r.objective).reduce(f64::min) {
|
|
common::metrics::training_metrics::set_hyperopt_best_objective(name, best);
|
|
common::metrics::training_metrics::set_hyperopt_trial_best_loss(name, best);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(objective)
|
|
}
|
|
}
|
|
|
|
/// Builder for `ArgminOptimizer` with fluent API
|
|
///
|
|
/// Provides a convenient way to configure optimizer settings.
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust
|
|
/// use ml::hyperopt::ArgminOptimizer;
|
|
///
|
|
/// let optimizer = ArgminOptimizer::builder()
|
|
/// .max_trials(50)
|
|
/// .n_initial(10)
|
|
/// .n_particles(30)
|
|
/// .seed(12345)
|
|
/// .build();
|
|
/// ```
|
|
#[derive(Debug, Clone)]
|
|
pub struct ArgminOptimizerBuilder {
|
|
max_trials: usize,
|
|
n_initial: usize,
|
|
n_particles: usize,
|
|
seed: Option<u64>,
|
|
max_iters_per_restart: usize,
|
|
model_name: Option<String>,
|
|
}
|
|
|
|
impl Default for ArgminOptimizerBuilder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl ArgminOptimizerBuilder {
|
|
/// Create a new builder with default settings
|
|
pub fn new() -> Self {
|
|
Self {
|
|
max_trials: 30,
|
|
n_initial: 5,
|
|
n_particles: 20,
|
|
seed: None,
|
|
max_iters_per_restart: 50,
|
|
model_name: None,
|
|
}
|
|
}
|
|
|
|
/// Set maximum number of trials
|
|
pub fn max_trials(mut self, max_trials: usize) -> Self {
|
|
self.max_trials = max_trials;
|
|
self
|
|
}
|
|
|
|
/// Set number of initial samples
|
|
pub fn n_initial(mut self, n_initial: usize) -> Self {
|
|
self.n_initial = n_initial;
|
|
self
|
|
}
|
|
|
|
/// Set number of particles in the swarm
|
|
pub fn n_particles(mut self, n_particles: usize) -> Self {
|
|
self.n_particles = n_particles;
|
|
self
|
|
}
|
|
|
|
/// Set random seed for reproducibility
|
|
pub fn seed(mut self, seed: u64) -> Self {
|
|
self.seed = Some(seed);
|
|
self
|
|
}
|
|
|
|
/// Set maximum iterations per restart
|
|
pub fn max_iters_per_restart(mut self, max_iters: usize) -> Self {
|
|
self.max_iters_per_restart = max_iters;
|
|
self
|
|
}
|
|
|
|
/// Set model name for per-trial Prometheus progress metrics
|
|
pub fn model_name<S: Into<String>>(mut self, name: S) -> Self {
|
|
self.model_name = Some(name.into());
|
|
self
|
|
}
|
|
|
|
/// Build the optimizer
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Panics if `max_trials <= n_initial` or if either is zero.
|
|
pub fn build(self) -> ArgminOptimizer {
|
|
assert!(
|
|
self.max_trials > self.n_initial,
|
|
"max_trials must be > n_initial"
|
|
);
|
|
assert!(self.n_initial > 0, "n_initial must be > 0");
|
|
assert!(self.n_particles > 0, "n_particles must be > 0");
|
|
|
|
ArgminOptimizer {
|
|
max_trials: self.max_trials,
|
|
n_initial: self.n_initial,
|
|
n_particles: self.n_particles,
|
|
seed: self.seed,
|
|
max_iters_per_restart: self.max_iters_per_restart,
|
|
model_name: self.model_name,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct TestParams {
|
|
x: f64,
|
|
y: f64,
|
|
}
|
|
|
|
impl ParameterSpace for TestParams {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
vec![(-5.0, 5.0), (-5.0, 5.0)]
|
|
}
|
|
|
|
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
|
Ok(Self { x: x[0], y: x[1] })
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
vec![self.x, self.y]
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
vec!["x", "y"]
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct TestMetrics {
|
|
loss: f64,
|
|
}
|
|
|
|
struct TestModel;
|
|
|
|
impl HyperparameterOptimizable for TestModel {
|
|
type Params = TestParams;
|
|
type Metrics = TestMetrics;
|
|
|
|
fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError> {
|
|
// Rosenbrock function: (1-x)^2 + 100(y-x^2)^2
|
|
// Minimum at (1, 1) with value 0
|
|
let loss = (1.0 - params.x).powi(2) + 100.0 * (params.y - params.x.powi(2)).powi(2);
|
|
Ok(TestMetrics { loss })
|
|
}
|
|
|
|
fn extract_objective(metrics: &Self::Metrics) -> f64 {
|
|
metrics.loss
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_optimizer_builder() {
|
|
let optimizer = ArgminOptimizer::builder()
|
|
.max_trials(20)
|
|
.n_initial(3)
|
|
.seed(42)
|
|
.build();
|
|
|
|
assert_eq!(optimizer.max_trials, 20);
|
|
assert_eq!(optimizer.n_initial, 3);
|
|
assert_eq!(optimizer.seed, Some(42));
|
|
}
|
|
|
|
#[test]
|
|
fn test_latin_hypercube_sampling() {
|
|
let mut rng = StdRng::seed_from_u64(42);
|
|
let bounds = vec![(-5.0, 5.0), (-5.0, 5.0)];
|
|
let samples = ArgminOptimizer::latin_hypercube_sampling(10, &bounds, &mut rng);
|
|
|
|
assert_eq!(samples.nrows(), 10);
|
|
assert_eq!(samples.ncols(), 2);
|
|
|
|
// Check all samples are within bounds
|
|
for i in 0..10 {
|
|
assert!(samples[[i, 0]] >= -5.0 && samples[[i, 0]] <= 5.0);
|
|
assert!(samples[[i, 1]] >= -5.0 && samples[[i, 1]] <= 5.0);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore] // Expensive test - run manually
|
|
fn test_optimizer_rosenbrock() {
|
|
let model = TestModel;
|
|
let optimizer = ArgminOptimizer::with_trials(15, 3);
|
|
|
|
let result = optimizer.optimize(model).unwrap();
|
|
|
|
// Should find minimum near (1, 1)
|
|
assert!(
|
|
result.best_objective < 1.0,
|
|
"Expected to find near-optimal solution"
|
|
);
|
|
assert!(
|
|
result.all_trials.len() <= 15,
|
|
"Should have at most 15 trials"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parameter_conversion_with_log_scale() {
|
|
// Test that log-scale parameters are converted correctly
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
struct LogScaleParams {
|
|
learning_rate: f64,
|
|
weight_decay: f64,
|
|
batch_size: usize,
|
|
}
|
|
|
|
impl ParameterSpace for LogScaleParams {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
vec![
|
|
(1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log scale)
|
|
(1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log scale)
|
|
(16.0, 256.0), // batch_size (linear)
|
|
]
|
|
}
|
|
|
|
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
|
Ok(Self {
|
|
learning_rate: x[0].exp(),
|
|
weight_decay: x[1].exp(),
|
|
batch_size: x[2].round() as usize,
|
|
})
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
vec![
|
|
self.learning_rate.ln(),
|
|
self.weight_decay.ln(),
|
|
self.batch_size as f64,
|
|
]
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
vec!["learning_rate", "weight_decay", "batch_size"]
|
|
}
|
|
}
|
|
|
|
// Test continuous space values (negative for log-scale)
|
|
let continuous = vec![-9.21, -11.51, 64.0]; // ln(1e-4), ln(1e-5), 64
|
|
|
|
// Convert to actual parameters
|
|
let params = LogScaleParams::from_continuous(&continuous).unwrap();
|
|
|
|
// Verify conversions (use relative tolerance for floating point)
|
|
// Allow 1% relative error due to log/exp floating point precision
|
|
assert!(
|
|
(params.learning_rate - 1e-4).abs() / 1e-4 < 0.01,
|
|
"Learning rate should be ~1e-4 (±1%), got {}",
|
|
params.learning_rate
|
|
);
|
|
assert!(
|
|
(params.weight_decay - 1e-5).abs() / 1e-5 < 0.01,
|
|
"Weight decay should be ~1e-5 (±1%), got {}",
|
|
params.weight_decay
|
|
);
|
|
assert_eq!(params.batch_size, 64, "Batch size should be 64");
|
|
|
|
// Verify round-trip
|
|
let recovered = params.to_continuous();
|
|
assert!((recovered[0] - continuous[0]).abs() < 1e-6);
|
|
assert!((recovered[1] - continuous[1]).abs() < 1e-6);
|
|
assert!((recovered[2] - continuous[2]).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parameter_logging_shows_converted_values() {
|
|
// This test verifies that parameters are logged AFTER conversion,
|
|
// so Debug output shows actual values (learning_rate ~1e-4, not -11)
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct LogParams {
|
|
value: f64,
|
|
}
|
|
|
|
impl ParameterSpace for LogParams {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
vec![(1e-5_f64.ln(), 1e-2_f64.ln())]
|
|
}
|
|
|
|
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
|
Ok(Self { value: x[0].exp() })
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
vec![self.value.ln()]
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
vec!["value"]
|
|
}
|
|
}
|
|
|
|
// Raw continuous value (log scale)
|
|
let continuous = vec![-11.51]; // ln(1e-5)
|
|
|
|
// Convert to actual parameter
|
|
let params = LogParams::from_continuous(&continuous).unwrap();
|
|
|
|
// Debug output should show converted value (in scientific notation or decimal)
|
|
let debug_str = format!("{:?}", params);
|
|
|
|
// Should show value close to 1e-5 (not the log value -11.5)
|
|
assert!(
|
|
debug_str.contains("e-5") || debug_str.contains("0.00001"),
|
|
"Debug output should show actual value ~1e-5, got: {}",
|
|
debug_str
|
|
);
|
|
|
|
// Should NOT contain raw log value
|
|
assert!(
|
|
!debug_str.contains("-11."),
|
|
"Debug output should NOT show raw log value -11.x"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_two_phase_optimizer_config() {
|
|
let optimizer = ArgminOptimizer::with_trials(30, 5);
|
|
assert_eq!(optimizer.max_trials, 30);
|
|
// Two-phase would run 15 + 15 trials
|
|
let phase_a = optimizer.max_trials / 2;
|
|
let phase_b = optimizer.max_trials - phase_a;
|
|
assert_eq!(phase_a, 15);
|
|
assert_eq!(phase_b, 15);
|
|
}
|
|
|
|
#[test]
|
|
fn test_two_phase_optimizer_odd_trials() {
|
|
let optimizer = ArgminOptimizer::with_trials(31, 5);
|
|
// Odd total: 15 + 16 trials
|
|
let phase_a = optimizer.max_trials / 2;
|
|
let phase_b = optimizer.max_trials - phase_a;
|
|
assert_eq!(phase_a, 15);
|
|
assert_eq!(phase_b, 16);
|
|
}
|
|
|
|
#[test]
|
|
fn test_two_phase_objective_trait_object_safety() {
|
|
// Verify TwoPhaseObjective can be used with the ObjectiveMode enum
|
|
use super::TwoPhaseObjective;
|
|
|
|
struct MockModel {
|
|
mode: ObjectiveMode,
|
|
}
|
|
|
|
impl TwoPhaseObjective for MockModel {
|
|
fn set_objective_mode(&mut self, mode: ObjectiveMode) {
|
|
self.mode = mode;
|
|
}
|
|
fn objective_mode(&self) -> ObjectiveMode {
|
|
self.mode
|
|
}
|
|
}
|
|
|
|
let mut model = MockModel {
|
|
mode: ObjectiveMode::Sharpe,
|
|
};
|
|
assert_eq!(model.objective_mode(), ObjectiveMode::Sharpe);
|
|
|
|
model.set_objective_mode(ObjectiveMode::EpisodeReward);
|
|
assert_eq!(model.objective_mode(), ObjectiveMode::EpisodeReward);
|
|
}
|
|
}
|