Add tolerations to the pod object

jupyterhub · Jul 4, 2018 · 1a31974 · 1a31974
1 parent f422420
commit 1a31974
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 2 deletions.
diff --git a/kubespawner/objects.py b/kubespawner/objects.py
@@ -52,7 +52,11 @@ def make_pod(
     extra_container_config=None,
     extra_pod_config=None,
     extra_containers=None,
+<<<<<<< HEAD
     scheduler_name=None
+=======
+    tolerations=None,
+>>>>>>> Add tolerations to the pod object
 ):
     """
     Make a k8s pod specification for running a user notebook.
@@ -233,6 +237,9 @@ def make_pod(
     if extra_containers:
         pod.spec.containers.extend(extra_containers)
 
+    if tolerations:
+        pod.spec.tolerations = tolerations
+
     pod.spec.init_containers = init_containers
     pod.spec.volumes = volumes
 

diff --git a/kubespawner/spawner.py b/kubespawner/spawner.py
@@ -823,6 +823,33 @@ def _hub_connect_port_default(self):
         """
     )
 
+    singleuser_tolerations = List(
+        None,
+        config=True,
+        help="""
+        List of tolerations that are to be assigned to the pod in order to be able to schedule the pod
+        on a node with the corresponding taints. See the official Kubernetes documentation for additional details 
+        https://kubernetes.io/docs/concepts/configuration/taint-and-toleration/
+        
+        Example:
+        
+            [
+                {
+                    'key': 'key',
+                    'operator': 'Equal',
+                    'value': 'value',
+                    'effect': 'NoSchedule'
+                },
+                {
+                    'key': 'key',
+                    'operator': 'Exists',
+                    'effect': 'NoSchedule'
+                }
+            ]
+         
+        """
+    )
+
     extra_resource_guarantees = Dict(
         {},
         config=True,
@@ -1195,6 +1222,7 @@ def get_pod_manifest(self):
             extra_container_config=self.extra_container_config,
             extra_pod_config=self.extra_pod_config,
             extra_containers=self.extra_containers,
+            tolerations=self.tolerations,
         )
 
     def get_pvc_manifest(self):

diff --git a/tests/test_objects.py b/tests/test_objects.py
@@ -360,7 +360,7 @@ def test_run_privileged_container():
                     "ports": [{
                         "name": "notebook-port",
                         "containerPort": 8888
-                    }],                    
+                    }],
                     "resources": {
                         "limits": {},
                         "requests": {}
@@ -919,7 +919,7 @@ def test_make_resources_all():
         }
     }
 
-    
+
 def test_make_pod_with_service_account():
     """
     Test specification of the simplest possible pod specification with non-default service account
@@ -963,3 +963,61 @@ def test_make_pod_with_service_account():
         "kind": "Pod",
         "apiVersion": "v1"
     }
+
+
+def test_make_pod_with_tolerations():
+    """
+    Test specification of the simplest possible pod specification with non-empty tolerations
+    """
+    tolerations = [
+        {
+            'key': 'key',
+            'operator': 'Equal',
+            'value': 'value',
+            'effect': 'NoSchedule'
+        },
+        {
+            'key': 'key',
+            'operator': 'Exists',
+            'effect': 'NoSchedule'
+        }
+    ]
+    assert api_client.sanitize_for_serialization(make_pod(
+        name='test',
+        image_spec='jupyter/singleuser:latest',
+        cmd=['jupyterhub-singleuser'],
+        port=8888,
+        image_pull_policy='IfNotPresent',
+        tolerations=tolerations
+    )) == {
+        "metadata": {
+            "name": "test",
+            "labels": {},
+            "annotations": {}
+        },
+        "spec": {
+            "securityContext": {},
+            "containers": [
+                {
+                    "env": [],
+                    "name": "notebook",
+                    "image": "jupyter/singleuser:latest",
+                    "imagePullPolicy": "IfNotPresent",
+                    "args": ["jupyterhub-singleuser"],
+                    "ports": [{
+                        "name": "notebook-port",
+                        "containerPort": 8888
+                    }],
+                    'volumeMounts': [],
+                    "resources": {
+                        "limits": {},
+                        "requests": {}
+                    }
+                }
+            ],
+            'volumes': [],
+            'tolerations': tolerations
+        },
+        "kind": "Pod",
+        "apiVersion": "v1"
+    }