diff --git a/kubeluigi/__init__.py b/kubeluigi/__init__.py index fa22023..6098656 100644 --- a/kubeluigi/__init__.py +++ b/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) \ No newline at end of file