From df5fcedecfb2710ce9bb130772bf3bdef3d062d8 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 17:47:13 +0100 Subject: [PATCH] feat: add tracing integration behind `tracing` feature flag Add optional structured logging via the tracing crate: - info-level spans on all optimize* methods (direction, n_trials/duration) - info events for trial completed, new best value, trial pruned - debug events for trial failed and parameter sampled - Zero overhead when the feature is disabled --- Cargo.toml | 2 + src/lib.rs | 25 ++++++++++ src/study.rs | 129 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/trial.rs | 13 +++++- 4 files changed, 168 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 64e40a2..902400a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,12 +23,14 @@ tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true optimizer-derive = { version = "0.1.0", path = "optimizer-derive", optional = true } serde = { version = "1", features = ["derive"], optional = true } serde_json = { version = "1", optional = true } +tracing = { version = "0.1", optional = true } [features] default = [] async = ["dep:tokio"] derive = ["dep:optimizer-derive"] serde = ["dep:serde", "dep:serde_json"] +tracing = ["dep:tracing"] [dev-dependencies] tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } diff --git a/src/lib.rs b/src/lib.rs index 7bedd32..3ca1885 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -183,6 +183,31 @@ //! - `async`: Enable async optimization methods (requires tokio) //! - `derive`: Enable `#[derive(Categorical)]` for enum parameters //! - `serde`: Enable `Serialize`/`Deserialize` on public types and `Study::save()`/`Study::load()` +//! - `tracing`: Emit structured log events via the [`tracing`](https://docs.rs/tracing) crate at key optimization points + +/// Emit a `tracing::info!` event when the `tracing` feature is enabled. +/// No-op otherwise. +#[cfg(feature = "tracing")] +macro_rules! trace_info { + ($($arg:tt)*) => { tracing::info!($($arg)*) }; +} + +#[cfg(not(feature = "tracing"))] +macro_rules! trace_info { + ($($arg:tt)*) => {}; +} + +/// Emit a `tracing::debug!` event when the `tracing` feature is enabled. +/// No-op otherwise. +#[cfg(feature = "tracing")] +macro_rules! trace_debug { + ($($arg:tt)*) => { tracing::debug!($($arg)*) }; +} + +#[cfg(not(feature = "tracing"))] +macro_rules! trace_debug { + ($($arg:tt)*) => {}; +} mod distribution; mod error; diff --git a/src/study.rs b/src/study.rs index 94e4447..97e2d87 100644 --- a/src/study.rs +++ b/src/study.rs @@ -340,6 +340,24 @@ where self.enqueued_params.lock().push_back(params); } + /// Returns the trial ID of the current best trial from the given slice. + #[cfg(feature = "tracing")] + fn best_id(&self, trials: &[CompletedTrial]) -> Option { + trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .max_by(|a, b| { + let ordering = a.value.partial_cmp(&b.value); + match self.direction { + Direction::Minimize => { + ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse) + } + Direction::Maximize => ordering.unwrap_or(core::cmp::Ordering::Equal), + } + }) + .map(|t| t.id) + } + /// Returns the number of enqueued parameter configurations. #[must_use] pub fn n_enqueued(&self) -> usize { @@ -789,18 +807,43 @@ where E: ToString + 'static, V: Default, { + #[cfg(feature = "tracing")] + let _span = + tracing::info_span!("optimize", n_trials, direction = ?self.direction).entered(); + for _ in 0..n_trials { let mut trial = self.create_trial(); match objective(&mut trial) { Ok(value) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); self.complete_trial(trial, value); + + #[cfg(feature = "tracing")] + { + tracing::info!(trial_id, "trial completed"); + let trials = self.completed_trials.read(); + if trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .count() + == 1 + || trials.last().map(|t| t.id) == self.best_id(&trials) + { + tracing::info!(trial_id, "new best value found"); + } + } } Err(e) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); if is_trial_pruned(&e) { self.prune_trial(trial); + trace_info!(trial_id, "trial pruned"); } else { self.fail_trial(trial, e.to_string()); + trace_debug!(trial_id, "trial failed"); } } } @@ -882,17 +925,25 @@ where Fut: Future>, E: ToString, { + #[cfg(feature = "tracing")] + let _span = + tracing::info_span!("optimize_async", n_trials, direction = ?self.direction).entered(); + for _ in 0..n_trials { let trial = self.create_trial(); + #[cfg(feature = "tracing")] + let trial_id = trial.id(); match objective(trial).await { Ok((trial, value)) => { self.complete_trial(trial, value); + trace_info!(trial_id, "trial completed"); } Err(e) => { // For async, we don't have the trial back on error // We'll just count this as a failed trial without recording it let _ = e.to_string(); + trace_debug!(trial_id, "trial failed"); } } } @@ -979,6 +1030,9 @@ where { use tokio::sync::Semaphore; + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize_parallel", n_trials, concurrency, direction = ?self.direction).entered(); + let semaphore = Arc::new(Semaphore::new(concurrency)); let objective = Arc::new(objective); @@ -1009,7 +1063,10 @@ where .map_err(|e| crate::Error::TaskError(e.to_string()))? { Ok((trial, value)) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); self.complete_trial(trial, value); + trace_info!(trial_id, "trial completed"); } Err(e) => { let _ = e.to_string(); @@ -1099,13 +1156,34 @@ where C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, E: ToString + 'static, { + #[cfg(feature = "tracing")] + let _span = + tracing::info_span!("optimize", n_trials, direction = ?self.direction).entered(); + for _ in 0..n_trials { let mut trial = self.create_trial(); match objective(&mut trial) { Ok(value) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); self.complete_trial(trial, value); + #[cfg(feature = "tracing")] + { + tracing::info!(trial_id, "trial completed"); + let trials = self.completed_trials.read(); + if trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .count() + == 1 + || trials.last().map(|t| t.id) == self.best_id(&trials) + { + tracing::info!(trial_id, "new best value found"); + } + } + // Get the just-completed trial for the callback let trials = self.completed_trials.read(); let Some(completed) = trials.last() else { @@ -1125,10 +1203,14 @@ where } } Err(e) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); if is_trial_pruned(&e) { self.prune_trial(trial); + trace_info!(trial_id, "trial pruned"); } else { self.fail_trial(trial, e.to_string()); + trace_debug!(trial_id, "trial failed"); } } } @@ -1192,19 +1274,29 @@ where E: ToString + 'static, V: Default, { + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize", duration_secs = duration.as_secs(), direction = ?self.direction).entered(); + let deadline = Instant::now() + duration; while Instant::now() < deadline { let mut trial = self.create_trial(); match objective(&mut trial) { Ok(value) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); self.complete_trial(trial, value); + trace_info!(trial_id, "trial completed"); } Err(e) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); if is_trial_pruned(&e) { self.prune_trial(trial); + trace_info!(trial_id, "trial pruned"); } else { self.fail_trial(trial, e.to_string()); + trace_debug!(trial_id, "trial failed"); } } } @@ -1287,14 +1379,34 @@ where C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, E: ToString + 'static, { + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize", duration_secs = duration.as_secs(), direction = ?self.direction).entered(); + let deadline = Instant::now() + duration; while Instant::now() < deadline { let mut trial = self.create_trial(); match objective(&mut trial) { Ok(value) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); self.complete_trial(trial, value); + #[cfg(feature = "tracing")] + { + tracing::info!(trial_id, "trial completed"); + let trials = self.completed_trials.read(); + if trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .count() + == 1 + || trials.last().map(|t| t.id) == self.best_id(&trials) + { + tracing::info!(trial_id, "new best value found"); + } + } + let trials = self.completed_trials.read(); let Some(completed) = trials.last() else { return Err(crate::Error::Internal( @@ -1310,10 +1422,14 @@ where } } Err(e) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); if is_trial_pruned(&e) { self.prune_trial(trial); + trace_info!(trial_id, "trial pruned"); } else { self.fail_trial(trial, e.to_string()); + trace_debug!(trial_id, "trial failed"); } } } @@ -1356,16 +1472,23 @@ where Fut: Future>, E: ToString, { + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize_until_async", duration_secs = duration.as_secs(), direction = ?self.direction).entered(); + let deadline = Instant::now() + duration; while Instant::now() < deadline { let trial = self.create_trial(); + #[cfg(feature = "tracing")] + let trial_id = trial.id(); match objective(trial).await { Ok((trial, value)) => { self.complete_trial(trial, value); + trace_info!(trial_id, "trial completed"); } Err(e) => { let _ = e.to_string(); + trace_debug!(trial_id, "trial failed"); } } } @@ -1415,6 +1538,9 @@ where { use tokio::sync::Semaphore; + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize_until_parallel", duration_secs = duration.as_secs(), concurrency, direction = ?self.direction).entered(); + let deadline = Instant::now() + duration; let semaphore = Arc::new(Semaphore::new(concurrency)); let objective = Arc::new(objective); @@ -1445,7 +1571,10 @@ where .map_err(|e| crate::Error::TaskError(e.to_string()))? { Ok((trial, value)) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); self.complete_trial(trial, value); + trace_info!(trial_id, "trial completed"); } Err(e) => { let _ = e.to_string(); diff --git a/src/trial.rs b/src/trial.rs index b4b90a0..cbb7f65 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -262,7 +262,11 @@ impl Trial { return false; }; let history_guard = history.read(); - pruner.should_prune(self.id, step, &self.intermediate_values, &history_guard) + let prune = pruner.should_prune(self.id, step, &self.intermediate_values, &history_guard); + if prune { + trace_info!(trial_id = self.id, step, "pruner recommends stopping"); + } + prune } /// Returns all intermediate values reported so far. @@ -367,6 +371,13 @@ impl Trial { let result = param.cast_param_value(&value)?; + trace_debug!( + trial_id = self.id, + param = %param.label(), + value = %value, + "parameter sampled" + ); + // Store distribution, value, and label self.distributions.insert(param_id, distribution); self.params.insert(param_id, value);