Move pod quota checking into separate class

binderhub/app.py

@@ -50,6 +50,7 @@
 from .log import log_request
 from .main import LegacyRedirectHandler, MainHandler, ParameterizedMainHandler
 from .metrics import MetricsHandler
+from .quota import KubernetesLaunchQuota, LaunchQuota
 from .ratelimit import RateLimiter
 from .registry import DockerRegistry
 from .repoproviders import (
@@ -304,6 +305,8 @@ def _valid_badge_base_url(self, proposal):
     pod_quota = Integer(
         None,
         help="""
+        DEPRECATED: Use c.LaunchQuota.total_quota
+
         The number of concurrent pods this hub has been designed to support.
 
         This quota is used as an indication for how much above or below the
@@ -319,6 +322,13 @@ def _valid_badge_base_url(self, proposal):
         config=True,
     )
 
+    @observe("pod_quota")
+    def _pod_quota_deprecated(self, change):
+        self.log.warning(
+            "BinderHub.pod_quota is deprecated, use LaunchQuota.total_quota"
+        )
+        self.config.LaunchQuota.total_quota = change.new
+
     per_repo_quota_higher = Integer(
         0,
         help="""
@@ -333,6 +343,17 @@ def _valid_badge_base_url(self, proposal):
         config=True,
     )
 
+    launch_quota_class = Type(
+        LaunchQuota,
+        default=KubernetesLaunchQuota,
+        help="""
+        The class used to check quotas for launched servers.
+
+        Must inherit from binderhub.quota.LaunchQuota
+        """,
+        config=True,
+    )
+
     log_tail_lines = Integer(
         100,
         help="""
@@ -791,6 +812,8 @@ def initialize(self, *args, **kwargs):
             with open(schema_file) as f:
                 self.event_log.register_schema(json.load(f))
 
+        launch_quota = self.launch_quota_class(parent=self, executor=self.executor)
+
         self.tornado_settings.update(
             {
                 "log_function": log_request,
@@ -814,6 +837,7 @@ def initialize(self, *args, **kwargs):
                 "per_repo_quota": self.per_repo_quota,
                 "per_repo_quota_higher": self.per_repo_quota_higher,
                 "repo_providers": self.repo_providers,
+                "launch_quota": launch_quota,
                 "rate_limiter": RateLimiter(parent=self),
                 "use_registry": self.use_registry,
                 "build_class": self.build_class,

binderhub/builder.py

@@ -2,7 +2,6 @@
 Handlers for working with version control services (i.e. GitHub) for builds.
 """
 
-import asyncio
 import hashlib
 import json
 import re
@@ -23,7 +22,7 @@
 
 from .base import BaseHandler
 from .build import Build, ProgressEvent
-from .utils import KUBE_REQUEST_TIMEOUT
+from .quota import LaunchQuotaExceeded
 
 # Separate buckets for builds and launches.
 # Builds and launches have very different characteristic times,
@@ -586,73 +585,29 @@ async def launch(self, provider):
         # Load the spec-specific configuration if it has been overridden
         repo_config = provider.repo_config(self.settings)
 
-        # the image name (without tag) is unique per repo
-        # use this to count the number of pods running with a given repo
-        # if we added annotations/labels with the repo name via KubeSpawner
-        # we could do this better
-        image_no_tag = self.image_name.rsplit(":", 1)[0]
-
-        # TODO: put busy users in a queue rather than fail?
-        # That would be hard to do without in-memory state.
-        repo_quota = repo_config.get("quota")
-        pod_quota = self.settings["pod_quota"]
-        if pod_quota is not None or repo_quota:
-            # Fetch info on currently running users *only* if quotas are set
-            matching_pods = 0
-
-            # TODO: run a watch to keep this up to date in the background
-            f = self.settings["executor"].submit(
-                self.settings["kubernetes_client"].list_namespaced_pod,
-                self.settings["build_namespace"],
-                label_selector="app=jupyterhub,component=singleuser-server",
-                _request_timeout=KUBE_REQUEST_TIMEOUT,
-                _preload_content=False,
+        launch_quota = self.settings["launch_quota"]
+        try:
+            quota_check = await launch_quota.check_repo_quota(
+                self.image_name, repo_config, self.repo_url
             )
-            resp = await asyncio.wrap_future(f)
-            pods = json.loads(resp.read())["items"]
-            total_pods = len(pods)
-
-            if pod_quota is not None and total_pods >= pod_quota:
-                # check overall quota first
-                LAUNCH_COUNT.labels(
-                    status="pod_quota",
-                    **self.repo_metric_labels,
-                ).inc()
-                app_log.error(f"BinderHub is full: {total_pods}/{pod_quota}")
-                await self.fail("Too many users on this BinderHub! Try again soon.")
-                return
-
-            for pod in pods:
-                for container in pod["spec"]["containers"]:
-                    # is the container running the same image as us?
-                    # if so, count one for the current repo.
-                    image = container["image"].rsplit(":", 1)[0]
-                    if image == image_no_tag:
-                        matching_pods += 1
-                        break
-
-            if repo_quota and matching_pods >= repo_quota:
-                LAUNCH_COUNT.labels(
-                    status="repo_quota",
-                    **self.repo_metric_labels,
-                ).inc()
-                app_log.error(
-                    f"{self.repo_url} has exceeded quota: {matching_pods}/{repo_quota} ({total_pods} total)"
-                )
-                await self.fail(
-                    f"Too many users running {self.repo_url}! Try again soon."
-                )
-                return
+        except LaunchQuotaExceeded as e:
+            LAUNCH_COUNT.labels(
+                status=e.status,
+                **self.repo_metric_labels,
+            ).inc()
+            await self.fail(e.message)
+            return
 
-            if matching_pods >= 0.5 * repo_quota:
+        if quota_check:
+            if quota_check.matching >= 0.5 * quota_check.quota:
                 log = app_log.warning
             else:
                 log = app_log.info
             log(
-                "Launching pod for %s: %s other pods running this repo (%s total)",
+                "Launching server for %s: %s other servers running this repo (%s total)",
                 self.repo_url,
-                matching_pods,
-                total_pods,
+                quota_check.matching,
+                quota_check.total,
             )
 
         await self.emit(

binderhub/quota.py

@@ -0,0 +1,161 @@
+"""
+Singleuser server quotas
+"""
+
+import asyncio
+import json
+import os
+from collections import namedtuple
+
+import kubernetes.config
+from kubernetes import client
+from traitlets import Any, Integer, Unicode, default
+from traitlets.config import LoggingConfigurable
+
+from .utils import KUBE_REQUEST_TIMEOUT
+
+
+class LaunchQuotaExceeded(Exception):
+    """Raised when a quota will be exceeded by a launch"""
+
+    def __init__(self, message, *, quota, used, status):
+        """
+        message: User-facing message
+        quota: Quota limit
+        used: Quota used
+        status: String indicating the type of quota
+        """
+        super().__init__()
+        self.message = message
+        self.quota = quota
+        self.used = used
+        self.status = status
+
+
+ServerQuotaCheck = namedtuple("ServerQuotaCheck", ["total", "matching", "quota"])
+
+
+class LaunchQuota(LoggingConfigurable):
+
+    executor = Any(
+        allow_none=True, help="Optional Executor to use for blocking operations"
+    )
+
+    total_quota = Integer(
+        None,
+        help="""
+        The number of concurrent singleuser servers that can be run.
+
+        None: no quota
+        0: the hub can't run any singleuser servers (e.g. in maintenance mode)
+        Positive integer: sets the quota
+        """,
+        allow_none=True,
+        config=True,
+    )
+
+    async def check_repo_quota(self, image_name, repo_config, repo_url):
+        """
+        Check whether launching a repository would exceed a quota.
+
+        Parameters
+        ----------
+        image_name: str
+        repo_config: dict
+        repo_url: str
+
+        Returns
+        -------
+        If quotas are disabled returns None
+        If quotas are exceeded raises LaunchQuotaExceeded
+        Otherwise returns:
+          - total servers
+          - matching servers running image_name
+          - quota
+        """
+        return None
+
+
+class KubernetesLaunchQuota(LaunchQuota):
+
+    api = Any(
+        help="Kubernetes API object to make requests (kubernetes.client.CoreV1Api())",
+    )
+
+    @default("api")
+    def _default_api(self):
+        try:
+            kubernetes.config.load_incluster_config()
+        except kubernetes.config.ConfigException:
+            kubernetes.config.load_kube_config()
+        return client.CoreV1Api()
+
+    namespace = Unicode(help="Kubernetes namespace to check", config=True)
+
+    @default("namespace")
+    def _default_namespace(self):
+        return os.getenv("BUILD_NAMESPACE", "default")
+
+    async def check_repo_quota(self, image_name, repo_config, repo_url):
+        # the image name (without tag) is unique per repo
+        # use this to count the number of pods running with a given repo
+        # if we added annotations/labels with the repo name via KubeSpawner
+        # we could do this better
+        image_no_tag = image_name.rsplit(":", 1)[0]
+
+        # TODO: put busy users in a queue rather than fail?
+        # That would be hard to do without in-memory state.
+        repo_quota = repo_config.get("quota")
+        pod_quota = self.total_quota
+
+        # Fetch info on currently running users *only* if quotas are set
+        if pod_quota is not None or repo_quota:
+            matching_pods = 0
+
+            # TODO: run a watch to keep this up to date in the background
+            f = self.executor.submit(
+                self.api.list_namespaced_pod,
+                self.namespace,
+                label_selector="app=jupyterhub,component=singleuser-server",
+                _request_timeout=KUBE_REQUEST_TIMEOUT,
+                _preload_content=False,
+            )
+            resp = await asyncio.wrap_future(f)
+            pods = json.loads(resp.read())["items"]
+            total_pods = len(pods)
+
+            if pod_quota is not None and total_pods >= pod_quota:
+                # check overall quota first
+                self.log.error(f"BinderHub is full: {total_pods}/{pod_quota}")
+                raise LaunchQuotaExceeded(
+                    "Too many users on this BinderHub! Try again soon.",
+                    quota=pod_quota,
+                    used=total_pods,
+                    status="pod_quota",
+                )
+
+            for pod in pods:
+                for container in pod["spec"]["containers"]:
+                    # is the container running the same image as us?
+                    # if so, count one for the current repo.
+                    image = container["image"].rsplit(":", 1)[0]
+                    if image == image_no_tag:
+                        matching_pods += 1
+                        break
+
+            if repo_quota and matching_pods >= repo_quota:
+                self.log.error(
+                    f"{repo_url} has exceeded quota: {matching_pods}/{repo_quota} ({total_pods} total)"
+                )
+                raise LaunchQuotaExceeded(
+                    f"Too many users running {repo_url}! Try again soon.",
+                    quota=repo_quota,
+                    used=matching_pods,
+                    status="repo_quota",
+                )
+
+            return ServerQuotaCheck(
+                total=total_pods, matching=matching_pods, quota=repo_quota
+            )
+
+        return None