Trust, but benchmark: How we let an AI agent optimize Elasticsearch

This article recounts the experience of developing a harness that automatically detects and implements optimizations in the Elasticsearch codebase. Elasticsearch is deployed across various workloads, such as heavy index building and real-time search and analytics. To ensure excellent performance across these workloads, the process requires both broad coverage and deep understanding of the codebase.

Traditionally, human engineers had to manually scrutinize each potential bottleneck, which is impractical given the limited engineering resources and the constantly evolving nature of the code.

However, with advancements in AI coding agents, performance optimization can now be approached semiautomatically. Unlike many software engineering challenges, optimizing code presents a cheap and objective verifier. When an AI model is tasked with making code faster, it provides a hard number indicating the outcome, backed by profiling tools that explain the results. This makes performance an ideal candidate for automation, provided there is trust in the numbers generated.

To ensure the accuracy of the optimizations, the authors devised a trustworthy measurement loop. The harness that drives the AI agent assumes the agent may be wrong a significant portion of the time but reliably identifies and validates the correct optimizations. This loop helps guide the AI agent to focus on the most promising areas for improvement.

The authors built a highly trustworthy measurement loop consisting of three main components: an exploration task, a human-driven exploitation task, and a validation run. The exploration task profiles a real workload and produces a ranked list of performance opportunities. A human engineer then selects one opportunity to be further explored through an exploitation task.

This task iterates against an approved microbenchmark, committing each accepted experiment. Finally, a validation run on the actual workload ensures the performance gains translate to real-world scenarios before a human reviews and opens a pull request.

To enable autonomous work, the authors identified three key properties a task should possess: an objective verdict for the AI agent to adhere to, a dense guiding signal to direct the agent towards relevant areas, and a bounded blast radius to minimize the impact of incorrect optimizations. Performance meets these criteria effectively, as benchmarks provide a clear verdict, profilers offer guidance, and rejected patches are easily reverted without affecting correctness.

The article emphasizes the importance of the AI agent's sensory apparatus, represented by the command line interface (CLI). Each command in the CLI should return a concise answer to a specific question related to the task and, when applicable, provide an explanation to aid the AI agent's reasoning process. This design ensures that the agent acts on clear signals rather than parsing and interpreting raw data.

The harness was engineered with four essential signals to assist the AI agent. Facet-decomposition allows the agent to break down complex performance-related signals into more manageable components. By decomposing facets, the agent can focus on specific aspects of performance optimization, streamlining the process and reducing the cognitive load on the human engineers.