feat(Backend + SDK): Update kfp backend and kubernetes sdk to support ImagePullPolicy

backend/src/v2/driver/driver.go

@@ -475,6 +475,24 @@ func extendPodSpecPatch(
 		podSpec.Containers[0].VolumeMounts = append(podSpec.Containers[0].VolumeMounts, volumeMounts...)
 	}
 
+	// Get image pull policy
+	pullPolicy := kubernetesExecutorConfig.GetImagePullPolicy()
+	if pullPolicy != "" {
+		policies := []string{"Always", "Never", "IfNotPresent"}
+		found := false
+		for _, value := range policies {
+			if value == pullPolicy {
+				found = true
+				break
+			}
+		}
+		if !found {
+			return fmt.Errorf("unsupported value: %s. ImagePullPolicy should be one of 'Always', 'Never' or 'IfNotPresent'", pullPolicy)
+		}
+		// We assume that the user container always gets executed first within a pod.
+		podSpec.Containers[0].ImagePullPolicy = k8score.PullPolicy(pullPolicy)
+	}
+
 	// Get node selector information
 	if kubernetesExecutorConfig.GetNodeSelector() != nil {
 		podSpec.NodeSelector = kubernetesExecutorConfig.GetNodeSelector().GetLabels()

backend/src/v2/driver/driver_test.go

@@ -1027,3 +1027,83 @@ func Test_extendPodSpecPatch_ActiveDeadlineSeconds(t *testing.T) {
 		})
 	}
 }
+
+func Test_extendPodSpecPatch_ImagePullPolicy(t *testing.T) {
+	tests := []struct {
+		name       string
+		k8sExecCfg *kubernetesplatform.KubernetesExecutorConfig
+		podSpec    *k8score.PodSpec
+		expected   *k8score.PodSpec
+	}{
+		{
+			"Valid - Always",
+			&kubernetesplatform.KubernetesExecutorConfig{
+				ImagePullPolicy: "Always",
+			},
+			&k8score.PodSpec{
+				Containers: []k8score.Container{
+					{
+						Name: "main",
+					},
+				},
+			},
+			&k8score.PodSpec{
+				Containers: []k8score.Container{
+					{
+						Name:            "main",
+						ImagePullPolicy: "Always",
+					},
+				},
+			},
+		},
+		{
+			"Valid - IfNotPresent",
+			&kubernetesplatform.KubernetesExecutorConfig{
+				ImagePullPolicy: "IfNotPresent",
+			},
+			&k8score.PodSpec{
+				Containers: []k8score.Container{
+					{
+						Name: "main",
+					},
+				},
+			},
+			&k8score.PodSpec{
+				Containers: []k8score.Container{
+					{
+						Name:            "main",
+						ImagePullPolicy: "IfNotPresent",
+					},
+				},
+			},
+		},
+		{
+			"Valid - Never",
+			&kubernetesplatform.KubernetesExecutorConfig{
+				ImagePullPolicy: "Never",
+			},
+			&k8score.PodSpec{
+				Containers: []k8score.Container{
+					{
+						Name: "main",
+					},
+				},
+			},
+			&k8score.PodSpec{
+				Containers: []k8score.Container{
+					{
+						Name:            "main",
+						ImagePullPolicy: "Never",
+					},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := extendPodSpecPatch(tt.podSpec, tt.k8sExecCfg, nil, nil)
+			assert.Nil(t, err)
+			assert.Equal(t, tt.expected, tt.podSpec)
+		})
+	}
+}

backend/third_party_licenses/apiserver.csv

@@ -61,7 +61,7 @@ github.com/klauspost/cpuid,https://github.com/klauspost/cpuid/blob/v1.3.1/LICENS
 github.com/klauspost/pgzip,https://github.com/klauspost/pgzip/blob/v1.2.5/LICENSE,MIT
 github.com/kubeflow/pipelines/api/v2alpha1/go,https://github.com/kubeflow/pipelines/blob/758c91f76784/api/LICENSE,Apache-2.0
 github.com/kubeflow/pipelines/backend,https://github.com/kubeflow/pipelines/blob/HEAD/LICENSE,Apache-2.0
-github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform,https://github.com/kubeflow/pipelines/blob/2983a7d49078/kubernetes_platform/LICENSE,Apache-2.0
+github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform,https://github.com/kubeflow/pipelines/blob/19a24e3e99db/kubernetes_platform/LICENSE,Apache-2.0
 github.com/kubeflow/pipelines/third_party/ml-metadata/go/ml_metadata,https://github.com/kubeflow/pipelines/blob/e1f0c010f800/third_party/ml-metadata/LICENSE,Apache-2.0
 github.com/lann/builder,https://github.com/lann/builder/blob/47ae307949d0/LICENSE,MIT
 github.com/lann/ps,https://github.com/lann/ps/blob/62de8c46ede0/LICENSE,MIT

backend/third_party_licenses/driver.csv

@@ -31,7 +31,7 @@ github.com/josharian/intern,https://github.com/josharian/intern/blob/v1.0.0/lice
 github.com/json-iterator/go,https://github.com/json-iterator/go/blob/v1.1.12/LICENSE,MIT
 github.com/kubeflow/pipelines/api/v2alpha1/go,https://github.com/kubeflow/pipelines/blob/758c91f76784/api/LICENSE,Apache-2.0
 github.com/kubeflow/pipelines/backend,https://github.com/kubeflow/pipelines/blob/HEAD/LICENSE,Apache-2.0
-github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform,https://github.com/kubeflow/pipelines/blob/2983a7d49078/kubernetes_platform/LICENSE,Apache-2.0
+github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform,https://github.com/kubeflow/pipelines/blob/19a24e3e99db/kubernetes_platform/LICENSE,Apache-2.0
 github.com/kubeflow/pipelines/third_party/ml-metadata/go/ml_metadata,https://github.com/kubeflow/pipelines/blob/e1f0c010f800/third_party/ml-metadata/LICENSE,Apache-2.0
 github.com/mailru/easyjson,https://github.com/mailru/easyjson/blob/v0.7.7/LICENSE,MIT
 github.com/modern-go/concurrent,https://github.com/modern-go/concurrent/blob/bacd9c7ef1dd/LICENSE,Apache-2.0

go.mod

@@ -31,7 +31,7 @@ require (
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.4 // indirect
 	github.com/kubeflow/pipelines/api v0.0.0-20230331215358-758c91f76784
-	github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240222213131-2983a7d49078
+	github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240305195700-19a24e3e99db
 	github.com/kubeflow/pipelines/third_party/ml-metadata v0.0.0-20230810215105-e1f0c010f800
 	github.com/lestrrat-go/strftime v1.0.4
 	github.com/mattn/go-sqlite3 v1.14.16

kubernetes_platform/python/README.md

@@ -4,6 +4,7 @@ The `kfp-kubernetes` Python library enables authoring [Kubeflow pipelines](https
 
 * [Secrets](https://kubernetes.io/docs/concepts/configuration/secret/)
 * [PersistentVolumeClaims](https://kubernetes.io/docs/concepts/storage/persistent-volumes/#persistentvolumeclaims)
+* [ImagePullPolicy](https://kubernetes.io/docs/concepts/containers/images/#image-pull-policy)
 
 See the [`kfp-kubernetes` reference documentation](https://kfp-kubernetes.readthedocs.io/).
 
@@ -203,3 +204,17 @@ def my_pipeline():
     kubernetes.set_timeout(task, 20)
 ```
 
+### ImagePullPolicy: One of "Always" "Never", "IfNotPresent".
+```python
+from kfp import dsl
+from kfp import kubernetes
+
+@dsl.component
+def simple_task():
+    print("hello-world")
+
+@dsl.pipeline
+def pipeline():
+    task = simple_task()
+    kubernetes.set_image_pull_policy(task, "Always")
+```

kubernetes_platform/python/kfp/kubernetes/__init__.py

@@ -22,6 +22,7 @@
     'CreatePVC',
     'DeletePVC',
     'mount_pvc',
+    'set_image_pull_policy',
     'use_field_path_as_env',
     'set_image_pull_secrets',
     'set_timeout',
@@ -31,11 +32,12 @@
     'use_secret_as_volume',
 ]
 
-from kfp.kubernetes.image import set_image_pull_secrets
-from kfp.kubernetes.config_map import use_config_map_as_volume
 from kfp.kubernetes.config_map import use_config_map_as_env
-from kfp.kubernetes.node_selector import add_node_selector
+from kfp.kubernetes.config_map import use_config_map_as_volume
 from kfp.kubernetes.field import use_field_path_as_env
+from kfp.kubernetes.image import set_image_pull_policy
+from kfp.kubernetes.image import set_image_pull_secrets
+from kfp.kubernetes.node_selector import add_node_selector
 from kfp.kubernetes.pod_metadata import add_pod_annotation
 from kfp.kubernetes.pod_metadata import add_pod_label
 from kfp.kubernetes.secret import use_secret_as_env

kubernetes_platform/python/kfp/kubernetes/image.py

@@ -38,11 +38,33 @@ def set_image_pull_secrets(
 
     # Assuming secret_names is a list of strings
     image_pull_secret = [
-        pb.ImagePullSecret(secret_name=secret_name) for secret_name in secret_names
+        pb.ImagePullSecret(secret_name=secret_name)
+        for secret_name in secret_names
     ]
 
     msg.image_pull_secret.extend(image_pull_secret)
 
     task.platform_config['kubernetes'] = json_format.MessageToDict(msg)
 
     return task
+
+
+def set_image_pull_policy(task: PipelineTask, policy: str) -> PipelineTask:
+    """Set image pull policy for the container.
+
+    Args:
+        task: Pipeline task.
+        policy: One of `Always`, `Never`, `IfNotPresent`.
+
+    Returns:
+        Task object with an added ImagePullPolicy specification.
+    """
+    if policy not in ['Always', 'Never', 'IfNotPresent']:
+        raise ValueError(
+            'Invalid imagePullPolicy. Must be one of `Always`, `Never`, `IfNotPresent`.'
+        )
+    msg = common.get_existing_kubernetes_config_as_message(task)
+    msg.image_pull_policy = policy
+    task.platform_config['kubernetes'] = json_format.MessageToDict(msg)
+
+    return task

kubernetes_platform/python/test/unit/test_image_pull_policy.py

@@ -0,0 +1,88 @@
+# Copyright 2024 The Kubeflow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.protobuf import json_format
+from kfp import dsl
+from kfp import kubernetes
+
+
+class TestImagePullPolicy:
+
+    def test_always(self):
+
+        @dsl.pipeline
+        def my_pipeline():
+            task = comp()
+            kubernetes.set_image_pull_policy(task, 'Always')
+
+        assert json_format.MessageToDict(my_pipeline.platform_spec) == {
+            'platforms': {
+                'kubernetes': {
+                    'deploymentSpec': {
+                        'executors': {
+                            'exec-comp': {
+                                'imagePullPolicy': 'Always'
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+    def test_if_not_present(self):
+
+        @dsl.pipeline
+        def my_pipeline():
+            task = comp()
+            kubernetes.set_image_pull_policy(task, 'IfNotPresent')
+
+        assert json_format.MessageToDict(my_pipeline.platform_spec) == {
+            'platforms': {
+                'kubernetes': {
+                    'deploymentSpec': {
+                        'executors': {
+                            'exec-comp': {
+                                'imagePullPolicy': 'IfNotPresent'
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+    def test_never(self):
+
+        @dsl.pipeline
+        def my_pipeline():
+            task = comp()
+            kubernetes.set_image_pull_policy(task, 'Never')
+
+        assert json_format.MessageToDict(my_pipeline.platform_spec) == {
+            'platforms': {
+                'kubernetes': {
+                    'deploymentSpec': {
+                        'executors': {
+                            'exec-comp': {
+                                'imagePullPolicy': 'Never'
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+
+@dsl.component
+def comp():
+    pass