add toleration method

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__)
@@ -24,6 +25,7 @@
 class KubernetesJobTask:
 
     volumes: List[AttachableVolume] = []
+    tolerations: List[V1Toleration]
 
     def _init_task_metadata(self):
         self.uu_name = self.name
@@ -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)