app/backend/scripts/bench_startup.py

"""How long it takes to get a run started, which is the thing Kedro is slow at.

A pipeline framework that boots the project per run pays that cost every time:
``kedro run`` on a pipeline that does nothing takes about a second, and a sweep
of five hundred configs therefore spends ten minutes doing nothing. Fluksio
answers that by not booting anything — the engine is already up and its workers
already have the code loaded, so submitting is one request.

This measures that claim against a running stack, so it can be checked rather
than asserted. Run it with the dev stack up::

    make bench-startup
    make bench-startup BENCH_ARGS="--runs 50 --kedro ../path/to/kedro/project"

The Kedro figure is optional and measured the same way — a null pipeline, timed
end to end — so the two numbers mean the same thing.
"""

from __future__ import annotations

import argparse
import json
import logging
import statistics
import subprocess
import sys
import time
import urllib.error
import urllib.request
from typing import Any

sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve().parents[1]))

from app.core.config import settings  # noqa: E402

log = logging.getLogger("bench")

#: A flow that computes nothing, so what is timed is the getting-started.
BENCH_FLOW = "bench_startup"
NODE_SOURCE = '"""Does nothing, on purpose."""\n\n\ndef process(n, params):\n    return {"out": n}\n'


class Api:
    def __init__(self, base: str, token: str) -> None:
        self.base = base.rstrip("/")
        self.token = token

    def call(self, method: str, path: str, body: Any = None) -> Any:
        request = urllib.request.Request(
            f"{self.base}{path}",
            method=method,
            data=json.dumps(body).encode() if body is not None else None,
            headers={
                "Authorization": f"Bearer {self.token}",
                "Content-Type": "application/json",
            },
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            return json.load(response) if response.status != 204 else None


def login(base: str) -> Api:
    data = (
        f"username={settings.FIRST_SUPERUSER}&password={settings.FIRST_SUPERUSER_PASSWORD}"
    ).encode()
    request = urllib.request.Request(
        f"{base.rstrip('/')}/api/v1/login/access-token",
        data=data,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        return Api(base, json.load(response)["access_token"])


def ensure_flow(api: Api) -> None:
    """A published batch flow with one node that returns what it was given."""
    definition = {
        "name": BENCH_FLOW,
        "title": "Startup benchmark",
        "version": 1,
        "mode": "batch",
        "outputs": ["out"],
        "inputs": [{"spec": {"name": "n", "dtype": "int"}, "initial": 1}],
        "nodes": [
            {
                "id": "noop",
                "type": "python",
                "requires": [{"name": "n", "dtype": "int"}],
                "provides": [{"name": "out", "dtype": "int"}],
            }
        ],
    }
    try:
        current = api.call("GET", f"/api/v1/flows/{BENCH_FLOW}")
        definition["version"] = current["definition"]["version"]
    except urllib.error.HTTPError:
        pass
    api.call("PUT", f"/api/v1/flows/{BENCH_FLOW}", definition)
    api.call(
        "PUT", f"/api/v1/flows/{BENCH_FLOW}/nodes/noop/source", {"code": NODE_SOURCE}
    )
    api.call(
        "POST",
        f"/api/v1/flows/{BENCH_FLOW}/publish",
        {"version": definition["version"]},
    )


def time_runs(api: Api, count: int) -> tuple[list[float], list[float]]:
    submits: list[float] = []
    totals: list[float] = []
    for index in range(count):
        start = time.perf_counter()
        run = api.call(
            "POST", f"/api/v1/runs/flows/{BENCH_FLOW}", {"params": {"n": index}}
        )
        submits.append((time.perf_counter() - start) * 1000)
        while True:
            got = api.call("GET", f"/api/v1/runs/{run['id']}")
            if got["status"] not in ("queued", "running"):
                break
            time.sleep(0.005)
        totals.append((time.perf_counter() - start) * 1000)
        if got["status"] != "ok":
            raise SystemExit(
                f"run {run['id']} ended {got['status']}: {got['status_reason']}"
            )
    return submits, totals


def time_kedro(project: str, pipeline: str, count: int) -> list[float]:
    """The same measurement for a Kedro project: a whole run, end to end."""
    timings: list[float] = []
    for _ in range(count):
        start = time.perf_counter()
        result = subprocess.run(
            ["kedro", "run", "--pipeline", pipeline],
            cwd=project,
            capture_output=True,
            check=False,
        )
        if result.returncode != 0:
            raise SystemExit(f"kedro run failed: {result.stderr.decode()[-400:]}")
        timings.append((time.perf_counter() - start) * 1000)
    return timings


def report(label: str, timings: list[float]) -> None:
    ordered = sorted(timings)
    log.info(
        f"  {label:<28} median {statistics.median(ordered):7.1f} ms"
        f"   min {ordered[0]:7.1f}   max {ordered[-1]:7.1f}"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default="http://api.localhost")
    parser.add_argument("--runs", type=int, default=20)
    parser.add_argument(
        "--kedro", default="", help="a Kedro project to compare against"
    )
    parser.add_argument("--kedro-pipeline", default="__default__")
    parser.add_argument("--kedro-runs", type=int, default=3)
    parser.add_argument(
        "--keep", action="store_true", help="leave the benchmark flow behind"
    )
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)

    api = login(args.base_url)
    ensure_flow(api)
    # One throwaway run first: the worker compiles the node on its first call,
    # and that cost belongs to the deployment rather than to a run.
    time_runs(api, 1)

    submits, totals = time_runs(api, args.runs)
    log.info(f"\nfluksio — {args.runs} runs of a flow that computes nothing")
    report("submit accepted", submits)
    report("submit -> result", totals)

    if args.kedro:
        kedro = time_kedro(args.kedro, args.kedro_pipeline, args.kedro_runs)
        log.info(
            f"\nkedro — {args.kedro_runs} runs of a pipeline that computes nothing"
        )
        report("kedro run", kedro)
        ratio = statistics.median(kedro) / statistics.median(totals)
        log.info(f"\n  a run costs {ratio:.0f}x less here, per run")

    if not args.keep:
        api.call("DELETE", f"/api/v1/flows/{BENCH_FLOW}")
        log.info(f"\nremoved the '{BENCH_FLOW}' flow")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())