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/3] 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/3] 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/3] 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"); + } +}