Benchmark for `fit` and `predct` at different `n_samples` for https://github.com/scikit-learn/scikit-learn/issues/34604

bench_pipeline_size_sweep.py

"""Sweep pipeline `fit` and `predict` runtime against the size of X.

This script times

    StandardScaler -> Nystroem -> LogisticRegression

as a function of how much data it is given, in two independent sweeps:

    fit      GridSearchCV(pipeline, PARAM_GRID, cv=3).fit(X[:n], y[:n]),
             for n from 2_000 to 100_000
    predict  one reference model fit at n=20_000, then predict on batches
             of 1 to 2_000 rows

Six arms are measured, so that "where do the fitted attributes live" can be
read against "what namespace and device is X in":

    numpy-plain        no Array API dispatch at all, today's plain sklearn
    numpy-dispatch     NumPy X with dispatch on
    mps-native-attrs   torch MPS X, attributes stay on the device
    mps-numpy-attrs    torch MPS X, attributes are NumPy
    cuda-native-attrs  torch CUDA X, attributes stay on the device
    cuda-numpy-attrs   torch CUDA X, attributes are NumPy

With NumPy X the attributes end up as NumPy either way, so the two NumPy arms
differ only by whether dispatch is on.

Whatever is not available on this machine is skipped, so the same script runs
on a laptop with MPS and on a box with CUDA. Each run writes one JSON file
named after the host; `plot_pipeline_size_sweep.py` merges them.

Needs the branch of https://github.com/scikit-learn/scikit-learn/pull/34667
checked out, plus the StandardScaler and Nystroem conversions.

    python bench_pipeline_size_sweep.py [--rounds 5] [--devices numpy torch-mps]
"""

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

# isort: skip_file

import os

os.environ.setdefault("SCIPY_ARRAY_API", "1")
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

import argparse
import importlib.util
import json
import pathlib
import platform
import socket
import statistics
import subprocess
import sys
import time
import warnings
from contextlib import nullcontext

import numpy as np

from sklearn import config_context
from sklearn.datasets import make_classification
from sklearn.exceptions import ConvergenceWarning
from sklearn.kernel_approximation import Nystroem
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

# MAX_ITER below is deliberately low to keep `fit` quick. We are timing, not
# fitting a good model, so the resulting warnings are just noise here.
warnings.filterwarnings("ignore", category=ConvergenceWarning)

# ---------------------------------------------------------------- knobs ----

DEVICES = ["numpy", "torch-mps", "torch-cuda"]

# The fit sweep: linearly spaced, with a finer step below 10_000.
FIT_SIZES = list(range(2_000, 10_000, 2_000)) + list(range(10_000, 100_001, 10_000))

# The predict sweep: log spaced, because the curve is a flat plateau where a
# fixed per-call cost dominates that bends into linear growth, and the knee is
# what we are looking for.
PREDICT_BATCHES = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1_000, 2_000]

# The predict sweep needs one model per arm; it is fit once, at this size.
REF_FIT_SIZE = 20_000

N_FEATURES = 50
N_COMPONENTS = 200
N_CLASSES = 3

# torch MPS has no float64, so this is not really a choice.
DTYPE = "float32"

MAX_ITER = 500

CV_FOLDS = 3
PARAM_GRID = {"clf__C": [0.1, 1.0], "kernel__gamma": [None, 0.1]}

N_ROUNDS = 5

# A cheap `predict` has to be repeated to rise above the clock, an expensive
# one must not be. MIN_REPS is what stops a single stalled call from becoming
# the reported value: without it a 500ms stall on a 7ms call satisfies
# MIN_TOTAL_TIME on its own and is returned as the median of one sample.
MIN_TOTAL_TIME = 0.05
MIN_REPS = 7
MAX_REPS = 200

WARMUP_FIT_SIZE = 2_000
WARMUP_PREDICT_SIZE = 100

# --------------------------------------------------------------- arms ----

# (name, device, array_api_dispatch, store_fitted_as_numpy, label)
ARMS = [
    ("numpy-plain", "numpy", False, False, "attrs numpy + X numpy (no dispatch)"),
    ("numpy-dispatch", "numpy", True, True, "attrs numpy + X numpy (dispatch)"),
    ("mps-native-attrs", "torch-mps", True, False, "attrs MPS + X MPS"),
    ("mps-numpy-attrs", "torch-mps", True, True, "attrs numpy + X MPS"),
    ("cuda-native-attrs", "torch-cuda", True, False, "attrs CUDA + X CUDA"),
    ("cuda-numpy-attrs", "torch-cuda", True, True, "attrs numpy + X CUDA"),
]


# ------------------------------------------------------------- backends ----


def available(device):
    """Whether `device` can be used here, and why not if it cannot.

    Importing and probing are both wrapped, because a library can be installed
    and still fail to load.
    """
    if device == "numpy":
        return True, ""

    if importlib.util.find_spec("torch") is None:
        return False, "torch is not installed"

    try:
        import torch

        if device == "torch-cpu":
            return True, ""
        if device == "torch-cuda":
            if not torch.cuda.is_available():
                return False, "torch reports no CUDA device"
        elif device == "torch-mps":
            if not torch.backends.mps.is_available():
                return False, "torch reports no MPS device"
        else:
            return False, f"unknown device {device!r}"
    except Exception as exc:
        return False, f"torch failed to load: {type(exc).__name__}: {exc}"

    return True, ""


def to_device(array, device):
    """Put a NumPy array into the namespace and on the device of `device`."""
    if device == "numpy":
        return array
    if device.startswith("torch"):
        import torch

        return torch.asarray(array, device=device.removeprefix("torch-"))
    raise ValueError(f"unknown device {device!r}")


def synchronize(device):
    """Block until queued work on `device` has finished."""
    if device == "torch-cuda":
        import torch

        torch.cuda.synchronize()
    elif device == "torch-mps":
        import torch

        torch.mps.synchronize()


def empty_cache(device):
    """Return cached device memory to the allocator."""
    if device == "torch-cuda":
        import torch

        torch.cuda.empty_cache()
    elif device == "torch-mps":
        import torch

        torch.mps.empty_cache()


# ------------------------------------------------------------- plumbing ----


def arm_config(dispatch, store_as_numpy):
    """Config context for one arm.

    The no-dispatch arm gets no config context at all, so that it is really
    today's plain scikit-learn rather than the dispatch path applied to NumPy.
    """
    if not dispatch:
        return nullcontext()
    return config_context(array_api_dispatch=True, store_fitted_as_numpy=store_as_numpy)


def measure_once(fn, device):
    """Time one call of `fn`, in seconds, along with whatever it returned.

    The synchronize is inside the timed region, so GPU numbers are end-to-end
    latencies rather than the time it took to queue the work.
    """
    synchronize(device)
    start = time.perf_counter()
    value = fn()
    synchronize(device)
    return time.perf_counter() - start, value


def measure_median(fn, device):
    """Median wall time of `fn`, repeating until `MIN_TOTAL_TIME` has elapsed.

    Cheap calls get enough repetitions to rise above the clock resolution while
    expensive ones do not take forever.
    """
    synchronize(device)
    times = []
    total = 0.0
    while True:
        start = time.perf_counter()
        fn()
        synchronize(device)
        elapsed = time.perf_counter() - start
        times.append(elapsed)
        total += elapsed
        if len(times) >= MAX_REPS:
            break
        if total >= MIN_TOTAL_TIME and len(times) >= MIN_REPS:
            break
    return statistics.median(times)


def build_pipeline():
    return Pipeline(
        [
            ("scale", StandardScaler()),
            ("kernel", Nystroem(n_components=N_COMPONENTS, random_state=0)),
            ("clf", LogisticRegression(max_iter=MAX_ITER)),
        ]
    )


def build_search():
    return GridSearchCV(build_pipeline(), PARAM_GRID, cv=CV_FOLDS)


def make_data(device, random_state=0):
    """Fit and predict data for one device.

    The two arms of a device share this, so it is generated once per device
    rather than once per arm.
    """
    n_predict = max(PREDICT_BATCHES)
    n_fit = max(FIT_SIZES)
    X, y = make_classification(
        n_samples=n_fit + n_predict,
        n_features=N_FEATURES,
        n_informative=max(N_FEATURES // 2, 5),
        n_redundant=0,
        n_classes=N_CLASSES,
        n_clusters_per_class=1,
        random_state=random_state,
    )
    X = X.astype(DTYPE)
    return {
        "X_fit": to_device(X[:n_fit], device),
        "y_fit": to_device(y[:n_fit], device),
        # Held out, so `predict` is not run on the training rows.
        "X_predict": to_device(X[n_fit:], device),
    }


def fmt_time(seconds):
    us = seconds * 1e6
    if us < 1000:
        return f"{us:.0f}us"
    if us < 1e6:
        return f"{us / 1000:.2f}ms"
    return f"{seconds:.2f}s"


def describe_environment(devices):
    import sklearn

    env = {
        "host": socket.gethostname(),
        "sklearn": sklearn.__version__,
        "numpy": np.__version__,
        "platform": platform.platform(),
        "python": platform.python_version(),
        "devices": list(devices),
    }
    try:
        import torch

        env["torch"] = torch.__version__
        if torch.cuda.is_available():
            env["cuda_devices"] = [
                torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())
            ]
            env["cuda_device_used"] = torch.cuda.get_device_name(0)
    except ImportError:
        pass
    if platform.system() == "Darwin":
        try:
            env["cpu"] = subprocess.run(
                ["sysctl", "-n", "machdep.cpu.brand_string"],
                capture_output=True,
                text=True,
                check=True,
            ).stdout.strip()
        except (OSError, subprocess.CalledProcessError):
            pass
    else:
        try:
            for line in pathlib.Path("/proc/cpuinfo").read_text().splitlines():
                if line.startswith("model name"):
                    env["cpu"] = line.split(":", 1)[1].strip()
                    break
        except OSError:
            pass
    return env


# ----------------------------------------------------------------- runs ----


class Arm:
    """One (device, config) combination, with its data and reference model."""

    def __init__(self, name, device, dispatch, store_as_numpy, label, data):
        self.name = name
        self.device = device
        self.dispatch = dispatch
        self.store_as_numpy = store_as_numpy
        self.label = label
        self.data = data
        self.reference = None

    def config(self):
        return arm_config(self.dispatch, self.store_as_numpy)

    def warmup(self):
        """One untimed pass, to absorb lazy allocation and kernel compilation.

        Without this the first timed round pays for it, and on MPS that is a
        factor of two on a whole grid search.
        """
        X = self.data["X_fit"][:WARMUP_FIT_SIZE]
        y = self.data["y_fit"][:WARMUP_FIT_SIZE]
        with self.config():
            fitted = build_search().fit(X, y)
            fitted.predict(self.data["X_predict"][:WARMUP_PREDICT_SIZE])
        synchronize(self.device)

    def fit_reference(self):
        """The single model the predict sweep runs against."""
        X = self.data["X_fit"][:REF_FIT_SIZE]
        y = self.data["y_fit"][:REF_FIT_SIZE]
        with self.config():
            self.reference = build_search().fit(X, y)
        synchronize(self.device)

    def time_fit(self, n_samples):
        X = self.data["X_fit"][:n_samples]
        y = self.data["y_fit"][:n_samples]
        with self.config():
            elapsed, _ = measure_once(lambda: build_search().fit(X, y), self.device)
        return elapsed

    def time_predict(self, batch_size):
        batch = self.data["X_predict"][:batch_size]
        with self.config():
            return measure_median(lambda: self.reference.predict(batch), self.device)


def release_fit_data(arms):
    """Drop the fit data once the reference models exist.

    The predict sweep times calls of a few milliseconds, and leaving the
    100_000 row training set resident makes those measurements sensitive to
    device memory pressure. Arms on the same device share one data dict, so the
    pops are deliberately idempotent.
    """
    for arm in arms:
        arm.data.pop("X_fit", None)
        arm.data.pop("y_fit", None)
    for device in {arm.device for arm in arms}:
        empty_cache(device)


def build_arms(devices):
    """The arms runnable on this machine, with per-device data."""
    data_by_device = {}
    arms = []
    for name, device, dispatch, store_as_numpy, label in ARMS:
        if device not in devices:
            continue
        if device not in data_by_device:
            print(f"generating data for {device}", flush=True)
            data_by_device[device] = make_data(device)
        arms.append(
            Arm(name, device, dispatch, store_as_numpy, label, data_by_device[device])
        )
    return arms


def sweep_fit(arms, rounds, records):
    """Time `fit` over `FIT_SIZES`, arms adjacent in time at every size.

    The arms at a given size run seconds apart, so a machine that throttles
    partway through the sweep cannot make one arm look worse than another.
    """
    for round_index in range(rounds):
        for n_samples in FIT_SIZES:
            cells = []
            for arm in arms:
                elapsed = arm.time_fit(n_samples)
                records.append(
                    {
                        "arm": arm.name,
                        "n_samples": n_samples,
                        "round": round_index,
                        "seconds": elapsed,
                    }
                )
                cells.append(f"{arm.name}={fmt_time(elapsed)}")
            print(
                f"  fit round {round_index} n={n_samples:<7d} " + "  ".join(cells),
                flush=True,
            )


def sweep_predict(arms, rounds, records):
    """Time `predict` over `PREDICT_BATCHES` against the reference models."""
    for round_index in range(rounds):
        for batch_size in PREDICT_BATCHES:
            cells = []
            for arm in arms:
                elapsed = arm.time_predict(batch_size)
                records.append(
                    {
                        "arm": arm.name,
                        "batch_size": batch_size,
                        "round": round_index,
                        "seconds": elapsed,
                    }
                )
                cells.append(f"{arm.name}={fmt_time(elapsed)}")
            print(
                f"  predict round {round_index} b={batch_size:<5d} " + "  ".join(cells),
                flush=True,
            )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--rounds", type=int, default=N_ROUNDS)
    parser.add_argument("--devices", nargs="*", default=None)
    parser.add_argument("--out", default=None)
    args = parser.parse_args()

    requested = args.devices or DEVICES
    devices = []
    for device in requested:
        usable, reason = available(device)
        if usable:
            devices.append(device)
        else:
            print(f"{device}: skipped, {reason}")
    if not devices:
        print("no usable devices")
        return 1

    print(f"devices: {', '.join(devices)}")
    print(
        f"fit sweep {FIT_SIZES[0]}..{FIT_SIZES[-1]} ({len(FIT_SIZES)} points), "
        f"predict sweep {PREDICT_BATCHES[0]}..{PREDICT_BATCHES[-1]} "
        f"({len(PREDICT_BATCHES)} points) against a model fit at {REF_FIT_SIZE}"
    )
    print(
        f"{N_FEATURES} features, {N_CLASSES} classes, {N_COMPONENTS} components, "
        f"{DTYPE}, GridSearchCV over {len(PARAM_GRID)} parameters, cv={CV_FOLDS}"
    )
    print(f"{args.rounds} rounds, arms adjacent in time at every size")

    arms = build_arms(devices)
    print(f"arms: {', '.join(arm.name for arm in arms)}")

    started = time.perf_counter()
    for arm in arms:
        print(f"warming up {arm.name}", flush=True)
        arm.warmup()

    fit_records = []
    sweep_fit(arms, args.rounds, fit_records)

    for arm in arms:
        print(f"fitting reference model for {arm.name}", flush=True)
        arm.fit_reference()
    release_fit_data(arms)

    predict_records = []
    sweep_predict(arms, args.rounds, predict_records)

    results = {
        "host": socket.gethostname(),
        "env": describe_environment(devices),
        "config": {
            "fit_sizes": FIT_SIZES,
            "predict_batches": PREDICT_BATCHES,
            "ref_fit_size": REF_FIT_SIZE,
            "n_features": N_FEATURES,
            "n_components": N_COMPONENTS,
            "n_classes": N_CLASSES,
            "dtype": DTYPE,
            "max_iter": MAX_ITER,
            "cv_folds": CV_FOLDS,
            "param_grid": {
                key: list(map(str, value)) for key, value in PARAM_GRID.items()
            },
            "rounds": args.rounds,
            "arms": {
                arm.name: {
                    "device": arm.device,
                    "array_api_dispatch": arm.dispatch,
                    "store_fitted_as_numpy": arm.store_as_numpy,
                    "label": arm.label,
                }
                for arm in arms
            },
        },
        "fit": fit_records,
        "predict": predict_records,
    }

    out = pathlib.Path(args.out or f"pipeline_size_sweep_{socket.gethostname()}.json")
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(results, indent=2))
    print(f"\nwrote {out} in {time.perf_counter() - started:.0f}s")
    return 0


if __name__ == "__main__":
    sys.exit(main())

plot_pipeline_size_sweep.py

"""Plot the sweeps written by `bench_pipeline_size_sweep.py`.

Reads every `pipeline_size_sweep_*.json` in the current directory, so a run on
a laptop with MPS and a run on a box with CUDA end up in the same figure. Each
line is the median across rounds, with a band from the minimum to the maximum,
because two lines that nearly touch are otherwise indistinguishable from two
lines that are actually the same.

Four figures are written, `fit` and `predict` against each of the two NumPy
references: the plain no-dispatch arm, which is what a user gets today, and the
dispatch-on arm, which is the same code path the GPU arms take.

    python plot_pipeline_size_sweep.py [--glob 'pipeline_size_sweep_*.json']
"""

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause
import argparse
import glob
import json
import pathlib
import statistics
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

# One color per arm. The two NumPy arms share a color because only one of them
# appears in any given figure.
COLORS = {
    "numpy-plain": "tab:gray",
    "numpy-dispatch": "tab:gray",
    "mps-native-attrs": "tab:blue",
    "mps-numpy-attrs": "tab:cyan",
    "cuda-native-attrs": "tab:red",
    "cuda-numpy-attrs": "tab:orange",
}

# The hostname a run happened to execute on says nothing to a reader, so the
# figures name the machines instead. Insertion order here is legend order; a
# host without an alias keeps its hostname and sorts after the known ones.
HOST_ALIASES = {"thead-mlt": "macbook", "tower": "linux"}

# Hosts are told apart by marker shape as well as line style. The line style
# alone is not enough: the NumPy arms are the only ones that appear on more
# than one machine, so they share a color, and a dashed and a solid swatch are
# hard to tell apart at legend size.
HOST_STYLES = [
    {"linestyle": "-", "marker": "o"},
    {"linestyle": "--", "marker": "D"},
    {"linestyle": ":", "marker": "v"},
    {"linestyle": "-.", "marker": "P"},
]

# The two arms of a device can land on top of each other to within a pixel, so
# the native-attributes arm is drawn wide and translucent and the
# numpy-attributes arm thin and opaque on top of it. That way an exact overlap
# still reads as two lines rather than one.
STYLES = {
    "numpy-plain": {"linewidth": 1.6, "alpha": 1.0, "zorder": 3},
    "numpy-dispatch": {"linewidth": 1.6, "alpha": 1.0, "zorder": 3},
    "mps-native-attrs": {"linewidth": 3.4, "alpha": 0.45, "zorder": 2},
    "mps-numpy-attrs": {"linewidth": 1.5, "alpha": 1.0, "zorder": 4},
    "cuda-native-attrs": {"linewidth": 3.4, "alpha": 0.45, "zorder": 2},
    "cuda-numpy-attrs": {"linewidth": 1.5, "alpha": 1.0, "zorder": 4},
}
DEFAULT_STYLE = {"linewidth": 1.6, "alpha": 1.0, "zorder": 3}

GPU_ARMS = [
    "mps-native-attrs",
    "mps-numpy-attrs",
    "cuda-native-attrs",
    "cuda-numpy-attrs",
]

# (file name suffix, NumPy reference arm, how to describe it)
REFERENCES = [
    ("plain", "numpy-plain", "NumPy reference: plain scikit-learn, no dispatch"),
    ("dispatch", "numpy-dispatch", "NumPy reference: NumPy with array API dispatch"),
]

SWEEPS = {
    "fit": {
        "x_key": "n_samples",
        "x_label": "n_samples in fit",
        "title": "GridSearchCV.fit runtime vs training set size",
        "x_scale": "linear",
        "y_scale": "linear",
        "y_label": "fit time (s)",
        "y_unit": 1.0,
    },
    "predict": {
        "x_key": "batch_size",
        "x_label": "rows passed to predict",
        "title": "predict runtime vs batch size",
        "x_scale": "log",
        "y_scale": "log",
        "y_label": "predict time (ms)",
        "y_unit": 1e3,
    },
}


def load(pattern):
    """Every result file matching `pattern`, keyed by host."""
    runs = {}
    for path in sorted(glob.glob(pattern)):
        payload = json.loads(pathlib.Path(path).read_text())
        host = payload.get("host") or pathlib.Path(path).stem
        runs[host] = payload
        print(f"read {path}: host {host}, arms {sorted(payload['config']['arms'])}")
    return runs


def display_host(host):
    return HOST_ALIASES.get(host, host)


def ordered_hosts(runs):
    """Hosts in the order their aliases are declared, unknown ones last."""
    known = [host for host in HOST_ALIASES if host in runs]
    return known + sorted(host for host in runs if host not in HOST_ALIASES)


def summarize(records, arm, x_key):
    """Median, minimum and maximum across rounds, per x value.

    The bench writes one record per round so that the statistic is chosen here
    rather than baked into the measurement.
    """
    by_x = {}
    for record in records:
        if record["arm"] != arm:
            continue
        by_x.setdefault(record[x_key], []).append(record["seconds"])
    xs = sorted(by_x)
    return (
        xs,
        [statistics.median(by_x[x]) for x in xs],
        [min(by_x[x]) for x in xs],
        [max(by_x[x]) for x in xs],
    )


def y_limits(runs, sweep_name):
    """Y limits covering every arm of `sweep_name`, both references included.

    Computed once per sweep and applied to both reference variants, so that the
    plain and the dispatch figure can be flipped between without the axis
    moving underneath.
    """
    spec = SWEEPS[sweep_name]
    arms = [reference for _, reference, _ in REFERENCES] + GPU_ARMS
    low_bound = float("inf")
    high_bound = float("-inf")
    for payload in runs.values():
        for arm in arms:
            if arm not in payload["config"]["arms"]:
                continue
            xs, _, low, high = summarize(payload[sweep_name], arm, spec["x_key"])
            if not xs:
                continue
            low_bound = min(low_bound, min(low) * spec["y_unit"])
            high_bound = max(high_bound, max(high) * spec["y_unit"])
    if low_bound > high_bound:
        return None
    if spec["y_scale"] == "log":
        return low_bound / 1.2, high_bound * 1.2
    # A runtime axis is easier to read anchored at zero.
    return 0.0, high_bound * 1.05


def draw(runs, sweep_name, reference_arm, reference_note, out_path, ylim=None):
    spec = SWEEPS[sweep_name]
    arms_to_draw = [reference_arm] + GPU_ARMS

    figure, axes = plt.subplots(figsize=(9, 6))
    hosts_in_order = ordered_hosts(runs)
    host_styles = {
        host: HOST_STYLES[index % len(HOST_STYLES)]
        for index, host in enumerate(hosts_in_order)
    }

    drawn = 0
    for host in hosts_in_order:
        payload = runs[host]
        available_arms = payload["config"]["arms"]
        for arm in arms_to_draw:
            if arm not in available_arms:
                continue
            xs, median, low, high = summarize(payload[sweep_name], arm, spec["x_key"])
            if not xs:
                continue
            unit = spec["y_unit"]
            median = [value * unit for value in median]
            low = [value * unit for value in low]
            high = [value * unit for value in high]
            color = COLORS.get(arm)
            style = STYLES.get(arm, DEFAULT_STYLE)
            axes.plot(
                xs,
                median,
                label=f"{display_host(host)}: {available_arms[arm]['label']}",
                color=color,
                markersize=4.5,
                **host_styles[host],
                **style,
            )
            axes.fill_between(
                xs,
                low,
                high,
                color=color,
                alpha=0.15,
                linewidth=0,
                zorder=1,
            )
            drawn += 1

    if not drawn:
        plt.close(figure)
        print(f"nothing to draw for {out_path}")
        return None

    axes.set_xscale(spec["x_scale"])
    axes.set_yscale(spec["y_scale"])
    if ylim is not None:
        axes.set_ylim(ylim)
    axes.set_xlabel(spec["x_label"])
    axes.set_ylabel(spec["y_label"])
    axes.set_title(spec["title"])
    axes.grid(True, which="both", alpha=0.25, linewidth=0.5)
    axes.legend(fontsize=8, loc="best")

    hosts = ", ".join(
        f"{display_host(host)} ({runs[host]['env'].get('cpu', 'unknown CPU')}"
        + (
            f", {runs[host]['env']['cuda_device_used']}"
            if "cuda_device_used" in runs[host]["env"]
            else ""
        )
        + ")"
        for host in hosts_in_order
    )
    reference = runs[hosts_in_order[0]]["config"]
    grid = (
        f"{reference['n_features']} features, "
        f"{reference['n_components']} Nystroem components, {reference['dtype']}, "
        f"GridSearchCV cv={reference['cv_folds']}"
    )
    if sweep_name == "predict":
        grid += f", model fit at n={reference['ref_fit_size']}"
    figure.text(0.5, 0.955, reference_note, ha="center", fontsize=8, color="0.3")
    figure.text(
        0.5,
        0.925,
        f"{grid}. Line is the median of {reference['rounds']} rounds, "
        f"band is min to max.",
        ha="center",
        fontsize=7.5,
        color="0.45",
    )
    figure.text(
        0.5,
        0.012,
        f"hosts: {hosts}.",
        ha="center",
        fontsize=7,
        color="0.45",
    )
    figure.tight_layout(rect=(0, 0.03, 1, 0.915))
    figure.savefig(out_path, dpi=150)
    plt.close(figure)
    print(f"wrote {out_path} ({drawn} lines)")
    return out_path


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--glob", default="pipeline_size_sweep_*.json")
    parser.add_argument("--prefix", default="pipeline_sweep")
    args = parser.parse_args()

    runs = load(args.glob)
    if not runs:
        print(f"no files matched {args.glob!r}")
        return 1

    for sweep_name in SWEEPS:
        ylim = y_limits(runs, sweep_name)
        for suffix, reference_arm, note in REFERENCES:
            draw(
                runs,
                sweep_name,
                reference_arm,
                note,
                f"{args.prefix}_{sweep_name}_{suffix}.png",
                ylim=ylim,
            )
    return 0


if __name__ == "__main__":
    sys.exit(main())