diff --git a/backend/src/v2/driver/driver.go b/backend/src/v2/driver/driver.go index ebb194f646e..56fb45b96f4 100644 --- a/backend/src/v2/driver/driver.go +++ b/backend/src/v2/driver/driver.go @@ -665,6 +665,33 @@ func extendPodSpecPatch( podSpec.Volumes = append(podSpec.Volumes, ephemeralVolume) podSpec.Containers[0].VolumeMounts = append(podSpec.Containers[0].VolumeMounts, ephemeralVolumeMount) } + + // EmptyDirMounts + for _, emptyDirVolumeSpec := range kubernetesExecutorConfig.GetEmptyDirMounts() { + var sizeLimitResource *k8sres.Quantity + if emptyDirVolumeSpec.GetSizeLimit() != "" { + r := k8sres.MustParse(emptyDirVolumeSpec.GetSizeLimit()) + sizeLimitResource = &r + } + + emptyDirVolume := k8score.Volume{ + Name: emptyDirVolumeSpec.GetVolumeName(), + VolumeSource: k8score.VolumeSource{ + EmptyDir: &k8score.EmptyDirVolumeSource{ + Medium: k8score.StorageMedium(emptyDirVolumeSpec.GetMedium()), + SizeLimit: sizeLimitResource, + }, + }, + } + emptyDirVolumeMount := k8score.VolumeMount{ + Name: emptyDirVolumeSpec.GetVolumeName(), + MountPath: emptyDirVolumeSpec.GetMountPath(), + } + + podSpec.Volumes = append(podSpec.Volumes, emptyDirVolume) + podSpec.Containers[0].VolumeMounts = append(podSpec.Containers[0].VolumeMounts, emptyDirVolumeMount) + } + return nil } diff --git a/backend/src/v2/driver/driver_test.go b/backend/src/v2/driver/driver_test.go index bea24890033..be64723ccfb 100644 --- a/backend/src/v2/driver/driver_test.go +++ b/backend/src/v2/driver/driver_test.go @@ -15,9 +15,10 @@ package driver import ( "encoding/json" + "testing" + k8sres "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "testing" "github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec" "github.com/kubeflow/pipelines/backend/src/v2/metadata" @@ -532,7 +533,7 @@ func Test_extendPodSpecPatch_Secret(t *testing.T) { { Name: "secret1", VolumeSource: k8score.VolumeSource{ - Secret: &k8score.SecretVolumeSource{SecretName: "secret1", Optional: &[]bool{false}[0],}, + Secret: &k8score.SecretVolumeSource{SecretName: "secret1", Optional: &[]bool{false}[0]}, }, }, }, @@ -730,7 +731,7 @@ func Test_extendPodSpecPatch_ConfigMap(t *testing.T) { VolumeSource: k8score.VolumeSource{ ConfigMap: &k8score.ConfigMapVolumeSource{ LocalObjectReference: k8score.LocalObjectReference{Name: "cm1"}, - Optional: &[]bool{false}[0],}, + Optional: &[]bool{false}[0]}, }, }, }, @@ -890,6 +891,165 @@ func Test_extendPodSpecPatch_ConfigMap(t *testing.T) { } } +func Test_extendPodSpecPatch_EmptyVolumeMount(t *testing.T) { + medium := "Memory" + sizeLimit := "1Gi" + var sizeLimitResource *k8sres.Quantity + r := k8sres.MustParse(sizeLimit) + sizeLimitResource = &r + + tests := []struct { + name string + k8sExecCfg *kubernetesplatform.KubernetesExecutorConfig + podSpec *k8score.PodSpec + expected *k8score.PodSpec + }{ + { + "Valid - emptydir mount with no medium or size limit", + &kubernetesplatform.KubernetesExecutorConfig{ + EmptyDirMounts: []*kubernetesplatform.EmptyDirMount{ + { + VolumeName: "emptydir1", + MountPath: "/data/path", + }, + }, + }, + &k8score.PodSpec{ + Containers: []k8score.Container{ + { + Name: "main", + }, + }, + }, + &k8score.PodSpec{ + Containers: []k8score.Container{ + { + Name: "main", + VolumeMounts: []k8score.VolumeMount{ + { + Name: "emptydir1", + MountPath: "/data/path", + }, + }, + }, + }, + Volumes: []k8score.Volume{ + { + Name: "emptydir1", + VolumeSource: k8score.VolumeSource{ + EmptyDir: &k8score.EmptyDirVolumeSource{}, + }, + }, + }, + }, + }, + { + "Valid - emptydir mount with medium and size limit", + &kubernetesplatform.KubernetesExecutorConfig{ + EmptyDirMounts: []*kubernetesplatform.EmptyDirMount{ + { + VolumeName: "emptydir1", + MountPath: "/data/path", + Medium: &medium, + SizeLimit: &sizeLimit, + }, + }, + }, + &k8score.PodSpec{ + Containers: []k8score.Container{ + { + Name: "main", + }, + }, + }, + &k8score.PodSpec{ + Containers: []k8score.Container{ + { + Name: "main", + VolumeMounts: []k8score.VolumeMount{ + { + Name: "emptydir1", + MountPath: "/data/path", + }, + }, + }, + }, + Volumes: []k8score.Volume{ + { + Name: "emptydir1", + VolumeSource: k8score.VolumeSource{ + EmptyDir: &k8score.EmptyDirVolumeSource{ + Medium: k8score.StorageMedium(medium), + SizeLimit: sizeLimitResource, + }, + }, + }, + }, + }, + }, + { + "Valid - multiple emptydir mounts", + &kubernetesplatform.KubernetesExecutorConfig{ + EmptyDirMounts: []*kubernetesplatform.EmptyDirMount{ + { + VolumeName: "emptydir1", + MountPath: "/data/path", + }, + { + VolumeName: "emptydir2", + MountPath: "/data/path2", + }, + }, + }, + &k8score.PodSpec{ + Containers: []k8score.Container{ + { + Name: "main", + }, + }, + }, + &k8score.PodSpec{ + Containers: []k8score.Container{ + { + Name: "main", + VolumeMounts: []k8score.VolumeMount{ + { + Name: "emptydir1", + MountPath: "/data/path", + }, + { + Name: "emptydir2", + MountPath: "/data/path2", + }, + }, + }, + }, + Volumes: []k8score.Volume{ + { + Name: "emptydir1", + VolumeSource: k8score.VolumeSource{ + EmptyDir: &k8score.EmptyDirVolumeSource{}, + }, + }, + { + Name: "emptydir2", + VolumeSource: k8score.VolumeSource{ + EmptyDir: &k8score.EmptyDirVolumeSource{}, + }, + }, + }, + }, + }, + } + 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) + }) + } +} + func Test_extendPodSpecPatch_ImagePullSecrets(t *testing.T) { tests := []struct { name string diff --git a/backend/third_party_licenses/apiserver.csv b/backend/third_party_licenses/apiserver.csv index b5dd1e1f6ed..ed15297596f 100644 --- a/backend/third_party_licenses/apiserver.csv +++ b/backend/third_party_licenses/apiserver.csv @@ -85,7 +85,7 @@ github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-exithandler/pkg/apis/exitha github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-kfptask/pkg/apis/kfptask,https://github.com/kubeflow/kfp-tekton/blob/a75d4b3711ff/tekton-catalog/tekton-kfptask/LICENSE,Apache-2.0 github.com/kubeflow/pipelines/api/v2alpha1/go,https://github.com/kubeflow/pipelines/blob/58ce09e07d03/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/8b2a099e8c9f/kubernetes_platform/LICENSE,Apache-2.0 +github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform,https://github.com/kubeflow/pipelines/blob/d911c8b73b49/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 diff --git a/backend/third_party_licenses/driver.csv b/backend/third_party_licenses/driver.csv index a2cf7c59854..fb514d83c70 100644 --- a/backend/third_party_licenses/driver.csv +++ b/backend/third_party_licenses/driver.csv @@ -33,7 +33,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/58ce09e07d03/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/8b2a099e8c9f/kubernetes_platform/LICENSE,Apache-2.0 +github.com/kubeflow/pipelines/kubernetes_platform/go/kubernetesplatform,https://github.com/kubeflow/pipelines/blob/d911c8b73b49/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 diff --git a/go.mod b/go.mod index 44e9ae0b7f4..1337ebb719b 100644 --- a/go.mod +++ b/go.mod @@ -32,7 +32,7 @@ require ( github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-exithandler v0.0.0-20231127195001-a75d4b3711ff github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-kfptask v0.0.0-20231127195001-a75d4b3711ff github.com/kubeflow/pipelines/api v0.0.0-20231027040853-58ce09e07d03 - github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240403164522-8b2a099e8c9f + github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240725205754-d911c8b73b49 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.19 diff --git a/go.sum b/go.sum index 013220bd44d..01d5672e4ad 100644 --- a/go.sum +++ b/go.sum @@ -2212,8 +2212,8 @@ github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-kfptask v0.0.0-202311271950 github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-kfptask v0.0.0-20231127195001-a75d4b3711ff/go.mod h1:lAFdPugzj3bcAXyN3+8y0NByidZ88zwGxMc+gdc8cHw= github.com/kubeflow/pipelines/api v0.0.0-20231027040853-58ce09e07d03 h1:reL3LbkRIozBkKSUYjtQFV2kVC1R4WHG9FrTClRT1FY= github.com/kubeflow/pipelines/api v0.0.0-20231027040853-58ce09e07d03/go.mod h1:T7TOQB36gGe97yUdfVAnYK5uuT0+uQbLNHDUHxYkmE4= -github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240403164522-8b2a099e8c9f h1:O5GmJN8tALpiqL0dUo4uhOkqHG8xOkNCgT7QI9q9GnE= -github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240403164522-8b2a099e8c9f/go.mod h1:CJkKr356RlpZP/gQRuHf3Myrn1qJtoUVe4EMCmtwarg= +github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240725205754-d911c8b73b49 h1:Xf1qun8x4ZJj/nLZpUSIaphDK04NKeV7WME31qIC8Xo= +github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240725205754-d911c8b73b49/go.mod h1:UhJrmpVSWawsAsSr1OzZfsEZTvQce+GvGwcZ58ULEhM= github.com/kubeflow/pipelines/third_party/ml-metadata v0.0.0-20230810215105-e1f0c010f800 h1:YAW+X9xCW8Yq5tQaBBQaLTNU9CJj8Nr7lx1+k66ZHJ0= github.com/kubeflow/pipelines/third_party/ml-metadata v0.0.0-20230810215105-e1f0c010f800/go.mod h1:chIDffBaVQ/asNl1pTTdbAymYcuBKf8BR3YtSP+3FEU= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= diff --git a/kubernetes_platform/python/kfp/kubernetes/__init__.py b/kubernetes_platform/python/kfp/kubernetes/__init__.py index fa149c31c09..91fb119212a 100644 --- a/kubernetes_platform/python/kfp/kubernetes/__init__.py +++ b/kubernetes_platform/python/kfp/kubernetes/__init__.py @@ -22,6 +22,7 @@ 'add_toleration', 'CreatePVC', 'DeletePVC', + 'empty_dir_mount', 'mount_pvc', 'set_image_pull_policy', 'use_field_path_as_env', @@ -49,3 +50,4 @@ from kfp.kubernetes.volume import CreatePVC from kfp.kubernetes.volume import DeletePVC from kfp.kubernetes.volume import mount_pvc +from kfp.kubernetes.empty_dir import empty_dir_mount diff --git a/kubernetes_platform/python/kfp/kubernetes/empty_dir.py b/kubernetes_platform/python/kfp/kubernetes/empty_dir.py new file mode 100644 index 00000000000..a8873d872ce --- /dev/null +++ b/kubernetes_platform/python/kfp/kubernetes/empty_dir.py @@ -0,0 +1,56 @@ +# 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 typing import Optional + +from google.protobuf import json_format +from kfp.dsl import PipelineTask +from kfp.kubernetes import common +from kfp.kubernetes import kubernetes_executor_config_pb2 as pb + + +def empty_dir_mount( + task: PipelineTask, + volume_name: str, + mount_path: str, + medium: Optional[str] = None, + size_limit: Optional[str] = None, +) -> PipelineTask: + """Mount an EmptyDir volume to the task's container. + + Args: + task: Pipeline task. + volume_name: Name of the EmptyDir volume. + mount_path: Path within the container at which the EmptyDir should be mounted. + medium: Storage medium to back the EmptyDir. Must be one of `Memory` or `HugePages`. Defaults to `None`. + size_limit: Maximum size of the EmptyDir. For example, `5Gi`. Defaults to `None`. + + Returns: + Task object with updated EmptyDir mount configuration. + """ + + msg = common.get_existing_kubernetes_config_as_message(task) + + empty_dir_mount = pb.EmptyDirMount( + volume_name=volume_name, + mount_path=mount_path, + medium=medium, + size_limit=size_limit, + ) + + msg.empty_dir_mounts.append(empty_dir_mount) + + task.platform_config['kubernetes'] = json_format.MessageToDict(msg) + + return task diff --git a/kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.py b/kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.py new file mode 100644 index 00000000000..a86c93d3cb7 --- /dev/null +++ b/kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.py @@ -0,0 +1,36 @@ +# 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 kfp import dsl +from kfp import kubernetes + + +@dsl.component +def comp(): + pass + +@dsl.pipeline +def my_pipeline(): + task = comp() + kubernetes.empty_dir_mount( + task, + volume_name='emptydir-vol-1', + mount_path='/mnt/my_vol_1', + medium='Memory', + size_limit='1Gi' + ) + +if __name__ == '__main__': + from kfp import compiler + compiler.Compiler().compile(my_pipeline, __file__.replace('.py', '.yaml')) diff --git a/kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.yaml b/kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.yaml new file mode 100644 index 00000000000..177fbd2a7cf --- /dev/null +++ b/kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.yaml @@ -0,0 +1,60 @@ +# PIPELINE DEFINITION +# Name: my-pipeline +components: + comp-comp: + executorLabel: exec-comp +deploymentSpec: + executors: + exec-comp: + container: + args: + - --executor_input + - '{{$}}' + - --function_to_execute + - comp + command: + - sh + - -c + - "\nif ! [ -x \"$(command -v pip)\" ]; then\n python3 -m ensurepip ||\ + \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\ + \ python3 -m pip install --quiet --no-warn-script-location 'kfp==2.7.0'\ + \ '--no-deps' 'typing-extensions>=3.7.4,<5; python_version<\"3.9\"' && \"\ + $0\" \"$@\"\n" + - sh + - -ec + - 'program_path=$(mktemp -d) + + + printf "%s" "$0" > "$program_path/ephemeral_component.py" + + _KFP_RUNTIME=true python3 -m kfp.dsl.executor_main --component_module_path "$program_path/ephemeral_component.py" "$@" + + ' + - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\ + \ *\n\ndef comp():\n pass\n\n" + image: python:3.9 +pipelineInfo: + name: my-pipeline +root: + dag: + tasks: + comp: + cachingOptions: + enableCache: true + componentRef: + name: comp-comp + taskInfo: + name: comp +schemaVersion: 2.1.0 +sdkVersion: kfp-2.7.0 +--- +platforms: + kubernetes: + deploymentSpec: + executors: + exec-comp: + emptyDirMounts: + - medium: Memory + mountPath: /mnt/my_vol_1 + sizeLimit: 1Gi + volumeName: emptydir-vol-1 diff --git a/kubernetes_platform/python/test/unit/test_empty_dir_mounts.py b/kubernetes_platform/python/test/unit/test_empty_dir_mounts.py new file mode 100644 index 00000000000..dcde2d140c4 --- /dev/null +++ b/kubernetes_platform/python/test/unit/test_empty_dir_mounts.py @@ -0,0 +1,136 @@ +# 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 TestEmptyDirMounts: + + def test_add_one(self): + + @dsl.pipeline + def my_pipeline(): + task = comp() + kubernetes.empty_dir_mount( + task, + volume_name='emptydir-vol-1', + mount_path='/mnt/my_vol_1', + medium='Memory', + size_limit='1Gi' + ) + + assert json_format.MessageToDict(my_pipeline.platform_spec) == { + 'platforms': { + 'kubernetes': { + 'deploymentSpec': { + 'executors': { + 'exec-comp': { + 'emptyDirMounts': [{ + 'medium': 'Memory', + 'mountPath': '/mnt/my_vol_1', + 'sizeLimit': '1Gi', + 'volumeName': 'emptydir-vol-1' + }] + } + } + } + } + } + } + + def test_add_two(self): + + @dsl.pipeline + def my_pipeline(): + task = comp() + kubernetes.empty_dir_mount( + task, + volume_name='emptydir-vol-1', + mount_path='/mnt/my_vol_1', + medium='Memory', + size_limit='1Gi' + ) + kubernetes.empty_dir_mount( + task, + volume_name='emptydir-vol-2', + mount_path='/mnt/my_vol_2' + ) + + assert json_format.MessageToDict(my_pipeline.platform_spec) == { + 'platforms': { + 'kubernetes': { + 'deploymentSpec': { + 'executors': { + 'exec-comp': { + 'emptyDirMounts': [{ + 'medium': 'Memory', + 'mountPath': '/mnt/my_vol_1', + 'sizeLimit': '1Gi', + 'volumeName': 'emptydir-vol-1' + }, + { + 'mountPath': '/mnt/my_vol_2', + 'volumeName': 'emptydir-vol-2' + }] + } + } + } + } + } + } + + def test_respects_other_configuration(self): + + @dsl.pipeline + def my_pipeline(): + task = comp() + + kubernetes.empty_dir_mount( + task, + volume_name='emptydir-vol-1', + mount_path='/mnt/my_vol_1', + medium='Memory', + size_limit='1Gi' + ) + + # this should exist too + kubernetes.set_image_pull_secrets(task, ['secret-name']) + + assert json_format.MessageToDict(my_pipeline.platform_spec) == { + 'platforms': { + 'kubernetes': { + 'deploymentSpec': { + 'executors': { + 'exec-comp': { + 'emptyDirMounts': [{ + 'medium': 'Memory', + 'mountPath': '/mnt/my_vol_1', + 'sizeLimit': '1Gi', + 'volumeName': 'emptydir-vol-1' + }], + 'imagePullSecret': [{ + 'secretName': 'secret-name' + }] + } + } + } + } + } + } + +@dsl.component +def comp(): + pass