Apply overrides before applying defaults to sidecar resource requirem…

…ents (flyteorg#197) Signed-off-by: Jeev B <jeev.balakrishnan@freenome.com>
eapolinario · Aug 12, 2021 · 81416b9 · 81416b9
1 parent e9a2b0e
commit 81416b9
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 16 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go
@@ -60,7 +60,7 @@ func ApplyResourceOverrides(ctx context.Context, resources v1.ResourceRequiremen
 
 	if _, found := resources.Requests[v1.ResourceMemory]; !found {
 		// use memory limit if set else default to config
-		if _, limitSet := resources.Limits[v1.ResourceCPU]; limitSet {
+		if _, limitSet := resources.Limits[v1.ResourceMemory]; limitSet {
 			resources.Requests[v1.ResourceMemory] = resources.Limits[v1.ResourceMemory]
 		} else {
 			resources.Requests[v1.ResourceMemory] = resource.MustParse(config.GetK8sPluginConfig().DefaultMemoryRequest)

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go
@@ -29,13 +29,6 @@ func TestApplyResourceOverrides_OverrideCpu(t *testing.T) {
 		},
 	})
 	assert.EqualValues(t, cpuRequest, overrides.Requests[v1.ResourceCPU])
-
-	// request equals limit if not set
-	overrides = ApplyResourceOverrides(context.Background(), v1.ResourceRequirements{
-		Limits: v1.ResourceList{
-			v1.ResourceCPU: cpuLimit,
-		},
-	})
 	assert.EqualValues(t, cpuLimit, overrides.Limits[v1.ResourceCPU])
 
 	// request equals limit if not set
@@ -45,6 +38,7 @@ func TestApplyResourceOverrides_OverrideCpu(t *testing.T) {
 		},
 	})
 	assert.EqualValues(t, cpuLimit, overrides.Requests[v1.ResourceCPU])
+	assert.EqualValues(t, cpuLimit, overrides.Limits[v1.ResourceCPU])
 }
 
 func TestApplyResourceOverrides_OverrideMemory(t *testing.T) {
@@ -70,15 +64,13 @@ func TestApplyResourceOverrides_OverrideMemory(t *testing.T) {
 	assert.EqualValues(t, memoryLimit, overrides.Limits[v1.ResourceMemory])
 
 	// request equals limit if not set
-	cpuLimit := resource.MustParse("2")
 	overrides = ApplyResourceOverrides(context.Background(), v1.ResourceRequirements{
 		Limits: v1.ResourceList{
 			v1.ResourceMemory: memoryLimit,
-			v1.ResourceCPU:    cpuLimit,
 		},
 	})
 	assert.EqualValues(t, memoryLimit, overrides.Requests[v1.ResourceMemory])
-	assert.EqualValues(t, cpuLimit, overrides.Requests[v1.ResourceCPU])
+	assert.EqualValues(t, memoryLimit, overrides.Limits[v1.ResourceMemory])
 }
 
 func TestApplyResourceOverrides_OverrideEphemeralStorage(t *testing.T) {

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go
@@ -197,8 +197,10 @@ func TestToK8sPod(t *testing.T) {
 		ResourceTolerations: map[v1.ResourceName][]v1.Toleration{
 			v1.ResourceStorage: {tolStorage},
 			ResourceNvidiaGPU:  {tolGPU},
-		}}),
-	)
+		},
+		DefaultCPURequest:    "1024m",
+		DefaultMemoryRequest: "1024Mi",
+	}))
 
 	op := &pluginsIOMock.OutputFilePaths{}
 	op.On("GetOutputPrefixPath").Return(storage.DataReference(""))
@@ -262,7 +264,9 @@ func TestToK8sPod(t *testing.T) {
 			DefaultNodeSelector: map[string]string{
 				"nodeId": "123",
 			},
-			SchedulerName: "myScheduler",
+			SchedulerName:        "myScheduler",
+			DefaultCPURequest:    "1024m",
+			DefaultMemoryRequest: "1024Mi",
 		}))
 
 		p, err := ToK8sPodSpec(ctx, x)

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/testdata/config.yaml b/flyteplugins/go/tasks/pluginmachinery/flytek8s/testdata/config.yaml
@@ -6,6 +6,8 @@ plugins:
   # All k8s plugins default configuration
   k8s:
     scheduler-name: flyte-scheduler
+    default-cpus: 1024m
+    default-memory: 1024Mi
     default-annotations:
       - annotationKey1: annotationValue1
       - annotationKey2: annotationValue2

diff --git a/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar.go b/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar.go
@@ -39,11 +39,11 @@ func validateAndFinalizePod(
 	for index, container := range pod.Spec.Containers {
 		if container.Name == primaryContainerName {
 			hasPrimaryContainer = true
-			container.Resources = *flytek8s.ApplyResourceOverrides(ctx, container.Resources)
 			if taskCtx.TaskExecutionMetadata().GetOverrides() != nil && taskCtx.TaskExecutionMetadata().GetOverrides().GetResources() != nil {
 				resOverrides := taskCtx.TaskExecutionMetadata().GetOverrides().GetResources()
 				flytek8s.MergeResources(*resOverrides, &container.Resources)
 			}
+			container.Resources = *flytek8s.ApplyResourceOverrides(ctx, container.Resources)
 		}
 		modifiedCommand, err := template.Render(ctx, container.Command, template.Parameters{
 			TaskExecMetadata: taskCtx.TaskExecutionMetadata(),

diff --git a/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar_test.go b/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar_test.go
@@ -478,7 +478,7 @@ func TestBuildSidecarResource(t *testing.T) {
 	}
 
 	// Assert resource requirements are correctly set
-	expectedCPURequest := resource.MustParse("1024m")
+	expectedCPURequest := resource.MustParse("2048m")
 	assert.Equal(t, expectedCPURequest.Value(), res.(*v1.Pod).Spec.Containers[0].Resources.Requests.Cpu().Value())
 	expectedMemRequest := resource.MustParse("1024Mi")
 	assert.Equal(t, expectedMemRequest.Value(), res.(*v1.Pod).Spec.Containers[0].Resources.Requests.Memory().Value())