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

@@ -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):
@@ -183,6 +183,73 @@ class Spawner(LoggingConfigurable):
         """
     ).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):

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'

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):
@@ -27,3 +29,39 @@ class Command(List):
         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]