From 9eb30f6ff60513d76fc7813ea37c220140f688c0 Mon Sep 17 00:00:00 2001 From: YuviPanda Date: Tue, 8 Nov 2016 17:17:10 -0800 Subject: [PATCH] Add resource limits / guarantees consistently to jupyterhub - Allows us to standardize this on the spawner base class, so there's a consistent interface for different spawners to implement this. - Specify the supported suffixes and various units we accept for memory and cpu units. - Standardize the way we expose resource limit / guarantees to single-user servers --- jupyterhub/spawner.py | 82 +++++++++++++++++++++++++++++- jupyterhub/tests/test_traitlets.py | 33 ++++++++++-- jupyterhub/traitlets.py | 44 ++++++++++++++-- 3 files changed, 150 insertions(+), 9 deletions(-) diff --git a/jupyterhub/spawner.py b/jupyterhub/spawner.py index 2b62bf44..21c29e71 100644 --- a/jupyterhub/spawner.py +++ b/jupyterhub/spawner.py @@ -24,7 +24,7 @@ from traitlets import ( validate, ) -from .traitlets import Command +from .traitlets import Command, MemorySpecification from .utils import random_port class Spawner(LoggingConfigurable): @@ -182,7 +182,74 @@ class Spawner(LoggingConfigurable): from having an effect on their server. """ ).tag(config=True) - + + mem_limit = MemorySpecification( + None, + allow_none=True, + help=""" + Maximum number of bytes a single-user server is allowed to use. + + Allows the following suffixes: + - K -> Kilobytes + - M -> Megabytes + - G -> Gigabytes + - T -> Terabytes + + If the single user server tries to allocate more memory than this, + it will fail. There is no guarantee that the single-user server + will be able to allocate this much memory - only that it can not + allocate more than this. + + This needs to be supported by your spawner for it to work. + """ + ).tag(config=True) + + cpu_limit = Float( + None, + allow_none=True, + help=""" + Maximum number of cpu-cores a single-user server is allowed to use. + + If this value is set to 0.5, allows use of 50% of one CPU. + If this value is set to 2, allows use of up to 2 CPUs. + + The single-user server will never be scheduled by the kernel to + use more cpu-cores than this. There is no guarantee that it can + access this many cpu-cores. + + This needs to be supported by your spawner for it to work. + """ + ).tag(config=True) + + mem_guarantee = MemorySpecification( + None, + allow_none=True, + help=""" + Minimum number of bytes a single-user server is guaranteed to have available. + + Allows the following suffixes: + - K -> Kilobytes + - M -> Megabytes + - G -> Gigabytes + - T -> Terabytes + + This needs to be supported by your spawner for it to work. + """ + ).tag(config=True) + + cpu_guarantee = Float( + None, + allow_none=True, + help=""" + Maximum number of cpu-cores a single-user server is allowed to use. + + If this value is set to 0.5, allows use of 50% of one CPU. + If this value is set to 2, allows use of up to 2 CPUs. + + Note that this needs to be supported by your spawner for it to work. + """ + ).tag(config=True) + def __init__(self, **kwargs): super(Spawner, self).__init__(**kwargs) if self.user.state: @@ -255,6 +322,17 @@ class Spawner(LoggingConfigurable): env[key] = value env['JPY_API_TOKEN'] = self.api_token + + # Put in limit and guarantee info if they exist. + if self.mem_limit: + env['LIMIT_MEM'] = str(self.mem_limit) + if self.mem_guarantee: + env['GUARANTEE_MEM'] = str(self.mem_guarantee) + if self.cpu_limit: + env['LIMIT_CPU'] = str(self.cpu_limit) + if self.cpu_guarantee: + env['GUARANTEE_CPU'] = str(self.cpu_guarantee) + return env def template_namespace(self): diff --git a/jupyterhub/tests/test_traitlets.py b/jupyterhub/tests/test_traitlets.py index 73eef270..2a920ddf 100644 --- a/jupyterhub/tests/test_traitlets.py +++ b/jupyterhub/tests/test_traitlets.py @@ -1,11 +1,12 @@ -from traitlets import HasTraits +import pytest +from traitlets import HasTraits, TraitError + +from jupyterhub.traitlets import URLPrefix, Command, MemorySpecification -from jupyterhub.traitlets import URLPrefix, Command def test_url_prefix(): class C(HasTraits): url = URLPrefix() - c = C() c.url = '/a/b/c/' assert c.url == '/a/b/c/' @@ -14,14 +15,38 @@ def test_url_prefix(): c.url = 'a/b/c/d' assert c.url == '/a/b/c/d/' + def test_command(): class C(HasTraits): cmd = Command('default command') cmd2 = Command(['default_cmd']) - c = C() assert c.cmd == ['default command'] assert c.cmd2 == ['default_cmd'] c.cmd = 'foo bar' assert c.cmd == ['foo bar'] + +def test_memoryspec(): + class C(HasTraits): + mem = MemorySpecification() + + c = C() + + c.mem = 1024 + assert c.mem == 1024 + + c.mem = '1024K' + assert c.mem == 1024 * 1024 + + c.mem = '1024M' + assert c.mem == 1024 * 1024 * 1024 + + c.mem = '1024G' + assert c.mem == 1024 * 1024 * 1024 * 1024 + + c.mem = '1024T' + assert c.mem == 1024 * 1024 * 1024 * 1024 * 1024 + + with pytest.raises(TraitError): + c.mem = '1024Gi' diff --git a/jupyterhub/traitlets.py b/jupyterhub/traitlets.py index 2daa9425..1f4fd7a8 100644 --- a/jupyterhub/traitlets.py +++ b/jupyterhub/traitlets.py @@ -1,8 +1,10 @@ -"""extra traitlets""" +""" +Traitlets that are used in JupyterHub +""" # Copyright (c) Jupyter Development Team. # Distributed under the terms of the Modified BSD License. -from traitlets import List, Unicode +from traitlets import List, Unicode, Integer, TraitError class URLPrefix(Unicode): def validate(self, obj, value): @@ -22,8 +24,44 @@ class Command(List): if isinstance(default_value, str): default_value = [default_value] super().__init__(Unicode(), default_value, **kwargs) - + def validate(self, obj, value): if isinstance(value, str): value = [value] return super().validate(obj, value) + + +class MemorySpecification(Integer): + """ + Allow easily specifying memory in units of 1024 with suffixes + + Suffixes allowed are: + - K -> Kilobyte + - M -> Megabyte + - G -> Gigabyte + - T -> Terabyte + """ + + UNIT_SUFFIXES = { + 'K': 1024, + 'M': 1024 * 1024, + 'G': 1024 * 1024 * 1024, + 'T': 1024 * 1024 * 1024 * 1024 + } + + def validate(self, obj, value): + """ + Validate that the passed in value is a valid memory specification + + It could either be a pure int, when it is taken as a byte value. + If it has one of the suffixes, it is converted into the appropriate + pure byte value. + """ + if isinstance(value, int): + return value + num = value[:-1] + suffix = value[-1] + if not num.isdigit() and suffix not in MemorySpecification.UNIT_SUFFIXES: + raise TraitError('{val} is not a valid memory specification. Must be an int or a string with suffix K, M, G, T'.format(val=value)) + else: + return int(num) * MemorySpecification.UNIT_SUFFIXES[suffix]