Skip to content

Commit 969f2f6

Browse files
authored
Merge pull request #1526 from manics/quota-no-k8s
Move pod quota checking into separate class
2 parents b886b2f + ce33719 commit 969f2f6

File tree

5 files changed

+297
-62
lines changed

5 files changed

+297
-62
lines changed

binderhub/app.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
from .log import log_request
5151
from .main import LegacyRedirectHandler, MainHandler, ParameterizedMainHandler
5252
from .metrics import MetricsHandler
53+
from .quota import KubernetesLaunchQuota, LaunchQuota
5354
from .ratelimit import RateLimiter
5455
from .registry import DockerRegistry
5556
from .repoproviders import (
@@ -304,6 +305,8 @@ def _valid_badge_base_url(self, proposal):
304305
pod_quota = Integer(
305306
None,
306307
help="""
308+
DEPRECATED: Use c.LaunchQuota.total_quota
309+
307310
The number of concurrent pods this hub has been designed to support.
308311
309312
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):
319322
config=True,
320323
)
321324

325+
@observe("pod_quota")
326+
def _pod_quota_deprecated(self, change):
327+
self.log.warning(
328+
"BinderHub.pod_quota is deprecated, use LaunchQuota.total_quota"
329+
)
330+
self.config.LaunchQuota.total_quota = change.new
331+
322332
per_repo_quota_higher = Integer(
323333
0,
324334
help="""
@@ -333,6 +343,17 @@ def _valid_badge_base_url(self, proposal):
333343
config=True,
334344
)
335345

346+
launch_quota_class = Type(
347+
LaunchQuota,
348+
default=KubernetesLaunchQuota,
349+
help="""
350+
The class used to check quotas for launched servers.
351+
352+
Must inherit from binderhub.quota.LaunchQuota
353+
""",
354+
config=True,
355+
)
356+
336357
log_tail_lines = Integer(
337358
100,
338359
help="""
@@ -791,6 +812,8 @@ def initialize(self, *args, **kwargs):
791812
with open(schema_file) as f:
792813
self.event_log.register_schema(json.load(f))
793814

815+
launch_quota = self.launch_quota_class(parent=self, executor=self.executor)
816+
794817
self.tornado_settings.update(
795818
{
796819
"log_function": log_request,
@@ -814,6 +837,7 @@ def initialize(self, *args, **kwargs):
814837
"per_repo_quota": self.per_repo_quota,
815838
"per_repo_quota_higher": self.per_repo_quota_higher,
816839
"repo_providers": self.repo_providers,
840+
"launch_quota": launch_quota,
817841
"rate_limiter": RateLimiter(parent=self),
818842
"use_registry": self.use_registry,
819843
"build_class": self.build_class,

binderhub/builder.py

Lines changed: 17 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
Handlers for working with version control services (i.e. GitHub) for builds.
33
"""
44

5-
import asyncio
65
import hashlib
76
import json
87
import re
@@ -23,7 +22,7 @@
2322

2423
from .base import BaseHandler
2524
from .build import Build, ProgressEvent
26-
from .utils import KUBE_REQUEST_TIMEOUT
25+
from .quota import LaunchQuotaExceeded
2726

2827
# Separate buckets for builds and launches.
2928
# Builds and launches have very different characteristic times,
@@ -591,73 +590,29 @@ async def launch(self, provider):
591590
# Load the spec-specific configuration if it has been overridden
592591
repo_config = provider.repo_config(self.settings)
593592

594-
# the image name (without tag) is unique per repo
595-
# use this to count the number of pods running with a given repo
596-
# if we added annotations/labels with the repo name via KubeSpawner
597-
# we could do this better
598-
image_no_tag = self.image_name.rsplit(":", 1)[0]
599-
600-
# TODO: put busy users in a queue rather than fail?
601-
# That would be hard to do without in-memory state.
602-
repo_quota = repo_config.get("quota")
603-
pod_quota = self.settings["pod_quota"]
604-
if pod_quota is not None or repo_quota:
605-
# Fetch info on currently running users *only* if quotas are set
606-
matching_pods = 0
607-
608-
# TODO: run a watch to keep this up to date in the background
609-
f = self.settings["executor"].submit(
610-
self.settings["kubernetes_client"].list_namespaced_pod,
611-
self.settings["build_namespace"],
612-
label_selector="app=jupyterhub,component=singleuser-server",
613-
_request_timeout=KUBE_REQUEST_TIMEOUT,
614-
_preload_content=False,
593+
launch_quota = self.settings["launch_quota"]
594+
try:
595+
quota_check = await launch_quota.check_repo_quota(
596+
self.image_name, repo_config, self.repo_url
615597
)
616-
resp = await asyncio.wrap_future(f)
617-
pods = json.loads(resp.read())["items"]
618-
total_pods = len(pods)
619-
620-
if pod_quota is not None and total_pods >= pod_quota:
621-
# check overall quota first
622-
LAUNCH_COUNT.labels(
623-
status="pod_quota",
624-
**self.repo_metric_labels,
625-
).inc()
626-
app_log.error(f"BinderHub is full: {total_pods}/{pod_quota}")
627-
await self.fail("Too many users on this BinderHub! Try again soon.")
628-
return
629-
630-
for pod in pods:
631-
for container in pod["spec"]["containers"]:
632-
# is the container running the same image as us?
633-
# if so, count one for the current repo.
634-
image = container["image"].rsplit(":", 1)[0]
635-
if image == image_no_tag:
636-
matching_pods += 1
637-
break
638-
639-
if repo_quota and matching_pods >= repo_quota:
640-
LAUNCH_COUNT.labels(
641-
status="repo_quota",
642-
**self.repo_metric_labels,
643-
).inc()
644-
app_log.error(
645-
f"{self.repo_url} has exceeded quota: {matching_pods}/{repo_quota} ({total_pods} total)"
646-
)
647-
await self.fail(
648-
f"Too many users running {self.repo_url}! Try again soon."
649-
)
650-
return
598+
except LaunchQuotaExceeded as e:
599+
LAUNCH_COUNT.labels(
600+
status=e.status,
601+
**self.repo_metric_labels,
602+
).inc()
603+
await self.fail(e.message)
604+
return
651605

652-
if matching_pods >= 0.5 * repo_quota:
606+
if quota_check:
607+
if quota_check.matching >= 0.5 * quota_check.quota:
653608
log = app_log.warning
654609
else:
655610
log = app_log.info
656611
log(
657-
"Launching pod for %s: %s other pods running this repo (%s total)",
612+
"Launching server for %s: %s other servers running this repo (%s total)",
658613
self.repo_url,
659-
matching_pods,
660-
total_pods,
614+
quota_check.matching,
615+
quota_check.total,
661616
)
662617

663618
await self.emit(

binderhub/quota.py

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
"""
2+
Singleuser server quotas
3+
"""
4+
5+
import asyncio
6+
import json
7+
import os
8+
from collections import namedtuple
9+
10+
import kubernetes.config
11+
from kubernetes import client
12+
from traitlets import Any, Integer, Unicode, default
13+
from traitlets.config import LoggingConfigurable
14+
15+
from .utils import KUBE_REQUEST_TIMEOUT
16+
17+
18+
class LaunchQuotaExceeded(Exception):
19+
"""Raised when a quota will be exceeded by a launch"""
20+
21+
def __init__(self, message, *, quota, used, status):
22+
"""
23+
message: User-facing message
24+
quota: Quota limit
25+
used: Quota used
26+
status: String indicating the type of quota
27+
"""
28+
super().__init__()
29+
self.message = message
30+
self.quota = quota
31+
self.used = used
32+
self.status = status
33+
34+
35+
ServerQuotaCheck = namedtuple("ServerQuotaCheck", ["total", "matching", "quota"])
36+
37+
38+
class LaunchQuota(LoggingConfigurable):
39+
40+
executor = Any(
41+
allow_none=True, help="Optional Executor to use for blocking operations"
42+
)
43+
44+
total_quota = Integer(
45+
None,
46+
help="""
47+
The number of concurrent singleuser servers that can be run.
48+
49+
None: no quota
50+
0: the hub can't run any singleuser servers (e.g. in maintenance mode)
51+
Positive integer: sets the quota
52+
""",
53+
allow_none=True,
54+
config=True,
55+
)
56+
57+
async def check_repo_quota(self, image_name, repo_config, repo_url):
58+
"""
59+
Check whether launching a repository would exceed a quota.
60+
61+
Parameters
62+
----------
63+
image_name: str
64+
repo_config: dict
65+
repo_url: str
66+
67+
Returns
68+
-------
69+
If quotas are disabled returns None
70+
If quotas are exceeded raises LaunchQuotaExceeded
71+
Otherwise returns:
72+
- total servers
73+
- matching servers running image_name
74+
- quota
75+
"""
76+
return None
77+
78+
79+
class KubernetesLaunchQuota(LaunchQuota):
80+
81+
api = Any(
82+
help="Kubernetes API object to make requests (kubernetes.client.CoreV1Api())",
83+
)
84+
85+
@default("api")
86+
def _default_api(self):
87+
try:
88+
kubernetes.config.load_incluster_config()
89+
except kubernetes.config.ConfigException:
90+
kubernetes.config.load_kube_config()
91+
return client.CoreV1Api()
92+
93+
namespace = Unicode(help="Kubernetes namespace to check", config=True)
94+
95+
@default("namespace")
96+
def _default_namespace(self):
97+
return os.getenv("BUILD_NAMESPACE", "default")
98+
99+
async def check_repo_quota(self, image_name, repo_config, repo_url):
100+
# the image name (without tag) is unique per repo
101+
# use this to count the number of pods running with a given repo
102+
# if we added annotations/labels with the repo name via KubeSpawner
103+
# we could do this better
104+
image_no_tag = image_name.rsplit(":", 1)[0]
105+
106+
# TODO: put busy users in a queue rather than fail?
107+
# That would be hard to do without in-memory state.
108+
repo_quota = repo_config.get("quota")
109+
pod_quota = self.total_quota
110+
111+
# Fetch info on currently running users *only* if quotas are set
112+
if pod_quota is not None or repo_quota:
113+
matching_pods = 0
114+
115+
# TODO: run a watch to keep this up to date in the background
116+
f = self.executor.submit(
117+
self.api.list_namespaced_pod,
118+
self.namespace,
119+
label_selector="app=jupyterhub,component=singleuser-server",
120+
_request_timeout=KUBE_REQUEST_TIMEOUT,
121+
_preload_content=False,
122+
)
123+
resp = await asyncio.wrap_future(f)
124+
pods = json.loads(resp.read())["items"]
125+
total_pods = len(pods)
126+
127+
if pod_quota is not None and total_pods >= pod_quota:
128+
# check overall quota first
129+
self.log.error(f"BinderHub is full: {total_pods}/{pod_quota}")
130+
raise LaunchQuotaExceeded(
131+
"Too many users on this BinderHub! Try again soon.",
132+
quota=pod_quota,
133+
used=total_pods,
134+
status="pod_quota",
135+
)
136+
137+
for pod in pods:
138+
for container in pod["spec"]["containers"]:
139+
# is the container running the same image as us?
140+
# if so, count one for the current repo.
141+
image = container["image"].rsplit(":", 1)[0]
142+
if image == image_no_tag:
143+
matching_pods += 1
144+
break
145+
146+
if repo_quota and matching_pods >= repo_quota:
147+
self.log.error(
148+
f"{repo_url} has exceeded quota: {matching_pods}/{repo_quota} ({total_pods} total)"
149+
)
150+
raise LaunchQuotaExceeded(
151+
f"Too many users running {repo_url}! Try again soon.",
152+
quota=repo_quota,
153+
used=matching_pods,
154+
status="repo_quota",
155+
)
156+
157+
return ServerQuotaCheck(
158+
total=total_pods, matching=matching_pods, quota=repo_quota
159+
)
160+
161+
return None

0 commit comments

Comments
 (0)