Skip to content

Commit

Permalink
Merge affinity from podtempalte and affinity-assistant
Browse files Browse the repository at this point in the history 
No matter user provide affinity in podtemplate or not the
affinity-assistant will overwrite the affinity in pod.

However if user set affinity in podtemplate merge
podtemplate's affinity with affinity-assistant's affinity
is better way.

So we inject the affinity of the affinity-assitant into the
affinity of the podtempalte as the affinity of the pod to
make sure that all the affinities work properly.

Related Issue: #5241

Signed-off-by: yuzhipeng <yuzp1996@qq.com>
  • Loading branch information
@yuzp1996
yuzp1996 committed Sep 8, 2022
1 parent 64984fc commit 723d147
Show file tree
Hide file tree
Showing 4 changed files with 361 additions and 38 deletions.
47 changes: 27 additions & 20 deletions pkg/internal/affinityassistant/transformer.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,32 +28,39 @@ import (
// NewTransformer returns a pod.Transformer that will pod affinity if needed
func NewTransformer(_ context.Context, annotations map[string]string) pod.Transformer {
return func(p *corev1.Pod) (*corev1.Pod, error) {
// Using node affinity on taskRuns sharing PVC workspace, with an Affinity Assistant
// is mutually exclusive with other affinity on taskRun pods. If other
// affinity is wanted, that should be added on the Affinity Assistant pod unless
// assistant is disabled. When Affinity Assistant is disabled, an affinityAssistantName is not set.
// Using node affinity on taskRuns sharing PVC workspace. When Affinity Assistant
// is disabled, an affinityAssistantName is not set.
if affinityAssistantName := annotations[workspace.AnnotationAffinityAssistantName]; affinityAssistantName != "" {
p.Spec.Affinity = nodeAffinityUsingAffinityAssistant(affinityAssistantName)
if p.Spec.Affinity == nil {
p.Spec.Affinity = &corev1.Affinity{}
}
mergeAffinityWithAffinityAssistant(p.Spec.Affinity, affinityAssistantName)
}
return p, nil
}
}

// nodeAffinityUsingAffinityAssistant achieves Node Affinity for taskRun pods
// sharing PVC workspace by setting PodAffinity so that taskRuns is
// scheduled to the Node were the Affinity Assistant pod is scheduled.
func nodeAffinityUsingAffinityAssistant(affinityAssistantName string) *corev1.Affinity {
return &corev1.Affinity{
PodAffinity: &corev1.PodAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{{
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelInstance: affinityAssistantName,
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
}},
func mergeAffinityWithAffinityAssistant(affinity *corev1.Affinity, affinityAssistantName string) {
podAffinityTerm := podAffinityTermUsingAffinityAssistant(affinityAssistantName)

if affinity.PodAffinity == nil {
affinity.PodAffinity = &corev1.PodAffinity{}
}

affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution =
append(affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution, *podAffinityTerm)
}

// podAffinityTermUsingAffinityAssistant achieves pod Affinity term for taskRun
// pods so that the taskRun is scheduled to the Node where the Affinity Assistant pod
// is scheduled.
func podAffinityTermUsingAffinityAssistant(affinityAssistantName string) *corev1.PodAffinityTerm {
return &corev1.PodAffinityTerm{LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelInstance: affinityAssistantName,
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
}
}
153 changes: 153 additions & 0 deletions pkg/internal/affinityassistant/transformer_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,3 +88,156 @@ func TestNewTransformer(t *testing.T) {
})
}
}

func TestNewTransformerWithNodeAffinity(t *testing.T) {

nodeAffinity := &corev1.NodeAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
NodeSelectorTerms: []corev1.NodeSelectorTerm{{
MatchFields: []corev1.NodeSelectorRequirement{{
Key: "kubernetes.io/hostname",
Operator: corev1.NodeSelectorOpNotIn,
Values: []string{"192.0.0.1"},
}},
}},
},
}

podAffinityTerm := &corev1.PodAffinityTerm{
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
"from": "podtemplate",
},
},
TopologyKey: "kubernetes.io/hostname",
}

for _, tc := range []struct {
description string
annotations map[string]string
pod *corev1.Pod
expected *corev1.Affinity
}{{
description: "no affinity annotation and pod contains nodeAffinity",
annotations: map[string]string{
"foo": "bar",
},
pod: &corev1.Pod{
Spec: corev1.PodSpec{
Affinity: &corev1.Affinity{
NodeAffinity: nodeAffinity,
},
},
},
expected: &corev1.Affinity{
NodeAffinity: nodeAffinity,
},
}, {
description: "affinity annotation and pod contains nodeAffinity",
annotations: map[string]string{
"foo": "bar",
workspace.AnnotationAffinityAssistantName: "baz",
},
pod: &corev1.Pod{
Spec: corev1.PodSpec{
Affinity: &corev1.Affinity{
NodeAffinity: nodeAffinity,
},
},
},

expected: &corev1.Affinity{
NodeAffinity: nodeAffinity,
PodAffinity: &corev1.PodAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{{
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelInstance: "baz",
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
}},
},
},
}, {
description: "affinity annotation with a different name and pod contains podAffinity",
annotations: map[string]string{
workspace.AnnotationAffinityAssistantName: "helloworld",
},
pod: &corev1.Pod{
Spec: corev1.PodSpec{
Affinity: &corev1.Affinity{PodAffinity: &corev1.PodAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm},
PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{{100, *podAffinityTerm}},
}},
},
},

expected: &corev1.Affinity{
PodAffinity: &corev1.PodAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm, {
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelInstance: "helloworld",
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
}},
PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{
{100, *podAffinityTerm},
},
},
},
}, {
description: "affinity annotation with a different name and pod contains podAffinity and nodeAffinity",
annotations: map[string]string{
workspace.AnnotationAffinityAssistantName: "helloworld",
},
pod: &corev1.Pod{
Spec: corev1.PodSpec{
Affinity: &corev1.Affinity{
PodAffinity: &corev1.PodAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{*podAffinityTerm},
PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{{100, *podAffinityTerm}},
},
NodeAffinity: nodeAffinity},
},
},

expected: &corev1.Affinity{
PodAffinity: &corev1.PodAffinity{
RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
*podAffinityTerm,
{
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelInstance: "helloworld",
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
},
},
PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{
{100, *podAffinityTerm},
},
},
NodeAffinity: nodeAffinity,
},
},
} {
t.Run(tc.description, func(t *testing.T) {
ctx := context.Background()
f := affinityassistant.NewTransformer(ctx, tc.annotations)
got, err := f(tc.pod)
if err != nil {
t.Fatalf("Transformer failed: %v", err)
}
if d := cmp.Diff(tc.expected, got.Spec.Affinity); d != "" {
t.Errorf("AffinityAssistant diff: %s", diff.PrintWantGot(d))
}
})
}
}
48 changes: 30 additions & 18 deletions pkg/reconciler/pipelinerun/affinity_assistant.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,19 +156,6 @@ func affinityAssistantStatefulSet(name string, pr *v1beta1.PipelineRun, claimNam
},
}}

// use podAntiAffinity to repel other affinity assistants
repelOtherAffinityAssistantsPodAffinityTerm := corev1.WeightedPodAffinityTerm{
Weight: 100,
PodAffinityTerm: corev1.PodAffinityTerm{
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
},
}

return &appsv1.StatefulSet{
TypeMeta: metav1.TypeMeta{
Kind: "StatefulSet",
Expand All @@ -195,11 +182,7 @@ func affinityAssistantStatefulSet(name string, pr *v1beta1.PipelineRun, claimNam
NodeSelector: tpl.NodeSelector,
ImagePullSecrets: tpl.ImagePullSecrets,

Affinity: &corev1.Affinity{
PodAntiAffinity: &corev1.PodAntiAffinity{
PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{repelOtherAffinityAssistantsPodAffinityTerm},
},
},
Affinity: getAssistantAffinityMergedWithPodTemplateAffinity(pr),
Volumes: []corev1.Volume{{
Name: "workspace",
VolumeSource: corev1.VolumeSource{
Expand Down Expand Up @@ -230,3 +213,32 @@ func (c *Reconciler) isAffinityAssistantDisabled(ctx context.Context) bool {
cfg := config.FromContextOrDefaults(ctx)
return cfg.FeatureFlags.DisableAffinityAssistant
}

// getAssistantAffinityMergedWithPodTemplateAffinity return the affinity that merged with PipelineRun PodTemplate affinity.
func getAssistantAffinityMergedWithPodTemplateAffinity(pr *v1beta1.PipelineRun) *corev1.Affinity {
// use podAntiAffinity to repel other affinity assistants
repelOtherAffinityAssistantsPodAffinityTerm := corev1.WeightedPodAffinityTerm{
Weight: 100,
PodAffinityTerm: corev1.PodAffinityTerm{
LabelSelector: &metav1.LabelSelector{
MatchLabels: map[string]string{
workspace.LabelComponent: workspace.ComponentNameAffinityAssistant,
},
},
TopologyKey: "kubernetes.io/hostname",
},
}

affinityAssistantsAffinity := &corev1.Affinity{}
if pr.Spec.PodTemplate != nil && pr.Spec.PodTemplate.Affinity != nil {
affinityAssistantsAffinity = pr.Spec.PodTemplate.Affinity
}
if affinityAssistantsAffinity.PodAntiAffinity == nil {
affinityAssistantsAffinity.PodAntiAffinity = &corev1.PodAntiAffinity{}
}
affinityAssistantsAffinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution =
append(affinityAssistantsAffinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
repelOtherAffinityAssistantsPodAffinityTerm)

return affinityAssistantsAffinity
}
Loading

0 comments on commit 723d147

Please sign in to comment.