diff --git a/src/lib.rs b/src/lib.rs index 1e4b5c8..58f6eb1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -210,7 +210,7 @@ pub use sampler::grid::GridSearchSampler; pub use sampler::random::RandomSampler; pub use sampler::tpe::TpeSampler; pub use study::Study; -pub use trial::Trial; +pub use trial::{AttrValue, Trial}; pub use types::{Direction, TrialState}; /// Convenient wildcard import for the most common types. @@ -236,6 +236,6 @@ pub mod prelude { pub use crate::sampler::random::RandomSampler; pub use crate::sampler::tpe::TpeSampler; pub use crate::study::Study; - pub use crate::trial::Trial; + pub use crate::trial::{AttrValue, Trial}; pub use crate::types::Direction; } diff --git a/src/pruner/hyperband.rs b/src/pruner/hyperband.rs index d15216c..98c6cca 100644 --- a/src/pruner/hyperband.rs +++ b/src/pruner/hyperband.rs @@ -278,6 +278,7 @@ mod tests { HashMap::new(), 0.0, values.to_vec(), + HashMap::new(), ) } diff --git a/src/pruner/percentile.rs b/src/pruner/percentile.rs index 1005f6c..7a5d6f6 100644 --- a/src/pruner/percentile.rs +++ b/src/pruner/percentile.rs @@ -204,6 +204,7 @@ mod tests { HashMap::new(), 0.0, values.to_vec(), + HashMap::new(), ) } diff --git a/src/pruner/successive_halving.rs b/src/pruner/successive_halving.rs index dec8a20..376ddea 100644 --- a/src/pruner/successive_halving.rs +++ b/src/pruner/successive_halving.rs @@ -240,6 +240,7 @@ mod tests { HashMap::new(), 0.0, values.to_vec(), + HashMap::new(), ) } diff --git a/src/pruner/wilcoxon.rs b/src/pruner/wilcoxon.rs index 27cead2..d148f20 100644 --- a/src/pruner/wilcoxon.rs +++ b/src/pruner/wilcoxon.rs @@ -311,6 +311,7 @@ mod tests { HashMap::new(), value, intermediate_values, + HashMap::new(), ) } diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 65698f2..1289f41 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -9,6 +9,7 @@ use std::collections::HashMap; use crate::distribution::Distribution; use crate::param::ParamValue; use crate::parameter::{ParamId, Parameter}; +use crate::trial::AttrValue; use crate::types::TrialState; /// A completed trial with its parameters, distributions, and objective value. @@ -32,6 +33,8 @@ pub struct CompletedTrial { pub intermediate_values: Vec<(u64, f64)>, /// The state of the trial (Complete, Pruned, or Failed). pub state: TrialState, + /// User-defined attributes stored during the trial. + pub user_attrs: HashMap, } impl CompletedTrial { @@ -51,10 +54,11 @@ impl CompletedTrial { value, intermediate_values: Vec::new(), state: TrialState::Complete, + user_attrs: HashMap::new(), } } - /// Creates a new completed trial with intermediate values. + /// Creates a new completed trial with intermediate values and user attributes. pub fn with_intermediate_values( id: u64, params: HashMap, @@ -62,6 +66,7 @@ impl CompletedTrial { param_labels: HashMap, value: V, intermediate_values: Vec<(u64, f64)>, + user_attrs: HashMap, ) -> Self { Self { id, @@ -71,6 +76,7 @@ impl CompletedTrial { value, intermediate_values, state: TrialState::Complete, + user_attrs, } } @@ -114,6 +120,18 @@ impl CompletedTrial { .expect("parameter type mismatch: stored value incompatible with parameter") }) } + + /// Gets a user attribute by key. + #[must_use] + pub fn user_attr(&self, key: &str) -> Option<&AttrValue> { + self.user_attrs.get(key) + } + + /// Returns all user attributes. + #[must_use] + pub fn user_attrs(&self) -> &HashMap { + &self.user_attrs + } } /// A pending (running) trial with its parameters and distributions, but no objective value yet. diff --git a/src/study.rs b/src/study.rs index 5705200..72dee4d 100644 --- a/src/study.rs +++ b/src/study.rs @@ -363,6 +363,7 @@ where trial.param_labels().clone(), value, trial.intermediate_values().to_vec(), + trial.user_attrs().clone(), ); completed.state = TrialState::Complete; self.completed_trials.write().push(completed); @@ -471,6 +472,7 @@ where trial.param_labels().clone(), V::default(), trial.intermediate_values().to_vec(), + trial.user_attrs().clone(), ); completed.state = TrialState::Pruned; self.completed_trials.write().push(completed); diff --git a/src/trial.rs b/src/trial.rs index bc918ca..b0dc05c 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -13,6 +13,49 @@ use crate::pruner::Pruner; use crate::sampler::{CompletedTrial, Sampler}; use crate::types::TrialState; +/// A user attribute value that can be stored on a trial. +#[derive(Clone, Debug, PartialEq)] +pub enum AttrValue { + /// A floating-point attribute. + Float(f64), + /// An integer attribute. + Int(i64), + /// A string attribute. + String(String), + /// A boolean attribute. + Bool(bool), +} + +impl From for AttrValue { + fn from(v: f64) -> Self { + Self::Float(v) + } +} + +impl From for AttrValue { + fn from(v: i64) -> Self { + Self::Int(v) + } +} + +impl From for AttrValue { + fn from(v: String) -> Self { + Self::String(v) + } +} + +impl From<&str> for AttrValue { + fn from(v: &str) -> Self { + Self::String(v.to_owned()) + } +} + +impl From for AttrValue { + fn from(v: bool) -> Self { + Self::Bool(v) + } +} + /// A trial represents a single evaluation of the objective function. /// /// Each trial has a unique ID and stores the sampled parameters along with @@ -41,6 +84,8 @@ pub struct Trial { intermediate_values: Vec<(u64, f64)>, /// The pruner used to decide whether to stop this trial early. pruner: Option>, + /// User-defined attributes for logging, debugging, and analysis. + user_attrs: HashMap, } impl core::fmt::Debug for Trial { @@ -55,6 +100,7 @@ impl core::fmt::Debug for Trial { .field("has_history", &self.history.is_some()) .field("intermediate_values", &self.intermediate_values) .field("has_pruner", &self.pruner.is_some()) + .field("user_attrs", &self.user_attrs) .finish() } } @@ -92,6 +138,7 @@ impl Trial { history: None, intermediate_values: Vec::new(), pruner: None, + user_attrs: HashMap::new(), } } @@ -121,6 +168,7 @@ impl Trial { history: Some(history), intermediate_values: Vec::new(), pruner: Some(pruner), + user_attrs: HashMap::new(), } } @@ -209,6 +257,23 @@ impl Trial { &self.intermediate_values } + /// Sets a user attribute on this trial. + pub fn set_user_attr(&mut self, key: impl Into, value: impl Into) { + self.user_attrs.insert(key.into(), value.into()); + } + + /// Gets a user attribute by key. + #[must_use] + pub fn user_attr(&self, key: &str) -> Option<&AttrValue> { + self.user_attrs.get(key) + } + + /// Returns all user attributes. + #[must_use] + pub fn user_attrs(&self) -> &HashMap { + &self.user_attrs + } + /// Sets the trial state to Complete. pub(crate) fn set_complete(&mut self) { self.state = TrialState::Complete; diff --git a/tests/median_pruner_tests.rs b/tests/median_pruner_tests.rs index 980675f..0f08c35 100644 --- a/tests/median_pruner_tests.rs +++ b/tests/median_pruner_tests.rs @@ -13,6 +13,7 @@ fn trial_with_values(id: u64, intermediate_values: Vec<(u64, f64)>) -> Completed HashMap::new(), 0.0, intermediate_values, + HashMap::new(), ) } diff --git a/tests/user_attr_tests.rs b/tests/user_attr_tests.rs new file mode 100644 index 0000000..866560b --- /dev/null +++ b/tests/user_attr_tests.rs @@ -0,0 +1,148 @@ +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::{AttrValue, Direction, Study}; + +#[test] +fn set_and_get_float_attr() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("score", 42.5); + assert_eq!(trial.user_attr("score"), Some(&AttrValue::Float(42.5))); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); +} + +#[test] +fn set_and_get_int_attr() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("epoch", 42_i64); + assert_eq!(trial.user_attr("epoch"), Some(&AttrValue::Int(42))); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); +} + +#[test] +fn set_and_get_string_attr() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("model", "resnet50"); + assert_eq!( + trial.user_attr("model"), + Some(&AttrValue::String("resnet50".to_owned())) + ); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); +} + +#[test] +fn set_and_get_bool_attr() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("converged", true); + assert_eq!(trial.user_attr("converged"), Some(&AttrValue::Bool(true))); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); +} + +#[test] +fn attrs_propagate_to_completed_trial() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("time_secs", 1.5); + trial.set_user_attr("tag", "baseline"); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert_eq!(best.user_attr("time_secs"), Some(&AttrValue::Float(1.5))); + assert_eq!( + best.user_attr("tag"), + Some(&AttrValue::String("baseline".to_owned())) + ); +} + +#[test] +fn overwrite_attr_replaces_value() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("key", "old"); + trial.set_user_attr("key", "new"); + assert_eq!( + trial.user_attr("key"), + Some(&AttrValue::String("new".to_owned())) + ); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert_eq!( + best.user_attr("key"), + Some(&AttrValue::String("new".to_owned())) + ); +} + +#[test] +fn missing_attr_returns_none() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + assert_eq!(trial.user_attr("nonexistent"), None); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert_eq!(best.user_attr("nonexistent"), None); +} + +#[test] +fn user_attrs_map_returns_all() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(1, |trial| { + let _ = x.suggest(trial)?; + trial.set_user_attr("a", 1.0); + trial.set_user_attr("b", true); + assert_eq!(trial.user_attrs().len(), 2); + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert_eq!(best.user_attrs().len(), 2); +}