better test for ephermeral storage

Signed-off-by: Future Outlier <[email protected]>
flyteorg · Jan 2, 2024 · 72b443d · 72b443d
1 parent bc92ed1
commit 72b443d
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 58 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go
@@ -24,7 +24,8 @@ import (
 
 var resourceRequirements = &v1.ResourceRequirements{
 	Limits: v1.ResourceList{
-		v1.ResourceCPU: resource.MustParse("1024m"),
+		v1.ResourceCPU:     resource.MustParse("1024m"),
+		v1.ResourceStorage: resource.MustParse("100M"),
 	},
 }
 

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go
@@ -33,8 +33,8 @@ func TestGetTolerationsForResources(t *testing.T) {
 		Effect:   v12.TaintEffectNoSchedule,
 	}
 
-	tolStorage := v12.Toleration{
-		Key:      "storage",
+	tolEphemeralStorage := v12.Toleration{
+		Key:      "ephemeral-storage",
 		Value:    "dedicated",
 		Operator: v12.TolerationOpExists,
 		Effect:   v12.TaintEffectNoSchedule,
@@ -55,8 +55,8 @@ func TestGetTolerationsForResources(t *testing.T) {
 			args{
 				v12.ResourceRequirements{
 					Limits: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
@@ -69,8 +69,8 @@ func TestGetTolerationsForResources(t *testing.T) {
 			args{
 				v12.ResourceRequirements{
 					Requests: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
@@ -83,12 +83,12 @@ func TestGetTolerationsForResources(t *testing.T) {
 			args{
 				v12.ResourceRequirements{
 					Limits: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 					Requests: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
@@ -101,84 +101,84 @@ func TestGetTolerationsForResources(t *testing.T) {
 			args{
 				v12.ResourceRequirements{
 					Limits: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
 			map[v12.ResourceName][]v12.Toleration{
-				v12.ResourceStorage: {tolStorage},
-				ResourceNvidiaGPU:   {tolGPU},
+				v12.ResourceEphemeralStorage: {tolEphemeralStorage},
+				ResourceNvidiaGPU:            {tolGPU},
 			},
 			nil,
-			[]v12.Toleration{tolStorage},
+			[]v12.Toleration{tolEphemeralStorage},
 		},
 		{
 			"tolerations-req",
 			args{
 				v12.ResourceRequirements{
 					Requests: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
 			map[v12.ResourceName][]v12.Toleration{
-				v12.ResourceStorage: {tolStorage},
-				ResourceNvidiaGPU:   {tolGPU},
+				v12.ResourceEphemeralStorage: {tolEphemeralStorage},
+				ResourceNvidiaGPU:            {tolGPU},
 			},
 			nil,
-			[]v12.Toleration{tolStorage},
+			[]v12.Toleration{tolEphemeralStorage},
 		},
 		{
 			"tolerations-both",
 			args{
 				v12.ResourceRequirements{
 					Limits: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 					Requests: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
 			map[v12.ResourceName][]v12.Toleration{
-				v12.ResourceStorage: {tolStorage},
-				ResourceNvidiaGPU:   {tolGPU},
+				v12.ResourceEphemeralStorage: {tolEphemeralStorage},
+				ResourceNvidiaGPU:            {tolGPU},
 			},
 			nil,
-			[]v12.Toleration{tolStorage},
+			[]v12.Toleration{tolEphemeralStorage},
 		},
 		{
 			"no-tolerations-both",
 			args{
 				v12.ResourceRequirements{
 					Limits: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
-						ResourceNvidiaGPU:   resource.MustParse("1"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
+						ResourceNvidiaGPU:            resource.MustParse("1"),
 					},
 					Requests: v12.ResourceList{
-						v12.ResourceCPU:     resource.MustParse("1024m"),
-						v12.ResourceStorage: resource.MustParse("100M"),
+						v12.ResourceCPU:              resource.MustParse("1024m"),
+						v12.ResourceEphemeralStorage: resource.MustParse("100M"),
 					},
 				},
 			},
 			map[v12.ResourceName][]v12.Toleration{
-				v12.ResourceStorage: {tolStorage},
-				ResourceNvidiaGPU:   {tolGPU},
+				v12.ResourceEphemeralStorage: {tolEphemeralStorage},
+				ResourceNvidiaGPU:            {tolGPU},
 			},
 			nil,
-			[]v12.Toleration{tolStorage, tolGPU},
+			[]v12.Toleration{tolEphemeralStorage, tolGPU},
 		},
 		{
 			"default-tolerations",
 			args{},
 			nil,
-			[]v12.Toleration{tolStorage},
-			[]v12.Toleration{tolStorage},
+			[]v12.Toleration{tolEphemeralStorage},
+			[]v12.Toleration{tolEphemeralStorage},
 		},
 	}
 	for _, tt := range tests {

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go
@@ -693,10 +693,12 @@ func TestApplyGPUNodeSelectors(t *testing.T) {
 func updatePod(t *testing.T) {
 	taskExecutionMetadata := dummyTaskExecutionMetadata(&v1.ResourceRequirements{
 		Limits: v1.ResourceList{
-			v1.ResourceCPU: resource.MustParse("1024m"),
+			v1.ResourceCPU:              resource.MustParse("1024m"),
+			v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 		},
 		Requests: v1.ResourceList{
-			v1.ResourceCPU: resource.MustParse("1024m"),
+			v1.ResourceCPU:              resource.MustParse("1024m"),
+			v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 		},
 	}, nil)
 
@@ -807,11 +809,13 @@ func toK8sPodInterruptible(t *testing.T) {
 
 	x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{
 		Limits: v1.ResourceList{
-			v1.ResourceCPU:    resource.MustParse("1024m"),
-			ResourceNvidiaGPU: resource.MustParse("1"),
+			v1.ResourceCPU:              resource.MustParse("1024m"),
+			v1.ResourceEphemeralStorage: resource.MustParse("100M"),
+			ResourceNvidiaGPU:           resource.MustParse("1"),
 		},
 		Requests: v1.ResourceList{
-			v1.ResourceCPU: resource.MustParse("1024m"),
+			v1.ResourceCPU:              resource.MustParse("1024m"),
+			v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 		},
 	}, nil)
 
@@ -849,9 +853,17 @@ func TestToK8sPod(t *testing.T) {
 		Effect:   v1.TaintEffectNoSchedule,
 	}
 
+	tolEphemeralStorage := v1.Toleration{
+		Key:      "ephemeral-storage",
+		Value:    "dedicated",
+		Operator: v1.TolerationOpExists,
+		Effect:   v1.TaintEffectNoSchedule,
+	}
+
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
 		ResourceTolerations: map[v1.ResourceName][]v1.Toleration{
-			ResourceNvidiaGPU: {tolGPU},
+			v1.ResourceEphemeralStorage: {tolEphemeralStorage},
+			ResourceNvidiaGPU:           {tolGPU},
 		},
 		DefaultCPURequest:    resource.MustParse("1024m"),
 		DefaultMemoryRequest: resource.MustParse("1024Mi"),
@@ -864,42 +876,48 @@ func TestToK8sPod(t *testing.T) {
 	t.Run("WithGPU", func(t *testing.T) {
 		x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{
 			Limits: v1.ResourceList{
-				v1.ResourceCPU:    resource.MustParse("1024m"),
-				ResourceNvidiaGPU: resource.MustParse("1"),
+				v1.ResourceCPU:              resource.MustParse("1024m"),
+				v1.ResourceEphemeralStorage: resource.MustParse("100M"),
+				ResourceNvidiaGPU:           resource.MustParse("1"),
 			},
 			Requests: v1.ResourceList{
-				v1.ResourceCPU: resource.MustParse("1024m"),
+				v1.ResourceCPU:              resource.MustParse("1024m"),
+				v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 			},
 		}, nil)
 
 		p, _, _, err := ToK8sPodSpec(ctx, x)
 		assert.NoError(t, err)
-		assert.Equal(t, len(p.Tolerations), 1)
+		assert.Equal(t, len(p.Tolerations), 2)
 	})
 
 	t.Run("NoGPU", func(t *testing.T) {
 		x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{
 			Limits: v1.ResourceList{
-				v1.ResourceCPU: resource.MustParse("1024m"),
+				v1.ResourceCPU:              resource.MustParse("1024m"),
+				v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 			},
 			Requests: v1.ResourceList{
-				v1.ResourceCPU: resource.MustParse("1024m"),
+				v1.ResourceCPU:              resource.MustParse("1024m"),
+				v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 			},
 		}, nil)
 
 		p, _, _, err := ToK8sPodSpec(ctx, x)
 		assert.NoError(t, err)
-		assert.Equal(t, len(p.Tolerations), 0)
+		assert.Equal(t, len(p.Tolerations), 1)
 		assert.Equal(t, "some-acceptable-name", p.Containers[0].Name)
 	})
 
 	t.Run("Default toleration, selector, scheduler", func(t *testing.T) {
 		x := dummyExecContext(dummyTaskTemplate(), &v1.ResourceRequirements{
 			Limits: v1.ResourceList{
-				v1.ResourceCPU: resource.MustParse("1024m"),
+				v1.ResourceCPU:              resource.MustParse("1024m"),
+				v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 			},
 			Requests: v1.ResourceList{
-				v1.ResourceCPU: resource.MustParse("1024m"),
+				v1.ResourceCPU:              resource.MustParse("1024m"),
+				v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 			},
 		}, nil)
 

diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go
@@ -205,6 +205,9 @@ func TestContainerTaskExecutor_BuildResource(t *testing.T) {
 			assert.NotEmpty(t, j.Spec.Containers)
 			assert.Equal(t, containerResourceRequirements.Limits[v1.ResourceCPU], j.Spec.Containers[0].Resources.Limits[v1.ResourceCPU])
 
+			ephemeralStorageRes := j.Spec.Containers[0].Resources.Limits[v1.ResourceEphemeralStorage]
+			assert.Equal(t, int64(0), (&ephemeralStorageRes).Value())
+
 			assert.Equal(t, command, j.Spec.Containers[0].Command)
 			assert.Equal(t, []string{"test-data-reference"}, j.Spec.Containers[0].Args)
 

diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/sidecar_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/sidecar_test.go
@@ -230,9 +230,16 @@ func TestBuildSidecarResource_TaskType2(t *testing.T) {
 		Effect:   v1.TaintEffectNoSchedule,
 	}
 
+	tolEphemeralStorage := v1.Toleration{
+		Key:      "ephemeral-storage",
+		Value:    "dedicated",
+		Operator: v1.TolerationOpExists,
+		Effect:   v1.TaintEffectNoSchedule,
+	}
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
 		ResourceTolerations: map[v1.ResourceName][]v1.Toleration{
-			ResourceNvidiaGPU: {tolGPU},
+			v1.ResourceStorage: {tolEphemeralStorage},
+			ResourceNvidiaGPU:  {tolGPU},
 		},
 		DefaultCPURequest:    resource.MustParse("1024m"),
 		DefaultMemoryRequest: resource.MustParse("1024Mi"),
@@ -333,9 +340,16 @@ func TestBuildSidecarResource_TaskType1(t *testing.T) {
 		Effect:   v1.TaintEffectNoSchedule,
 	}
 
+	tolEphemeralStorage := v1.Toleration{
+		Key:      "ephemeral-storage",
+		Value:    "dedicated",
+		Operator: v1.TolerationOpExists,
+		Effect:   v1.TaintEffectNoSchedule,
+	}
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
 		ResourceTolerations: map[v1.ResourceName][]v1.Toleration{
-			ResourceNvidiaGPU: {tolGPU},
+			v1.ResourceStorage: {tolEphemeralStorage},
+			ResourceNvidiaGPU:  {tolGPU},
 		},
 		DefaultCPURequest:    resource.MustParse("1024m"),
 		DefaultMemoryRequest: resource.MustParse("1024Mi"),
@@ -399,7 +413,8 @@ func TestBuildSideResource_TaskType1_InvalidSpec(t *testing.T) {
 
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
 		ResourceTolerations: map[v1.ResourceName][]v1.Toleration{
-			ResourceNvidiaGPU: {},
+			v1.ResourceStorage: {},
+			ResourceNvidiaGPU:  {},
 		},
 		DefaultCPURequest:    resource.MustParse("1024m"),
 		DefaultMemoryRequest: resource.MustParse("1024Mi"),
@@ -442,9 +457,16 @@ func TestBuildSidecarResource(t *testing.T) {
 		Effect:   v1.TaintEffectNoSchedule,
 	}
 
+	tolEphemeralStorage := v1.Toleration{
+		Key:      "ephemeral-storage",
+		Value:    "dedicated",
+		Operator: v1.TolerationOpExists,
+		Effect:   v1.TaintEffectNoSchedule,
+	}
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
 		ResourceTolerations: map[v1.ResourceName][]v1.Toleration{
-			ResourceNvidiaGPU: {tolGPU},
+			v1.ResourceStorage: {tolEphemeralStorage},
+			ResourceNvidiaGPU:  {tolGPU},
 		},
 		DefaultCPURequest:    resource.MustParse("1024m"),
 		DefaultMemoryRequest: resource.MustParse("1024Mi"),

diff --git a/flyteplugins/go/tasks/plugins/presto/helpers_test.go b/flyteplugins/go/tasks/plugins/presto/helpers_test.go
@@ -42,7 +42,8 @@ func GetPrestoQueryTaskTemplate() idlCore.TaskTemplate {
 
 var resourceRequirements = &v1.ResourceRequirements{
 	Limits: v1.ResourceList{
-		v1.ResourceCPU: resource.MustParse("1024m"),
+		v1.ResourceCPU:              resource.MustParse("1024m"),
+		v1.ResourceEphemeralStorage: resource.MustParse("100M"),
 	},
 }