DataBench — Benchmark Smarter with AI Agents

Automates end-to-end data platform benchmarking: job submission, status tracking, result collection, cost analysis, and report generation — through natural language or a visual dashboard.

Quick Start

git clone ssh://git.amazon.com/pkg/DataBench
cd DataBench
pip install -r requirements.txt

Run the Dashboard

PYTHONPATH=.. streamlit run ui/app.py

Opens at http://localhost:8501 with:

• Dashboard — ClickBench-inspired ranking of all benchmark runs (GPU vs CPU, EMR vs OSS, Parquet vs Iceberg)

• Report Builder — Upload driver logs or CSVs, compare runs, export JSON

• Submit Jobs — Pick a template, configure, launch TPC-DS benchmarks

• Monitor — Track running EMR jobs in real-time

Run Tests

python -m pytest tests/ -v

Use as an MCP Server in Claude Code

DataBench ships an MCP server (databench.mcpserver.server) that exposes the benchmark tools — benchmark-list-templates, benchmark-submit, benchmark-status, benchmark-collect, benchmark-cost, benchmark-compare, benchmark-report, and more — directly to Claude Code.

First install the package so databench is importable:

pip install -e .

Then register the server (--scope user makes it available in all your projects):

claude mcp add databench --scope user \
  -e PYTHONPATH=/path/to/sourcecode/DataBench \
  -- python -m databench.mcpserver.server

Use the absolute path to the python interpreter where you ran pip install -e . (find it with which python) so Claude Code launches the server with the right environment. Verify it connected:

claude mcp get databench   # Status: ✔ Connected

Restart your Claude Code session and the databench tools will be available (check with /mcp inside Claude, or claude mcp list). To remove it:

claude mcp remove databench -s user

Talk to it in plain English

Once the MCP server is connected, you drive the whole workflow — provisioning nodegroups, running benchmarks, and generating reports — by just asking Claude Code in natural language. Claude picks the right templates, creates the EKS managed nodegroups (RAID0 local disk, cluster-autoscaler wiring, pod-template node pinning), submits the jobs, polls them, and produces the report. Real examples:

Provision infrastructure

• "Create managed nodegroups in the loadtest-mcp EKS cluster for r7g.4xlarge and r8g.4xlarge, scaling from 1 to 8, each with 4×64GB EBS volumes striped as RAID0 mounted at /var/data."

• "Create r7gd.4xlarge and r8gd.4xlarge nodegroups and mount the local NVMe as the Spark spill dir."

• "Install the cluster autoscaler so it can scale only these benchmark nodegroups up and down — don't let it touch the ops nodegroup or Karpenter."

• "Bump all four benchmark nodegroups to max 8 nodes."

• "Remove the 8xlarge nodegroups but keep the templates."

Run benchmarks

• "Run TPC-DS 3TB comparing r7g vs r8g on 8 nodes for EMR on EKS 7.12 Spark performance."

• "Benchmark TPC-DS 3TB across r7g.4xl, r8g.4xl, r7gd.4xl and r8gd.4xl, then compare the results."

• "Show me the template setup first, then run r7g.8xlarge vs r8g.8xlarge on 4 nodes — keep total CPU and memory the same as the 4xlarge run."

• "Split the big executor into 6 smaller pods per node without changing total CPU or memory, then re-run."

Check status & get reports

• "What's the status of the r7gd benchmark jobs?"

• "Compare all of r7g, r7gd, r8g and r8gd and give me a downloadable report."

• "Generate the Spark cost-performance comparison report."

Claude handles the mechanics behind these — picking benchmark-list-templates, benchmark-submit, benchmark-status, benchmark-compare, and benchmark-report, plus the eksctl/kubectl/aws steps for nodegroup lifecycle — and asks for confirmation before anything that costs money.

Related MCP server: Databricks Unity Catalog MCP Server

What's Inside

DataBench/
├── models/
│   └── benchmark_result.py    # BenchmarkResult, QueryResult, ComparisonResult
├── tools/
│   ├── collect.py             # Parse driver logs, CSV, JSON → BenchmarkResult
│   ├── compare.py             # Speedup calculator: compare_runs() → ComparisonResult
│   ├── cost.py                # Instance pricing lookup + cost-per-run calculation
│   ├── report.py              # xlsx report generator (Summary, Query Details, Configs)
│   └── status.py              # Cross-platform job status poller
├── adapters/
│   ├── emr_eks.py             # EMR on EKS: submit jobs, check status
│   ├── emr_ec2.py             # EMR on EC2: submit steps
│   └── athena.py              # Athena: submit TPC-DS queries
├── configs/                   # 9 parameterized benchmark templates
├── ui/
│   └── app.py                 # Streamlit dashboard
├── tests/
│   ├── test_collect.py
│   ├── test_compare.py
│   ├── test_compare_real_data.py
│   ├── test_cost.py
│   ├── test_report.py
│   ├── test_athena_submit.py
│   └── sample_benchmark_result.json
├── conftest.py                # pytest import fix (DataBench → databench)
└── README.md

Benchmark Results (baked into dashboard)

Interface Contract

Every tool produces/consumes BenchmarkResult:

from databench.models import BenchmarkResult

result = BenchmarkResult.from_json(open("result.json").read())
print(result.total_median())        # 534.0 seconds
print(result.median_time("q1"))     # 4.5 seconds
print(result.query_names())         # ["q1", "q2", ...]

Compare & Report

from databench.tools.compare import compare_runs
from databench.tools.report import generate_report

# Compare GPU vs CPU
comparison = compare_runs([gpu_result, cpu_result], baseline_run_id="cpu-run")
print(comparison.aggregates)  # speedups, wins, cost savings

# Generate xlsx report
report = generate_report([gpu_result, cpu_result], baseline_run_id="cpu-run")
print(report["file_path"])    # databench-report-gpu-vs-cpu.xlsx

Cost Lookup

from databench.tools.cost import calculate_cost

cost = calculate_cost("g6.4xlarge", node_count=8, seconds=534.0, region="us-east-1")
print(cost)  # {"price_per_hour": 1.323, "cluster_cost_per_hour": 10.584, "cost_per_run": 1.57}

Config Templates

from databench.configs import list_templates, load_template

# List all templates
for t in list_templates():
    print(f"{t['template_id']:30s} GPU={t['gpu']}")

# Load with variable resolution
cfg = load_template("gpu-parquet-g6-4xl", bucket="my-bucket", region="us-east-1")

AWS Credentials

The dashboard uses your AWS CLI profile. Set it in the sidebar or:

ada credentials update --profile aws-emr-bda-admin --provider isengard --once

Team

• Karthik Prabhakar (subbakk) — EMR adapters, orchestration, config templates, UI

• Pathik Shah (pathshah) — Report generation, comparison, cost, Athena adapter

Links

• Architecture Doc (Quip)

• Taskei Epic