Repository: DataBiosphere/toil

Issue:
Preemptable Toil Kubernetes pods should tolerate the `cloud.google.com/gke-preemptible="true"` taint
A Kubernetes cluster might want to taint preemptable nodes. The only vaguely standard taint I have found for this is `cloud.google.com/gke-preemptible="true"`. I'm going to put it on some autoscaling nodes on the GI Kubernetes cluster, so it would be good if Toil jobs that were preemptable knew to tolerate it.

Note the spelling; Toil says "preemptable", while cloud providers seem to have settled on "preemptible".

┆Issue is synchronized with this [Jira Task](https://ucsc-cgl.atlassian.net/browse/TOIL-1102)
┆friendlyId: TOIL-1102


Discussion:

Patch:
diff --git a/src/toil/batchSystems/kubernetes.py b/src/toil/batchSystems/kubernetes.py
index 6f561c7b..362b18ac 100644
--- a/src/toil/batchSystems/kubernetes.py
+++ b/src/toil/batchSystems/kubernetes.py
@@ -332,48 +332,101 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
         self.user_script = userScript
 
     # setEnv is provided by BatchSystemSupport, updates self.environment
-
-    def _create_affinity(self, preemptable: bool) -> kubernetes.client.V1Affinity:
-        """
-        Make a V1Affinity that places pods appropriately depending on if they
-        tolerate preemptable nodes or not.
+    
+    @staticmethod
+    def _apply_placement_constraints(preemptable: bool, pod_spec: kubernetes.client.V1PodSpec) -> None:
         """
+        Set .affinity and/or .tolerations on the given pod spec, so that it
+        runs on the right kind of nodes, according to whether it is allowed to
+        be preempted.
 
-        # Describe preemptable nodes
+        Preemptable jobs will be able to run on preemptable or non-preemptable
+        nodes, and will prefer preemptable nodes if available.
 
-        # There's no labeling standard for knowing which nodes are
-        # preemptable across different cloud providers/Kubernetes clusters,
-        # so we use the labels that EKS uses. Toil-managed Kubernetes
-        # clusters also use this label. If we come to support more kinds of
-        # preemptable nodes, we will need to add more labels to avoid here.
-        preemptable_label = "eks.amazonaws.com/capacityType"
-        preemptable_value = "SPOT"
+        Non-preemptable jobs will not be allowed to run on nodes that are
+        marked as preemptable.
 
-        non_spot = [kubernetes.client.V1NodeSelectorRequirement(key=preemptable_label,
-                                                                operator='NotIn',
-                                                                values=[preemptable_value])]
-        unspecified = [kubernetes.client.V1NodeSelectorRequirement(key=preemptable_label,
-                                                                   operator='DoesNotExist')]
-        # These are OR'd
-        node_selector_terms = [kubernetes.client.V1NodeSelectorTerm(match_expressions=non_spot),
-                               kubernetes.client.V1NodeSelectorTerm(match_expressions=unspecified)]
-        node_selector = kubernetes.client.V1NodeSelector(node_selector_terms=node_selector_terms)
+        Understands the labeling scheme used by EKS, and the taint scheme used
+        by GCE. The Toil-managed Kubernetes setup will mimic at least one of
+        these.
+        """
 
+        # We consider nodes preemptable if they have any of these label or taint values.
+        # We tolerate all effects of specified taints.
+        # Amazon just uses a label, while Google
+        # <https://cloud.google.com/kubernetes-engine/docs/how-to/preemptible-vms>
+        # uses a label and a taint.
+        PREEMPTABLE_SCHEMES = {'labels': [('eks.amazonaws.com/capacityType', ['SPOT']),
+                                          ('cloud.google.com/gke-preemptible', ['true'])],
+                               'taints': [('cloud.google.com/gke-preemptible', ['true'])]}
+
+        # We will compose a node selector with these requirements to require or prefer.
+        # These requirements will be AND-ed and turn into a single term.
+        node_selector_requirements: List[kubernetes.client.V1NodeSelectorRequirement] = []
+        # These terms will be OR'd, along with a term made of those ANDed requirements
+        node_selector_terms: List[kubernetes.client.V1NodeSelectorTerm] = []
+        # And this list of tolerations to apply
+        tolerations: List[kubernetes.client.V1Toleration] = []
 
         if preemptable:
-            # We can put this job anywhere. But we would be smart to prefer
-            # preemptable nodes first, if available, so we don't block any
-            # non-preemptable jobs.
-            node_preference = kubernetes.client.V1PreferredSchedulingTerm(weight=1, preference=node_selector)
-
-            node_affinity = kubernetes.client.V1NodeAffinity(preferred_during_scheduling_ignored_during_execution=[node_preference])
+            # We want to seek preemptable labels and tolerate preemptable taints.
+            for label, values in PREEMPTABLE_SCHEMES['labels']:
+                is_spot = kubernetes.client.V1NodeSelectorRequirement(key=label,
+                                                                      operator='In',
+                                                                      values=values)
+                # We want to OR all the labels, so they all need to become separate terms.
+                node_selector_terms.append(kubernetes.client.V1NodeSelectorTerm(
+                    match_expressions=[is_spot]
+                ))
+            for taint, values in PREEMPTABLE_SCHEMES['taints']:
+                for value in values:
+                    # Each toleration can tolerate one value
+                    spot_ok = kubernetes.client.V1Toleration(key=taint,
+                                                             value=value)
+                    tolerations.append(spot_ok)
         else:
-            # We need to add some selector stuff to keep the job off of
-            # nodes that might be preempted.
-            node_affinity = kubernetes.client.V1NodeAffinity(required_during_scheduling_ignored_during_execution=node_selector)
+            # We want to prohibit preemptable labels
+            for label, values in PREEMPTABLE_SCHEMES['labels']:
+                # So we need to say that each preemptable label either doesn't
+                # have any of the preemptable values, or doesn't exist.
+                # Although the docs don't really say so, NotIn also matches
+                # cases where the label doesn't exist. This is suggested by
+                # <https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#set-based-requirement>
+                # So we create a NotIn for each preemptable label and AND them
+                # all together.
+                not_spot = kubernetes.client.V1NodeSelectorRequirement(key=label,
+                                                                       operator='NotIn',
+                                                                       values=values)
+                node_selector_requirements.append(not_spot)
+
+        # Now combine everything
+        if node_selector_requirements:
+            # We have requirements that want to be a term
+            node_selector_terms.append(kubernetes.client.V1NodeSelectorTerm(
+                match_expressions=node_selector_requirements
+            ))
+
+        if node_selector_terms:
+            # Make the terms into a node selector
+            node_selector = kubernetes.client.V1NodeSelector(node_selector_terms=node_selector_terms)
+            if preemptable:
+                # Node selector sense is a preference, so we wrap it
+                node_preference = kubernetes.client.V1PreferredSchedulingTerm(weight=1, preference=node_selector)
+                # And make an affinity around a preference
+                node_affinity = kubernetes.client.V1NodeAffinity(
+                    preferred_during_scheduling_ignored_during_execution=[node_preference]
+                )
+            else:
+                # Node selector sense is a requirement, so make an affinity around a requirement
+                node_affinity = kubernetes.client.V1NodeAffinity(
+                    required_during_scheduling_ignored_during_execution=node_selector
+                )
+            # Apply the affinity
+            pod_spec.affinity = node_affinity
 
-        # Make the node affinity into an overall affinity
-        return kubernetes.client.V1Affinity(node_affinity=node_affinity)
+        if tolerations:
+            # Apply the tolerations
+            pod_spec.tolerations = tolerations
 
     def _create_pod_spec(
             self,
@@ -460,7 +513,7 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
                                                volumes=volumes,
                                                restart_policy="Never")
         # Tell the spec where to land
-        pod_spec.affinity = self._create_affinity(job_desc.preemptable)
+        self._apply_placement_constraints(job_desc.preemptable, pod_spec)
 
         if self.service_account:
             # Apply service account if set


Test patch:
diff --git a/src/toil/test/__init__.py b/src/toil/test/__init__.py
index 6306e282..e4236910 100644
--- a/src/toil/test/__init__.py
+++ b/src/toil/test/__init__.py
@@ -458,9 +458,18 @@ def needs_tes(test_item: MT) -> MT:
     return test_item
 
 
-def needs_kubernetes(test_item: MT) -> MT:
+def needs_kubernetes_installed(test_item: MT) -> MT:
     """Use as a decorator before test classes or methods to run only if Kubernetes is installed."""
     test_item = _mark_test('kubernetes', test_item)
+    try:
+        import kubernetes
+    except ImportError:
+        return unittest.skip("Install Toil with the 'kubernetes' extra to include this test.")(test_item)
+    return test_item
+
+def needs_kubernetes(test_item: MT) -> MT:
+    """Use as a decorator before test classes or methods to run only if Kubernetes is installed and configured."""
+    test_item = needs_kubernetes_installed(test_item)
     try:
         import kubernetes
         try:
@@ -472,7 +481,8 @@ def needs_kubernetes(test_item: MT) -> MT:
                 return unittest.skip("Configure Kubernetes (~/.kube/config, $KUBECONFIG, "
                                      "or current pod) to include this test.")(test_item)
     except ImportError:
-        return unittest.skip("Install Toil with the 'kubernetes' extra to include this test.")(test_item)
+        # We should already be skipping this test
+        pass
     return test_item
 
 
diff --git a/src/toil/test/batchSystems/batchSystemTest.py b/src/toil/test/batchSystems/batchSystemTest.py
index 02a0c5b3..9fef0f5a 100644
--- a/src/toil/test/batchSystems/batchSystemTest.py
+++ b/src/toil/test/batchSystems/batchSystemTest.py
@@ -18,6 +18,7 @@ import os
 import subprocess
 import sys
 import tempfile
+import textwrap
 import time
 from abc import ABCMeta, abstractmethod
 from fractions import Fraction
@@ -49,6 +50,7 @@ from toil.test import (ToilTest,
                        needs_fetchable_appliance,
                        needs_gridengine,
                        needs_htcondor,
+                       needs_kubernetes_installed,
                        needs_kubernetes,
                        needs_lsf,
                        needs_mesos,
@@ -456,6 +458,60 @@ class KubernetesBatchSystemTest(hidden.AbstractBatchSystemTest):
         return KubernetesBatchSystem(config=self.config,
                                      maxCores=numCores, maxMemory=1e9, maxDisk=2001)
 
+@needs_kubernetes_installed
+class KubernetesBatchSystemBenchTest(ToilTest):
+    """
+    Kubernetes batch system unit tests that don't need to actually talk to a cluster.
+    """
+
+    def test_preemptability_constraints(self):
+        """
+        Make sure we generate the right preemptability constraints.
+        """
+
+        # Make sure we can print diffs of these long strings
+        self.maxDiff = 10000
+
+        from kubernetes.client import V1PodSpec
+        from toil.batchSystems.kubernetes import KubernetesBatchSystem
+
+        normal_spec = V1PodSpec(containers=[])
+        KubernetesBatchSystem._apply_placement_constraints(False, normal_spec)
+        self.assertEqual(str(normal_spec.affinity), textwrap.dedent("""
+        {'preferred_during_scheduling_ignored_during_execution': None,
+         'required_during_scheduling_ignored_during_execution': {'node_selector_terms': [{'match_expressions': [{'key': 'eks.amazonaws.com/capacityType',
+                                                                                                                 'operator': 'NotIn',
+                                                                                                                 'values': ['SPOT']},
+                                                                                                                {'key': 'cloud.google.com/gke-preemptible',
+                                                                                                                 'operator': 'NotIn',
+                                                                                                                 'values': ['true']}],
+                                                                                          'match_fields': None}]}}
+        """).strip())
+        self.assertEqual(str(normal_spec.tolerations), "None")
+
+        spot_spec = V1PodSpec(containers=[])
+        KubernetesBatchSystem._apply_placement_constraints(True, spot_spec)
+        self.assertEqual(str(spot_spec.affinity), textwrap.dedent("""
+        {'preferred_during_scheduling_ignored_during_execution': [{'preference': {'node_selector_terms': [{'match_expressions': [{'key': 'eks.amazonaws.com/capacityType',
+                                                                                                                                  'operator': 'In',
+                                                                                                                                  'values': ['SPOT']}],
+                                                                                                           'match_fields': None},
+                                                                                                          {'match_expressions': [{'key': 'cloud.google.com/gke-preemptible',
+                                                                                                                                  'operator': 'In',
+                                                                                                                                  'values': ['true']}],
+                                                                                                           'match_fields': None}]},
+                                                                   'weight': 1}],
+         'required_during_scheduling_ignored_during_execution': None}
+        """).strip())
+        self.assertEqual(str(spot_spec.tolerations), textwrap.dedent("""
+        [{'effect': None,
+         'key': 'cloud.google.com/gke-preemptible',
+         'operator': None,
+         'toleration_seconds': None,
+         'value': 'true'}]
+        """).strip())
+
+
 @needs_tes
 @needs_fetchable_appliance
 class TESBatchSystemTest(hidden.AbstractBatchSystemTest):