feat(controller): Allow to configure main container resources (#4656)

Signed-off-by: terrytangyuan <terrytangyuan@gmail.com>
argoproj · Dec 9, 2020 · b1e2c20 · b1e2c20
1 parent 4f9fab9
commit b1e2c20
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 0 deletions.
diff --git a/config/config.go b/config/config.go
@@ -31,6 +31,9 @@ type Config struct {
 	// DEPRECATED: use `executor.resources` in configmap instead
 	ExecutorResources *apiv1.ResourceRequirements `json:"executorResources,omitempty"`
 
+	// MainContainer holds container customization for the main container
+	MainContainer *apiv1.Container `json:"mainContainer,omitempty"`
+
 	// KubeConfig specifies a kube config file for the wait & init containers
 	KubeConfig *KubeConfig `json:"kubeConfig,omitempty"`
 

diff --git a/workflow/controller/workflowpod.go b/workflow/controller/workflowpod.go
@@ -151,6 +151,12 @@ func (woc *wfOperationCtx) createWorkflowPod(nodeName string, mainCtr apiv1.Cont
 	wfSpec := woc.execWf.Spec.DeepCopy()
 
 	mainCtr.Name = common.MainContainerName
+	// Allow customization of main container resources.
+	if isResourcesSpecified(woc.controller.Config.MainContainer) &&
+		// Container resources in workflow spec takes precedence over the main container's configuration in controller.
+		!(isResourcesSpecified(tmpl.Container) && tmpl.Container.Name == "main") {
+		mainCtr.Resources = woc.controller.Config.MainContainer.Resources
+	}
 
 	var activeDeadlineSeconds *int64
 	wfDeadline := woc.getWorkflowDeadline()

diff --git a/workflow/controller/workflowpod_test.go b/workflow/controller/workflowpod_test.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	apiv1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/utils/pointer"
 	"sigs.k8s.io/yaml"
@@ -1146,6 +1147,33 @@ func TestPodSpecPatch(t *testing.T) {
 
 }
 
+func TestMainContainerCustomization(t *testing.T) {
+	mainCtrSpec := &apiv1.Container{
+		Name: "main",
+		Resources: apiv1.ResourceRequirements{
+			Limits: apiv1.ResourceList{
+				apiv1.ResourceCPU:    resource.MustParse("0.200"),
+				apiv1.ResourceMemory: resource.MustParse("512Mi"),
+			},
+		},
+	}
+	// podSpecPatch in workflow spec takes precedence over the main container's
+	// configuration in controller so here we respect what's specified in podSpecPatch.
+	wf := unmarshalWF(helloWorldWfWithPatch)
+	woc := newWoc(*wf)
+	woc.controller.Config.MainContainer = mainCtrSpec
+	mainCtr := woc.execWf.Spec.Templates[0].Container
+	pod, _ := woc.createWorkflowPod(wf.Name, *mainCtr, &wf.Spec.Templates[0], &createWorkflowPodOpts{})
+	assert.Equal(t, "0.800", pod.Spec.Containers[1].Resources.Limits.Cpu().AsDec().String())
+
+	wf = unmarshalWF(helloWorldWf)
+	woc = newWoc(*wf)
+	woc.controller.Config.MainContainer = mainCtrSpec
+	mainCtr = woc.execWf.Spec.Templates[0].Container
+	pod, _ = woc.createWorkflowPod(wf.Name, *mainCtr, &wf.Spec.Templates[0], &createWorkflowPodOpts{})
+	assert.Equal(t, "0.200", pod.Spec.Containers[1].Resources.Limits.Cpu().AsDec().String())
+}
+
 var helloWindowsWf = `
 apiVersion: argoproj.io/v1alpha1
 kind: Workflow