ci: benchmark the object store WAL against S3 from an EC2 runner

Add a workflow_dispatch workflow that loads the o11ybench agent observability
dataset into a fresh GreptimeDB container per cell and compares three WAL
configurations that differ only in [wal]: Raft Engine with sync_write on the
instance store, and the object store WAL in durable and enqueued mode, with
the data files on S3 in the runner's region for all three. Rounds alternate
the configuration order. A summary reports rows per second, D/R and E/R as
the geometric mean over rounds, append acknowledgement p50 and p99, and the
object store WAL counters. A smoke mode runs the same chain on ubuntu-latest
against the MinIO fixture.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

.github/scripts/oswal-s3-bench-run.sh

@@ -0,0 +1,134 @@
+#!/usr/bin/env bash
+# Host side of the object store WAL S3 benchmark (oswal-s3-bench.yml).
+#   wal-disk DIR  mount the instance store at DIR when the host has one
+#   generate      generate the dataset once with the o11ybench runtime
+#   cells         run every cell of $PLAN, each on a fresh GreptimeDB container
+set -euo pipefail
+
+: "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}"
+# Same isolation as the o11ybench runtime's own docker run.
+runtime=(docker run --rm --user "$(id -u):$(id -g)" --label "o11ybench.owner=$OWNER"
+    --security-opt no-new-privileges --cap-drop ALL)
+
+wal_disk() {
+    local root=$1 device
+    device=$(lsblk -dnpo NAME,MODEL | awk '/Instance Storage/ {print $1; exit}')
+    sudo -n mkdir -p "$root"
+    if [[ -n "$device" ]]; then
+        sudo -n umount "$device" 2>/dev/null || true
+        sudo -n mkfs.ext4 -q -F "$device"
+        sudo -n mount -o noatime "$device" "$root"
+    fi
+    sudo -n chown "$(id -u):$(id -g)" "$root"
+    findmnt -no SOURCE,FSTYPE -T "$root"
+}
+
+generate() {
+    local args
+    mapfile -t args < <(jq -r '.generator_args[]' "$PLAN")
+    mkdir -p "$CORPUS_DIR"
+    "${runtime[@]}" --name "$OWNER-generator" \
+        --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus" \
+        "$RUNTIME_IMAGE" scripts/generate_agent_observations.py "${args[@]}" \
+        --output /work/corpus/agent_observations.jsonl \
+        --summary-output /work/corpus/summary.json --evidence-output /work/corpus/evidence.json \
+        > "$RESULTS/generator.log" 2>&1
+    jq -e --argjson rows "$(jq .rows "$PLAN")" '.line_count == $rows' "$CORPUS_DIR/summary.json" >/dev/null
+}
+
+# One cell in its own process, so errexit holds inside it.
+cell() {
+    : "${GREPTIMEDB_IMAGE:?}" "${RUN_PREFIX:?}" "${WAL_ROOT:?}" "${S3_BUCKET:?}" "${S3_REGION:?}"
+    : "${AWS_ACCESS_KEY_ID:?}" "${AWS_SECRET_ACCESS_KEY:?}"
+    # Globals, not locals: the EXIT trap runs after this function returns.
+    spec=$1 cid= wal_dir= version= load_seconds=
+    local concurrency config name
+    id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec"); config=$(jq -r .config <<< "$spec")
+    dir="$RESULTS/cells/$id"; name="$OWNER-$id"
+    mkdir -p "$dir"
+    # The prefix holds data/ and the WAL's default prefix wal/ side by side.
+    local settings=(STORAGE__TYPE=S3 "STORAGE__BUCKET=$S3_BUCKET" "STORAGE__REGION=$S3_REGION"
+        "STORAGE__ROOT=$RUN_PREFIX/$id")
+    [[ -z "${S3_CONTAINER_ENDPOINT:-}" ]] || settings+=("STORAGE__ENDPOINT=$S3_CONTAINER_ENDPOINT")
+    mapfile -t -O "${#settings[@]}" settings < <(jq -r '.wal | to_entries[] | "WAL__\(.key)=\(.value)"' <<< "$spec")
+    local docker_env=() mounts=() setting
+    for setting in "${settings[@]}"; do
+        docker_env+=(-e "GREPTIMEDB_STANDALONE__$setting")
+        printf 'GREPTIMEDB_STANDALONE__%s\n' "$setting" >> "$dir/config.env"
+    done
+    # Pass the two keys by name: their values never reach argv, a file or the artifact.
+    export GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID
+    export GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY
+    docker_env+=(-e GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID -e GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY)
+    printf 'GREPTIMEDB_STANDALONE__STORAGE__%s=<redacted>\n' ACCESS_KEY_ID SECRET_ACCESS_KEY >> "$dir/config.env"
+    if jq -e '.wal.DIR' <<< "$spec" >/dev/null; then
+        wal_dir="$WAL_ROOT/$id"; mkdir -p "$wal_dir"
+        mounts=(--mount "type=bind,src=$wal_dir,dst=$(jq -r .wal.DIR <<< "$spec")")
+        wal_device=$(findmnt -no SOURCE -T "$wal_dir")
+    else
+        wal_device="s3://$S3_BUCKET/$RUN_PREFIX/$id/wal"
+    fi
+    finish() {
+        local rc=$?
+        trap - EXIT
+        set +e
+        if [[ -n "$cid" ]]; then
+            docker stop -t 120 "$cid" >/dev/null
+            docker logs "$cid" > "$dir/server.log" 2>&1
+            docker rm -fv "$cid" >/dev/null
+        fi
+        [[ -z "$wal_dir" ]] || sudo -n rm -rf "$wal_dir"
+        jq -n --argjson spec "$spec" --arg device "$wal_device" --arg version "$version" \
+            --arg load_seconds "$load_seconds" --argjson exit_code "$rc" \
+            '$spec + {wal_device: $device, version: $version, exit_code: $exit_code,
+              load_wall_seconds: ($load_seconds | tonumber? // null)}' > "$dir/cell.json"
+        exit "$rc"
+    }
+    trap finish EXIT
+    cid=$(docker run -d --init --name "$name" --label "o11ybench.owner=$OWNER" \
+        --add-host host.docker.internal:host-gateway \
+        -p 127.0.0.1:4000:4000 -p 127.0.0.1:4001:4001 "${docker_env[@]}" "${mounts[@]}" \
+        "$GREPTIMEDB_IMAGE" standalone start --http-addr=0.0.0.0:4000 --grpc-bind-addr=0.0.0.0:4001)
+    local deadline=$((SECONDS + 300))
+    until curl --noproxy '*' -fsS --max-time 5 http://127.0.0.1:4000/health >/dev/null 2>&1; do
+        [[ $(docker inspect --format '{{.State.Running}}' "$cid") == true ]] || { echo "$id: server exited" >&2; exit 1; }
+        ((SECONDS < deadline)) || { echo "$id: server not healthy in 300s" >&2; exit 1; }
+        sleep 2
+    done
+    version=$(docker exec "$cid" greptime --version | sed -n 's/^version: //p')
+    [[ -n "$version" ]]
+    local started ended
+    started=$(date +%s.%N)
+    # Only the load phase, run exactly as run-agent-observability.sh runs it.
+    "${runtime[@]}" --name "$name-load" --network host \
+        --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus,readonly" --mount "type=bind,src=$dir,dst=/results" \
+        -e GRPC_ENDPOINT=127.0.0.1:4001 -e "DATA_PROFILE=$(jq -r .data_profile "$PLAN")" -e RESULT_NAME=greptimedb-1 \
+        -e "ENDPOINT=http://127.0.0.1:4000/v1/sql?db=public&format=greptimedb_v1" \
+        -e "TARGET_IMAGE=$GREPTIMEDB_IMAGE" -e "TARGET_VERSION=$version" -e "TARGET_DIGEST=${GREPTIMEDB_IMAGE##*@}" \
+        -e "TARGET_ID=$name" -e TARGET_HOST=127.0.0.1 -e PREPARE_WORKERS= -e "MAX_IN_FLIGHT=$concurrency" \
+        -e "GREPTIMEDB_BATCH_SIZE=$(jq -r .batch_rows "$PLAN")" --entrypoint /bin/bash \
+        "$RUNTIME_IMAGE" scripts/agent-observability/greptimedb.sh load > "$dir/load.log" 2>&1
+    ended=$(date +%s.%N)
+    load_seconds=$(awk -v a="$started" -v b="$ended" 'BEGIN {printf "%.3f", b - a}')
+    jq -e --argjson rows "$(jq .rows "$PLAN")" '.count_before == 0 and .count_after == $rows' \
+        "$dir/greptimedb-1/load-result.json" >/dev/null
+    curl --noproxy '*' -fsS --max-time 30 http://127.0.0.1:4000/metrics > "$dir/metrics.all"
+    grep -E '^(# (HELP|TYPE) )?greptime_logstore_object_store_wal_' "$dir/metrics.all" > "$dir/metrics.prom" || true
+    rm "$dir/metrics.all"
+}
+
+case "${1:-}" in
+    wal-disk) wal_disk "${2:?}" ;;
+    generate) generate ;;
+    cell) cell "${2:?}" ;;
+    cells)
+        mapfile -t specs < <(jq -c '.cells[]' "$PLAN")
+        failed=0
+        for spec in "${specs[@]}"; do
+            echo "::group::cell $(jq -r .id <<< "$spec")"
+            bash "${BASH_SOURCE[0]}" cell "$spec" || { echo "::error::cell $(jq -r .id <<< "$spec") failed"; failed=1; }
+            echo '::endgroup::'
+        done
+        exit "$failed" ;;
+    *) echo "usage: $0 wal-disk DIR | generate | cells" >&2; exit 2 ;;
+esac

.github/scripts/oswal-s3-bench-summary.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+# Copyright 2023 Greptime Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Summarize the object store WAL S3 benchmark as Markdown and summary.json."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import re
+from pathlib import Path
+
+PREFIX = "greptime_logstore_object_store_wal_"
+COUNTERS = ("created_objects_total", "stalled_appends_total", "poisoned_total",
+            "create_failures_total", "create_conflicts_total")
+SAMPLE = re.compile(r'([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{([^}]*)\})?\s+(\S+)')
+
+
+def parse_metrics(text: str) -> dict:
+    """Map each unlabeled sample to its value and each histogram to its buckets."""
+    samples, buckets = {}, {}
+    for line in text.splitlines():
+        match = SAMPLE.fullmatch(line.strip())
+        if not match or not match.group(1).startswith(PREFIX):
+            continue
+        name, labels, value = match.group(1)[len(PREFIX):], match.group(2) or "", float(match.group(3))
+        le = re.search(r'le="([^"]+)"', labels)
+        if name.endswith("_bucket") and le:
+            buckets.setdefault(name[:-len("_bucket")], []).append((float(le.group(1)), value))
+        else:
+            samples[name] = value
+    return {"samples": samples, "buckets": {name: sorted(items) for name, items in buckets.items()}}
+
+
+def quantile(q: float, buckets: list) -> float | None:
+    """Prometheus histogram_quantile over cumulative (upper bound, count) buckets."""
+    if not buckets or buckets[-1][1] == 0:
+        return None
+    rank = q * buckets[-1][1]
+    lower, below = 0.0, 0.0
+    for upper, count in buckets:
+        if count >= rank:
+            if math.isinf(upper):
+                return lower
+            return lower + (upper - lower) * (rank - below) / (count - below) if count > below else upper
+        lower, below = upper, count
+    return lower
+
+
+def geomean(values: list) -> float | None:
+    return math.exp(sum(map(math.log, values)) / len(values)) if values and all(v > 0 for v in values) else None
+
+
+def load_cell(path: Path) -> dict:
+    cell = json.loads((path / "cell.json").read_text())
+    result_path = path / "greptimedb-1" / "load-result.json"
+    result = json.loads(result_path.read_text()) if result_path.exists() else {}
+    metrics_path = path / "metrics.prom"
+    metrics = parse_metrics(metrics_path.read_text() if metrics_path.exists() else "")
+    ack = metrics["buckets"].get("append_ack_seconds", [])
+    load_seconds = result.get("load_seconds")
+    row = {
+        "id": cell["id"], "round": cell["round"], "concurrency": cell["concurrency"],
+        "config": cell["config"], "exit_code": cell["exit_code"], "wal_device": cell["wal_device"],
+        "version": cell.get("version"), "load_wall_seconds": cell.get("load_wall_seconds"),
+        "rows": result.get("count_after"), "load_seconds": load_seconds,
+        "rows_per_second": result.get("rows_per_second") if cell["exit_code"] == 0 else None,
+        "append_ack_p50_seconds": quantile(0.5, ack), "append_ack_p99_seconds": quantile(0.99, ack),
+    }
+    for counter in COUNTERS:
+        row[counter] = metrics["samples"].get(counter)
+    created = row["created_objects_total"]
+    row["creates_per_minute"] = created * 60 / load_seconds if created is not None and load_seconds else None
+    return row
+
+
+def ratios(rows: list) -> list:
+    """D/R and E/R per round and their geometric mean over the rounds, per concurrency."""
+    speed = {(r["round"], r["concurrency"], r["config"]): r["rows_per_second"] for r in rows}
+    result = []
+    for concurrency in sorted({r["concurrency"] for r in rows}):
+        rounds = sorted({r["round"] for r in rows if r["concurrency"] == concurrency})
+        for config in ("D", "E"):
+            per_round = []
+            for round_number in rounds:
+                base, value = speed.get((round_number, concurrency, "R")), speed.get((round_number, concurrency, config))
+                per_round.append(value / base if base and value else None)
+            complete = [value for value in per_round if value is not None]
+            result.append({"concurrency": concurrency, "ratio": f"{config}/R", "per_round": per_round,
+                           "geomean": geomean(complete) if len(complete) == len(per_round) else None})
+    return result
+
+
+def fmt(value, digits=0, scale=1.0) -> str:
+    return "—" if value is None else f"{value * scale:,.{digits}f}"
+
+
+def render(run: dict, rows: list, pairs: list) -> str:
+    lines = [f"## Object store WAL against S3 ({'smoke, MinIO' if run.get('smoke') else run.get('instance_type')})", ""]
+    lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({rows[0]['version'] if rows else '—'}); "
+              f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""]
+    lines += ["| Cell | Rows/s | Load s | Wall s | Rows | Ack p50 ms | Ack p99 ms | Creates | Creates/min "
+              "| Stalled | Poisoned | Create failures | Conflicts | WAL on |",
+              "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |"]
+    for r in rows:
+        cell = r["id"] if r["exit_code"] == 0 else f"{r['id']} (failed, exit {r['exit_code']})"
+        lines.append(" | ".join([
+            f"| {cell}", fmt(r["rows_per_second"]), fmt(r["load_seconds"], 1), fmt(r["load_wall_seconds"], 1),
+            fmt(r["rows"]), fmt(r["append_ack_p50_seconds"], 1, 1000), fmt(r["append_ack_p99_seconds"], 1, 1000),
+            *(fmt(r[c]) for c in COUNTERS[:1]), fmt(r["creates_per_minute"], 1),
+            *(fmt(r[c]) for c in COUNTERS[1:]), f"`{r['wal_device']}` |"]))
+    lines += ["", "| Concurrency | Ratio | Per round | Geometric mean |", "| ---: | --- | --- | ---: |"]
+    for p in pairs:
+        lines.append(f"| {p['concurrency']} | {p['ratio']} | {', '.join(fmt(v, 3) for v in p['per_round'])} "
+                     f"| {fmt(p['geomean'], 3)} |")
+    return "\n".join(lines) + "\n"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--root", type=Path, required=True)
+    args = parser.parse_args()
+    args.root.mkdir(parents=True, exist_ok=True)
+    run_path = args.root / "run.json"
+    run = json.loads(run_path.read_text()) if run_path.exists() else {}
+    plan_path = args.root / "plan.json"
+    order = [cell["id"] for cell in json.loads(plan_path.read_text())["cells"]] if plan_path.exists() else []
+    cells = {path.name: load_cell(path) for path in (args.root / "cells").glob("*") if (path / "cell.json").exists()}
+    rows = [cells[i] for i in order if i in cells] + [cells[i] for i in sorted(cells) if i not in order]
+    pairs = ratios(rows)
+    (args.root / "summary.json").write_text(json.dumps({"run": run, "cells": rows, "ratios": pairs}, indent=2) + "\n")
+    print(render(run, rows, pairs) if rows else "No benchmark cell completed.\n")
+
+
+if __name__ == "__main__":
+    main()

.github/scripts/oswal-s3-bench.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python3
+# Copyright 2023 Greptime Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["boto3>=1.35,<2"]
+# ///
+
+"""Plan, image tag and bucket helpers of the object store WAL S3 benchmark.
+
+`plan` and `tag` need only the standard library; `probe` and `cleanup` need
+boto3 and run under `uv run`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import sys
+import time
+import urllib.request
+
+SMOKE_ROWS = 50_000
+# The o11ybench data profiles; a matching row count generates the canonical dataset.
+DATA_PROFILES = {5_000: "S", 2_000_000: "D", 10_000_000: "P", 100_000_000: "M"}
+# Identical [storage] for every configuration; only [wal] differs.
+WAL_CONFIGS = {
+    "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"},
+    "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"},
+    "E": {"PROVIDER": "experimental_object_store", "ACK_MODE": "enqueued"},
+}
+ORDERS = (("R", "D", "E"), ("E", "D", "R"))
+NIGHTLY = re.compile(r"v\d+\.\d+\.\d+(?:-[0-9A-Za-z.]+)?-nightly-(\d{8})")
+TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}")
+
+
+def append_output(name: str, value: str) -> None:
+    if path := os.environ.get("GITHUB_OUTPUT"):
+        with open(path, "a") as output:
+            output.write(f"{name}={value}\n")
+
+
+def positive(name: str, value: str, upper: int) -> int:
+    if not re.fullmatch(r"[1-9][0-9]*", value) or int(value) > upper:
+        raise ValueError(f"{name} must be an integer in 1..{upper}, got {value!r}")
+    return int(value)
+
+
+def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool) -> dict:
+    levels = [positive("concurrency", level.strip(), 256) for level in concurrency.split(",")]
+    if len(set(levels)) != len(levels):
+        raise ValueError(f"concurrency repeats a level: {concurrency!r}")
+    row_count = SMOKE_ROWS if smoke else positive("rows", rows, 100_000_000)
+    round_count = 1 if smoke else positive("rounds", rounds, 10)
+    profile = DATA_PROFILES.get(row_count)
+    cells = []
+    for round_number in range(1, round_count + 1):
+        # Alternate the order between rounds so no configuration always runs first.
+        for level in levels:
+            for config in ORDERS[(round_number - 1) % 2]:
+                cells.append({
+                    "id": f"{round_number}-{level}-{config}", "round": round_number,
+                    "concurrency": level, "config": config, "wal": WAL_CONFIGS[config],
+                })
+    return {
+        "smoke": smoke, "rows": row_count, "rounds": round_count,
+        "batch_rows": positive("batch_rows", batch_rows, 100_000), "concurrency": levels,
+        "data_profile": profile or "custom",
+        "generator_args": ["--data-profile", profile] if profile else ["--line-count", str(row_count)],
+        "cells": cells,
+    }
+
+
+def newest_nightly(tags: list[str]) -> str:
+    dated = [(match.group(1), tag) for tag in tags if (match := NIGHTLY.fullmatch(tag))]
+    if not dated:
+        raise ValueError("no v*-nightly-* tag found")
+    return max(dated)[1]
+
+
+def resolve_tag(requested: str) -> str:
+    if requested != "nightly":
+        if not TAG.fullmatch(requested):
+            raise ValueError(f"invalid image tag: {requested!r}")
+        return requested
+    url = ("https://hub.docker.com/v2/repositories/greptime/greptimedb/tags"
+           "?page_size=100&name=nightly&ordering=last_updated")
+    with urllib.request.urlopen(url, timeout=30) as response:
+        return newest_nightly([tag["name"] for tag in json.load(response)["results"]])
+
+
+def s3_client(args):
+    import boto3
+    from botocore.config import Config
+
+    return boto3.client("s3", region_name=args.region, endpoint_url=args.endpoint or None,
+                        config=Config(retries={"mode": "standard", "max_attempts": 5}))
+
+
+def probe(args) -> None:
+    client, key = s3_client(args), f"{args.prefix}/probe/object"
+    deadline = time.monotonic() + 60
+    while True:
+        try:
+            client.put_object(Bucket=args.bucket, Key=key, Body=b"probe")
+            break
+        except Exception:
+            # A fresh MinIO creates its bucket after the container starts.
+            if time.monotonic() > deadline:
+                raise
+            time.sleep(3)
+    listed = client.list_objects_v2(Bucket=args.bucket, Prefix=f"{args.prefix}/probe/")
+    if [item["Key"] for item in listed.get("Contents", [])] != [key]:
+        raise RuntimeError(f"listing {args.prefix}/probe/ did not return exactly {key}")
+    client.delete_object(Bucket=args.bucket, Key=key)
+    print(f"put, list and delete succeeded under s3://{args.bucket}/{args.prefix}/probe/")
+
+
+def cleanup(args) -> None:
+    client, deleted = s3_client(args), 0
+    for page in client.get_paginator("list_objects_v2").paginate(Bucket=args.bucket, Prefix=f"{args.prefix}/"):
+        keys = [{"Key": item["Key"]} for item in page.get("Contents", [])]
+        if keys:
+            result = client.delete_objects(Bucket=args.bucket, Delete={"Objects": keys, "Quiet": True})
+            if result.get("Errors"):
+                raise RuntimeError(f"delete failed for {len(result['Errors'])} objects")
+            deleted += len(keys)
+    print(f"deleted {deleted} objects under s3://{args.bucket}/{args.prefix}/")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    plan = commands.add_parser("plan")
+    for name in ("concurrency", "rows", "batch-rows", "rounds"):
+        plan.add_argument(f"--{name}", required=True)
+    plan.add_argument("--smoke", choices=("true", "false"), required=True)
+    plan.add_argument("--output")
+    commands.add_parser("tag").add_argument("requested")
+    for name in ("probe", "cleanup"):
+        bucket = commands.add_parser(name)
+        bucket.add_argument("--bucket", required=True)
+        bucket.add_argument("--region", required=True)
+        bucket.add_argument("--endpoint", default="")
+        bucket.add_argument("--prefix", required=True)
+    args = parser.parse_args()
+    if args.command == "plan":
+        result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true")
+        if args.output:
+            with open(args.output, "w") as output:
+                json.dump(result, output, indent=2)
+        print(f"{len(result['cells'])} cells, {result['rows']} rows, profile {result['data_profile']}")
+    elif args.command == "tag":
+        tag = resolve_tag(args.requested)
+        append_output("greptimedb_tag", tag)
+        print(f"GreptimeDB image tag: {tag}")
+    else:
+        if not re.fullmatch(r"oswal-bench/[0-9]+-[0-9]+", args.prefix):
+            raise ValueError(f"refusing a prefix outside oswal-bench/<run>: {args.prefix!r}")
+        (probe if args.command == "probe" else cleanup)(args)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())