measure-preload-content

#!/usr/bin/env python3
"""Measure the cost of _preload_content=True across kubernetes versions

This has grown over time, and is *super* expensive with kubernetes 12
and lots of pods.

CSV output is produced on stdout, so call with e.g.

    ./measure-preload-content | tee measurements.csv
"""
import json
import os
import sys
import time
import timeit
from functools import lru_cache, partial
from subprocess import run
from tempfile import TemporaryDirectory

kube_ctx = "prod"
kube_ns = "prod"


@lru_cache()
def kube_api():
    """Cached function to return a single kube API client"""
    import kubernetes.client
    import kubernetes.config

    kubernetes.config.load_kube_config(context=kube_ctx)
    return kubernetes.client.CoreV1Api()


def list_pods(_preload_content, label_selector):
    """List pods, return the number of pods returned"""
    kube = kube_api()
    r = kube.list_namespaced_pod(
        kube_ns,
        label_selector=label_selector,
        _preload_content=_preload_content,
    )
    if _preload_content:
        return len(r.items)
    else:
        return len(json.loads(r.read())["items"])


def list_user_pods(preload):
    """List user pods in a jupyterhub deployment (lots)"""
    return list_pods(
        _preload_content=preload,
        label_selector="app=jupyterhub,component=singleuser-server",
    )


def list_placeholder_pods(preload):
    """List placeholder pods in a jupyterhub deployment (several, fewer than users)"""
    return list_pods(
        _preload_content=preload,
        label_selector="app=jupyterhub,component=user-placeholder",
    )


def list_build_pods(preload):
    """List build pods in a binderhub deployment (a few)"""
    return list_pods(
        _preload_content=preload,
        label_selector="component=binderhub-build",
    )


def list_hub_pods(preload):
    """List hub pods in a jupyterhub deployment (always 1)"""
    return list_pods(
        _preload_content=preload,
        label_selector="app=jupyterhub,component=hub",
    )


def measure(f, repeat=7):
    """Measure how long it takes to call f()

    - collects `repeat` data points
    - For each data point, run f enough times
      that each measurement takes at least 200ms,
      but at least twice.
      (lifted from IPython's %timeit)

    Returns:
      (number, times, counts):
          number of calls per measurement
          list of times (in seconds) per f() call): (measured time / number)
          and number of pods returned for each call
    """
    tic = time.perf_counter()
    f()
    toc = time.perf_counter()
    # always run it at least 2 times
    number = 2
    t = toc - tic
    # make sure we take at least 200ms
    # lifted from IPython's %timeit
    while t * number < 0.2:
        number *= 10
    times = []
    counts = []
    print(f"Measuring {f} with r={repeat},n={number}", file=sys.stderr)
    for i in range(repeat):
        tic = time.perf_counter()
        count = 0
        for j in range(number):
            count = max(f(), count)
        toc = time.perf_counter()
        times.append((toc - tic) / number)
        counts.append(count)
    return number, times, counts


def run_one():
    """Run one collection of data

    This function is called within the test environment
    """
    import kubernetes

    # ensure cached kube client is created
    kube = kube_api()

    for list_f in [list_user_pods, list_placeholder_pods, list_hub_pods]:
        for preload in (True, False):
            number, times, counts = measure(partial(list_f, preload))
            for t, c in zip(times, counts):
                print(
                    f"{list_f.__name__},{preload},{kubernetes.__version__},{number},{t},{c}"
                )


def setup_env(kube_version, td):
    """Set up a virtual environment in which to run a test"""
    env_dir = os.path.join(td, f"kube-{kube_version}")
    run([sys.executable, "-m", "venv", env_dir], stdout=sys.stderr)
    run(
        [
            os.path.join(env_dir, "bin", "python3"),
            "-mpip",
            "install",
            f"kubernetes=={kube_version}.*",
        ],
        stdout=sys.stderr,
    )
    return env_dir


def run_all():
    """Run all the data collection

    - create a test environment for each major kubernetes version
    - run `run_one()` in that env
    """
    with TemporaryDirectory() as td:
        print("fname,preload,kube_version,number,per_call,pods")
        for kube_version in [7, 8, 9, 10, 11, 12]:
            env = setup_env(kube_version, td)
            run([os.path.join(env, "bin", "python3"), "-u", __file__, "one"])


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "command",
        help="command to run",
        choices=["one", "all"],
        default="all",
        nargs="?",
    )
    opts = parser.parse_args()
    if opts.command == "all":
        run_all()
    elif opts.command == "one":
        run_one()

measurements.csv

preload-content.ipynb