From e4ffb4ce9fd3c72ee2f4bb200a35ca762f465a9f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 7 Feb 2026 01:10:26 +0000 Subject: [PATCH 1/6] Initial plan From e8795ba55a0f8db6319ef663bfa9d51451bd46d2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 7 Feb 2026 01:21:31 +0000 Subject: [PATCH 2/6] Add autotuner tool for parameter optimization Co-authored-by: harsha-simhadri <5590673+harsha-simhadri@users.noreply.github.com> --- Cargo.lock | 1 + diskann-tools/AUTOTUNER.md | 192 ++++++++++ diskann-tools/Cargo.toml | 1 + diskann-tools/src/bin/autotuner.rs | 554 +++++++++++++++++++++++++++++ 4 files changed, 748 insertions(+) create mode 100644 diskann-tools/AUTOTUNER.md create mode 100644 diskann-tools/src/bin/autotuner.rs 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-tools/AUTOTUNER.md b/diskann-tools/AUTOTUNER.md new file mode 100644 index 000000000..b2b5a81f0 --- /dev/null +++ b/diskann-tools/AUTOTUNER.md @@ -0,0 +1,192 @@ +# 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 helps optimize DiskANN indexes by automatically: +- Sweeping over key parameters: R (max_degree), l_build, l_search (search_l), and quantization bytes (num_pq_chunks) +- Running benchmarks for each parameter combination +- Analyzing results to find the best configuration based on QPS, latency, or recall + +## Installation + +Build the autotuner tool: + +```bash +cargo build --release --package diskann-tools --bin autotuner +``` + +## Usage + +### 1. Generate an Example Sweep Configuration + +First, generate an example sweep configuration file: + +```bash +cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json +``` + +This creates a JSON file with example parameter ranges: + +```json +{ + "max_degree": [16, 32, 64], + "l_build": [50, 75, 100], + "search_l": [10, 20, 30, 40, 50, 75, 100], + "num_pq_chunks": [8, 16, 32] +} +``` + +### 2. Prepare a Base Configuration + +Create a base benchmark configuration file (or use an existing one from `diskann-benchmark/example/`). The autotuner will use this as a template and modify the parameters specified in the sweep configuration. + +Example base configuration (`base_config.json`): + +```json +{ + "search_directories": ["test_data/disk_index_search"], + "jobs": [{ + "type": "async-index-build", + "content": { + "source": { + "index-source": "Build", + "data_type": "float32", + "data": "data.fbin", + "distance": "squared_l2", + "max_degree": 32, + "l_build": 50, + "alpha": 1.2, + "backedge_ratio": 1.0, + "num_threads": 1, + "start_point_strategy": "medoid" + }, + "search_phase": { + "search-type": "topk", + "queries": "queries.fbin", + "groundtruth": "groundtruth.bin", + "reps": 5, + "num_threads": [1], + "runs": [{ + "search_n": 10, + "search_l": [20, 30, 40], + "recall_k": 10 + }] + } + } + }] +} +``` + +### 3. Run the Parameter Sweep + +Run the autotuner to sweep over parameters and find the best configuration: + +```bash +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 +``` + +#### Options: + +- `--base-config`: Path to the base benchmark configuration (template) +- `--sweep-config`: Path to the sweep configuration (parameter ranges) +- `--output-dir`: Directory where results will be saved +- `--criterion`: Optimization criterion: + - `qps`: Maximize queries per second (default) + - `latency`: Minimize latency + - `recall`: Maximize recall +- `--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 the benchmark command + +## Output + +The autotuner generates the following files in the output directory: + +- `config_.json`: Generated configuration for each parameter combination +- `results_.json`: Benchmark results for each configuration +- `sweep_summary.json`: Summary of all results with the best configuration highlighted + +### Example Summary Output: + +```json +{ + "criterion": "qps", + "target_recall": 0.95, + "total_configs": 9, + "successful_configs": 9, + "best_config": { + "config_id": "0005_R32_L75", + "parameters": { + "max_degree": 32, + "l_build": 75, + "search_l": [10, 20, 30, 40, 50, 75, 100] + }, + "metrics": { + "qps": [12345.6, 11234.5, ...], + "recall": [0.98, 0.96, ...] + } + }, + "all_results": [...] +} +``` + +## Examples + +### Optimize for Maximum QPS at 95% 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_qps \ + --criterion qps \ + --target-recall 0.95 +``` + +### Optimize for Minimum Latency at 99% 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_latency \ + --criterion latency \ + --target-recall 0.99 +``` + +### 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 +``` + +## Parameter Descriptions + +- **max_degree (R)**: The maximum degree of the graph. Higher values increase index size but can improve recall. +- **l_build**: The search queue length during index construction. Higher values improve index quality but increase build time. +- **search_l**: The search queue length during queries. Higher values improve recall but reduce throughput. +- **num_pq_chunks**: Number of product quantization chunks (for quantized indexes). Affects compression ratio and search accuracy. + +## Tips + +1. **Start with a coarse sweep**: Begin with a small set of parameter values to get a rough idea of the performance landscape. +2. **Refine the sweep**: Once you identify promising regions, create a more fine-grained sweep around those values. +3. **Consider the trade-offs**: Higher QPS often comes at the cost of lower recall, and vice versa. +4. **Test data size**: Use representative data sizes for your production workload. + +## Notes + +- The autotuner generates one configuration per combination of build parameters (max_degree, l_build, num_pq_chunks). +- All search_l values are tested for each build configuration, allowing the tool to find the best search_l given the build parameters. +- For quantized indexes (PQ, SQ, etc.), make sure to use the appropriate base configuration template. 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..cf281e41d --- /dev/null +++ b/diskann-tools/src/bin/autotuner.rs @@ -0,0 +1,554 @@ +/* + * 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 +//! (R/max_degree, l_build, l_search/search_l, quantization bytes where applicable) +//! and identify the best configuration based on specified optimization criteria. + +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, + }, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct SweepConfig { + /// Max degree (R) values to sweep + #[serde(skip_serializing_if = "Option::is_none")] + max_degree: Option>, + + /// l_build values to sweep + #[serde(skip_serializing_if = "Option::is_none")] + l_build: Option>, + + /// search_l values to sweep + #[serde(skip_serializing_if = "Option::is_none")] + search_l: Option>, + + /// num_pq_chunks values to sweep (for quantized indexes) + #[serde(skip_serializing_if = "Option::is_none")] + num_pq_chunks: Option>, +} + +#[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 } => generate_example(&output), + } +} + +fn generate_example(output: &Path) -> Result<()> { + let example = SweepConfig { + max_degree: Some(vec![16, 32, 64]), + l_build: Some(vec![50, 75, 100]), + search_l: Some(vec![10, 20, 30, 40, 50, 75, 100]), + num_pq_chunks: Some(vec![8, 16, 32]), + }; + + 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:"); + println!(" - max_degree: {:?}", example.max_degree.unwrap()); + println!(" - l_build: {:?}", example.l_build.unwrap()); + println!(" - search_l: {:?}", example.search_l.unwrap()); + println!(" - num_pq_chunks: {:?}", example.num_pq_chunks.unwrap()); + + 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)); + run_benchmark( + benchmark_cmd, + benchmark_args, + &config_file, + &output_file, + )?; + + // 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"); + } + } + + // 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(); + + // Get job type to determine which parameters to sweep + let job_type = base_config + .get("jobs") + .and_then(|j| j.as_array()) + .and_then(|arr| arr.first()) + .and_then(|job| job.get("type")) + .and_then(|t| t.as_str()) + .unwrap_or("unknown"); + + let is_pq = job_type.contains("pq"); + + // Get default values from base config + let default_max_degree = sweep_spec.max_degree.as_ref().map(|v| v[0]).unwrap_or(32); + let default_l_build = sweep_spec.l_build.as_ref().map(|v| v[0]).unwrap_or(50); + let default_search_l = sweep_spec.search_l.clone().unwrap_or_else(|| vec![20, 30, 40]); + + // Generate all combinations of build parameters + let max_degrees = sweep_spec.max_degree.as_ref().map_or(vec![default_max_degree], |v| v.clone()); + let l_builds = sweep_spec.l_build.as_ref().map_or(vec![default_l_build], |v| v.clone()); + let pq_chunks = if is_pq { + sweep_spec.num_pq_chunks.as_ref().map_or(vec![], |v| v.clone()) + } else { + vec![0] // dummy value for non-PQ + }; + + let mut config_id = 0; + for &max_degree in &max_degrees { + for &l_build in &l_builds { + for &num_pq_chunks in &pq_chunks { + if !is_pq && num_pq_chunks != 0 { + continue; + } + + // Create modified config + let mut config = base_config.clone(); + let mut params = HashMap::new(); + + // Update max_degree and l_build + if let Some(jobs) = config.get_mut("jobs").and_then(|j| j.as_array_mut()) { + for job in jobs.iter_mut() { + update_build_params(job, max_degree, l_build, is_pq, num_pq_chunks)?; + + // Update search_l values - try two different paths + let search_phase = job.get_mut("content") + .and_then(|c| c.get_mut("search_phase")); + + if search_phase.is_none() { + if let Some(search_phase) = job.get_mut("content") + .and_then(|c| c.get_mut("index_operation")) + .and_then(|io| io.get_mut("search_phase")) + { + if let Some(runs) = search_phase.get_mut("runs").and_then(|r| r.as_array_mut()) { + for run in runs.iter_mut() { + run["search_l"] = serde_json::to_value(&default_search_l)?; + } + } + } + } else if let Some(search_phase) = search_phase { + if let Some(runs) = search_phase.get_mut("runs").and_then(|r| r.as_array_mut()) { + for run in runs.iter_mut() { + run["search_l"] = serde_json::to_value(&default_search_l)?; + } + } + } + } + } + + params.insert("max_degree".to_string(), serde_json::json!(max_degree)); + params.insert("l_build".to_string(), serde_json::json!(l_build)); + params.insert("search_l".to_string(), serde_json::json!(default_search_l)); + if is_pq { + params.insert("num_pq_chunks".to_string(), serde_json::json!(num_pq_chunks)); + } + + let id = format!( + "{:04}_R{}_L{}{}", + config_id, + max_degree, + l_build, + if is_pq { format!("_PQ{}", num_pq_chunks) } else { String::new() } + ); + configs.push((id, config, params)); + config_id += 1; + } + } + } + + Ok(configs) +} + +fn update_build_params( + job: &mut Value, + max_degree: u32, + l_build: u32, + is_pq: bool, + num_pq_chunks: u32, +) -> Result<()> { + // Try different paths based on job type + if let Some(source) = job.get_mut("content").and_then(|c| c.get_mut("source")) { + // async-index-build format + if source.get("index-source").is_some() { + source["max_degree"] = serde_json::json!(max_degree); + source["l_build"] = serde_json::json!(l_build); + } + // disk-index format + if source.get("disk-index-source").is_some() { + source["max_degree"] = serde_json::json!(max_degree); + source["l_build"] = serde_json::json!(l_build); + if is_pq { + source["num_pq_chunks"] = serde_json::json!(num_pq_chunks); + } + } + } + + // async-index-build-pq format + if is_pq { + if let Some(content) = job.get_mut("content") { + if let Some(index_op) = content.get_mut("index_operation").and_then(|io| io.get_mut("source")) { + index_op["max_degree"] = serde_json::json!(max_degree); + index_op["l_build"] = serde_json::json!(l_build); + } + content["num_pq_chunks"] = serde_json::json!(num_pq_chunks); + } + } + + 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())) +} From c49d374ced96b2f20689a0076c689f63009aa2f3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 7 Feb 2026 01:28:17 +0000 Subject: [PATCH 3/6] Add tests and documentation for autotuner Co-authored-by: harsha-simhadri <5590673+harsha-simhadri@users.noreply.github.com> --- diskann-benchmark/README.md | 29 +++++ diskann-tools/src/bin/autotuner.rs | 175 +++++++++++++++++++++++++++++ 2 files changed, 204 insertions(+) diff --git a/diskann-benchmark/README.md b/diskann-benchmark/README.md index 6e072b757..53fb2829e 100644 --- a/diskann-benchmark/README.md +++ b/diskann-benchmark/README.md @@ -491,3 +491,32 @@ 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). + +See [diskann-tools/AUTOTUNER.md](../diskann-tools/AUTOTUNER.md) for detailed documentation on using the autotuner. + +### 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 autotuner sweeps over: +- **max_degree (R)**: Graph degree parameter +- **l_build**: Search queue length during index construction +- **search_l**: Search queue length during queries +- **num_pq_chunks**: Product quantization chunks (for quantized indexes) + +Results are saved to the output directory with a summary highlighting the best configuration. + diff --git a/diskann-tools/src/bin/autotuner.rs b/diskann-tools/src/bin/autotuner.rs index cf281e41d..633098b17 100644 --- a/diskann-tools/src/bin/autotuner.rs +++ b/diskann-tools/src/bin/autotuner.rs @@ -552,3 +552,178 @@ fn save_json(path: &Path, data: &T) -> Result<()> { 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_sweep_config_serialization() { + let config = SweepConfig { + max_degree: Some(vec![16, 32, 64]), + l_build: Some(vec![50, 100]), + search_l: Some(vec![10, 20, 30]), + num_pq_chunks: Some(vec![8, 16]), + }; + + let json = serde_json::to_string(&config).unwrap(); + let deserialized: SweepConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(config.max_degree, deserialized.max_degree); + assert_eq!(config.l_build, deserialized.l_build); + assert_eq!(config.search_l, deserialized.search_l); + assert_eq!(config.num_pq_chunks, deserialized.num_pq_chunks); + } + + #[test] + fn test_generate_configurations() { + let base_config = serde_json::json!({ + "search_directories": ["test_data"], + "jobs": [{ + "type": "async-index-build", + "content": { + "source": { + "index-source": "Build", + "data_type": "float32", + "data": "data.fbin", + "distance": "squared_l2", + "max_degree": 32, + "l_build": 50, + "alpha": 1.2, + "num_threads": 1 + }, + "search_phase": { + "search-type": "topk", + "queries": "queries.fbin", + "groundtruth": "gt.bin", + "reps": 1, + "num_threads": [1], + "runs": [{ + "search_n": 10, + "search_l": [20], + "recall_k": 10 + }] + } + } + }] + }); + + let sweep_config = SweepConfig { + max_degree: Some(vec![16, 32]), + l_build: Some(vec![50, 100]), + search_l: Some(vec![10, 20, 30]), + num_pq_chunks: None, + }; + + let configs = generate_configurations(&base_config, &sweep_config).unwrap(); + + // Should generate 2 max_degree * 2 l_build = 4 configurations + assert_eq!(configs.len(), 4); + + // Check that parameters are correctly set + for (config_id, config, params) in configs { + assert!(config_id.contains("R")); + assert!(config_id.contains("L")); + assert!(params.contains_key("max_degree")); + assert!(params.contains_key("l_build")); + assert!(params.contains_key("search_l")); + + // Verify the config has updated values + let job = &config["jobs"][0]; + let source = &job["content"]["source"]; + let max_degree = source["max_degree"].as_u64().unwrap() as u32; + let l_build = source["l_build"].as_u64().unwrap() as u32; + + assert!(max_degree == 16 || max_degree == 32); + assert!(l_build == 50 || l_build == 100); + } + } + + #[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"); + } +} From 42d0e9013946be9c59d8148aff2b28a11c4d354d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 05:37:28 +0000 Subject: [PATCH 4/6] Refactor autotuner to use path-based configuration system The new design addresses maintainability concerns by: - Using JSON paths instead of hardcoded structure assumptions - Making the tool adaptable to any benchmark configuration format - Allowing users to specify which parameters to sweep via paths - Staying in sync with benchmark framework changes automatically Users now specify parameters as path-value pairs (e.g., "jobs.0.content.source.max_degree": [16, 32, 64]) instead of relying on fixed parameter names. This makes the tool robust to changes in the benchmark framework's JSON schema. Co-authored-by: harsha-simhadri <5590673+harsha-simhadri@users.noreply.github.com> --- diskann-tools/AUTOTUNER.md | 270 ++++++++++++------ diskann-tools/src/bin/autotuner.rs | 440 ++++++++++++++--------------- 2 files changed, 396 insertions(+), 314 deletions(-) diff --git a/diskann-tools/AUTOTUNER.md b/diskann-tools/AUTOTUNER.md index b2b5a81f0..7c7a1073a 100644 --- a/diskann-tools/AUTOTUNER.md +++ b/diskann-tools/AUTOTUNER.md @@ -4,10 +4,13 @@ The autotuner is a tool that builds on top of the DiskANN benchmark framework to ## Overview -The autotuner helps optimize DiskANN indexes by automatically: -- Sweeping over key parameters: R (max_degree), l_build, l_search (search_l), and quantization bytes (num_pq_chunks) -- Running benchmarks for each parameter combination -- Analyzing results to find the best configuration based on 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, 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 @@ -17,103 +20,157 @@ Build the autotuner tool: cargo build --release --package diskann-tools --bin autotuner ``` -## Usage +## Quick Start ### 1. Generate an Example Sweep Configuration -First, generate an example sweep configuration file: - ```bash cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json ``` -This creates a JSON file with example parameter ranges: +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 -```json -{ - "max_degree": [16, 32, 64], - "l_build": [50, 75, 100], - "search_l": [10, 20, 30, 40, 50, 75, 100], - "num_pq_chunks": [8, 16, 32] -} +# For disk indexes +cargo run --release --package diskann-tools --bin autotuner -- example --output sweep_config.json --benchmark-type disk ``` -### 2. Prepare a Base Configuration - -Create a base benchmark configuration file (or use an existing one from `diskann-benchmark/example/`). The autotuner will use this as a template and modify the parameters specified in the sweep configuration. +### 2. Understanding the Sweep Configuration -Example base configuration (`base_config.json`): +The generated sweep configuration uses JSON paths to specify which parameters to sweep: ```json { - "search_directories": ["test_data/disk_index_search"], - "jobs": [{ - "type": "async-index-build", - "content": { - "source": { - "index-source": "Build", - "data_type": "float32", - "data": "data.fbin", - "distance": "squared_l2", - "max_degree": 32, - "l_build": 50, - "alpha": 1.2, - "backedge_ratio": 1.0, - "num_threads": 1, - "start_point_strategy": "medoid" - }, - "search_phase": { - "search-type": "topk", - "queries": "queries.fbin", - "groundtruth": "groundtruth.bin", - "reps": 5, - "num_threads": [1], - "runs": [{ - "search_n": 10, - "search_l": [20, 30, 40], - "recall_k": 10 - }] - } + "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] + ] } - }] + ] } ``` -### 3. Run the Parameter Sweep +**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. -Run the autotuner to sweep over parameters and find the best configuration: +### 4. Run the Parameter Sweep ```bash cargo run --release --package diskann-tools --bin autotuner -- sweep \ - --base-config base_config.json \ + --base-config diskann-benchmark/example/async.json \ --sweep-config sweep_config.json \ --output-dir ./autotuner_results \ --criterion qps \ --target-recall 0.95 ``` -#### Options: +## 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 -- `--base-config`: Path to the base benchmark configuration (template) -- `--sweep-config`: Path to the sweep configuration (parameter ranges) -- `--output-dir`: Directory where results will be saved -- `--criterion`: Optimization criterion: - - `qps`: Maximize queries per second (default) - - `latency`: Minimize latency - - `recall`: Maximize recall -- `--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 the 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_.json`: Generated configuration for each parameter combination -- `results_.json`: Benchmark results for each configuration -- `sweep_summary.json`: Summary of all results with the best configuration highlighted +- `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 Output: +### Example Summary ```json { @@ -122,24 +179,23 @@ The autotuner generates the following files in the output directory: "total_configs": 9, "successful_configs": 9, "best_config": { - "config_id": "0005_R32_L75", + "config_id": "0005", "parameters": { - "max_degree": 32, - "l_build": 75, - "search_l": [10, 20, 30, 40, 50, 75, 100] + "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": [...] + } } ``` ## Examples -### Optimize for Maximum QPS at 95% Recall +### Example 1: Optimize for Maximum QPS ```bash cargo run --release --package diskann-tools --bin autotuner -- sweep \ @@ -150,7 +206,7 @@ cargo run --release --package diskann-tools --bin autotuner -- sweep \ --target-recall 0.95 ``` -### Optimize for Minimum Latency at 99% Recall +### Example 2: Optimize for Minimum Latency ```bash cargo run --release --package diskann-tools --bin autotuner -- sweep \ @@ -161,7 +217,7 @@ cargo run --release --package diskann-tools --bin autotuner -- sweep \ --target-recall 0.99 ``` -### Optimize for Maximum Recall +### Example 3: Optimize for Maximum Recall ```bash cargo run --release --package diskann-tools --bin autotuner -- sweep \ @@ -171,22 +227,60 @@ cargo run --release --package diskann-tools --bin autotuner -- sweep \ --criterion recall ``` -## Parameter Descriptions +### Example 4: Sweep Over Product-Quantized Index Parameters -- **max_degree (R)**: The maximum degree of the graph. Higher values increase index size but can improve recall. -- **l_build**: The search queue length during index construction. Higher values improve index quality but increase build time. -- **search_l**: The search queue length during queries. Higher values improve recall but reduce throughput. -- **num_pq_chunks**: Number of product quantization chunks (for quantized indexes). Affects compression ratio and search accuracy. +```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**: Begin with a small set of parameter values to get a rough idea of the performance landscape. -2. **Refine the sweep**: Once you identify promising regions, create a more fine-grained sweep around those values. -3. **Consider the trade-offs**: Higher QPS often comes at the cost of lower recall, and vice versa. -4. **Test data size**: Use representative data sizes for your production workload. +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 -## Notes +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 -- The autotuner generates one configuration per combination of build parameters (max_degree, l_build, num_pq_chunks). -- All search_l values are tested for each build configuration, allowing the tool to find the best search_l given the build parameters. -- For quantized indexes (PQ, SQ, etc.), make sure to use the appropriate base configuration template. +Users control what gets swept by specifying paths, not by relying on tool-specific parameter names. diff --git a/diskann-tools/src/bin/autotuner.rs b/diskann-tools/src/bin/autotuner.rs index 633098b17..944bc5d56 100644 --- a/diskann-tools/src/bin/autotuner.rs +++ b/diskann-tools/src/bin/autotuner.rs @@ -6,8 +6,10 @@ //! Autotuner for DiskANN benchmarks. //! //! This tool builds on top of the benchmark framework to sweep over a subset of parameters -//! (R/max_degree, l_build, l_search/search_l, quantization bytes where applicable) //! 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}; @@ -62,26 +64,26 @@ enum Commands { /// 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 SweepConfig { - /// Max degree (R) values to sweep - #[serde(skip_serializing_if = "Option::is_none")] - max_degree: Option>, - - /// l_build values to sweep - #[serde(skip_serializing_if = "Option::is_none")] - l_build: Option>, - - /// search_l values to sweep - #[serde(skip_serializing_if = "Option::is_none")] - search_l: Option>, +struct ParameterSweep { + /// JSON path to the parameter (e.g., "jobs.0.content.source.max_degree") + path: String, + /// Values to sweep over + values: Vec, +} - /// num_pq_chunks values to sweep (for quantized indexes) - #[serde(skip_serializing_if = "Option::is_none")] - num_pq_chunks: Option>, +#[derive(Debug, Serialize, Deserialize, Clone)] +struct SweepConfig { + /// List of parameters to sweep with their paths and values + parameters: Vec, } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -122,16 +124,76 @@ fn main() -> Result<()> { &benchmark_cmd, benchmark_args.as_deref(), ), - Commands::Example { output } => generate_example(&output), + Commands::Example { output, benchmark_type } => generate_example(&output, benchmark_type.as_deref()), } } -fn generate_example(output: &Path) -> Result<()> { - let example = SweepConfig { - max_degree: Some(vec![16, 32, 64]), - l_build: Some(vec![50, 75, 100]), - search_l: Some(vec![10, 20, 30, 40, 50, 75, 100]), - num_pq_chunks: Some(vec![8, 16, 32]), +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) @@ -139,11 +201,14 @@ fn generate_example(output: &Path) -> Result<()> { serde_json::to_writer_pretty(file, &example)?; println!("Example sweep configuration written to: {}", output.display()); - println!("\nThis configuration will sweep over:"); - println!(" - max_degree: {:?}", example.max_degree.unwrap()); - println!(" - l_build: {:?}", example.l_build.unwrap()); - println!(" - search_l: {:?}", example.search_l.unwrap()); - println!(" - num_pq_chunks: {:?}", example.num_pq_chunks.unwrap()); + 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(()) } @@ -196,24 +261,24 @@ fn run_sweep( // Run benchmark let output_file = output_dir.join(format!("results_{}.json", config_id)); - run_benchmark( - benchmark_cmd, - benchmark_args, - &config_file, - &output_file, - )?; - - // 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"); + 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); + } } } @@ -258,128 +323,92 @@ fn generate_configurations( ) -> Result)>> { let mut configs = Vec::new(); - // Get job type to determine which parameters to sweep - let job_type = base_config - .get("jobs") - .and_then(|j| j.as_array()) - .and_then(|arr| arr.first()) - .and_then(|job| job.get("type")) - .and_then(|t| t.as_str()) - .unwrap_or("unknown"); - - let is_pq = job_type.contains("pq"); - - // Get default values from base config - let default_max_degree = sweep_spec.max_degree.as_ref().map(|v| v[0]).unwrap_or(32); - let default_l_build = sweep_spec.l_build.as_ref().map(|v| v[0]).unwrap_or(50); - let default_search_l = sweep_spec.search_l.clone().unwrap_or_else(|| vec![20, 30, 40]); + // Generate all combinations of parameter values + let combinations = generate_parameter_combinations(&sweep_spec.parameters); - // Generate all combinations of build parameters - let max_degrees = sweep_spec.max_degree.as_ref().map_or(vec![default_max_degree], |v| v.clone()); - let l_builds = sweep_spec.l_build.as_ref().map_or(vec![default_l_build], |v| v.clone()); - let pq_chunks = if is_pq { - sweep_spec.num_pq_chunks.as_ref().map_or(vec![], |v| v.clone()) - } else { - vec![0] // dummy value for non-PQ - }; + 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()); + } - let mut config_id = 0; - for &max_degree in &max_degrees { - for &l_build in &l_builds { - for &num_pq_chunks in &pq_chunks { - if !is_pq && num_pq_chunks != 0 { - continue; - } + // Generate config ID from parameter values + let config_id = format!("{:04}", idx); + configs.push((config_id, config, params)); + } - // Create modified config - let mut config = base_config.clone(); - let mut params = HashMap::new(); - - // Update max_degree and l_build - if let Some(jobs) = config.get_mut("jobs").and_then(|j| j.as_array_mut()) { - for job in jobs.iter_mut() { - update_build_params(job, max_degree, l_build, is_pq, num_pq_chunks)?; - - // Update search_l values - try two different paths - let search_phase = job.get_mut("content") - .and_then(|c| c.get_mut("search_phase")); - - if search_phase.is_none() { - if let Some(search_phase) = job.get_mut("content") - .and_then(|c| c.get_mut("index_operation")) - .and_then(|io| io.get_mut("search_phase")) - { - if let Some(runs) = search_phase.get_mut("runs").and_then(|r| r.as_array_mut()) { - for run in runs.iter_mut() { - run["search_l"] = serde_json::to_value(&default_search_l)?; - } - } - } - } else if let Some(search_phase) = search_phase { - if let Some(runs) = search_phase.get_mut("runs").and_then(|r| r.as_array_mut()) { - for run in runs.iter_mut() { - run["search_l"] = serde_json::to_value(&default_search_l)?; - } - } - } - } - } + Ok(configs) +} - params.insert("max_degree".to_string(), serde_json::json!(max_degree)); - params.insert("l_build".to_string(), serde_json::json!(l_build)); - params.insert("search_l".to_string(), serde_json::json!(default_search_l)); - if is_pq { - params.insert("num_pq_chunks".to_string(), serde_json::json!(num_pq_chunks)); - } +/// Generate all combinations of parameter values using Cartesian product +fn generate_parameter_combinations(parameters: &[ParameterSweep]) -> Vec> { + if parameters.is_empty() { + return vec![vec![]]; + } - let id = format!( - "{:04}_R{}_L{}{}", - config_id, - max_degree, - l_build, - if is_pq { format!("_PQ{}", num_pq_chunks) } else { String::new() } - ); - configs.push((id, config, params)); - config_id += 1; + 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; } - - Ok(configs) + + result } -fn update_build_params( - job: &mut Value, - max_degree: u32, - l_build: u32, - is_pq: bool, - num_pq_chunks: u32, -) -> Result<()> { - // Try different paths based on job type - if let Some(source) = job.get_mut("content").and_then(|c| c.get_mut("source")) { - // async-index-build format - if source.get("index-source").is_some() { - source["max_degree"] = serde_json::json!(max_degree); - source["l_build"] = serde_json::json!(l_build); - } - // disk-index format - if source.get("disk-index-source").is_some() { - source["max_degree"] = serde_json::json!(max_degree); - source["l_build"] = serde_json::json!(l_build); - if is_pq { - source["num_pq_chunks"] = serde_json::json!(num_pq_chunks); - } - } +/// 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; - // async-index-build-pq format - if is_pq { - if let Some(content) = job.get_mut("content") { - if let Some(index_op) = content.get_mut("index_operation").and_then(|io| io.get_mut("source")) { - index_op["max_degree"] = serde_json::json!(max_degree); - index_op["l_build"] = serde_json::json!(l_build); + // 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); } - content["num_pq_chunks"] = serde_json::json!(num_pq_chunks); } } @@ -558,85 +587,44 @@ mod tests { use super::*; #[test] - fn test_sweep_config_serialization() { - let config = SweepConfig { - max_degree: Some(vec![16, 32, 64]), - l_build: Some(vec![50, 100]), - search_l: Some(vec![10, 20, 30]), - num_pq_chunks: Some(vec![8, 16]), - }; - - let json = serde_json::to_string(&config).unwrap(); - let deserialized: SweepConfig = serde_json::from_str(&json).unwrap(); + fn test_set_json_path_simple() { + let mut json = serde_json::json!({ + "a": { + "b": 42 + } + }); - assert_eq!(config.max_degree, deserialized.max_degree); - assert_eq!(config.l_build, deserialized.l_build); - assert_eq!(config.search_l, deserialized.search_l); - assert_eq!(config.num_pq_chunks, deserialized.num_pq_chunks); + set_json_path(&mut json, "a.b", serde_json::json!(100)).unwrap(); + assert_eq!(json["a"]["b"], 100); } #[test] - fn test_generate_configurations() { - let base_config = serde_json::json!({ - "search_directories": ["test_data"], - "jobs": [{ - "type": "async-index-build", - "content": { - "source": { - "index-source": "Build", - "data_type": "float32", - "data": "data.fbin", - "distance": "squared_l2", - "max_degree": 32, - "l_build": 50, - "alpha": 1.2, - "num_threads": 1 - }, - "search_phase": { - "search-type": "topk", - "queries": "queries.fbin", - "groundtruth": "gt.bin", - "reps": 1, - "num_threads": [1], - "runs": [{ - "search_n": 10, - "search_l": [20], - "recall_k": 10 - }] - } - } - }] + fn test_set_json_path_array() { + let mut json = serde_json::json!({ + "jobs": [ + {"content": {"source": {"max_degree": 32}}} + ] }); - let sweep_config = SweepConfig { - max_degree: Some(vec![16, 32]), - l_build: Some(vec![50, 100]), - search_l: Some(vec![10, 20, 30]), - num_pq_chunks: None, - }; - - let configs = generate_configurations(&base_config, &sweep_config).unwrap(); - - // Should generate 2 max_degree * 2 l_build = 4 configurations - assert_eq!(configs.len(), 4); - - // Check that parameters are correctly set - for (config_id, config, params) in configs { - assert!(config_id.contains("R")); - assert!(config_id.contains("L")); - assert!(params.contains_key("max_degree")); - assert!(params.contains_key("l_build")); - assert!(params.contains_key("search_l")); + 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); + } - // Verify the config has updated values - let job = &config["jobs"][0]; - let source = &job["content"]["source"]; - let max_degree = source["max_degree"].as_u64().unwrap() as u32; - let l_build = source["l_build"].as_u64().unwrap() as u32; + #[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)], + }, + ]; - assert!(max_degree == 16 || max_degree == 32); - assert!(l_build == 50 || l_build == 100); - } + let combinations = generate_parameter_combinations(¶meters); + assert_eq!(combinations.len(), 4); // 2 * 2 = 4 } #[test] From 14e99502caf8a1cdccc1d1681d3378ce87aca3b5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 05:37:51 +0000 Subject: [PATCH 5/6] Update benchmark README to reflect path-based autotuner design Co-authored-by: harsha-simhadri <5590673+harsha-simhadri@users.noreply.github.com> --- diskann-benchmark/README.md | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/diskann-benchmark/README.md b/diskann-benchmark/README.md index 53fb2829e..853638534 100644 --- a/diskann-benchmark/README.md +++ b/diskann-benchmark/README.md @@ -495,7 +495,9 @@ Refer to implementations within the benchmarking framework for what some of this 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). -See [diskann-tools/AUTOTUNER.md](../diskann-tools/AUTOTUNER.md) for detailed documentation on using the autotuner. +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 @@ -512,11 +514,15 @@ cargo run --release --package diskann-tools --bin autotuner -- sweep \ --target-recall 0.95 ``` -The autotuner sweeps over: -- **max_degree (R)**: Graph degree parameter -- **l_build**: Search queue length during index construction -- **search_l**: Search queue length during queries -- **num_pq_chunks**: Product quantization chunks (for quantized indexes) +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]} + ] +} +``` -Results are saved to the output directory with a summary highlighting the best configuration. +This design makes the autotuner adaptable to any benchmark configuration format without requiring code changes when the benchmark framework evolves. From ab1760196fc004ed73e0edaec55c8e0452dfebbe Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 05:38:42 +0000 Subject: [PATCH 6/6] Clarify config_id format and all_results field in documentation Co-authored-by: harsha-simhadri <5590673+harsha-simhadri@users.noreply.github.com> --- diskann-tools/AUTOTUNER.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/diskann-tools/AUTOTUNER.md b/diskann-tools/AUTOTUNER.md index 7c7a1073a..a711f2a05 100644 --- a/diskann-tools/AUTOTUNER.md +++ b/diskann-tools/AUTOTUNER.md @@ -172,6 +172,8 @@ The autotuner generates the following files in the output directory: ### Example Summary +The `sweep_summary.json` file contains all sweep results. Here's an example (truncated for brevity): + ```json { "criterion": "qps", @@ -189,10 +191,15 @@ The autotuner generates the following files in the output directory: "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