Skip to content

Commit

Permalink
feat(Backend + SDK): Update kfp backend and kubernetes sdk to support…
Browse files Browse the repository at this point in the history
… ImagePullPolicy (#10417)

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

Signed-off-by: Revital Sur <eres@il.ibm.com>

* Fix format.

Signed-off-by: Revital Sur <eres@il.ibm.com>

* Update apiserver.csv.

Signed-off-by: Revital Sur <eres@il.ibm.com>

* Update licenses.

Signed-off-by: Revital Sur <eres@il.ibm.com>

---------

Signed-off-by: Revital Sur <eres@il.ibm.com>
  • Loading branch information
revit13 authored Mar 5, 2024
1 parent b7ea6e7 commit 83cabab
Show file tree
Hide file tree
Showing 10 changed files with 234 additions and 9 deletions.
18 changes: 18 additions & 0 deletions backend/src/v2/driver/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
80 changes: 80 additions & 0 deletions backend/src/v2/driver/driver_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
}
}
2 changes: 1 addition & 1 deletion backend/third_party_licenses/apiserver.csv
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion backend/third_party_licenses/driver.csv
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions go.sum

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions kubernetes_platform/python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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/).

Expand Down Expand Up @@ -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")
```
8 changes: 5 additions & 3 deletions kubernetes_platform/python/kfp/kubernetes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
'CreatePVC',
'DeletePVC',
'mount_pvc',
'set_image_pull_policy',
'use_field_path_as_env',
'set_image_pull_secrets',
'set_timeout',
Expand All @@ -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
Expand Down
24 changes: 23 additions & 1 deletion kubernetes_platform/python/kfp/kubernetes/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
88 changes: 88 additions & 0 deletions kubernetes_platform/python/test/unit/test_image_pull_policy.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 83cabab

Please sign in to comment.