From 5fe0a75f787872b52ef05cdf2eee6937ef9fae3a Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 17:33:48 +0100 Subject: [PATCH] feat: add serde serialization support behind `serde` feature flag Add Serialize/Deserialize derives to all public data types (ParamValue, Distribution, Direction, TrialState, ParamId, AttrValue, CompletedTrial) gated behind a `serde` feature flag. Introduce StudySnapshot struct and Study::save()/Study::load() for persisting and restoring study state as human-readable JSON. --- Cargo.toml | 4 + src/distribution.rs | 4 + src/lib.rs | 5 ++ src/param.rs | 1 + src/parameter.rs | 1 + src/sampler/mod.rs | 1 + src/study.rs | 74 +++++++++++++++++ src/trial.rs | 1 + src/types.rs | 2 + tests/serde_tests.rs | 184 +++++++++++++++++++++++++++++++++++++++++++ 10 files changed, 277 insertions(+) create mode 100644 tests/serde_tests.rs diff --git a/Cargo.toml b/Cargo.toml index 1c74d31..64e40a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,15 +21,19 @@ thiserror = "2" parking_lot = "0.12" 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 } [features] default = [] async = ["dep:tokio"] derive = ["dep:optimizer-derive"] +serde = ["dep:serde", "dep:serde_json"] [dev-dependencies] tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } optimizer-derive = { version = "0.1.0", path = "optimizer-derive" } +serde_json = "1" [[example]] name = "async_api_optimization" diff --git a/src/distribution.rs b/src/distribution.rs index fd64f80..a597a8a 100644 --- a/src/distribution.rs +++ b/src/distribution.rs @@ -2,6 +2,7 @@ /// Distribution for floating-point parameters. #[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct FloatDistribution { /// Lower bound (inclusive). pub low: f64, @@ -15,6 +16,7 @@ pub struct FloatDistribution { /// Distribution for integer parameters. #[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct IntDistribution { /// Lower bound (inclusive). pub low: i64, @@ -28,6 +30,7 @@ pub struct IntDistribution { /// Distribution for categorical parameters. #[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct CategoricalDistribution { /// Number of choices available. pub n_choices: usize, @@ -35,6 +38,7 @@ pub struct CategoricalDistribution { /// Enum wrapping all parameter distribution types. #[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub enum Distribution { /// A floating-point distribution. Float(FloatDistribution), diff --git a/src/lib.rs b/src/lib.rs index 58f6eb1..7bedd32 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -182,6 +182,7 @@ //! //! - `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()` mod distribution; mod error; @@ -210,6 +211,8 @@ pub use sampler::grid::GridSearchSampler; pub use sampler::random::RandomSampler; pub use sampler::tpe::TpeSampler; pub use study::Study; +#[cfg(feature = "serde")] +pub use study::StudySnapshot; pub use trial::{AttrValue, Trial}; pub use types::{Direction, TrialState}; @@ -236,6 +239,8 @@ pub mod prelude { pub use crate::sampler::random::RandomSampler; pub use crate::sampler::tpe::TpeSampler; pub use crate::study::Study; + #[cfg(feature = "serde")] + pub use crate::study::StudySnapshot; pub use crate::trial::{AttrValue, Trial}; pub use crate::types::Direction; } diff --git a/src/param.rs b/src/param.rs index b5c0009..6dd3e60 100644 --- a/src/param.rs +++ b/src/param.rs @@ -6,6 +6,7 @@ /// For categorical parameters, the `Categorical` variant stores /// the index into the choices array. #[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub enum ParamValue { /// A floating-point parameter value. Float(f64), diff --git a/src/parameter.rs b/src/parameter.rs index 3fa820c..34f5fd3 100644 --- a/src/parameter.rs +++ b/src/parameter.rs @@ -38,6 +38,7 @@ static NEXT_PARAM_ID: AtomicU64 = AtomicU64::new(0); /// Each parameter is assigned a unique `ParamId` at creation time. Cloning a parameter /// copies its `ParamId`, so clones refer to the same logical parameter. #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct ParamId(u64); impl ParamId { diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 1289f41..6f4c07f 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -18,6 +18,7 @@ use crate::types::TrialState; /// parameter values, their distributions, and the objective value returned /// by the objective function. #[derive(Clone, Debug)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct CompletedTrial { /// The unique identifier for this trial. pub id: u64, diff --git a/src/study.rs b/src/study.rs index 894ac5e..ec50734 100644 --- a/src/study.rs +++ b/src/study.rs @@ -1650,6 +1650,80 @@ impl Study { } } +/// A serializable snapshot of a study's state. +/// +/// Since [`Study`] contains non-serializable fields (samplers, atomics, etc.), +/// this struct captures the essential state needed to save and restore a study. +/// +/// # Schema versioning +/// +/// The `version` field enables future schema evolution without breaking existing files. +/// The current version is `1`. +/// +/// # Sampler state +/// +/// Sampler state is **not** included in the snapshot. After loading, the study +/// uses a default `RandomSampler`. Call [`Study::set_sampler`] to restore +/// the desired sampler configuration. +#[cfg(feature = "serde")] +#[derive(serde::Serialize, serde::Deserialize)] +pub struct StudySnapshot { + /// Schema version for forward compatibility. + pub version: u32, + /// The optimization direction. + pub direction: Direction, + /// All completed (and pruned) trials. + pub trials: Vec>, + /// The next trial ID to assign. + pub next_trial_id: u64, + /// Optional metadata (creation timestamp, sampler description, etc.). + pub metadata: HashMap, +} + +#[cfg(feature = "serde")] +impl Study { + /// Saves the study state to a JSON file. + /// + /// # Errors + /// + /// Returns an I/O error if the file cannot be created or written. + pub fn save(&self, path: impl AsRef) -> std::io::Result<()> { + let snapshot = StudySnapshot { + version: 1, + direction: self.direction, + trials: self.trials(), + next_trial_id: self.next_trial_id.load(Ordering::Relaxed), + metadata: HashMap::new(), + }; + let file = std::fs::File::create(path)?; + serde_json::to_writer_pretty(file, &snapshot).map_err(std::io::Error::other) + } +} + +#[cfg(feature = "serde")] +impl Study { + /// Loads a study from a JSON file. + /// + /// The loaded study uses a `RandomSampler` by default. Call + /// [`set_sampler()`](Self::set_sampler) to restore the original sampler + /// configuration. + /// + /// # Errors + /// + /// Returns an I/O error if the file cannot be read or parsed. + pub fn load(path: impl AsRef) -> std::io::Result { + let file = std::fs::File::open(path)?; + let snapshot: StudySnapshot = serde_json::from_reader(file) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + let study = Study::new(snapshot.direction); + *study.completed_trials.write() = snapshot.trials; + study + .next_trial_id + .store(snapshot.next_trial_id, Ordering::Relaxed); + Ok(study) + } +} + /// Returns `true` if the error represents a pruned trial. /// /// Checks via `Any` downcasting whether `e` is `Error::TrialPruned` or diff --git a/src/trial.rs b/src/trial.rs index 3cc33c2..b4b90a0 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -15,6 +15,7 @@ use crate::types::TrialState; /// A user attribute value that can be stored on a trial. #[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub enum AttrValue { /// A floating-point attribute. Float(f64), diff --git a/src/types.rs b/src/types.rs index b5e4d96..9e25026 100644 --- a/src/types.rs +++ b/src/types.rs @@ -2,6 +2,7 @@ /// The direction of optimization. #[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub enum Direction { /// Minimize the objective value. Minimize, @@ -11,6 +12,7 @@ pub enum Direction { /// The state of a trial in its lifecycle. #[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub enum TrialState { /// The trial is currently running. Running, diff --git a/tests/serde_tests.rs b/tests/serde_tests.rs new file mode 100644 index 0000000..cb0e812 --- /dev/null +++ b/tests/serde_tests.rs @@ -0,0 +1,184 @@ +#![cfg(feature = "serde")] + +use std::collections::HashMap; + +use optimizer::parameter::{FloatParam, IntParam, Parameter}; +use optimizer::sampler::CompletedTrial; +use optimizer::{Direction, ParamValue, Study, StudySnapshot, TrialState}; + +#[test] +fn round_trip_save_load() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(-10.0, 10.0).name("x"); + let n = IntParam::new(1, 100).name("n"); + + study + .optimize(5, |trial| { + let x_val = x.suggest(trial)?; + let n_val = n.suggest(trial)?; + Ok::<_, optimizer::Error>(x_val * x_val + n_val as f64) + }) + .unwrap(); + + let dir = tempdir(); + let path = dir.join("study.json"); + + study.save(&path).unwrap(); + let loaded: Study = Study::load(&path).unwrap(); + + assert_eq!(loaded.direction(), study.direction()); + assert_eq!(loaded.n_trials(), study.n_trials()); + + let orig_trials = study.trials(); + let loaded_trials = loaded.trials(); + + for (orig, loaded) in orig_trials.iter().zip(loaded_trials.iter()) { + assert_eq!(orig.id, loaded.id); + assert!((orig.value - loaded.value).abs() < 1e-10); + assert_eq!(orig.state, loaded.state); + assert_eq!(orig.params.len(), loaded.params.len()); + assert_eq!(orig.distributions, loaded.distributions); + assert_eq!(orig.param_labels, loaded.param_labels); + } + + // Clean up + std::fs::remove_dir_all(&dir).ok(); +} + +#[test] +fn json_output_is_human_readable() { + let study: Study = Study::new(Direction::Maximize); + let x = FloatParam::new(0.0, 1.0).name("x"); + + study + .optimize(2, |trial| { + let v = x.suggest(trial)?; + Ok::<_, optimizer::Error>(v) + }) + .unwrap(); + + let dir = tempdir(); + let path = dir.join("study.json"); + study.save(&path).unwrap(); + + let contents = std::fs::read_to_string(&path).unwrap(); + + // Verify it's pretty-printed JSON with recognizable fields + assert!(contents.contains("\"version\"")); + assert!(contents.contains("\"direction\"")); + assert!(contents.contains("\"trials\"")); + assert!(contents.contains("\"next_trial_id\"")); + assert!(contents.contains("\"Maximize\"")); + + std::fs::remove_dir_all(&dir).ok(); +} + +#[test] +fn round_trip_empty_study() { + let study: Study = Study::new(Direction::Minimize); + + let dir = tempdir(); + let path = dir.join("empty.json"); + + study.save(&path).unwrap(); + let loaded: Study = Study::load(&path).unwrap(); + + assert_eq!(loaded.direction(), Direction::Minimize); + assert_eq!(loaded.n_trials(), 0); + + std::fs::remove_dir_all(&dir).ok(); +} + +#[test] +fn snapshot_version_field_is_present() { + let study: Study = Study::new(Direction::Minimize); + + let dir = tempdir(); + let path = dir.join("version.json"); + study.save(&path).unwrap(); + + let contents = std::fs::read_to_string(&path).unwrap(); + let snapshot: StudySnapshot = serde_json::from_str(&contents).unwrap(); + + assert_eq!(snapshot.version, 1); + + std::fs::remove_dir_all(&dir).ok(); +} + +#[test] +fn completed_trial_serde_round_trip() { + let trial = CompletedTrial::new(42, HashMap::new(), HashMap::new(), HashMap::new(), 2.78); + + let json = serde_json::to_string(&trial).unwrap(); + let deserialized: CompletedTrial = serde_json::from_str(&json).unwrap(); + + assert_eq!(deserialized.id, 42); + assert_eq!(deserialized.value, 2.78); + assert_eq!(deserialized.state, TrialState::Complete); +} + +#[test] +fn param_value_serde_round_trip() { + let values = vec![ + ParamValue::Float(1.23), + ParamValue::Int(42), + ParamValue::Categorical(2), + ]; + + for val in &values { + let json = serde_json::to_string(val).unwrap(); + let deserialized: ParamValue = serde_json::from_str(&json).unwrap(); + assert_eq!(&deserialized, val); + } +} + +#[test] +fn direction_serde_round_trip() { + let min_json = serde_json::to_string(&Direction::Minimize).unwrap(); + let max_json = serde_json::to_string(&Direction::Maximize).unwrap(); + + assert_eq!( + serde_json::from_str::(&min_json).unwrap(), + Direction::Minimize + ); + assert_eq!( + serde_json::from_str::(&max_json).unwrap(), + Direction::Maximize + ); +} + +#[test] +fn round_trip_preserves_trial_id_counter() { + let study: Study = Study::new(Direction::Minimize); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(10, |trial| { + let v = x.suggest(trial)?; + Ok::<_, optimizer::Error>(v) + }) + .unwrap(); + + let dir = tempdir(); + let path = dir.join("counter.json"); + study.save(&path).unwrap(); + + let loaded: Study = Study::load(&path).unwrap(); + + // Creating a new trial should use an ID >= 10 + let trial = loaded.create_trial(); + assert!(trial.id() >= 10); + + std::fs::remove_dir_all(&dir).ok(); +} + +/// Helper to create a unique temporary directory. +fn tempdir() -> std::path::PathBuf { + use std::sync::atomic::{AtomicU64, Ordering}; + static COUNTER: AtomicU64 = AtomicU64::new(0); + let id = COUNTER.fetch_add(1, Ordering::Relaxed); + let dir = + std::env::temp_dir().join(format!("optimizer_serde_test_{}_{id}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + dir +}