Add emptyDir sizeLimit support

Signed-off-by: Jacob Salway <jacob.salway@rokt.com>
kubeflow · Apr 18, 2024 · fc1c352 · fc1c352
1 parent 5e9c0b4
commit fc1c352
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/pkg/controller/sparkapplication/submission.go b/pkg/controller/sparkapplication/submission.go
@@ -516,6 +516,9 @@ func buildLocalVolumeOptions(prefix string, volume v1.Volume, volumeMount v1.Vol
 		}
 	case volume.EmptyDir != nil:
 		options = append(options, fmt.Sprintf(VolumeMountPathTemplate, "emptyDir", volume.Name, volumeMount.MountPath))
+		if volume.EmptyDir.SizeLimit != nil {
+			options = append(options, fmt.Sprintf(VolumeMountOptionTemplate, "emptyDir", volume.Name, "sizeLimit", volume.EmptyDir.SizeLimit.String()))
+		}
 	case volume.PersistentVolumeClaim != nil:
 		options = append(options, fmt.Sprintf(VolumeMountPathTemplate, "persistentVolumeClaim", volume.Name, volumeMount.MountPath))
 		options = append(options, fmt.Sprintf(VolumeMountOptionTemplate, "persistentVolumeClaim", volume.Name, "claimName", volume.PersistentVolumeClaim.ClaimName))

diff --git a/pkg/controller/sparkapplication/submission_test.go b/pkg/controller/sparkapplication/submission_test.go
@@ -29,6 +29,7 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/kubeflow/spark-operator/pkg/apis/sparkoperator.k8s.io/v1beta2"
@@ -366,6 +367,61 @@ func TestAddLocalDir_Driver_Executor(t *testing.T) {
 	assert.Equal(t, fmt.Sprintf(VolumeMountOptionPathTemplate, "executor", "hostPath", volumes[0].Name, "path", volumes[0].HostPath.Path), localDirOptions[3])
 }
 
+func TestAddEmptyDir_Driver_Executor_WithSizeLimit(t *testing.T) {
+	sizeLimit := resource.MustParse("5Gi")
+	volumes := []corev1.Volume{
+		{
+			Name: "spark-local-dir-1",
+			VolumeSource: corev1.VolumeSource{
+				EmptyDir: &corev1.EmptyDirVolumeSource{
+					SizeLimit: &sizeLimit,
+				},
+			},
+		},
+	}
+
+	volumeMounts := []corev1.VolumeMount{
+		{
+			Name:      "spark-local-dir-1",
+			MountPath: "/tmp/mnt-1",
+		},
+	}
+
+	app := &v1beta2.SparkApplication{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "spark-test",
+			UID:  "spark-test-1",
+		},
+		Spec: v1beta2.SparkApplicationSpec{
+			Volumes: volumes,
+			Driver: v1beta2.DriverSpec{
+				SparkPodSpec: v1beta2.SparkPodSpec{
+					VolumeMounts: volumeMounts,
+				},
+			},
+			Executor: v1beta2.ExecutorSpec{
+				SparkPodSpec: v1beta2.SparkPodSpec{
+					VolumeMounts: volumeMounts,
+				},
+			},
+		},
+	}
+
+	localDirOptions, err := addLocalDirConfOptions(app)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, 0, len(app.Spec.Volumes))
+	assert.Equal(t, 0, len(app.Spec.Driver.VolumeMounts))
+	assert.Equal(t, 0, len(app.Spec.Executor.VolumeMounts))
+	assert.Equal(t, 4, len(localDirOptions))
+	assert.Equal(t, fmt.Sprintf(VolumeMountPathTemplate, "driver", "emptyDir", volumes[0].Name, volumeMounts[0].MountPath), localDirOptions[0])
+	assert.Equal(t, fmt.Sprintf(VolumeMountOptionPathTemplate, "driver", "emptyDir", volumes[0].Name, "sizeLimit", volumes[0].EmptyDir.SizeLimit.String()), localDirOptions[1])
+	assert.Equal(t, fmt.Sprintf(VolumeMountPathTemplate, "executor", "emptyDir", volumes[0].Name, volumeMounts[0].MountPath), localDirOptions[2])
+	assert.Equal(t, fmt.Sprintf(VolumeMountOptionPathTemplate, "executor", "emptyDir", volumes[0].Name, "sizeLimit", volumes[0].EmptyDir.SizeLimit.String()), localDirOptions[3])
+}
+
 func TestPopulateLabels_Driver_Executor(t *testing.T) {
 	const (
 		AppLabelKey        = "app-label-key"