fix: Apply workflow level PodSpecPatch in agent pod. Fixes argoproj#1…

…2387 (argoproj#12440) Signed-off-by: oninowang <[email protected]> Co-authored-by: jswxstw <[email protected]>
jmeridth · Jan 3, 2024 · 198818b · 198818b
1 parent bedd47f
commit 198818b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/workflow/controller/agent.go b/workflow/controller/agent.go
@@ -18,6 +18,7 @@ import (
 	wfv1 "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
 	"github.com/argoproj/argo-workflows/v3/util/env"
 	"github.com/argoproj/argo-workflows/v3/workflow/common"
+	"github.com/argoproj/argo-workflows/v3/workflow/util"
 )
 
 func (woc *wfOperationCtx) getAgentPodName() string {
@@ -243,6 +244,14 @@ func (woc *wfOperationCtx) createAgentPod(ctx context.Context) (*apiv1.Pod, erro
 	addSchedulingConstraints(pod, woc.execWf.Spec.DeepCopy(), tmpl)
 	woc.addMetadata(pod, tmpl)
 
+	if woc.execWf.Spec.HasPodSpecPatch() {
+		patchedPodSpec, err := util.ApplyPodSpecPatch(pod.Spec, woc.execWf.Spec.PodSpecPatch)
+		if err != nil {
+			return nil, errors.Wrap(err, "", "Error applying PodSpecPatch")
+		}
+		pod.Spec = *patchedPodSpec
+	}
+
 	if woc.controller.Config.InstanceID != "" {
 		pod.ObjectMeta.Labels[common.LabelKeyControllerInstanceID] = woc.controller.Config.InstanceID
 	}
@@ -293,17 +302,17 @@ func (woc *wfOperationCtx) getExecutorPlugins(ctx context.Context) ([]apiv1.Cont
 }
 
 func addresses(containers []apiv1.Container) []string {
-	var addresses []string
+	var pluginAddresses []string
 	for _, c := range containers {
-		addresses = append(addresses, fmt.Sprintf("http://localhost:%d", c.Ports[0].ContainerPort))
+		pluginAddresses = append(pluginAddresses, fmt.Sprintf("http://localhost:%d", c.Ports[0].ContainerPort))
 	}
-	return addresses
+	return pluginAddresses
 }
 
 func names(containers []apiv1.Container) []string {
-	var addresses []string
+	var pluginNames []string
 	for _, c := range containers {
-		addresses = append(addresses, c.Name)
+		pluginNames = append(pluginNames, c.Name)
 	}
-	return addresses
+	return pluginNames
 }
diff --git a/workflow/controller/agent_test.go b/workflow/controller/agent_test.go
@@ -20,6 +20,8 @@ metadata:
   name: http-template
   namespace: default
 spec:
+  podSpecPatch: |
+    nodeName: virtual-node
   entrypoint: main
   templates:
     - name: main
@@ -107,6 +109,7 @@ status:
 		for _, pod := range pods.Items {
 			assert.NotNil(t, pod)
 			assert.True(t, strings.HasSuffix(pod.Name, "-agent"))
+			assert.Equal(t, "virtual-node", pod.Spec.NodeName)
 		}
 	})
 	t.Run("CreateTaskSetWithInstanceID", func(t *testing.T) {
@@ -133,6 +136,7 @@ status:
 			assert.NotNil(t, pod)
 			assert.True(t, strings.HasSuffix(pod.Name, "-agent"))
 			assert.Equal(t, "testID", pod.ObjectMeta.Labels[common.LabelKeyControllerInstanceID])
+			assert.Equal(t, "virtual-node", pod.Spec.NodeName)
 		}
 	})
 }