1725 lines
66 KiB
Rust
1725 lines
66 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.
|
||
// Pass the shared trial_counter so LHS evaluations already counted above
|
||
// are included in the budget — fixes the "0 forever" bug where PSO ran
|
||
// unconstrained because the observer had its own separate counter.
|
||
let observer = crate::TrialBudgetObserver::new(self.max_trials, Arc::clone(&trial_counter));
|
||
|
||
// 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);
|
||
|
||
// Algorithm selection by budget:
|
||
// - Budget < 2× swarm → LHS (space-filling, optimal for small budgets)
|
||
// - Budget ≥ 2× swarm → PSO (swarm optimization, needs full iterations)
|
||
// PSO with a partial swarm is worse than random — 75% of particles
|
||
// have no cost value, so gbest is computed from incomplete data.
|
||
if remaining_trials < self.n_particles * 2 {
|
||
info!(
|
||
"Using LHS for remaining {} trials (budget < {} particles × 2 = need PSO budget ≥ {})",
|
||
remaining_trials, self.n_particles, self.n_particles * 2
|
||
);
|
||
let extra_samples = Self::latin_hypercube_sampling(remaining_trials, &bounds, &mut rng);
|
||
let mut model_guard = cost_fn.model.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?;
|
||
for i in 0..remaining_trials {
|
||
let continuous_vec: Vec<f64> = extra_samples.row(i).to_vec();
|
||
Self::evaluate_point(
|
||
&continuous_vec,
|
||
&mut *model_guard,
|
||
&trial_results,
|
||
&trial_counter,
|
||
¶m_names,
|
||
)?;
|
||
}
|
||
drop(model_guard);
|
||
}
|
||
|
||
let max_iters_by_budget =
|
||
((remaining_trials as f64) / (self.n_particles as f64)).ceil() as usize;
|
||
let max_iters = if remaining_trials < self.n_particles * 2 {
|
||
0 // LHS already handled above
|
||
} else {
|
||
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 with tuned parameters.
|
||
// Default argmin: inertia=0.72, cognitive=1.19, social=1.19
|
||
// Problem: social >> inertia → premature convergence to global best.
|
||
// Fix: higher inertia (0.9) for exploration, lower social (0.8) to
|
||
// reduce collapse toward the single LHS best point.
|
||
let solver = ParticleSwarm::new((lower_bounds, upper_bounds), self.n_particles)
|
||
.with_inertia_factor(0.9)
|
||
.map_err(|e| anyhow::anyhow!("PSO inertia: {e}"))?
|
||
.with_cognitive_factor(1.5)
|
||
.map_err(|e| anyhow::anyhow!("PSO cognitive: {e}"))?
|
||
.with_social_factor(0.8)
|
||
.map_err(|e| anyhow::anyhow!("PSO social: {e}"))?;
|
||
|
||
// 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 — NaN/divergence is a valid PSO outcome,
|
||
// not a fatal error. Score failed trials with maximum penalty so PSO
|
||
// learns to avoid that region of the search space.
|
||
let (objective, trial_metrics) = match model.train_with_params(params.clone()) {
|
||
Ok(metrics) => {
|
||
let raw_objective = M::extract_objective(&metrics);
|
||
let obj = if raw_objective.is_finite() {
|
||
raw_objective
|
||
} else {
|
||
tracing::error!("Trial {} produced NON-FINITE objective ({:?}), using penalty 1e6", trial_num, raw_objective);
|
||
1e6
|
||
};
|
||
(obj, M::extract_metrics(&metrics))
|
||
}
|
||
Err(e) => {
|
||
tracing::error!("Trial {} training FAILED: {:#}", trial_num, e);
|
||
eprintln!("!!! FULL ERROR CHAIN: {:?}", e);
|
||
(1e6, None)
|
||
}
|
||
};
|
||
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() {
|
||
tracing::error!("LHS evaluation FAILED (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 sharing the same counter as LHS/evaluate_point.
|
||
let observer = crate::TrialBudgetObserver::new(effective_max_trials, Arc::clone(&trial_counter));
|
||
|
||
// 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)
|
||
.with_inertia_factor(0.9)
|
||
.map_err(|e| anyhow::anyhow!("PSO inertia: {e}"))?
|
||
.with_cognitive_factor(1.5)
|
||
.map_err(|e| anyhow::anyhow!("PSO cognitive: {e}"))?
|
||
.with_social_factor(0.8)
|
||
.map_err(|e| anyhow::anyhow!("PSO social: {e}"))?;
|
||
|
||
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 the shared trial counter FIRST — this is the canonical counter
|
||
// used by both the observer (for budget enforcement) and the results log.
|
||
let trial_num = self.trial_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
|
||
|
||
// Check budget immediately after incrementing — abort if we've exceeded the limit.
|
||
if trial_num > self.max_trials {
|
||
warn!(
|
||
"Trial budget exhausted: {}/{} trials. Stopping PSO.",
|
||
trial_num, self.max_trials
|
||
);
|
||
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();
|
||
|
||
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) => {
|
||
tracing::error!("Trial {} training FAILED: {:#}", 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 the shared trial counter FIRST — shared with LHS phase and observer.
|
||
let trial_num = self.trial_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
|
||
|
||
// Check budget immediately after incrementing.
|
||
if trial_num > self.max_trials {
|
||
warn!(
|
||
"Trial budget exhausted: {}/{} trials. Stopping PSO.",
|
||
trial_num, self.max_trials
|
||
);
|
||
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();
|
||
|
||
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) => {
|
||
tracing::error!("Trial {} training FAILED: {:#}", 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);
|
||
}
|
||
}
|