Merge pull request #102 from foxish/plumb-new-options

Added new methods for specifying other resource limits/requests
jupyterhub · Nov 27, 2017 · 054f6d6 · 054f6d6
2 parents 06a2e09 + ce68bf4
commit 054f6d6
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 0 deletions.
diff --git a/kubespawner/objects.py b/kubespawner/objects.py
@@ -40,6 +40,8 @@ def make_pod(
     cpu_guarantee=None,
     mem_limit=None,
     mem_guarantee=None,
+    extra_resource_limits=None,
+    extra_resource_guarantees=None,
     lifecycle_hooks=None,
     init_containers=None,
     service_account=None,
@@ -197,12 +199,19 @@ def make_pod(
         notebook_container.resources.requests['cpu'] = cpu_guarantee
     if mem_guarantee:
         notebook_container.resources.requests['memory'] = mem_guarantee
+    if extra_resource_guarantees:
+        for k in extra_resource_guarantees:
+            notebook_container.resources.requests[k] = extra_resource_guarantees[k]
 
     notebook_container.resources.limits = {}
     if cpu_limit:
         notebook_container.resources.limits['cpu'] = cpu_limit
     if mem_limit:
         notebook_container.resources.limits['memory'] = mem_limit
+    if extra_resource_limits:
+        for k in extra_resource_limits:
+            notebook_container.resources.limits[k] = extra_resource_limits[k]
+
     notebook_container.volume_mounts = volume_mounts + hack_volume_mounts
     pod.spec.containers.append(notebook_container)
 

diff --git a/kubespawner/spawner.py b/kubespawner/spawner.py
@@ -668,6 +668,28 @@ def _hub_connect_port_default(self):
         """
     )
 
+    extra_resource_guarantees = Dict(
+        {},
+        config=True,
+        help="""
+        The dictionary used to request arbitrary resources.
+        Default is None and means no additional resources are requested.
+        For example, to request 3 Nvidia GPUs
+            `{"nvidia.com/gpu": "3"}`
+        """
+    )
+
+    extra_resource_limits = Dict(
+        {},
+        config=True,
+        help="""
+        The dictionary used to limit arbitrary resources.
+        Default is None and means no additional resources are limited.
+        For example, to add a limit of 3 Nvidia GPUs
+            `{"nvidia.com/gpu": "3"}`
+        """
+    )
+
     def _expand_user_properties(self, template):
         # Make sure username and servername match the restrictions for DNS labels
         safe_chars = set(string.ascii_lowercase + string.digits)
@@ -763,6 +785,8 @@ def get_pod_manifest(self):
             cpu_guarantee=self.cpu_guarantee,
             mem_limit=self.mem_limit,
             mem_guarantee=self.mem_guarantee,
+            extra_resource_limits=self.extra_resource_limits,
+            extra_resource_guarantees=self.extra_resource_guarantees,
             lifecycle_hooks=self.singleuser_lifecycle_hooks,
             init_containers=self.singleuser_init_containers,
             service_account=self.singleuser_service_account,

diff --git a/tests/test_objects.py b/tests/test_objects.py
@@ -631,6 +631,65 @@ def test_make_pod_with_extra_containers():
         "apiVersion": "v1"
     }
 
+def test_make_pod_with_extra_resources():
+    """
+    Test specification of extra resources (like GPUs)
+    """
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cpu_limit=2,
+        cpu_guarantee=1,
+        extra_resource_limits={"nvidia.com/gpu": "5", "k8s.io/new-resource": "1"},
+        extra_resource_guarantees={"nvidia.com/gpu": "3"},
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        mem_limit='1Gi',
+        mem_guarantee='512Mi',
+        image_pull_policy='IfNotPresent',
+        image_pull_secret="myregistrykey",
+        node_selector={"disk": "ssd"}
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+        },
+        "spec": {
+            "securityContext": {},
+            "imagePullSecrets": [{"name": "myregistrykey"}],
+            "nodeSelector": {"disk": "ssd"},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [{'name': 'no-api-access-please', 'mountPath': '/var/run/secrets/kubernetes.io/serviceaccount', 'readOnly': True}],
+                    "resources": {
+                        "limits": {
+                            "cpu": 2,
+                            "memory": '1Gi',
+                            "nvidia.com/gpu": "5",
+                            "k8s.io/new-resource": "1"
+                        },
+                        "requests": {
+                            "cpu": 1,
+                            "memory": '512Mi',
+                            "nvidia.com/gpu": "3"
+                        }
+                    }
+                }
+            ],
+            'volumes': [{'name': 'no-api-access-please', 'emptyDir': {}}],
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }
 
 def test_make_pvc_simple():
     """