From 2651e61b2cdd788fe6fa60fd85c4f13259ae1923 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 15:46:56 +0100 Subject: [PATCH] feat: add Pruner trait and NopPruner default implementation Introduce the pruning extension point for the trial pruning system. The Pruner trait allows deciding whether to stop a trial early based on intermediate values. NopPruner (never prunes) is the default. --- src/lib.rs | 3 +++ src/pruner/mod.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++ src/pruner/nop.rs | 17 ++++++++++++++ src/study.rs | 58 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 138 insertions(+) create mode 100644 src/pruner/mod.rs create mode 100644 src/pruner/nop.rs diff --git a/src/lib.rs b/src/lib.rs index 498a2c7..ddbf426 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -188,6 +188,7 @@ mod error; mod kde; mod param; pub mod parameter; +pub mod pruner; pub mod sampler; mod study; mod trial; @@ -200,6 +201,7 @@ pub use param::ParamValue; pub use parameter::{ BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamId, Parameter, }; +pub use pruner::{NopPruner, Pruner}; pub use sampler::CompletedTrial; pub use sampler::grid::GridSearchSampler; pub use sampler::random::RandomSampler; @@ -222,6 +224,7 @@ pub mod prelude { pub use crate::parameter::{ BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter, }; + pub use crate::pruner::{NopPruner, Pruner}; pub use crate::sampler::CompletedTrial; pub use crate::sampler::grid::GridSearchSampler; pub use crate::sampler::random::RandomSampler; diff --git a/src/pruner/mod.rs b/src/pruner/mod.rs new file mode 100644 index 0000000..06ea9f4 --- /dev/null +++ b/src/pruner/mod.rs @@ -0,0 +1,60 @@ +//! Pruner trait and implementations for trial pruning. +//! +//! Pruners decide whether to stop (prune) a trial early based on its +//! intermediate values compared to other trials. This is useful for +//! discarding unpromising trials before they complete, saving compute. + +mod nop; + +pub use nop::NopPruner; + +use crate::sampler::CompletedTrial; + +/// Trait for pluggable trial pruning strategies. +/// +/// Pruners are consulted after each intermediate value is reported to +/// decide whether the trial should be stopped early. The trait requires +/// `Send + Sync` to support concurrent and async optimization. +/// +/// # Implementing a custom pruner +/// +/// ``` +/// use optimizer::pruner::Pruner; +/// use optimizer::sampler::CompletedTrial; +/// +/// struct MyPruner { +/// threshold: f64, +/// } +/// +/// impl Pruner for MyPruner { +/// fn should_prune( +/// &self, +/// _trial_id: u64, +/// _step: u64, +/// intermediate_values: &[(u64, f64)], +/// _completed_trials: &[CompletedTrial], +/// ) -> bool { +/// // Prune if the latest value exceeds the threshold +/// intermediate_values +/// .last() +/// .is_some_and(|&(_, v)| v > self.threshold) +/// } +/// } +/// ``` +pub trait Pruner: Send + Sync { + /// Decide whether to prune a trial at the given step. + /// + /// # Arguments + /// + /// * `trial_id` - The current trial's ID. + /// * `step` - The step at which the intermediate value was reported. + /// * `intermediate_values` - All `(step, value)` pairs reported so far for this trial. + /// * `completed_trials` - History of all completed trials (for comparison). + fn should_prune( + &self, + trial_id: u64, + step: u64, + intermediate_values: &[(u64, f64)], + completed_trials: &[CompletedTrial], + ) -> bool; +} diff --git a/src/pruner/nop.rs b/src/pruner/nop.rs new file mode 100644 index 0000000..9ccfad9 --- /dev/null +++ b/src/pruner/nop.rs @@ -0,0 +1,17 @@ +use super::Pruner; +use crate::sampler::CompletedTrial; + +/// A pruner that never prunes. This is the default when no pruner is configured. +pub struct NopPruner; + +impl Pruner for NopPruner { + fn should_prune( + &self, + _trial_id: u64, + _step: u64, + _intermediate_values: &[(u64, f64)], + _completed_trials: &[CompletedTrial], + ) -> bool { + false + } +} diff --git a/src/study.rs b/src/study.rs index 79718a3..bdb85fe 100644 --- a/src/study.rs +++ b/src/study.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use parking_lot::RwLock; +use crate::pruner::{NopPruner, Pruner}; use crate::sampler::random::RandomSampler; use crate::sampler::{CompletedTrial, Sampler}; use crate::trial::Trial; @@ -40,6 +41,8 @@ where direction: Direction, /// The sampler used to generate parameter values. sampler: Arc, + /// The pruner used to decide whether to stop trials early. + pruner: Arc, /// Completed trials (wrapped in Arc for sharing with Trial). completed_trials: Arc>>>, /// Counter for generating unique trial IDs. @@ -109,6 +112,7 @@ where Self { direction, sampler, + pruner: Arc::new(NopPruner), completed_trials, next_trial_id: AtomicU64::new(0), trial_factory, @@ -158,6 +162,46 @@ where /// let mut study: Study = Study::new(Direction::Minimize); /// study.set_sampler(TpeSampler::new()); /// ``` + /// Creates a new study with a custom sampler and pruner. + /// + /// # Arguments + /// + /// * `direction` - Whether to minimize or maximize the objective function. + /// * `sampler` - The sampler to use for parameter sampling. + /// * `pruner` - The pruner to use for trial pruning. + /// + /// # Examples + /// + /// ``` + /// use optimizer::pruner::NopPruner; + /// use optimizer::sampler::random::RandomSampler; + /// use optimizer::{Direction, Study}; + /// + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler_and_pruner(Direction::Minimize, sampler, NopPruner); + /// ``` + pub fn with_sampler_and_pruner( + direction: Direction, + sampler: impl Sampler + 'static, + pruner: impl Pruner + 'static, + ) -> Self + where + V: 'static, + { + let sampler: Arc = Arc::new(sampler); + let completed_trials = Arc::new(RwLock::new(Vec::new())); + let trial_factory = Self::make_trial_factory(&sampler, &completed_trials); + + Self { + direction, + sampler, + pruner: Arc::new(pruner), + completed_trials, + next_trial_id: AtomicU64::new(0), + trial_factory, + } + } + pub fn set_sampler(&mut self, sampler: impl Sampler + 'static) where V: 'static, @@ -166,6 +210,20 @@ where self.trial_factory = Self::make_trial_factory(&self.sampler, &self.completed_trials); } + /// Sets a new pruner for the study. + /// + /// # Arguments + /// + /// * `pruner` - The pruner to use for trial pruning. + pub fn set_pruner(&mut self, pruner: impl Pruner + 'static) { + self.pruner = Arc::new(pruner); + } + + /// Returns a reference to the study's pruner. + pub fn pruner(&self) -> &dyn Pruner { + &*self.pruner + } + /// Generates the next unique trial ID. pub(crate) fn next_trial_id(&self) -> u64 { self.next_trial_id.fetch_add(1, Ordering::SeqCst)