Merge pull request #22 from optimizely/toleration-feature

Toleration feature

kubeluigi/__init__.py

@@ -16,6 +16,7 @@
 
 from kubernetes.client import ApiClient
 from kubernetes.client.models.v1_job import V1Job
+from kubernetes.client import V1Toleration
 from kubernetes.client.models.v1_pod_spec import V1PodSpec
 
 logger = logging.getLogger(__name__)
@@ -27,6 +28,7 @@ class KubernetesJobTask:
 
     def _init_task_metadata(self):
         self.uu_name = self.name
+        self._tolerations: List[V1Toleration] = []
 
     def _init_kubernetes(self):
         self.kubernetes_client = kubernetes_client()
@@ -79,7 +81,7 @@ def build_job_definition(self) -> V1Job:
         schema = self.spec_schema()
         schema_with_volumes = self._attach_volumes_to_spec(schema)
         pod_template_spec = pod_spec_from_dict(
-            self.uu_name, schema_with_volumes, self.labels, self.restart_policy
+            self.uu_name, schema_with_volumes, self.labels, self.restart_policy, tolerations=self._tolerations
         )
 
         job = job_definition(
@@ -147,3 +149,7 @@ def add_volume(self, volume):
         adds a volume to the task
         """
         return self.volumes.append(volume)
+
+    def add_toleration(self, key, value, effect='NoSchedule', operator='Equals'):
+        toleration = V1Toleration(key=key, value=value, effect=effect, operator=operator)
+        return self._tolerations.append(toleration)