kube_orphans.py

"""
Check for orphan single-user pods

Compares JupyterHub API list of running servers
to list of running pods in kubernetes
in order to identify discrepancies

see https://github.com/jupyterhub/kubespawner/pull/742
"""

import json
import os
import shlex
from subprocess import run
from urllib.parse import urlencode

import requests

hub_api_token = os.getenv("JUPYTERHUB_API_TOKEN")


def get_running_servers(hub_url):
    """
    Get running users from jupyterhub API
    """
    hub_url = hub_url.rstrip("/")
    users_url = hub_url + "/hub/api/users"
    s = requests.Session()
    s.headers = {
        "Authorization": f"Bearer {hub_api_token}",
        "Accept": "application/jupyterhub-pagination+json",
    }
    running = {}
    params = {"state": "active"}
    next_params = {"offset": "0"}
    while next_params:
        params.update(next_params)
        r = s.get(users_url + "?" + urlencode(params))
        r.raise_for_status()
        page = r.json()
        for user in page["items"]:
            for server_name, server in user["servers"].items():
                running[f"{user['name']}/{server_name}"] = server
        next_params = page["_pagination"]["next"]
    return running


def get_user_pods(namespace=None):
    cmd = [
        "kubectl",
        "get",
        "pod",
        "-o",
        "json",
        "-l",
        "component=singleuser-server",
    ]
    if namespace:
        cmd.extend(["--namespace", namespace])

    p = run(cmd, capture_output=True, check=True, text=True)
    pods = json.loads(p.stdout)["items"]
    user_pods = {}
    for pod in pods:
        annotations = pod["metadata"]["annotations"]

        username = annotations['hub.jupyter.org/username']
        servername = annotations.get('hub.jupyter.org/servername', "")
        key = f"{username}/{servername}"
        user_pods[key] = pod
    return user_pods


def main():
    if not hub_api_token:
        raise ValueError("Need $JUPYTERHUB_API_TOKEN")

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("hub_url", help="URL of your Hub")
    parser.add_argument(
        "--namespace",
        "-n",
        default=None,
        help="Namespace for your Hub (if not default)",
    )
    args = parser.parse_args()
    user_pods = get_user_pods(args.namespace)
    servers = get_running_servers(args.hub_url)
    orphan_pods = set(user_pods).difference(servers)
    server_not_found = set(servers).difference(user_pods)
    print(f"Found {len(servers)} active servers (according to JupyterHub)")
    print(f"Found {len(user_pods)} active pods (according to kubernetes)")
    for server in server_not_found:
        print(
            f"Found no pod for {server}. This is suspicious, don't trust the orphan output!"
        )

    pod_names = []
    for server_name in orphan_pods:
        pod = user_pods[server_name]
        pod_name = pod["metadata"]["name"]
        pod_names.append(pod_name)
        print(f"Found orphaned pod {pod_name} for {server_name}")

    print(
        "Delete command to remove orphan pods (double check that these should be deleted!!!):"
    )
    kubectl = ["kubectl"]
    if args.namespace:
        kubectl += ["--namespace", args.namespace]
    print(
        "    " + shlex.join(
            kubectl + [
                "delete",
                "pod",
            ]
            + pod_names
        )
    )


if __name__ == "__main__":
    main()