diff --git a/Cargo.lock b/Cargo.lock index 426fe8795..b623daa2d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -682,6 +682,7 @@ dependencies = [ "rayon", "rstest", "serde", + "serde_json", "tracing", "tracing-subscriber", "vfs", diff --git a/diskann-benchmark/README.md b/diskann-benchmark/README.md index 6e072b757..853638534 100644 --- a/diskann-benchmark/README.md +++ b/diskann-benchmark/README.md @@ -491,3 +491,38 @@ error reporting in the event of a dispatch fail much easier for the user to unde Refer to implementations within the benchmarking framework for what some of this may look like. +## Autotuner Tool + +The `autotuner` tool builds on top of the benchmark framework to automatically sweep over parameter combinations and identify the best configuration based on optimization criteria (QPS, latency, or recall). + +The autotuner uses a **path-based configuration system** that doesn't hardcode JSON structure, making it robust to changes in the benchmark framework. You specify which parameters to sweep by providing JSON paths. + +See [diskann-tools/AUTOTUNER.md](../diskann-tools/AUTOTUNER.md) for detailed documentation. + +### Quick Start + +```sh +# Generate an example sweep configuration +cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json + +# Run parameter sweep to find optimal configuration +cargo run --release --package diskann-tools --bin autotuner -- sweep \ + --base-config base_config.json \ + --sweep-config sweep_config.json \ + --output-dir ./autotuner_results \ + --criterion qps \ + --target-recall 0.95 +``` + +The sweep configuration uses JSON paths to specify parameters: +```json +{ + "parameters": [ + {"path": "jobs.0.content.source.max_degree", "values": [16, 32, 64]}, + {"path": "jobs.0.content.source.l_build", "values": [50, 75, 100]} + ] +} +``` + +This design makes the autotuner adaptable to any benchmark configuration format without requiring code changes when the benchmark framework evolves. + diff --git a/diskann-tools/AUTOTUNER.md b/diskann-tools/AUTOTUNER.md new file mode 100644 index 000000000..a711f2a05 --- /dev/null +++ b/diskann-tools/AUTOTUNER.md @@ -0,0 +1,293 @@ +# DiskANN Autotuner + +The autotuner is a tool that builds on top of the DiskANN benchmark framework to automatically sweep over parameter combinations and identify the best configuration based on specified optimization criteria. + +## Overview + +The autotuner uses a **path-based configuration system** that doesn't hardcode JSON structure, making it robust to changes in the benchmark framework. You specify which parameters to sweep by providing JSON paths, making the tool adaptable to any benchmark configuration format. + +Key features: +- **Framework-agnostic**: Works with any benchmark JSON structure +- **Flexible parameter sweeping**: Specify any JSON path to override +- **Multiple optimization criteria**: QPS, latency, or recall +- **Automatic result analysis**: Identifies best configuration based on your criteria + +## Installation + +Build the autotuner tool: + +```bash +cargo build --release --package diskann-tools --bin autotuner +``` + +## Quick Start + +### 1. Generate an Example Sweep Configuration + +```bash +cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json +``` + +For specific benchmark types: +```bash +# For product-quantized indexes +cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json --benchmark-type pq + +# For disk indexes +cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json --benchmark-type disk +``` + +### 2. Understanding the Sweep Configuration + +The generated sweep configuration uses JSON paths to specify which parameters to sweep: + +```json +{ + "parameters": [ + { + "path": "jobs.0.content.source.max_degree", + "values": [16, 32, 64] + }, + { + "path": "jobs.0.content.source.l_build", + "values": [50, 75, 100] + }, + { + "path": "jobs.0.content.search_phase.runs.0.search_l", + "values": [ + [10, 20, 30, 40, 50], + [20, 40, 60, 80, 100], + [30, 60, 90, 120, 150] + ] + } + ] +} +``` + +**Path syntax:** +- Use dot notation: `jobs.0.content.source.max_degree` +- Array indices are numbers: `jobs.0` refers to first job +- The autotuner will generate all combinations of parameter values + +### 3. Prepare a Base Configuration + +Use an existing benchmark configuration file as your template (from `diskann-benchmark/example/`), or create your own. + +### 4. Run the Parameter Sweep + +```bash +cargo run --release --package diskann-tools --bin autotuner -- sweep \ + --base-config diskann-benchmark/example/async.json \ + --sweep-config sweep_config.json \ + --output-dir ./autotuner_results \ + --criterion qps \ + --target-recall 0.95 +``` + +## Configuration Guide + +### Finding the Right JSON Paths + +To find the correct paths for your benchmark configuration: + +1. **Examine your base configuration** - Look at the JSON structure +2. **Identify the parameters** - Find where parameters like `max_degree`, `l_build`, `search_l` are located +3. **Write the path** - Use dot notation with array indices + +**Example paths for common parameters:** + +| Benchmark Type | Parameter | Path | +|----------------|-----------|------| +| async-index-build | max_degree | `jobs.0.content.source.max_degree` | +| async-index-build | l_build | `jobs.0.content.source.l_build` | +| async-index-build | search_l | `jobs.0.content.search_phase.runs.0.search_l` | +| async-index-build-pq | max_degree | `jobs.0.content.index_operation.source.max_degree` | +| async-index-build-pq | l_build | `jobs.0.content.index_operation.source.l_build` | +| async-index-build-pq | num_pq_chunks | `jobs.0.content.num_pq_chunks` | +| disk-index | max_degree | `jobs.0.content.source.max_degree` | +| disk-index | l_build | `jobs.0.content.source.l_build` | +| disk-index | search_list | `jobs.0.content.search_phase.search_list` | + +### Creating Custom Sweep Configurations + +You can sweep over any parameter, not just the standard ones: + +```json +{ + "parameters": [ + { + "path": "jobs.0.content.source.alpha", + "values": [1.0, 1.2, 1.5] + }, + { + "path": "jobs.0.content.source.backedge_ratio", + "values": [0.8, 1.0, 1.2] + }, + { + "path": "jobs.0.content.search_phase.num_threads", + "values": [[1], [2], [4], [8]] + } + ] +} +``` + +## Command Line Options + +### `autotuner sweep` + +Run a parameter sweep to find the optimal configuration. + +```bash +autotuner sweep [OPTIONS] --base-config --sweep-config --output-dir +``` + +**Options:** +- `-b, --base-config ` - Base benchmark configuration JSON file (template) +- `-s, --sweep-config ` - Parameter sweep specification JSON file +- `-o, --output-dir ` - Output directory for results +- `-c, --criterion ` - Optimization criterion: `qps`, `latency`, or `recall` (default: `qps`) +- `-t, --target-recall ` - Target recall threshold for qps/latency optimization (default: 0.95) +- `--benchmark-cmd ` - Path to diskann-benchmark binary (default: `cargo`) +- `--benchmark-args ` - Additional arguments for benchmark command + +### `autotuner example` + +Generate an example sweep configuration. + +```bash +autotuner example --output [--benchmark-type ] +``` + +**Options:** +- `-o, --output ` - Output file for example configuration +- `--benchmark-type ` - Generate example for specific type: `pq`, `disk`, or default + +## Output + +The autotuner generates the following files in the output directory: + +- `config_XXXX.json` - Generated configuration for each parameter combination +- `results_XXXX.json` - Benchmark results for each configuration +- `sweep_summary.json` - Summary of all results with the best configuration highlighted + +### Example Summary + +The `sweep_summary.json` file contains all sweep results. Here's an example (truncated for brevity): + +```json +{ + "criterion": "qps", + "target_recall": 0.95, + "total_configs": 9, + "successful_configs": 9, + "best_config": { + "config_id": "0005", + "parameters": { + "jobs.0.content.source.max_degree": 32, + "jobs.0.content.source.l_build": 75, + "jobs.0.content.search_phase.runs.0.search_l": [10, 20, 30, 40, 50] + }, + "metrics": { + "qps": [12345.6, 11234.5, ...], + "recall": [0.98, 0.96, ...] + } + }, + "all_results": [ + // Array of all sweep results, same format as best_config + ] +} +``` + +**Note**: The `config_id` is a sequential number (e.g., "0005") representing the configuration index in the sweep. To see which parameter values correspond to which ID, check the `parameters` field. + +## Examples + +### Example 1: Optimize for Maximum QPS + +```bash +cargo run --release --package diskann-tools --bin autotuner -- sweep \ + --base-config diskann-benchmark/example/async.json \ + --sweep-config sweep_config.json \ + --output-dir ./results_qps \ + --criterion qps \ + --target-recall 0.95 +``` + +### Example 2: Optimize for Minimum Latency + +```bash +cargo run --release --package diskann-tools --bin autotuner -- sweep \ + --base-config diskann-benchmark/example/async.json \ + --sweep-config sweep_config.json \ + --output-dir ./results_latency \ + --criterion latency \ + --target-recall 0.99 +``` + +### Example 3: Optimize for Maximum Recall + +```bash +cargo run --release --package diskann-tools --bin autotuner -- sweep \ + --base-config diskann-benchmark/example/async.json \ + --sweep-config sweep_config.json \ + --output-dir ./results_recall \ + --criterion recall +``` + +### Example 4: Sweep Over Product-Quantized Index Parameters + +```bash +# 1. Generate PQ-specific example +cargo run --release --package diskann-tools --bin autotuner -- example \ + --output pq_sweep.json \ + --benchmark-type pq + +# 2. Run sweep +cargo run --release --package diskann-tools --bin autotuner -- sweep \ + --base-config diskann-benchmark/example/product.json \ + --sweep-config pq_sweep.json \ + --output-dir ./results_pq \ + --criterion qps \ + --target-recall 0.95 +``` + +## Adapting to Benchmark Framework Changes + +The path-based design makes the autotuner robust to changes in the benchmark framework: + +1. **If parameter locations change** - Simply update the paths in your sweep configuration +2. **If new parameters are added** - Add new path-value pairs to your sweep configuration +3. **If benchmark output format changes** - Only the result parsing logic needs updating (not the sweep logic) + +The autotuner itself doesn't hardcode any assumptions about the benchmark JSON structure, so it remains compatible as long as: +- The benchmark accepts JSON configuration files +- The benchmark produces JSON output with QPS/recall metrics + +## Tips + +1. **Start with a coarse sweep**: Use a small number of parameter values to explore the space quickly +2. **Refine iteratively**: Once you identify promising regions, create a finer-grained sweep +3. **Use appropriate benchmarks**: Make sure your base configuration is appropriate for your workload +4. **Check paths carefully**: Invalid paths will cause the sweep to fail - verify with your base config +5. **Monitor resource usage**: Each configuration runs a full benchmark, which can be time-consuming + +## Troubleshooting + +**Problem**: "Path 'X' not found in JSON" +**Solution**: Check that the path exists in your base configuration file. Use dot notation with array indices (e.g., `jobs.0.content`) + +**Problem**: Benchmark failures during sweep +**Solution**: Test your base configuration manually first to ensure it works before running the sweep + +**Problem**: No results meet criteria +**Solution**: Lower your target recall threshold or adjust your parameter ranges + +## Design Philosophy + +The autotuner is designed to be: +- **Maintainable**: No hardcoded assumptions about JSON structure +- **Flexible**: Works with any benchmark configuration format +- **Robust**: Survives changes to the benchmark framework +- **Extensible**: Easy to add new optimization criteria or parameter types + +Users control what gets swept by specifying paths, not by relying on tool-specific parameter names. diff --git a/diskann-tools/Cargo.toml b/diskann-tools/Cargo.toml index ae987dca9..74e86d309 100644 --- a/diskann-tools/Cargo.toml +++ b/diskann-tools/Cargo.toml @@ -24,6 +24,7 @@ ordered-float = "4.2.0" rand_distr.workspace = true rand.workspace = true serde = { workspace = true, features = ["derive"] } +serde_json.workspace = true bincode.workspace = true opentelemetry.workspace = true diskann-quantization = { workspace = true } diff --git a/diskann-tools/src/bin/autotuner.rs b/diskann-tools/src/bin/autotuner.rs new file mode 100644 index 000000000..944bc5d56 --- /dev/null +++ b/diskann-tools/src/bin/autotuner.rs @@ -0,0 +1,717 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Autotuner for DiskANN benchmarks. +//! +//! This tool builds on top of the benchmark framework to sweep over a subset of parameters +//! and identify the best configuration based on specified optimization criteria. +//! +//! The autotuner uses a path-based configuration system that doesn't hardcode JSON structure, +//! making it robust to changes in the benchmark framework. + +use anyhow::{Context, Result}; +use clap::{Parser, Subcommand}; +use serde::{Deserialize, Serialize}; +use serde_json::{self, Value}; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +#[derive(Parser, Debug)] +#[command(name = "autotuner")] +#[command(about = "Autotuner for DiskANN benchmarks - sweeps parameters to find optimal configuration")] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand, Debug)] +enum Commands { + /// Run parameter sweep to find optimal configuration + Sweep { + /// Base configuration file (JSON) to use as template + #[arg(short, long)] + base_config: PathBuf, + + /// Parameter sweep specification file (JSON) + #[arg(short, long)] + sweep_config: PathBuf, + + /// Output directory for sweep results + #[arg(short, long)] + output_dir: PathBuf, + + /// Optimization criterion: "qps", "latency", or "recall" + #[arg(short = 'c', long, default_value = "qps")] + criterion: String, + + /// Target recall threshold (for qps/latency optimization) + #[arg(short, long)] + target_recall: Option, + + /// Path to diskann-benchmark binary + #[arg(long, default_value = "cargo")] + benchmark_cmd: String, + + /// Additional args for benchmark command + #[arg(long)] + benchmark_args: Option, + }, + + /// Generate example sweep configuration + Example { + /// Output file for example sweep configuration + #[arg(short, long)] + output: PathBuf, + + /// Generate example for specific benchmark type + #[arg(long)] + benchmark_type: Option, + }, +} + +/// Path-based parameter specification for sweeping +#[derive(Debug, Serialize, Deserialize, Clone)] +struct ParameterSweep { + /// JSON path to the parameter (e.g., "jobs.0.content.source.max_degree") + path: String, + /// Values to sweep over + values: Vec, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct SweepConfig { + /// List of parameters to sweep with their paths and values + parameters: Vec, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct SweepResult { + config_id: String, + parameters: HashMap, + metrics: BenchmarkMetrics, + output_file: PathBuf, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct BenchmarkMetrics { + qps: Vec, + recall: Vec, + latency_p50: Option>, + latency_p90: Option>, + latency_p99: Option>, +} + +fn main() -> Result<()> { + let cli = Cli::parse(); + + match cli.command { + Commands::Sweep { + base_config, + sweep_config, + output_dir, + criterion, + target_recall, + benchmark_cmd, + benchmark_args, + } => run_sweep( + &base_config, + &sweep_config, + &output_dir, + &criterion, + target_recall, + &benchmark_cmd, + benchmark_args.as_deref(), + ), + Commands::Example { output, benchmark_type } => generate_example(&output, benchmark_type.as_deref()), + } +} + +fn generate_example(output: &Path, benchmark_type: Option<&str>) -> Result<()> { + let example = match benchmark_type { + Some("pq") | Some("product-quantization") => SweepConfig { + parameters: vec![ + ParameterSweep { + path: "jobs.0.content.index_operation.source.max_degree".to_string(), + values: vec![16, 32, 64].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ParameterSweep { + path: "jobs.0.content.index_operation.source.l_build".to_string(), + values: vec![50, 75, 100].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ParameterSweep { + path: "jobs.0.content.index_operation.search_phase.runs.0.search_l".to_string(), + values: vec![ + serde_json::json!([10, 20, 30, 40, 50]), + serde_json::json!([20, 40, 60, 80, 100]), + serde_json::json!([30, 60, 90, 120, 150]), + ], + }, + ParameterSweep { + path: "jobs.0.content.num_pq_chunks".to_string(), + values: vec![8, 16, 32].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ], + }, + Some("disk") | Some("disk-index") => SweepConfig { + parameters: vec![ + ParameterSweep { + path: "jobs.0.content.source.max_degree".to_string(), + values: vec![16, 32, 64].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ParameterSweep { + path: "jobs.0.content.source.l_build".to_string(), + values: vec![50, 75, 100].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ParameterSweep { + path: "jobs.0.content.search_phase.search_list".to_string(), + values: vec![ + serde_json::json!([10, 20, 40]), + serde_json::json!([20, 40, 80]), + serde_json::json!([30, 60, 120]), + ], + }, + ], + }, + _ => SweepConfig { + parameters: vec![ + ParameterSweep { + path: "jobs.0.content.source.max_degree".to_string(), + values: vec![16, 32, 64].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ParameterSweep { + path: "jobs.0.content.source.l_build".to_string(), + values: vec![50, 75, 100].into_iter().map(|v| serde_json::json!(v)).collect(), + }, + ParameterSweep { + path: "jobs.0.content.search_phase.runs.0.search_l".to_string(), + values: vec![ + serde_json::json!([10, 20, 30, 40, 50]), + serde_json::json!([20, 40, 60, 80, 100]), + serde_json::json!([30, 60, 90, 120, 150]), + ], + }, + ], + }, + }; + + let file = std::fs::File::create(output) + .with_context(|| format!("Failed to create output file: {}", output.display()))?; + serde_json::to_writer_pretty(file, &example)?; + + println!("Example sweep configuration written to: {}", output.display()); + println!("\nThis configuration will sweep over {} parameters:", example.parameters.len()); + for param in &example.parameters { + println!(" - {}: {} values", param.path, param.values.len()); + } + println!("\nTo use this configuration:"); + println!(" 1. Adjust the paths if your benchmark JSON structure is different"); + println!(" 2. Modify the values to sweep over"); + println!(" 3. Run: autotuner sweep --base-config --sweep-config --output-dir "); + + Ok(()) +} + +fn run_sweep( + base_config: &Path, + sweep_config: &Path, + output_dir: &Path, + criterion: &str, + target_recall: Option, + benchmark_cmd: &str, + benchmark_args: Option<&str>, +) -> Result<()> { + println!("Starting parameter sweep..."); + println!("Base config: {}", base_config.display()); + println!("Sweep config: {}", sweep_config.display()); + println!("Output dir: {}", output_dir.display()); + println!("Optimization criterion: {}", criterion); + if let Some(recall) = target_recall { + println!("Target recall: {}", recall); + } + + // Create output directory + std::fs::create_dir_all(output_dir) + .with_context(|| format!("Failed to create output directory: {}", output_dir.display()))?; + + // Load base configuration + let base_config_data: Value = load_json(base_config)?; + + // Load sweep configuration + let sweep_spec: SweepConfig = load_json(sweep_config)?; + + // Generate configurations + let configs = generate_configurations(&base_config_data, &sweep_spec)?; + println!("Generated {} configurations to evaluate", configs.len()); + + // Run benchmarks for each configuration + let mut results = Vec::new(); + for (idx, (config_id, config_data, params)) in configs.iter().enumerate() { + println!( + "\n[{}/{}] Running configuration: {}", + idx + 1, + configs.len(), + config_id + ); + + // Save configuration to file + let config_file = output_dir.join(format!("config_{}.json", config_id)); + save_json(&config_file, config_data)?; + + // Run benchmark + let output_file = output_dir.join(format!("results_{}.json", config_id)); + match run_benchmark(benchmark_cmd, benchmark_args, &config_file, &output_file) { + Ok(_) => { + // Parse results + if let Ok(metrics) = parse_benchmark_results(&output_file) { + results.push(SweepResult { + config_id: config_id.clone(), + parameters: params.clone(), + metrics, + output_file: output_file.clone(), + }); + println!(" ✓ Completed"); + } else { + println!(" ✗ Failed to parse results"); + } + } + Err(e) => { + println!(" ✗ Benchmark failed: {}", e); + } + } + } + + // Analyze results and find best configuration + if results.is_empty() { + anyhow::bail!("No successful benchmark runs to analyze"); + } + + let best_config = find_best_configuration(&results, criterion, target_recall)?; + + // Save summary + let summary_file = output_dir.join("sweep_summary.json"); + let summary = serde_json::json!({ + "criterion": criterion, + "target_recall": target_recall, + "total_configs": configs.len(), + "successful_configs": results.len(), + "best_config": best_config, + "all_results": results, + }); + save_json(&summary_file, &summary)?; + + println!("\n========================================"); + println!("Sweep completed!"); + println!("========================================"); + println!("Best configuration: {}", best_config.config_id); + println!("Parameters:"); + for (key, value) in &best_config.parameters { + println!(" {}: {}", key, value); + } + println!("Metrics:"); + println!(" QPS: {:?}", best_config.metrics.qps); + println!(" Recall: {:?}", best_config.metrics.recall); + println!("\nSummary saved to: {}", summary_file.display()); + + Ok(()) +} + +fn generate_configurations( + base_config: &Value, + sweep_spec: &SweepConfig, +) -> Result)>> { + let mut configs = Vec::new(); + + // Generate all combinations of parameter values + let combinations = generate_parameter_combinations(&sweep_spec.parameters); + + for (idx, param_values) in combinations.iter().enumerate() { + let mut config = base_config.clone(); + let mut params = HashMap::new(); + + // Apply each parameter override + for (param_spec, value) in sweep_spec.parameters.iter().zip(param_values.iter()) { + set_json_path(&mut config, ¶m_spec.path, value.clone())?; + params.insert(param_spec.path.clone(), value.clone()); + } + + // Generate config ID from parameter values + let config_id = format!("{:04}", idx); + configs.push((config_id, config, params)); + } + + Ok(configs) +} + +/// Generate all combinations of parameter values using Cartesian product +fn generate_parameter_combinations(parameters: &[ParameterSweep]) -> Vec> { + if parameters.is_empty() { + return vec![vec![]]; + } + + let mut result = vec![vec![]]; + + for param in parameters { + let mut new_result = Vec::new(); + for combination in &result { + for value in ¶m.values { + let mut new_combination = combination.clone(); + new_combination.push(value.clone()); + new_result.push(new_combination); + } + } + result = new_result; + } + + result +} + +/// Set a value in JSON using a dot-separated path +fn set_json_path(json: &mut Value, path: &str, value: Value) -> Result<()> { + let parts: Vec<&str> = path.split('.').collect(); + + if parts.is_empty() { + anyhow::bail!("Empty path"); + } + + let mut current = json; + + // Navigate to the parent of the target field + for (i, &part) in parts.iter().enumerate() { + let is_last = i == parts.len() - 1; + + // Check if this is an array index + if let Ok(index) = part.parse::() { + if let Some(array) = current.as_array_mut() { + if index >= array.len() { + anyhow::bail!("Array index {} out of bounds at path {}", index, path); + } + if is_last { + array[index] = value; + return Ok(()); + } else { + current = &mut array[index]; + } + } else { + anyhow::bail!("Expected array at path element '{}' in path {}", part, path); + } + } else { + // Object key + if let Some(object) = current.as_object_mut() { + if is_last { + object.insert(part.to_string(), value); + return Ok(()); + } else { + current = object.get_mut(part) + .ok_or_else(|| anyhow::anyhow!("Path '{}' not found in JSON at '{}'", path, part))?; + } + } else { + anyhow::bail!("Expected object at path element '{}' in path {}", part, path); + } + } + } + + Ok(()) +} + +fn run_benchmark( + benchmark_cmd: &str, + benchmark_args: Option<&str>, + config_file: &Path, + output_file: &Path, +) -> Result<()> { + let mut cmd = if benchmark_cmd == "cargo" { + let mut c = std::process::Command::new("cargo"); + c.arg("run"); + c.arg("--release"); + c.arg("--package"); + c.arg("diskann-benchmark"); + c.arg("--"); + c + } else { + std::process::Command::new(benchmark_cmd) + }; + + if let Some(args) = benchmark_args { + for arg in args.split_whitespace() { + cmd.arg(arg); + } + } + + cmd.arg("run"); + cmd.arg("--input-file"); + cmd.arg(config_file); + cmd.arg("--output-file"); + cmd.arg(output_file); + + let output = cmd.output() + .with_context(|| format!("Failed to run benchmark command: {}", benchmark_cmd))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!("Benchmark failed: {}", stderr); + } + + Ok(()) +} + +fn parse_benchmark_results(output_file: &Path) -> Result { + let data: Vec = load_json(output_file)?; + + // Extract metrics from first job result + let result = data.first() + .ok_or_else(|| anyhow::anyhow!("No results in output file"))?; + + let search_results = result + .get("results") + .and_then(|r: &Value| r.get("search")) + .and_then(|s: &Value| s.get("Topk")) + .and_then(|t: &Value| t.as_array()) + .ok_or_else(|| anyhow::anyhow!("Invalid result format"))?; + + let mut qps_values = Vec::new(); + let mut recall_values = Vec::new(); + + for search_result in search_results { + // Extract QPS (can be array or single value) + if let Some(qps) = search_result.get("qps") { + if let Some(qps_arr) = qps.as_array() { + for val in qps_arr { + if let Some(q) = val.as_f64() { + qps_values.push(q); + } + } + } else if let Some(q) = qps.as_f64() { + qps_values.push(q); + } + } + + // Extract recall + if let Some(recall) = search_result.get("recall").and_then(|r: &Value| r.get("average")) { + if let Some(r) = recall.as_f64() { + recall_values.push(r); + } + } + } + + Ok(BenchmarkMetrics { + qps: qps_values, + recall: recall_values, + latency_p50: None, + latency_p90: None, + latency_p99: None, + }) +} + +fn find_best_configuration( + results: &[SweepResult], + criterion: &str, + target_recall: Option, +) -> Result { + if results.is_empty() { + anyhow::bail!("No results to analyze"); + } + + let target_recall = target_recall.unwrap_or(0.95); + + let best = match criterion { + "qps" => { + // Find configuration with max QPS at target recall + results + .iter() + .filter_map(|r| { + // Find max QPS where recall >= target + let max_qps = r.metrics.qps.iter() + .zip(&r.metrics.recall) + .filter(|(_, &recall)| recall >= target_recall) + .map(|(&qps, _)| qps) + .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + max_qps.map(|qps| (r, qps)) + }) + .max_by(|(_, qps_a), (_, qps_b)| qps_a.partial_cmp(qps_b).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(r, _)| r) + } + "recall" => { + // Find configuration with max recall + results + .iter() + .max_by(|a, b| { + let max_recall_a = a.metrics.recall.iter().cloned().fold(0.0, f64::max); + let max_recall_b = b.metrics.recall.iter().cloned().fold(0.0, f64::max); + max_recall_a.partial_cmp(&max_recall_b).unwrap_or(std::cmp::Ordering::Equal) + }) + } + "latency" => { + // Find configuration with min latency at target recall + // For now, use inverse of QPS as proxy for latency + results + .iter() + .filter_map(|r| { + let min_latency = r.metrics.qps.iter() + .zip(&r.metrics.recall) + .filter(|(_, &recall)| recall >= target_recall) + .map(|(&qps, _)| 1.0 / qps) + .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + min_latency.map(|latency| (r, latency)) + }) + .min_by(|(_, lat_a), (_, lat_b)| lat_a.partial_cmp(lat_b).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(r, _)| r) + } + _ => anyhow::bail!("Unknown criterion: {}", criterion), + }; + + best.cloned() + .ok_or_else(|| anyhow::anyhow!("No configuration meets the specified criteria")) +} + +fn load_json(path: &Path) -> Result { + let file = std::fs::File::open(path) + .with_context(|| format!("Failed to open file: {}", path.display()))?; + let reader = std::io::BufReader::new(file); + serde_json::from_reader(reader) + .with_context(|| format!("Failed to parse JSON from: {}", path.display())) +} + +fn save_json(path: &Path, data: &T) -> Result<()> { + let file = std::fs::File::create(path) + .with_context(|| format!("Failed to create file: {}", path.display()))?; + serde_json::to_writer_pretty(file, data) + .with_context(|| format!("Failed to write JSON to: {}", path.display())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_set_json_path_simple() { + let mut json = serde_json::json!({ + "a": { + "b": 42 + } + }); + + set_json_path(&mut json, "a.b", serde_json::json!(100)).unwrap(); + assert_eq!(json["a"]["b"], 100); + } + + #[test] + fn test_set_json_path_array() { + let mut json = serde_json::json!({ + "jobs": [ + {"content": {"source": {"max_degree": 32}}} + ] + }); + + set_json_path(&mut json, "jobs.0.content.source.max_degree", serde_json::json!(64)).unwrap(); + assert_eq!(json["jobs"][0]["content"]["source"]["max_degree"], 64); + } + + #[test] + fn test_generate_parameter_combinations() { + let parameters = vec![ + ParameterSweep { + path: "a".to_string(), + values: vec![serde_json::json!(1), serde_json::json!(2)], + }, + ParameterSweep { + path: "b".to_string(), + values: vec![serde_json::json!(10), serde_json::json!(20)], + }, + ]; + + let combinations = generate_parameter_combinations(¶meters); + assert_eq!(combinations.len(), 4); // 2 * 2 = 4 + } + + #[test] + fn test_benchmark_metrics() { + let metrics = BenchmarkMetrics { + qps: vec![1000.0, 2000.0, 1500.0], + recall: vec![0.95, 0.97, 0.99], + latency_p50: None, + latency_p90: None, + latency_p99: None, + }; + + // Test serialization + let json = serde_json::to_string(&metrics).unwrap(); + let deserialized: BenchmarkMetrics = serde_json::from_str(&json).unwrap(); + + assert_eq!(metrics.qps, deserialized.qps); + assert_eq!(metrics.recall, deserialized.recall); + } + + #[test] + fn test_find_best_configuration_qps() { + let results = vec![ + SweepResult { + config_id: "config1".to_string(), + parameters: HashMap::new(), + metrics: BenchmarkMetrics { + qps: vec![1000.0, 1500.0], + recall: vec![0.94, 0.96], + latency_p50: None, + latency_p90: None, + latency_p99: None, + }, + output_file: PathBuf::from("output1.json"), + }, + SweepResult { + config_id: "config2".to_string(), + parameters: HashMap::new(), + metrics: BenchmarkMetrics { + qps: vec![2000.0, 2500.0], + recall: vec![0.96, 0.98], + latency_p50: None, + latency_p90: None, + latency_p99: None, + }, + output_file: PathBuf::from("output2.json"), + }, + ]; + + let best = find_best_configuration(&results, "qps", Some(0.95)).unwrap(); + + // config2 has higher QPS at recall >= 0.95 + assert_eq!(best.config_id, "config2"); + } + + #[test] + fn test_find_best_configuration_recall() { + let results = vec![ + SweepResult { + config_id: "config1".to_string(), + parameters: HashMap::new(), + metrics: BenchmarkMetrics { + qps: vec![1000.0], + recall: vec![0.95], + latency_p50: None, + latency_p90: None, + latency_p99: None, + }, + output_file: PathBuf::from("output1.json"), + }, + SweepResult { + config_id: "config2".to_string(), + parameters: HashMap::new(), + metrics: BenchmarkMetrics { + qps: vec![800.0], + recall: vec![0.99], + latency_p50: None, + latency_p90: None, + latency_p99: None, + }, + output_file: PathBuf::from("output2.json"), + }, + ]; + + let best = find_best_configuration(&results, "recall", None).unwrap(); + + // config2 has higher recall + assert_eq!(best.config_id, "config2"); + } +}