From 7877302152f8893e1155725e5de7055ae0ba8ce2 Mon Sep 17 00:00:00 2001 From: Austin Zhao Date: Tue, 5 Jul 2022 15:30:09 -0400 Subject: [PATCH] [TEP-0104] Update Pod with Task-level Requirements --- pkg/pod/pod.go | 25 ++++++ pkg/pod/pod_test.go | 209 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 234 insertions(+) diff --git a/pkg/pod/pod.go b/pkg/pod/pod.go index 1f0d3eb856b..b5470dee8e5 100644 --- a/pkg/pod/pod.go +++ b/pkg/pod/pod.go @@ -141,6 +141,11 @@ func (b *Builder) Build(ctx context.Context, taskRun *v1beta1.TaskRun, taskSpec if err != nil { return nil, err } + if alphaAPIEnabled && taskRun.Spec.ComputeResources != nil { + if err = applyTaskLevelComputeResources(steps, taskRun.Spec.ComputeResources, taskSpec); err != nil { + return nil, err + } + } sidecars, err := v1beta1.MergeSidecarsWithOverrides(taskSpec.Sidecars, taskRun.Spec.SidecarOverrides) if err != nil { return nil, err @@ -354,6 +359,26 @@ func (b *Builder) Build(ctx context.Context, taskRun *v1beta1.TaskRun, taskSpec return newPod, nil } +// applyTaskLevelComputeResources applies the task-level compute resource requirements to the Pod. +func applyTaskLevelComputeResources(steps []v1beta1.Step, computeResources *corev1.ResourceRequirements, taskSpec v1beta1.TaskSpec) error { + if taskSpec.StepTemplate != nil { + taskSpec.StepTemplate.Resources = corev1.ResourceRequirements{} + } + + // TODO: apply LimitRanges + + for i := range steps { + if i == 0 { + steps[i].Resources.Requests = computeResources.Requests + } else { + steps[i].Resources.Requests = nil + } + steps[i].Resources.Limits = computeResources.Limits + } + + return nil +} + // makeLabels constructs the labels we will propagate from TaskRuns to Pods. func makeLabels(s *v1beta1.TaskRun) map[string]string { labels := make(map[string]string, len(s.ObjectMeta.Labels)+1) diff --git a/pkg/pod/pod_test.go b/pkg/pod/pod_test.go index c0a23c4298a..5823245c941 100644 --- a/pkg/pod/pod_test.go +++ b/pkg/pod/pod_test.go @@ -2014,6 +2014,215 @@ debug-fail-continue-heredoc-randomly-generated-mz4c7 } } +func TestPodBuild_TaskLevelResourceRequirements(t *testing.T) { + testcases := []struct { + desc string + ts v1beta1.TaskSpec + trs v1beta1.TaskRunSpec + expectedComputeResources []corev1.ResourceRequirements + }{{ + desc: "only with requests", + ts: v1beta1.TaskSpec{ + Steps: []v1beta1.Step{{ + Name: "1st-step", + Image: "image", + Command: []string{"cmd"}, + }, { + Name: "2nd-step", + Image: "image", + Command: []string{"cmd"}, + }}, + }, + trs: v1beta1.TaskRunSpec{ + ComputeResources: &corev1.ResourceRequirements{ + Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")}, + }, + }, + expectedComputeResources: []corev1.ResourceRequirements{{ + Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")}, + }, { + Requests: nil, + }}, + }, { + desc: "only with limits", + ts: v1beta1.TaskSpec{ + Steps: []v1beta1.Step{{ + Name: "1st-step", + Image: "image", + Command: []string{"cmd"}, + }, { + Name: "2nd-step", + Image: "image", + Command: []string{"cmd"}, + }}, + }, + trs: v1beta1.TaskRunSpec{ + ComputeResources: &corev1.ResourceRequirements{ + Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")}, + }, + }, + expectedComputeResources: []corev1.ResourceRequirements{{ + Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")}, + }, { + Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("500m")}, + }}, + }, { + desc: "both with requests and limits", + ts: v1beta1.TaskSpec{ + Steps: []v1beta1.Step{{ + Name: "1st-step", + Image: "image", + Command: []string{"cmd"}, + }, { + Name: "2nd-step", + Image: "image", + Command: []string{"cmd"}, + }}, + }, + trs: v1beta1.TaskRunSpec{ + ComputeResources: &corev1.ResourceRequirements{ + Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")}, + Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")}, + }, + }, + expectedComputeResources: []corev1.ResourceRequirements{{ + Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")}, + Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")}, + }, { + Limits: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("2")}, + }}, + }, { + desc: "overwrite StepTemplate resources requirements", + ts: v1beta1.TaskSpec{ + Steps: []v1beta1.Step{{ + Name: "1st-step", + Image: "image", + Command: []string{"cmd"}, + }, { + Name: "2nd-step", + Image: "image", + Command: []string{"cmd"}, + }}, + StepTemplate: &v1beta1.StepTemplate{ + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("500m"), + corev1.ResourceMemory: resource.MustParse("500Mi"), + }, + }, + }, + }, + trs: v1beta1.TaskRunSpec{ + ComputeResources: &corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("2"), + corev1.ResourceMemory: resource.MustParse("2Gi"), + }, + }, + }, + expectedComputeResources: []corev1.ResourceRequirements{{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("2"), + corev1.ResourceMemory: resource.MustParse("2Gi"), + }, + }, { + Requests: nil, + }}, + }, { + desc: "with sidecar resource requirements", + ts: v1beta1.TaskSpec{ + Steps: []v1beta1.Step{{ + Name: "1st-step", + Image: "image", + Command: []string{"cmd"}, + }}, + Sidecars: []v1beta1.Sidecar{{ + Name: "sidecar", + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("750m"), + }, + Limits: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("1.5"), + }, + }, + }}, + }, + trs: v1beta1.TaskRunSpec{ + ComputeResources: &corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("2"), + corev1.ResourceMemory: resource.MustParse("2Gi"), + }, + }, + }, + expectedComputeResources: []corev1.ResourceRequirements{{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("2"), + corev1.ResourceMemory: resource.MustParse("2Gi"), + }, + }, { + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("750m"), + }, + Limits: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("1.5"), + }, + }}, + }} + + for _, tc := range testcases { + t.Run(tc.desc, func(t *testing.T) { + names.TestingSeed() + store := config.NewStore(logtesting.TestLogger(t)) + enableAlphaAPI := map[string]string{"enable-api-fields": "alpha"} + store.OnConfigChanged( + &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{Name: config.GetFeatureFlagsConfigName(), Namespace: system.Namespace()}, + Data: enableAlphaAPI, + }, + ) + + kubeclient := fakek8s.NewSimpleClientset( + &corev1.ServiceAccount{ObjectMeta: metav1.ObjectMeta{Name: "default", Namespace: "default"}}, + ) + builder := Builder{ + Images: images, + KubeClient: kubeclient, + } + tr := &v1beta1.TaskRun{ + ObjectMeta: metav1.ObjectMeta{ + Name: "foo-taskrun", + Namespace: "default", + }, + Spec: tc.trs, + } + + gotPod, err := builder.Build(store.ToContext(context.Background()), tr, tc.ts) + if err != nil { + t.Fatalf("builder.Build: %v", err) + } + + if err := verifyTaskLevelComputeResources(tc.expectedComputeResources, gotPod.Spec.Containers); err != nil { + t.Errorf("verifyTaskLevelComputeResources: %v", err) + } + }) + } +} + +// verifyTaskLevelComputeResources verifies that the given TaskRun's containers have the expected compute resources. +func verifyTaskLevelComputeResources(expectedComputeResources []corev1.ResourceRequirements, containers []corev1.Container) error { + if len(expectedComputeResources) != len(containers) { + return fmt.Errorf("expected %d compute resource requirements, got %d", len(expectedComputeResources), len(containers)) + } + for i, r := range expectedComputeResources { + if d := cmp.Diff(r, containers[i].Resources); d != "" { + return fmt.Errorf("container \"#%d\" resource requirements don't match %s", i, diff.PrintWantGot(d)) + } + } + return nil +} + func TestMakeLabels(t *testing.T) { taskRunName := "task-run-name" want := map[string]string{