Repository: kubeflow/pipelines

Issue:
SDK feature request: support soft constraint for scheduling a component to a preemptible node pool
From a pipeline, it would be useful to have an option to schedule onto a pre-emptible node if such a node pool is available, otherwise fall back to using non-preemptible.
(Under the hood this can apparently be done with node affinity and `preferredDuringSchedulingIgnoredDuringExecution`).

Ning's suggestion: expose config in the `gcp.use_preemptible_nodepool()` call to specify whether the constraint should be hard or soft.

/assign @gaoning777 (as per discussion)


Discussion:
/assign @gaoning777 

Resolving patch:
diff --git a/sdk/python/kfp/gcp.py b/sdk/python/kfp/gcp.py
--- a/sdk/python/kfp/gcp.py
+++ b/sdk/python/kfp/gcp.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from kubernetes.client import V1Toleration
+from kubernetes.client import V1Toleration, V1Affinity, V1NodeAffinity, \
+  V1NodeSelector, V1NodeSelectorTerm, V1NodeSelectorRequirement, V1PreferredSchedulingTerm
 
 def use_gcp_secret(secret_name='user-gcp-sa', secret_file_path_in_volume=None, volume_name=None, secret_volume_mount_path='/secret/gcp-credentials'):
     """An operator that configures the container to use GCP service account.
@@ -102,13 +103,31 @@ def _set_tpu_spec(task):
 def use_preemptible_nodepool(toleration: V1Toleration = V1Toleration(effect='NoSchedule',
                                                              key='preemptible',
                                                              operator='Equal',
-                                                             value='true')):
+                                                             value='true'),
+                              hard_constraint: bool = False):
   """An operator that configures the GKE preemptible in a container op.
+  Args:
+    toleration (V1Toleration): toleration to pods, default is the preemptible label.
+    hard_constraint (bool): the constraint of scheduling the pods on preemptible
+        nodepools is hard. (Default: False)
   """
 
   def _set_preemptible(task):
     task.add_toleration(toleration)
-    task.add_node_selector_constraint("cloud.google.com/gke-preemptible", "true")
+    node_selector_term = V1NodeSelectorTerm(match_expressions=[
+        V1NodeSelectorRequirement(key='cloud.google.com/gke-preemptible',
+                                  operator='In',
+                                  values=['true'])]
+    )
+    if hard_constraint:
+      node_affinity = V1NodeAffinity(required_during_scheduling_ignored_during_execution=
+                        V1NodeSelector(node_selector_terms=[node_selector_term]))
+    else:
+      node_affinity = V1NodeAffinity(preferred_during_scheduling_ignored_during_execution=
+                        V1PreferredSchedulingTerm(preference=node_selector_term,
+                                                  weight=50))
+    affinity = V1Affinity(node_affinity=node_affinity)
+    task.add_affinity(affinity=affinity)
     return task
 
   return _set_preemptible


Test patch:
diff --git a/sdk/python/tests/compiler/testdata/preemptible_tpu_gpu.yaml b/sdk/python/tests/compiler/testdata/preemptible_tpu_gpu.yaml
--- a/sdk/python/tests/compiler/testdata/preemptible_tpu_gpu.yaml
+++ b/sdk/python/tests/compiler/testdata/preemptible_tpu_gpu.yaml
@@ -10,7 +10,17 @@ spec:
   entrypoint: pipeline-flip-coin
   serviceAccountName: pipeline-runner
   templates:
-    - container:
+    - affinity:
+        nodeAffinity:
+          preferredDuringSchedulingIgnoredDuringExecution:
+            preference:
+              matchExpressions:
+                - key: cloud.google.com/gke-preemptible
+                  operator: In
+                  values:
+                    - 'true'
+            weight: 50
+      container:
         args:
           - python -c "import random; result = 'heads' if random.randint(0,1) == 0 else
             'tails'; print(result)" | tee /tmp/output
@@ -22,8 +32,6 @@ spec:
           limits:
             nvidia.com/gpu: 1
       name: flip
-      nodeSelector:
-        cloud.google.com/gke-preemptible: 'true'
       outputs:
         artifacts:
           - name: flip-output