foxhunt/ml/src/hyperopt/observer.rs

use argmin::core::observers::Observe;
use argmin::core::{Error, State};
use std::sync::{Arc, Mutex};

/// Observer that enforces strict trial budget limits
///
/// Argmin's PSO does not reliably respect max_iters when set to small values (e.g., 1).
/// This observer tracks evaluations across all iterations and terminates when budget is hit.
#[derive(Clone, Debug)]
pub struct TrialBudgetObserver {
    max_trials: usize,
    trials_used: Arc<Mutex<usize>>,
}

impl TrialBudgetObserver {
    pub fn new(max_trials: usize) -> Self {
        Self {
            max_trials,
            trials_used: Arc::new(Mutex::new(0)),
        }
    }

    pub fn increment_trial(&self) {
        let mut count = self.trials_used.lock().unwrap();
        *count += 1;
    }

    pub fn should_terminate(&self) -> bool {
        let count = self.trials_used.lock().unwrap();
        *count >= self.max_trials
    }

    pub fn get_trials_used(&self) -> usize {
        *self.trials_used.lock().unwrap()
    }
}

impl<I> Observe<I> for TrialBudgetObserver
where
    I: State,
{
    fn observe_iter(&mut self, _state: &I, _kv: &argmin::core::KV) -> Result<(), Error> {
        // Check if we've exceeded budget
        if self.should_terminate() {
            tracing::warn!(
                "Trial budget exhausted: {}/{} trials used. Terminating optimization.",
                self.get_trials_used(),
                self.max_trials
            );
            // Return error to signal termination
            return Err(Error::msg(format!(
                "Trial budget exhausted: {}/{} trials",
                self.get_trials_used(),
                self.max_trials
            )));
        }

        Ok(())
    }
}