Skip to content

Commit 1063e9c

Browse files
gregsheremetagregsheremeta
committed
feat(Backend + SDK): Update kfp backend and kubernetes sdk to support EmptyDir
Update kfp backend and kubernetes sdk to support mounting EmptyDir volumes to task pods. Inspired by #10427 Fixes: #10656 Signed-off-by: Greg Sheremeta <gshereme@redhat.com>
1 parent d911c8b commit 1063e9c

File tree

9 files changed

+484
-6
lines changed

9 files changed

+484
-6
lines changed

‎backend/src/v2/driver/driver.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import (
3636
"google.golang.org/protobuf/encoding/protojson"
3737
"google.golang.org/protobuf/types/known/structpb"
3838
k8score "k8s.io/api/core/v1"
39+
"k8s.io/apimachinery/pkg/api/resource"
3940
k8sres "k8s.io/apimachinery/pkg/api/resource"
4041
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
4142
"k8s.io/client-go/kubernetes"
@@ -665,6 +666,33 @@ func extendPodSpecPatch(
665666
podSpec.Volumes = append(podSpec.Volumes, ephemeralVolume)
666667
podSpec.Containers[0].VolumeMounts = append(podSpec.Containers[0].VolumeMounts, ephemeralVolumeMount)
667668
}
669+
670+
// EmptyDirMounts
671+
for _, emptyDirVolumeSpec := range kubernetesExecutorConfig.GetEmptyDirMounts() {
672+
var sizeLimitResource *resource.Quantity
673+
if emptyDirVolumeSpec.GetSizeLimit() != "" {
674+
r := k8sres.MustParse(emptyDirVolumeSpec.GetSizeLimit())
675+
sizeLimitResource = &r
676+
}
677+
678+
emptyDirVolume := k8score.Volume{
679+
Name: emptyDirVolumeSpec.GetVolumeName(),
680+
VolumeSource: k8score.VolumeSource{
681+
EmptyDir: &k8score.EmptyDirVolumeSource{
682+
Medium: k8score.StorageMedium(emptyDirVolumeSpec.GetMedium()),
683+
SizeLimit: sizeLimitResource,
684+
},
685+
},
686+
}
687+
emptyDirVolumeMount := k8score.VolumeMount{
688+
Name: emptyDirVolumeSpec.GetVolumeName(),
689+
MountPath: emptyDirVolumeSpec.GetMountPath(),
690+
}
691+
692+
podSpec.Volumes = append(podSpec.Volumes, emptyDirVolume)
693+
podSpec.Containers[0].VolumeMounts = append(podSpec.Containers[0].VolumeMounts, emptyDirVolumeMount)
694+
}
695+
668696
return nil
669697
}
670698

‎backend/src/v2/driver/driver_test.go

Lines changed: 163 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,10 @@ package driver
1515

1616
import (
1717
"encoding/json"
18+
"testing"
19+
1820
k8sres "k8s.io/apimachinery/pkg/api/resource"
1921
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
20-
"testing"
2122

2223
"github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec"
2324
"github.com/kubeflow/pipelines/backend/src/v2/metadata"
@@ -532,7 +533,7 @@ func Test_extendPodSpecPatch_Secret(t *testing.T) {
532533
{
533534
Name: "secret1",
534535
VolumeSource: k8score.VolumeSource{
535-
Secret: &k8score.SecretVolumeSource{SecretName: "secret1", Optional: &[]bool{false}[0],},
536+
Secret: &k8score.SecretVolumeSource{SecretName: "secret1", Optional: &[]bool{false}[0]},
536537
},
537538
},
538539
},
@@ -730,7 +731,7 @@ func Test_extendPodSpecPatch_ConfigMap(t *testing.T) {
730731
VolumeSource: k8score.VolumeSource{
731732
ConfigMap: &k8score.ConfigMapVolumeSource{
732733
LocalObjectReference: k8score.LocalObjectReference{Name: "cm1"},
733-
Optional: &[]bool{false}[0],},
734+
Optional: &[]bool{false}[0]},
734735
},
735736
},
736737
},
@@ -890,6 +891,165 @@ func Test_extendPodSpecPatch_ConfigMap(t *testing.T) {
890891
}
891892
}
892893

894+
func Test_extendPodSpecPatch_EmptyVolumeMount(t *testing.T) {
895+
medium := "Memory"
896+
sizeLimit := "1Gi"
897+
var sizeLimitResource *k8sres.Quantity
898+
r := k8sres.MustParse(sizeLimit)
899+
sizeLimitResource = &r
900+
901+
tests := []struct {
902+
name string
903+
k8sExecCfg *kubernetesplatform.KubernetesExecutorConfig
904+
podSpec *k8score.PodSpec
905+
expected *k8score.PodSpec
906+
}{
907+
{
908+
"Valid - emptydir mount with no medium or size limit",
909+
&kubernetesplatform.KubernetesExecutorConfig{
910+
EmptyDirMounts: []*kubernetesplatform.EmptyDirMount{
911+
{
912+
VolumeName: "emptydir1",
913+
MountPath: "/data/path",
914+
},
915+
},
916+
},
917+
&k8score.PodSpec{
918+
Containers: []k8score.Container{
919+
{
920+
Name: "main",
921+
},
922+
},
923+
},
924+
&k8score.PodSpec{
925+
Containers: []k8score.Container{
926+
{
927+
Name: "main",
928+
VolumeMounts: []k8score.VolumeMount{
929+
{
930+
Name: "emptydir1",
931+
MountPath: "/data/path",
932+
},
933+
},
934+
},
935+
},
936+
Volumes: []k8score.Volume{
937+
{
938+
Name: "emptydir1",
939+
VolumeSource: k8score.VolumeSource{
940+
EmptyDir: &k8score.EmptyDirVolumeSource{},
941+
},
942+
},
943+
},
944+
},
945+
},
946+
{
947+
"Valid - emptydir mount with medium and size limit",
948+
&kubernetesplatform.KubernetesExecutorConfig{
949+
EmptyDirMounts: []*kubernetesplatform.EmptyDirMount{
950+
{
951+
VolumeName: "emptydir1",
952+
MountPath: "/data/path",
953+
Medium: &medium,
954+
SizeLimit: &sizeLimit,
955+
},
956+
},
957+
},
958+
&k8score.PodSpec{
959+
Containers: []k8score.Container{
960+
{
961+
Name: "main",
962+
},
963+
},
964+
},
965+
&k8score.PodSpec{
966+
Containers: []k8score.Container{
967+
{
968+
Name: "main",
969+
VolumeMounts: []k8score.VolumeMount{
970+
{
971+
Name: "emptydir1",
972+
MountPath: "/data/path",
973+
},
974+
},
975+
},
976+
},
977+
Volumes: []k8score.Volume{
978+
{
979+
Name: "emptydir1",
980+
VolumeSource: k8score.VolumeSource{
981+
EmptyDir: &k8score.EmptyDirVolumeSource{
982+
Medium: k8score.StorageMedium(medium),
983+
SizeLimit: sizeLimitResource,
984+
},
985+
},
986+
},
987+
},
988+
},
989+
},
990+
{
991+
"Valid - multiple emptydir mounts",
992+
&kubernetesplatform.KubernetesExecutorConfig{
993+
EmptyDirMounts: []*kubernetesplatform.EmptyDirMount{
994+
{
995+
VolumeName: "emptydir1",
996+
MountPath: "/data/path",
997+
},
998+
{
999+
VolumeName: "emptydir2",
1000+
MountPath: "/data/path2",
1001+
},
1002+
},
1003+
},
1004+
&k8score.PodSpec{
1005+
Containers: []k8score.Container{
1006+
{
1007+
Name: "main",
1008+
},
1009+
},
1010+
},
1011+
&k8score.PodSpec{
1012+
Containers: []k8score.Container{
1013+
{
1014+
Name: "main",
1015+
VolumeMounts: []k8score.VolumeMount{
1016+
{
1017+
Name: "emptydir1",
1018+
MountPath: "/data/path",
1019+
},
1020+
{
1021+
Name: "emptydir2",
1022+
MountPath: "/data/path2",
1023+
},
1024+
},
1025+
},
1026+
},
1027+
Volumes: []k8score.Volume{
1028+
{
1029+
Name: "emptydir1",
1030+
VolumeSource: k8score.VolumeSource{
1031+
EmptyDir: &k8score.EmptyDirVolumeSource{},
1032+
},
1033+
},
1034+
{
1035+
Name: "emptydir2",
1036+
VolumeSource: k8score.VolumeSource{
1037+
EmptyDir: &k8score.EmptyDirVolumeSource{},
1038+
},
1039+
},
1040+
},
1041+
},
1042+
},
1043+
}
1044+
for _, tt := range tests {
1045+
t.Run(tt.name, func(t *testing.T) {
1046+
err := extendPodSpecPatch(tt.podSpec, tt.k8sExecCfg, nil, nil)
1047+
assert.Nil(t, err)
1048+
assert.Equal(t, tt.expected, tt.podSpec)
1049+
})
1050+
}
1051+
}
1052+
8931053
func Test_extendPodSpecPatch_ImagePullSecrets(t *testing.T) {
8941054
tests := []struct {
8951055
name string

‎go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ require (
3232
github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-exithandler v0.0.0-20231127195001-a75d4b3711ff
3333
github.com/kubeflow/kfp-tekton/tekton-catalog/tekton-kfptask v0.0.0-20231127195001-a75d4b3711ff
3434
github.com/kubeflow/pipelines/api v0.0.0-20231027040853-58ce09e07d03
35-
github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240403164522-8b2a099e8c9f
35+
github.com/kubeflow/pipelines/kubernetes_platform v0.0.0-20240725205754-d911c8b73b49
3636
github.com/kubeflow/pipelines/third_party/ml-metadata v0.0.0-20230810215105-e1f0c010f800
3737
github.com/lestrrat-go/strftime v1.0.4
3838
github.com/mattn/go-sqlite3 v1.14.19

‎go.sum

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎kubernetes_platform/python/kfp/kubernetes/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
'add_toleration',
2323
'CreatePVC',
2424
'DeletePVC',
25+
'empty_dir_mount',
2526
'mount_pvc',
2627
'set_image_pull_policy',
2728
'use_field_path_as_env',
@@ -49,3 +50,4 @@
4950
from kfp.kubernetes.volume import CreatePVC
5051
from kfp.kubernetes.volume import DeletePVC
5152
from kfp.kubernetes.volume import mount_pvc
53+
from kfp.kubernetes.empty_dir import empty_dir_mount

‎kubernetes_platform/python/kfp/kubernetes/empty_dir.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Copyright 2024 The Kubeflow Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from typing import Optional
16+
17+
from google.protobuf import json_format
18+
from kfp.dsl import PipelineTask
19+
from kfp.kubernetes import common
20+
from kfp.kubernetes import kubernetes_executor_config_pb2 as pb
21+
22+
23+
def empty_dir_mount(
24+
task: PipelineTask,
25+
volume_name: str,
26+
mount_path: str,
27+
medium: Optional[str] = None,
28+
size_limit: Optional[str] = None,
29+
) -> PipelineTask:
30+
"""Mount an EmptyDir volume to the task's container.
31+
32+
Args:
33+
task: Pipeline task.
34+
volume_name: Name of the EmptyDir volume.
35+
mount_path: Path within the container at which the EmptyDir should be mounted.
36+
medium: Storage medium to back the EmptyDir. Must be one of `Memory` or `HugePages`. Defaults to `None`.
37+
size_limit: Maximum size of the EmptyDir. For example, `5Gi`. Defaults to `None`.
38+
39+
Returns:
40+
Task object with updated EmptyDir mount configuration.
41+
"""
42+
43+
msg = common.get_existing_kubernetes_config_as_message(task)
44+
45+
empty_dir_mount = pb.EmptyDirMount(
46+
volume_name=volume_name,
47+
mount_path=mount_path,
48+
medium=medium,
49+
size_limit=size_limit,
50+
)
51+
52+
msg.empty_dir_mounts.append(empty_dir_mount)
53+
54+
task.platform_config['kubernetes'] = json_format.MessageToDict(msg)
55+
56+
return task

‎kubernetes_platform/python/test/snapshot/data/empty_dir_mounts.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Copyright 2024 The Kubeflow Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from kfp import dsl
16+
from kfp import kubernetes
17+
18+
19+
@dsl.component
20+
def comp():
21+
pass
22+
23+
@dsl.pipeline
24+
def my_pipeline():
25+
task = comp()
26+
kubernetes.empty_dir_mount(
27+
task,
28+
volume_name='emptydir-vol-1',
29+
mount_path='/mnt/my_vol_1',
30+
medium='Memory',
31+
size_limit='1Gi'
32+
)
33+
34+
if __name__ == '__main__':
35+
from kfp import compiler
36+
compiler.Compiler().compile(my_pipeline, __file__.replace('.py', '.yaml'))

0 commit comments

Comments
 (0)