Add ability to specify affinity rules at both the workflow and templa…

…te level (resolves argoproj#701)

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Changelog
 
 ## 2.0.0-beta2 (Unreleased)
++ Add ability to specify affinity rules at both the workflow and template level
 - Fix issue preventing the referencing of artifacts in a container with retries
 - Fix issue preventing the use of volumes in a sidecar
 

pkg/apis/workflow/v1alpha1/types.go

@@ -82,6 +82,10 @@ type WorkflowSpec struct {
  // a nodeSelector specified in the template.
  NodeSelector map[string]string `json:"nodeSelector,omitempty"`
 
+ // Affinity sets the scheduling constraints for all pods in the workflow.
+ // Can be overridden by an affinity specified in the template
+ Affinity *apiv1.Affinity `json:"affinity,omitempty"`
+
  // OnExit is a template reference which is invoked at the end of the
  // workflow, irrespective of the success, failure, or error of the
  // primary workflow.
@@ -100,6 +104,10 @@ type Template struct {
  // run on the selected node(s). Overrides the selector set at the workflow level.
  NodeSelector map[string]string `json:"nodeSelector,omitempty"`
 
+ // Affinity sets the pod's scheduling constraints
+ // Overrides the affinity set at the workflow level (if any)
+ Affinity *apiv1.Affinity `json:"affinity,omitempty"`
+
  // Deamon will allow a workflow to proceed to the next step so long as the container reaches readiness
  Daemon *bool `json:"daemon,omitempty"`
 

pkg/apis/workflow/v1alpha1/zz_generated.deepcopy.go

@@ -566,6 +566,15 @@ func (in *Template) DeepCopyInto(out *Template) {
  (*out)[key] = val
  }
  }
+ if in.Affinity != nil {
+ in, out := &in.Affinity, &out.Affinity
+ if *in == nil {
+ *out = nil
+ } else {
+ *out = new(v1.Affinity)
+ (*in).DeepCopyInto(*out)
+ }
+ }
  if in.Daemon != nil {
  in, out := &in.Daemon, &out.Daemon
  if *in == nil {
@@ -773,6 +782,15 @@ func (in *WorkflowSpec) DeepCopyInto(out *WorkflowSpec) {
  (*out)[key] = val
  }
  }
+ if in.Affinity != nil {
+ in, out := &in.Affinity, &out.Affinity
+ if *in == nil {
+ *out = nil
+ } else {
+ *out = new(v1.Affinity)
+ (*in).DeepCopyInto(*out)
+ }
+ }
  return
 }
 

workflow/controller/workflowpod.go

@@ -159,7 +159,7 @@ func (woc *wfOperationCtx) createWorkflowPod(nodeName string, mainCtr apiv1.Cont
  pod.Spec.InitContainers = []apiv1.Container{initCtr}
  }
 
- woc.addNodeSelectors(&pod, tmpl)
+ woc.addSchedulingConstraints(&pod, tmpl)
 
  err := woc.addVolumeReferences(&pod, tmpl)
  if err != nil {
@@ -260,15 +260,19 @@ func (woc *wfOperationCtx) newExecContainer(name string, privileged bool) *apiv1
  return &exec
 }
 
-// addNodeSelectors applies any node selectors, either set in the workflow or the template, to the pod
-func (woc *wfOperationCtx) addNodeSelectors(pod *apiv1.Pod, tmpl *wfv1.Template) {
+// addSchedulingConstraints applies any node selectors or affinity rules to the pod, either set in the workflow or the template
+func (woc *wfOperationCtx) addSchedulingConstraints(pod *apiv1.Pod, tmpl *wfv1.Template) {
+ // Set nodeSelector (if specified)
  if len(tmpl.NodeSelector) > 0 {
  pod.Spec.NodeSelector = tmpl.NodeSelector
- return
- }
- if len(woc.wf.Spec.NodeSelector) > 0 {
+ } else if len(woc.wf.Spec.NodeSelector) > 0 {
  pod.Spec.NodeSelector = woc.wf.Spec.NodeSelector
- return
+ }
+ // Set affinity (if specified)
+ if tmpl.Affinity != nil {
+ pod.Spec.Affinity = tmpl.Affinity
+ } else if woc.wf.Spec.Affinity != nil {
+ pod.Spec.Affinity = woc.wf.Spec.Affinity
  }
 }
 

workflow/controller/workflowpod_test.go

@@ -7,6 +7,7 @@ import (
  "github.com/ghodss/yaml"
  log "github.com/sirupsen/logrus"
  "github.com/stretchr/testify/assert"
+ apiv1 "k8s.io/api/core/v1"
  metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -90,3 +91,34 @@ func TestServiceAccount(t *testing.T) {
  assert.Nil(t, err)
  assert.Equal(t, pod.Spec.ServiceAccountName, "foo")
 }
+
+// TestAffinity verifies the ability to carry forward affinity rules
+func TestAffinity(t *testing.T) {
+ woc := newWoc()
+ woc.wf.Spec.Affinity = &apiv1.Affinity{
+ NodeAffinity: &apiv1.NodeAffinity{
+ RequiredDuringSchedulingIgnoredDuringExecution: &apiv1.NodeSelector{
+ NodeSelectorTerms: []apiv1.NodeSelectorTerm{
+ {
+ MatchExpressions: []apiv1.NodeSelectorRequirement{
+ {
+ Key: "kubernetes.io/e2e-az-name",
+ Operator: apiv1.NodeSelectorOpIn,
+ Values: []string{
+ "e2e-az1",
+ "e2e-az2",
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ }
+ err := woc.executeContainer(woc.wf.Spec.Entrypoint, &woc.wf.Spec.Templates[0])
+ assert.Nil(t, err)
+ podName := getPodName(woc.wf)
+ pod, err := woc.controller.kubeclientset.CoreV1().Pods("").Get(podName, metav1.GetOptions{})
+ assert.Nil(t, err)
+ assert.NotNil(t, pod.Spec.Affinity)
+}