gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add examples for async API parameter optimization and ML hyperparameter tuning

Browse Source
This commit is contained in:
Manuel Raimann
2026-01-31 11:56:50 +01:00
parent b9a90ebec3
commit af8a9f7638
4 changed files with 612 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+16
View File
@@ -61,6 +61,21 @@ jobs:
cargo test --verbose --features "${{ matrix.features }}"
fi
examples:
name: Examples
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- name: Install Rust
run: |
rustup override set stable
rustup update stable
- uses: Swatinem/rust-cache@v2
- name: Run sync example (ml_hyperparameter_tuning)
run: cargo run --example ml_hyperparameter_tuning
- name: Run async example (async_api_optimization)
run: cargo run --example async_api_optimization --features async
docs:
name: Docs
runs-on: ubuntu-latest
@@ -243,6 +258,7 @@ jobs:
- fmt
- clippy
- test
- examples
- docs
- feature-check
- coverage
Cargo.toml
+10 -1
View File
@@ -23,4 +23,13 @@ default = []
async = ["dep:tokio"]
[dev-dependencies]
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
[[example]]
name = "async_api_optimization"
path = "examples/async_api_optimization.rs"
required-features = ["async"]
[[example]]
name = "ml_hyperparameter_tuning"
path = "examples/ml_hyperparameter_tuning.rs"
examples/async_api_optimization.rs
+317
View File
@@ -0,0 +1,317 @@
//! Async API Parameter Optimization Example
//!
//! This example shows how to use async/parallel optimization to tune
//! configuration parameters for a web service. Each evaluation simulates
//! an async operation (like deploying and load-testing a service).
//!
//! # Key Concepts Demonstrated
//!
//! - Async optimization with `optimize_parallel_with_sampler`
//! - Running multiple trials concurrently for faster optimization
//! - Boolean and categorical parameter types
//! - Measuring speedup from parallelism
//!
//! # When to Use Async Optimization
//!
//! Use async/parallel optimization when your objective function involves:
//! - Network requests (API calls, database queries)
//! - File I/O operations
//! - External service calls
//! - Any operation where you're waiting for I/O rather than computing
//!
//! With parallelism, you can evaluate multiple configurations simultaneously,
//! significantly reducing total optimization time.
//!
//! Run with: `cargo run --example async_api_optimization --features async`
use std::time::{Duration, Instant};
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, ParamValue, Study, Trial};
// ============================================================================
// Configuration: Service parameters we want to tune
// ============================================================================
/// Configuration for a web service.
///
/// In a real application, these parameters would control:
/// - Memory allocation (cache sizes)
/// - Connection management (pool sizes, timeouts)
/// - Request handling (batching, compression)
/// - Protocol options (HTTP version, load balancing)
struct ServiceConfig {
cache_size_mb: i64,
connection_pool_size: i64,
request_timeout_ms: i64,
retry_count: i64,
batch_size: i64,
compression_level: i64,
use_http2: bool,
load_balancing: String,
}
// ============================================================================
// Objective Function: Evaluate a service configuration
// ============================================================================
/// Simulates deploying and load-testing a service configuration.
///
/// In a real scenario, this function might:
/// 1. Deploy the configuration to a staging environment
/// 2. Run load tests against the service
/// 3. Collect metrics (latency, throughput, error rate)
/// 4. Return a composite score
///
/// The async sleep simulates the I/O time of these operations.
/// This is where parallel execution helps - while one trial is waiting
/// for I/O, other trials can run.
#[allow(clippy::too_many_arguments)]
async fn evaluate_service(config: &ServiceConfig) -> f64 {
// Simulate async I/O (deployment, load testing, metric collection)
tokio::time::sleep(Duration::from_millis(50)).await;
// Calculate a score based on how close we are to optimal values
// Lower score = better configuration
let mut score = 0.0;
// Cache size: too small = cache misses, too large = wasted memory
// Optimal around 512MB
let cache_optimal = 512.0;
score += ((config.cache_size_mb as f64 - cache_optimal) / 256.0).powi(2);
// Connection pool: too small = contention, too large = resource waste
// Optimal around 100
let pool_optimal = 100.0;
score += ((config.connection_pool_size as f64 - pool_optimal) / 50.0).powi(2);
// Timeout: too short = false failures, too long = slow recovery
// Optimal around 5000ms
let timeout_optimal = 5000.0;
score += ((config.request_timeout_ms as f64 - timeout_optimal) / 2000.0).powi(2);
// Retries: too few = fragile, too many = amplifies failures
// Optimal around 3
let retry_optimal = 3.0;
score += ((config.retry_count as f64 - retry_optimal) / 2.0).powi(2);
// Batch size: trade-off between latency and throughput
// Optimal around 64
let batch_optimal = 64.0;
score += ((config.batch_size as f64 - batch_optimal) / 32.0).powi(2);
// Compression level: trade-off between CPU and bandwidth
// Optimal around 6
let compression_optimal = 6.0;
score += ((config.compression_level as f64 - compression_optimal) / 3.0).powi(2);
// HTTP/2 is generally better for our use case
if !config.use_http2 {
score += 0.5;
}
// Load balancing strategy affects performance
score += match config.load_balancing.as_str() {
"round_robin" => 0.0, // Best for our use case
"least_connections" => 0.1, // Good alternative
"ip_hash" => 0.2, // OK for session affinity
"random" => 0.3, // Not ideal
_ => 1.0,
};
// Add noise to simulate real-world variability
let noise = (config.cache_size_mb as f64 * 0.1).sin() * 0.05;
score + noise
}
/// The async objective function for each trial.
///
/// For async optimization, the objective function must:
/// 1. Take ownership of the Trial (not a mutable reference)
/// 2. Return a Future
/// 3. Return both the Trial and the result value as a tuple
///
/// This ownership pattern allows the trial to be used across await points.
async fn objective(mut trial: Trial) -> optimizer::Result<(Trial, f64)> {
// Sample configuration parameters
// Stepped integers: only sample multiples of the step value
let cache_size_mb = trial.suggest_int_step("cache_size_mb", 64, 1024, 64)?;
let connection_pool_size = trial.suggest_int_step("connection_pool_size", 10, 200, 10)?;
let request_timeout_ms = trial.suggest_int_step("request_timeout_ms", 1000, 10000, 500)?;
// Regular integer
let retry_count = trial.suggest_int("retry_count", 0, 5)?;
// Log-scale integer: good for parameters like batch sizes
// that might vary from 1 to 256
let batch_size = trial.suggest_int_log("batch_size", 1, 256)?;
// Regular integer for compression level
let compression_level = trial.suggest_int("compression_level", 0, 9)?;
// Boolean: internally uses categorical with [false, true]
let use_http2 = trial.suggest_bool("use_http2")?;
// Categorical: choose from a list of options
let load_balancing = trial.suggest_categorical(
"load_balancing",
&["round_robin", "least_connections", "random", "ip_hash"],
)?;
// Build configuration
let config = ServiceConfig {
cache_size_mb,
connection_pool_size,
request_timeout_ms,
retry_count,
batch_size,
compression_level,
use_http2,
load_balancing: load_balancing.to_string(),
};
// Evaluate (this is the async part)
let score = evaluate_service(&config).await;
// Return both the trial and the score
Ok((trial, score))
}
// ============================================================================
// Helper Functions
// ============================================================================
/// Formats a parameter value for display.
fn format_param(name: &str, value: &ParamValue) -> String {
match (name, value) {
(_, ParamValue::Float(v)) => format!("{v:.4}"),
(_, ParamValue::Int(v)) => format!("{v}"),
("use_http2", ParamValue::Categorical(idx)) => {
if *idx == 1 { "true" } else { "false" }.to_string()
}
("load_balancing", ParamValue::Categorical(idx)) => {
["round_robin", "least_connections", "random", "ip_hash"]
.get(*idx)
.unwrap_or(&"unknown")
.to_string()
}
(_, ParamValue::Categorical(idx)) => format!("category_{idx}"),
}
}
/// Prints the results of the optimization.
fn print_results(study: &Study<f64>, elapsed: Duration, n_trials: usize) {
println!("\n{}", "=".repeat(60));
println!("\nOptimization completed!");
println!("Total trials: {}", study.n_trials());
println!("Time elapsed: {elapsed:.2?}");
// Calculate speedup from parallelism
// Each trial takes ~50ms, so sequential would take n_trials * 50ms
let sequential_time = n_trials as f64 * 0.050;
let actual_time = elapsed.as_secs_f64();
println!(
"Effective parallelism: {:.1}x speedup",
sequential_time / actual_time
);
}
/// Prints the best configuration found.
fn print_best_config(study: &Study<f64>) -> optimizer::Result<()> {
let best = study.best_trial()?;
println!("\nBest configuration found:");
println!(" Score: {:.6}", best.value);
println!("\n Parameters:");
// Print parameters in a logical order
let param_order = [
"cache_size_mb",
"connection_pool_size",
"request_timeout_ms",
"retry_count",
"batch_size",
"compression_level",
"use_http2",
"load_balancing",
];
for name in param_order {
if let Some(value) = best.params.get(name) {
let display = format_param(name, value);
println!(" {name}: {display}");
}
}
Ok(())
}
/// Prints the top N trials.
fn print_top_trials(study: &Study<f64>, n: usize) {
println!("\nTop {n} trials:");
let mut trials = study.trials();
trials.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap());
for (i, trial) in trials.iter().take(n).enumerate() {
println!(
" {}. Trial #{}: score = {:.6}",
i + 1,
trial.id,
trial.value
);
}
}
// ============================================================================
// Main: Set up and run the async optimization
// ============================================================================
#[tokio::main]
async fn main() -> optimizer::Result<()> {
println!("=== Async API Parameter Optimization Example ===\n");
// Step 1: Create a TPE sampler
let sampler = TpeSampler::builder()
.n_startup_trials(8)
.gamma(0.2)
.seed(123)
.build()
.expect("Failed to build TPE sampler");
// Step 2: Create a study to minimize the score
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
// Step 3: Configure optimization
let n_trials = 40;
let concurrency = 4; // Run 4 trials in parallel
println!("Starting parallel optimization with {concurrency} concurrent evaluations...\n");
let start = Instant::now();
// Step 4: Run parallel async optimization
//
// optimize_parallel_with_sampler:
// - Runs up to `concurrency` trials simultaneously
// - Each trial calls the objective function
// - Uses a semaphore to limit concurrent evaluations
// - Collects results as trials complete
//
// The "_with_sampler" suffix means the TPE sampler gets access to
// trial history for informed sampling.
study
.optimize_parallel_with_sampler(n_trials, concurrency, objective)
.await?;
let elapsed = start.elapsed();
// Step 5: Print results
print_results(&study, elapsed, n_trials);
print_best_config(&study)?;
print_top_trials(&study, 5);
Ok(())
}
examples/ml_hyperparameter_tuning.rs
+269
View File
@@ -0,0 +1,269 @@
//! Machine Learning Hyperparameter Tuning Example
//!
//! This example shows how to use the optimizer library to find the best
//! hyperparameters for a machine learning model. We simulate a gradient
//! boosting model (like XGBoost or LightGBM) and search for optimal settings.
//!
//! # Key Concepts Demonstrated
//!
//! - Creating a Study with a TPE (Tree-Parzen Estimator) sampler
//! - Defining an objective function that the optimizer will minimize
//! - Using different parameter types: floats, integers, log-scale, stepped
//! - Using callbacks to monitor progress and implement early stopping
//!
//! # How It Works
//!
//! 1. Create a `Study` - this manages the optimization process
//! 2. Define an objective function that takes a `Trial` and returns a score
//! 3. Inside the objective, use `trial.suggest_*()` to sample parameters
//! 4. The optimizer runs many trials, learning which parameter regions work best
//! 5. After optimization, retrieve the best parameters found
//!
//! Run with: `cargo run --example ml_hyperparameter_tuning`
use std::ops::ControlFlow;
use optimizer::sampler::CompletedTrial;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, ParamValue, Study, Trial};
// ============================================================================
// Configuration: Hyperparameters we want to tune
// ============================================================================
/// Holds all the hyperparameters for our model.
///
/// In a real application, you would pass these to your ML framework
/// (e.g., XGBoost, LightGBM, scikit-learn).
struct ModelConfig {
learning_rate: f64,
max_depth: i64,
n_estimators: i64,
subsample: f64,
colsample_bytree: f64,
min_child_weight: i64,
reg_alpha: f64,
reg_lambda: f64,
}
// ============================================================================
// Objective Function: What we want to optimize
// ============================================================================
/// Simulates training a model and returns the validation loss.
///
/// In a real scenario, this function would:
/// 1. Create a model with the given hyperparameters
/// 2. Train it on your training data
/// 3. Evaluate it on validation data
/// 4. Return the validation metric (e.g., RMSE, log loss, accuracy)
///
/// The optimizer will try to MINIMIZE this value (we set Direction::Minimize).
#[allow(clippy::too_many_arguments)]
fn evaluate_model(config: &ModelConfig) -> f64 {
// Simulated optimal hyperparameters:
// learning_rate ~ 0.05, max_depth ~ 6, n_estimators ~ 200
// subsample ~ 0.8, colsample_bytree ~ 0.8, min_child_weight ~ 3
// reg_alpha ~ 0.1, reg_lambda ~ 1.0
let mut loss = 0.15; // Base loss
// Each term penalizes deviation from the optimal value
loss += (config.learning_rate - 0.05).powi(2) * 100.0;
loss += ((config.max_depth - 6) as f64).powi(2) * 0.01;
loss += ((config.n_estimators - 200) as f64).powi(2) * 0.00001;
loss += (config.subsample - 0.8).powi(2) * 10.0;
loss += (config.colsample_bytree - 0.8).powi(2) * 10.0;
loss += ((config.min_child_weight - 3) as f64).powi(2) * 0.05;
loss += (config.reg_alpha - 0.1).powi(2) * 5.0;
loss += (config.reg_lambda - 1.0).powi(2) * 2.0;
// Add some noise to simulate real-world variability
let noise = (config.learning_rate * 1000.0).sin() * 0.01;
loss + noise
}
/// The objective function that the optimizer calls for each trial.
///
/// This function:
/// 1. Uses `trial.suggest_*()` methods to sample hyperparameter values
/// 2. Builds a model configuration from those values
/// 3. Evaluates the model and returns the loss
///
/// The optimizer learns from the results to suggest better parameters
/// in future trials.
fn objective(trial: &mut Trial) -> optimizer::Result<f64> {
// Sample hyperparameters using different strategies:
// Log-scale: Good for parameters spanning multiple orders of magnitude
// The learning rate might be 0.001, 0.01, or 0.1 - log-scale samples evenly across these
let learning_rate = trial.suggest_float_log("learning_rate", 0.001, 0.3)?;
// Regular integer: Uniformly samples from the range [3, 12]
let max_depth = trial.suggest_int("max_depth", 3, 12)?;
// Stepped integer: Only samples multiples of 50 (50, 100, 150, ..., 500)
// Useful when you only want to test specific values
let n_estimators = trial.suggest_int_step("n_estimators", 50, 500, 50)?;
// Regular float: Uniformly samples from [0.5, 1.0]
let subsample = trial.suggest_float("subsample", 0.5, 1.0)?;
let colsample_bytree = trial.suggest_float("colsample_bytree", 0.5, 1.0)?;
// More parameters
let min_child_weight = trial.suggest_int("min_child_weight", 1, 10)?;
let reg_alpha = trial.suggest_float_log("reg_alpha", 1e-3, 10.0)?;
let reg_lambda = trial.suggest_float_log("reg_lambda", 1e-3, 10.0)?;
// Build configuration and evaluate
let config = ModelConfig {
learning_rate,
max_depth,
n_estimators,
subsample,
colsample_bytree,
min_child_weight,
reg_alpha,
reg_lambda,
};
let loss = evaluate_model(&config);
Ok(loss)
}
// ============================================================================
// Callback Function: Monitor progress and implement early stopping
// ============================================================================
/// Called after each successful trial completes.
///
/// Use callbacks to:
/// - Log progress to console or file
/// - Save checkpoints
/// - Implement early stopping when a good solution is found
/// - Track metrics over time
///
/// Return `ControlFlow::Continue(())` to keep optimizing.
/// Return `ControlFlow::Break(())` to stop early.
fn on_trial_complete(study: &Study<f64>, trial: &CompletedTrial<f64>) -> ControlFlow<()> {
// Helper to extract parameter values
let get_float = |name: &str| -> f64 {
match trial.params.get(name) {
Some(ParamValue::Float(v)) => *v,
_ => 0.0,
}
};
let get_int = |name: &str| -> i64 {
match trial.params.get(name) {
Some(ParamValue::Int(v)) => *v,
_ => 0,
}
};
// Print progress
println!(
"{:>5} {:>10.5} {:>10} {:>12} {:>10.3} {:>12.3} {:>8} {:>10.4} {:>10.4} {:>12.6}",
study.n_trials(),
get_float("learning_rate"),
get_int("max_depth"),
get_int("n_estimators"),
get_float("subsample"),
get_float("colsample_bytree"),
get_int("min_child_weight"),
get_float("reg_alpha"),
get_float("reg_lambda"),
trial.value,
);
// Early stopping: if we find an excellent solution, stop early
if trial.value < 0.16 {
println!("\nEarly stopping: found excellent solution!");
return ControlFlow::Break(());
}
ControlFlow::Continue(())
}
// ============================================================================
// Main: Set up and run the optimization
// ============================================================================
fn main() -> optimizer::Result<()> {
println!("=== ML Hyperparameter Tuning Example ===\n");
// Step 1: Create a sampler
//
// TPE (Tree-Parzen Estimator) is a Bayesian optimization algorithm.
// It learns from previous trials to suggest better parameters.
// - n_startup_trials: Number of random trials before TPE kicks in
// - gamma: What fraction of trials are considered "good" (lower = more selective)
// - seed: For reproducibility
let sampler = TpeSampler::builder()
.n_startup_trials(10)
.gamma(0.25)
.seed(42)
.build()
.expect("Failed to build TPE sampler");
// Step 2: Create a study
//
// The study manages the optimization process. We want to MINIMIZE
// the loss (lower is better). Use Direction::Maximize for metrics
// where higher is better (like accuracy).
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
// Print header
println!("Starting hyperparameter optimization...\n");
println!(
"{:>5} {:>10} {:>10} {:>12} {:>10} {:>12} {:>8} {:>10} {:>10} {:>12}",
"Trial",
"LR",
"MaxDepth",
"Estimators",
"Subsample",
"ColSample",
"MinCW",
"Alpha",
"Lambda",
"Loss"
);
println!("{}", "-".repeat(110));
// Step 3: Run optimization
//
// optimize_with_callback_sampler runs the objective function for up to
// n_trials iterations. After each trial, it calls the callback.
// The "_sampler" suffix means the TPE sampler gets access to trial
// history for informed sampling.
let n_trials = 50;
study.optimize_with_callback_sampler(n_trials, objective, on_trial_complete)?;
// Step 4: Get the best result
println!("\n{}", "=".repeat(110));
println!("\nOptimization completed!");
println!("Total trials: {}", study.n_trials());
let best = study.best_trial()?;
println!("\nBest trial:");
println!(" Loss: {:.6}", best.value);
println!(" Parameters:");
for (name, value) in &best.params {
match value {
ParamValue::Float(v) => println!(" {name}: {v:.6}"),
ParamValue::Int(v) => println!(" {name}: {v}"),
ParamValue::Categorical(v) => println!(" {name}: category {v}"),
}
}
// Step 5: Use the best parameters (in a real app)
//
// Now you would take best.params and use them to train your final model
// on the full dataset.
Ok(())
}