feat: add enqueue_trial() for pre-specified parameter evaluation
Add Study::enqueue() to push specific parameter configurations onto a FIFO queue. The next call to ask()/create_trial() or the next iteration in optimize() pops the front entry and injects it into the trial so that suggest_param() returns the pre-filled value instead of sampling. Parameters missing from an enqueued map fall back to normal sampling, and once the queue is drained regular sampler-driven trials resume.
This commit is contained in:
+62
-2
@@ -6,11 +6,14 @@ use core::future::Future;
|
||||
use core::ops::ControlFlow;
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
use core::time::Duration;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
|
||||
use crate::param::ParamValue;
|
||||
use crate::parameter::ParamId;
|
||||
use crate::pruner::{NopPruner, Pruner};
|
||||
use crate::sampler::random::RandomSampler;
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
@@ -53,6 +56,8 @@ where
|
||||
/// Set automatically for `Study<f64>` so that `create_trial()` and all
|
||||
/// optimization methods use the sampler without requiring `_with_sampler` suffixes.
|
||||
trial_factory: Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>,
|
||||
/// Queue of parameter configurations to evaluate next.
|
||||
enqueued_params: Arc<Mutex<VecDeque<HashMap<ParamId, ParamValue>>>>,
|
||||
}
|
||||
|
||||
impl<V> Study<V>
|
||||
@@ -170,6 +175,7 @@ where
|
||||
completed_trials,
|
||||
next_trial_id: AtomicU64::new(0),
|
||||
trial_factory,
|
||||
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -261,6 +267,7 @@ where
|
||||
completed_trials,
|
||||
next_trial_id: AtomicU64::new(0),
|
||||
trial_factory,
|
||||
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -292,6 +299,52 @@ where
|
||||
&*self.pruner
|
||||
}
|
||||
|
||||
/// Enqueues a specific parameter configuration to be evaluated next.
|
||||
///
|
||||
/// The next call to [`ask()`](Self::ask) or the next trial in [`optimize()`](Self::optimize)
|
||||
/// will use these exact parameters instead of sampling from the sampler.
|
||||
///
|
||||
/// Multiple configurations can be enqueued; they are evaluated in FIFO order.
|
||||
/// If an enqueued configuration is missing a parameter that the objective calls
|
||||
/// `suggest()` on, that parameter falls back to normal sampling.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `params` - A map from parameter IDs to the values to use.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use std::collections::HashMap;
|
||||
///
|
||||
/// use optimizer::parameter::{FloatParam, IntParam, Parameter};
|
||||
/// use optimizer::{Direction, ParamValue, Study};
|
||||
///
|
||||
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// let x = FloatParam::new(0.0, 10.0);
|
||||
/// let y = IntParam::new(1, 100);
|
||||
///
|
||||
/// // Evaluate these specific configurations first
|
||||
/// study.enqueue(HashMap::from([
|
||||
/// (x.id(), ParamValue::Float(0.001)),
|
||||
/// (y.id(), ParamValue::Int(3)),
|
||||
/// ]));
|
||||
///
|
||||
/// // Next trial will use x=0.001, y=3
|
||||
/// let mut trial = study.ask();
|
||||
/// assert_eq!(x.suggest(&mut trial).unwrap(), 0.001);
|
||||
/// assert_eq!(y.suggest(&mut trial).unwrap(), 3);
|
||||
/// ```
|
||||
pub fn enqueue(&self, params: HashMap<ParamId, ParamValue>) {
|
||||
self.enqueued_params.lock().push_back(params);
|
||||
}
|
||||
|
||||
/// Returns the number of enqueued parameter configurations.
|
||||
#[must_use]
|
||||
pub fn n_enqueued(&self) -> usize {
|
||||
self.enqueued_params.lock().len()
|
||||
}
|
||||
|
||||
/// Generates the next unique trial ID.
|
||||
pub(crate) fn next_trial_id(&self) -> u64 {
|
||||
self.next_trial_id.fetch_add(1, Ordering::SeqCst)
|
||||
@@ -321,11 +374,18 @@ where
|
||||
/// ```
|
||||
pub fn create_trial(&self) -> Trial {
|
||||
let id = self.next_trial_id();
|
||||
if let Some(factory) = &self.trial_factory {
|
||||
let mut trial = if let Some(factory) = &self.trial_factory {
|
||||
factory(id)
|
||||
} else {
|
||||
Trial::new(id)
|
||||
};
|
||||
|
||||
// If there are enqueued params, inject them into this trial
|
||||
if let Some(fixed_params) = self.enqueued_params.lock().pop_front() {
|
||||
trial.set_fixed_params(fixed_params);
|
||||
}
|
||||
|
||||
trial
|
||||
}
|
||||
|
||||
/// Records a completed trial with its objective value.
|
||||
|
||||
+21
-2
@@ -86,6 +86,8 @@ pub struct Trial {
|
||||
pruner: Option<Arc<dyn Pruner>>,
|
||||
/// User-defined attributes for logging, debugging, and analysis.
|
||||
user_attrs: HashMap<String, AttrValue>,
|
||||
/// Pre-filled parameter values from enqueue (used instead of sampling).
|
||||
fixed_params: HashMap<ParamId, ParamValue>,
|
||||
}
|
||||
|
||||
impl core::fmt::Debug for Trial {
|
||||
@@ -101,6 +103,7 @@ impl core::fmt::Debug for Trial {
|
||||
.field("intermediate_values", &self.intermediate_values)
|
||||
.field("has_pruner", &self.pruner.is_some())
|
||||
.field("user_attrs", &self.user_attrs)
|
||||
.field("fixed_params", &self.fixed_params)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -139,6 +142,7 @@ impl Trial {
|
||||
intermediate_values: Vec::new(),
|
||||
pruner: None,
|
||||
user_attrs: HashMap::new(),
|
||||
fixed_params: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -169,9 +173,18 @@ impl Trial {
|
||||
intermediate_values: Vec::new(),
|
||||
pruner: Some(pruner),
|
||||
user_attrs: HashMap::new(),
|
||||
fixed_params: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets pre-filled parameters on this trial.
|
||||
///
|
||||
/// When `suggest_param` is called for a parameter that has a fixed value,
|
||||
/// the fixed value is used instead of sampling.
|
||||
pub(crate) fn set_fixed_params(&mut self, params: HashMap<ParamId, ParamValue>) {
|
||||
self.fixed_params = params;
|
||||
}
|
||||
|
||||
/// Samples a value from the given distribution using the sampler.
|
||||
///
|
||||
/// If the trial has a sampler, it delegates to the sampler's sample method
|
||||
@@ -343,8 +356,14 @@ impl Trial {
|
||||
});
|
||||
}
|
||||
|
||||
// Sample using the sampler
|
||||
let value = self.sample_value(&distribution);
|
||||
// Check for a pre-filled (enqueued) value for this parameter
|
||||
let value = if let Some(fixed_value) = self.fixed_params.remove(¶m_id) {
|
||||
fixed_value
|
||||
} else {
|
||||
// Sample using the sampler
|
||||
self.sample_value(&distribution)
|
||||
};
|
||||
|
||||
let result = param.cast_param_value(&value)?;
|
||||
|
||||
// Store distribution, value, and label
|
||||
|
||||
@@ -1700,3 +1700,195 @@ fn test_ask_and_tell_with_custom_value_type() {
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
assert_eq!(study.best_value().unwrap(), 40);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests: enqueue trials
|
||||
// =============================================================================
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use optimizer::ParamValue;
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_params_evaluated_first() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
let y = IntParam::new(1, 100);
|
||||
|
||||
// Enqueue a specific configuration
|
||||
study.enqueue(HashMap::from([
|
||||
(x.id(), ParamValue::Float(5.0)),
|
||||
(y.id(), ParamValue::Int(42)),
|
||||
]));
|
||||
|
||||
// The first trial should use the enqueued params
|
||||
let mut trial = study.ask();
|
||||
let x_val = x.suggest(&mut trial).unwrap();
|
||||
let y_val = y.suggest(&mut trial).unwrap();
|
||||
|
||||
assert_eq!(x_val, 5.0);
|
||||
assert_eq!(y_val, 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_fifo_order() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Enqueue two configs
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
|
||||
// First trial gets first enqueued value
|
||||
let mut trial1 = study.ask();
|
||||
assert_eq!(x.suggest(&mut trial1).unwrap(), 1.0);
|
||||
|
||||
// Second trial gets second enqueued value
|
||||
let mut trial2 = study.ask();
|
||||
assert_eq!(x.suggest(&mut trial2).unwrap(), 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_then_normal_sampling_resumes() {
|
||||
let sampler = RandomSampler::with_seed(42);
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Enqueue one config
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(5.0))]));
|
||||
|
||||
// First trial uses enqueued value
|
||||
let mut trial1 = study.ask();
|
||||
assert_eq!(x.suggest(&mut trial1).unwrap(), 5.0);
|
||||
study.tell(trial1, Ok::<_, &str>(25.0));
|
||||
|
||||
// Second trial uses normal sampling (not 5.0)
|
||||
let mut trial2 = study.ask();
|
||||
let x_val = x.suggest(&mut trial2).unwrap();
|
||||
// The sampled value should be in [0, 10] but extremely unlikely to be exactly 5.0
|
||||
assert!((0.0..=10.0).contains(&x_val));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_with_optimize() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Enqueue two specific configs
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
|
||||
let mut values = Vec::new();
|
||||
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x_val = x.suggest(trial)?;
|
||||
values.push(x_val);
|
||||
Ok::<_, Error>(x_val * x_val)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// First two trials should use enqueued values
|
||||
assert_eq!(values[0], 1.0);
|
||||
assert_eq!(values[1], 2.0);
|
||||
// All 5 trials should have completed
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_partial_params_fall_back_to_sampling() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
let y = IntParam::new(1, 100);
|
||||
|
||||
// Enqueue only x, not y
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(3.0))]));
|
||||
|
||||
let mut trial = study.ask();
|
||||
let x_val = x.suggest(&mut trial).unwrap();
|
||||
let y_val = y.suggest(&mut trial).unwrap();
|
||||
|
||||
// x should be the enqueued value
|
||||
assert_eq!(x_val, 3.0);
|
||||
// y should be sampled (within range)
|
||||
assert!((1..=100).contains(&y_val));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_trials_appear_in_completed_trials() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(7.0))]));
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let x_val = x.suggest(trial)?;
|
||||
Ok::<_, Error>(x_val)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let trials = study.trials();
|
||||
assert_eq!(trials.len(), 1);
|
||||
assert_eq!(trials[0].value, 7.0);
|
||||
assert_eq!(
|
||||
*trials[0].params.get(&x.id()).unwrap(),
|
||||
ParamValue::Float(7.0)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_with_ask_and_tell() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(4.0))]));
|
||||
|
||||
let mut trial = study.ask();
|
||||
let x_val = x.suggest(&mut trial).unwrap();
|
||||
assert_eq!(x_val, 4.0);
|
||||
|
||||
study.tell(trial, Ok::<_, &str>(x_val * x_val));
|
||||
assert_eq!(study.n_trials(), 1);
|
||||
assert_eq!(study.best_value().unwrap(), 16.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_n_enqueued() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
assert_eq!(study.n_enqueued(), 0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
assert_eq!(study.n_enqueued(), 1);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
assert_eq!(study.n_enqueued(), 2);
|
||||
|
||||
// Creating a trial dequeues one
|
||||
let _ = study.ask();
|
||||
assert_eq!(study.n_enqueued(), 1);
|
||||
|
||||
let _ = study.ask();
|
||||
assert_eq!(study.n_enqueued(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_counted_in_n_trials() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x_val = x.suggest(trial)?;
|
||||
Ok::<_, Error>(x_val)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// All 5 trials count, including the 2 enqueued ones
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user