Quota Policies

check-quotas.py

#!/usr/bin/env python
import subprocess
import json
import argparse
import math
import sys

NFS_TEMPLATE = """
jupyterhub-home-nfs:
  quotaEnforcer:
    extraConfig:
      secret-quota-overrides: |
        c.QuotaManager.quota_overrides = {{
{body}
        }}
"""

MiB = 1024**2
GiB = MiB * 1024

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("cluster")
    parser.add_argument("hub")
    parser.add_argument("-q", "--quota-gib", default=10, type=int)
    parser.add_argument("-s", "--shared-quota-gib", default=100, type=int)
    parser.add_argument("-b", "--buffer-gib", default=2, type=int)

    parser.add_argument("--output", choices=["pretty", "json", "template"])
    parser.add_argument("--shrink", action="store_true")
    parser.add_argument("--shared-quota", action="store_true")
    args = parser.parse_args()

    quota_bytes = args.quota_gib * GiB
    buffer_bytes = args.buffer_gib * GiB
    # If we're shrinking the quota, any value thats below the margin needs to be excepted
    # Otherwise, users already know about this'
    margin_bytes = buffer_bytes if args.shrink else 0

    quota_usage = json.loads(
        subprocess.check_output(
            ["deployer", "exec", "get-quota-usage", args.cluster, args.hub, "--json"],
            text=True,
        )
    )

    user_overrides_gib = {}
    for name, usage_bytes in quota_usage.items():
        if usage_bytes + margin_bytes <= quota_bytes:
            continue
        user_overrides_gib[name] = math.ceil((usage_bytes + buffer_bytes) / GiB)

    for key in "_shared", "_shared-public", "_shared-readwrite":
        if key in quota_usage and not args.shared_quota:
            # Update shared, taking largest of default shared-quota or the computed value)
            user_overrides_gib[key] = max(
                args.shared_quota_gib, user_overrides_gib.get(key, -1)
            )

    match args.output:
        case "json":
            json.dump(user_overrides_gib, sys.stdout, indent=4, sort_keys=True)
        case "pretty":
            for key, value in user_overrides_gib.items():
                current_value = math.ceil(quota_usage[key] / GiB)
                print(f"{key}:{value}GiB (current {current_value}GiB)")
            print(f"\n\n{len(user_overrides_gib) / len(quota_usage):.0%}")
        case "template":
            body_lines = [
                f"          {name!r}: {value!r}"
                for name, value in user_overrides_gib.items()
            ]
            body = ",\n".join(body_lines)
            print(NFS_TEMPLATE.format(body=body))

format-changes.sh

wl-paste | tr "'" '"' | jq '. | to_entries[] | "\(.key)[\(.value)GiB], "' -r | sd "\[0GiB\]" "[NA]" | xargs echo

gen-checks.py

python -c '
import yaml
import sys

with open(sys.argv[1], "r") as f:
  data = yaml.safe_load(f.read())
cluster = data["name"]
for line in data["hubs"]:
        hub = line["name"]
        print(f"check-quotas {cluster} {hub} --output template -q 10")
'  2>/dev/null cluster.yaml