Skip to content
This repository was archived by the owner on Oct 9, 2023. It is now read-only.

Commit a87ca40

Browse files
fg91Fabio Grätz
fg91
and
Fabio Grätz
authored
Pass K8sPluginConfig to spark driver and executor pods #patch (#271)
* Pass default tolerations to spark driver and executor Signed-off-by: fg91 <fabiograetz@googlemail.com> * Test passing default tolerations to spark driver and executor Signed-off-by: fg91 <fabiograetz@googlemail.com> * Pass scheduler name to driver and executor SparkPodSpec Signed-off-by: fg91 <fabiograetz@googlemail.com> * Carry DefaultNodeSelector from k8s plugin config to SparkPodSpec Signed-off-by: fg91 <fabiograetz@googlemail.com> * Carry over EnableHostNetworkingPod Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Test carrying over of default env vars Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Carry over DefaultEnvVarsFromEnv Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Carry over DefaultAffinity Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Doc behaviour of default and interruptible NodeSelector and Tolerations Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Don't carry over default env vars from env and fix test Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Lint Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> * Apply node selector requirement to pod affinity Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> Signed-off-by: fg91 <fabiograetz@googlemail.com> Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com> Co-authored-by: Fabio Grätz <fabiogratz@googlemail.com>
1 parent 932e97b commit a87ca40

3 files changed

Lines changed: 188 additions & 33 deletions

File tree

‎go/tasks/pluginmachinery/flytek8s/pod_helper.go‎

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ const SIGKILL = 137
2626
const defaultContainerTemplateName = "default"
2727
const primaryContainerTemplateName = "primary"
2828

29-
// ApplyInterruptibleNodeAffinity configures the node-affinity for the pod using the configuration specified.
30-
func ApplyInterruptibleNodeAffinity(interruptible bool, podSpec *v1.PodSpec) {
29+
// ApplyInterruptibleNodeSelectorRequirement configures the node selector requirement of the node-affinity using the configuration specified.
30+
func ApplyInterruptibleNodeSelectorRequirement(interruptible bool, affinity *v1.Affinity) {
3131
// Determine node selector terms to add to node affinity
3232
var nodeSelectorRequirement v1.NodeSelectorRequirement
3333
if interruptible {
@@ -42,24 +42,31 @@ func ApplyInterruptibleNodeAffinity(interruptible bool, podSpec *v1.PodSpec) {
4242
nodeSelectorRequirement = *config.GetK8sPluginConfig().NonInterruptibleNodeSelectorRequirement
4343
}
4444

45-
if podSpec.Affinity == nil {
46-
podSpec.Affinity = &v1.Affinity{}
45+
if affinity.NodeAffinity == nil {
46+
affinity.NodeAffinity = &v1.NodeAffinity{}
4747
}
48-
if podSpec.Affinity.NodeAffinity == nil {
49-
podSpec.Affinity.NodeAffinity = &v1.NodeAffinity{}
48+
if affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
49+
affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
5050
}
51-
if podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
52-
podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
53-
}
54-
if len(podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms) > 0 {
55-
nodeSelectorTerms := podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
51+
if len(affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms) > 0 {
52+
nodeSelectorTerms := affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms
5653
for i := range nodeSelectorTerms {
5754
nst := &nodeSelectorTerms[i]
5855
nst.MatchExpressions = append(nst.MatchExpressions, nodeSelectorRequirement)
5956
}
6057
} else {
61-
podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = []v1.NodeSelectorTerm{v1.NodeSelectorTerm{MatchExpressions: []v1.NodeSelectorRequirement{nodeSelectorRequirement}}}
58+
affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = []v1.NodeSelectorTerm{v1.NodeSelectorTerm{MatchExpressions: []v1.NodeSelectorRequirement{nodeSelectorRequirement}}}
59+
}
60+
61+
}
62+
63+
// ApplyInterruptibleNodeAffinity configures the node-affinity for the pod using the configuration specified.
64+
func ApplyInterruptibleNodeAffinity(interruptible bool, podSpec *v1.PodSpec) {
65+
if podSpec.Affinity == nil {
66+
podSpec.Affinity = &v1.Affinity{}
6267
}
68+
69+
ApplyInterruptibleNodeSelectorRequirement(interruptible, podSpec.Affinity)
6370
}
6471

6572
// UpdatePod updates the base pod spec used to execute tasks. This is configured with plugins and task metadata-specific options

‎go/tasks/plugins/k8s/spark/spark.go‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo
9090
for _, envVar := range envVars {
9191
sparkEnvVars[envVar.Name] = envVar.Value
9292
}
93+
9394
sparkEnvVars["FLYTE_MAX_ATTEMPTS"] = strconv.Itoa(int(taskCtx.TaskExecutionMetadata().GetMaxAttempts()))
9495

9596
serviceAccountName := flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata())
@@ -99,24 +100,34 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo
99100
}
100101
driverSpec := sparkOp.DriverSpec{
101102
SparkPodSpec: sparkOp.SparkPodSpec{
103+
Affinity: config.GetK8sPluginConfig().DefaultAffinity,
102104
Annotations: annotations,
103105
Labels: labels,
104106
EnvVars: sparkEnvVars,
105107
Image: &container.Image,
106108
SecurityContenxt: config.GetK8sPluginConfig().DefaultPodSecurityContext.DeepCopy(),
107109
DNSConfig: config.GetK8sPluginConfig().DefaultPodDNSConfig.DeepCopy(),
110+
Tolerations: config.GetK8sPluginConfig().DefaultTolerations,
111+
SchedulerName: &config.GetK8sPluginConfig().SchedulerName,
112+
NodeSelector: config.GetK8sPluginConfig().DefaultNodeSelector,
113+
HostNetwork: config.GetK8sPluginConfig().EnableHostNetworkingPod,
108114
},
109115
ServiceAccount: &serviceAccountName,
110116
}
111117

112118
executorSpec := sparkOp.ExecutorSpec{
113119
SparkPodSpec: sparkOp.SparkPodSpec{
120+
Affinity: config.GetK8sPluginConfig().DefaultAffinity.DeepCopy(),
114121
Annotations: annotations,
115122
Labels: labels,
116123
Image: &container.Image,
117124
EnvVars: sparkEnvVars,
118125
SecurityContenxt: config.GetK8sPluginConfig().DefaultPodSecurityContext.DeepCopy(),
119126
DNSConfig: config.GetK8sPluginConfig().DefaultPodDNSConfig.DeepCopy(),
127+
Tolerations: config.GetK8sPluginConfig().DefaultTolerations,
128+
SchedulerName: &config.GetK8sPluginConfig().SchedulerName,
129+
NodeSelector: config.GetK8sPluginConfig().DefaultNodeSelector,
130+
HostNetwork: config.GetK8sPluginConfig().EnableHostNetworkingPod,
120131
},
121132
}
122133

@@ -225,11 +236,16 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo
225236
j.Spec.MainClass = &sparkJob.MainClass
226237
}
227238

228-
// Add Tolerations/NodeSelector to only Executor pods.
239+
// Add Interruptible Tolerations/NodeSelector to only Executor pods.
240+
// The Interruptible NodeSelector takes precedence over the DefaultNodeSelector
229241
if taskCtx.TaskExecutionMetadata().IsInterruptible() {
230-
j.Spec.Executor.Tolerations = config.GetK8sPluginConfig().InterruptibleTolerations
242+
j.Spec.Executor.Tolerations = append(j.Spec.Executor.Tolerations, config.GetK8sPluginConfig().InterruptibleTolerations...)
231243
j.Spec.Executor.NodeSelector = config.GetK8sPluginConfig().InterruptibleNodeSelector
232244
}
245+
246+
// Add interruptible/non-interruptible node selector requirements to executor pod
247+
flytek8s.ApplyInterruptibleNodeSelectorRequirement(taskCtx.TaskExecutionMetadata().IsInterruptible(), j.Spec.Executor.Affinity)
248+
233249
return j, nil
234250
}
235251

‎go/tasks/plugins/k8s/spark/spark_test.go‎

Lines changed: 151 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package spark
33
import (
44
"context"
55
"fmt"
6+
"os"
67
"strconv"
78
"testing"
89

@@ -353,7 +354,67 @@ func TestBuildResourceSpark(t *testing.T) {
353354
dnsOptVal1 := "1"
354355
dnsOptVal2 := "1"
355356
dnsOptVal3 := "3"
357+
358+
// Set scheduler
359+
schedulerName := "custom-scheduler"
360+
361+
// Node selectors
362+
defaultNodeSelector := map[string]string{
363+
"x/default": "true",
364+
}
365+
interruptibleNodeSelector := map[string]string{
366+
"x/interruptible": "true",
367+
}
368+
369+
defaultPodHostNetwork := true
370+
371+
// Default env vars passed explicitly and default env vars derived from environment
372+
defaultEnvVars := make(map[string]string)
373+
defaultEnvVars["foo"] = "bar"
374+
375+
defaultEnvVarsFromEnv := make(map[string]string)
376+
targetKeyFromEnv := "TEST_VAR_FROM_ENV_KEY"
377+
targetValueFromEnv := "TEST_VAR_FROM_ENV_VALUE"
378+
os.Setenv(targetKeyFromEnv, targetValueFromEnv)
379+
defer os.Unsetenv(targetKeyFromEnv)
380+
defaultEnvVarsFromEnv["fooEnv"] = targetKeyFromEnv
381+
382+
// Default affinity/anti-affinity
383+
defaultAffinity := &corev1.Affinity{
384+
NodeAffinity: &corev1.NodeAffinity{
385+
RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
386+
NodeSelectorTerms: []corev1.NodeSelectorTerm{
387+
{
388+
MatchExpressions: []corev1.NodeSelectorRequirement{
389+
{
390+
Key: "x/default",
391+
Operator: corev1.NodeSelectorOpIn,
392+
Values: []string{"true"},
393+
},
394+
},
395+
},
396+
},
397+
},
398+
},
399+
}
400+
401+
// interruptible/non-interruptible nodeselector requirement
402+
interruptibleNodeSelectorRequirement := &corev1.NodeSelectorRequirement{
403+
Key: "x/interruptible",
404+
Operator: corev1.NodeSelectorOpIn,
405+
Values: []string{"true"},
406+
}
407+
408+
nonInterruptibleNodeSelectorRequirement := &corev1.NodeSelectorRequirement{
409+
Key: "x/non-interruptible",
410+
Operator: corev1.NodeSelectorOpIn,
411+
Values: []string{"true"},
412+
}
413+
414+
// NonInterruptibleNodeSelectorRequirement
415+
356416
assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
417+
DefaultAffinity: defaultAffinity,
357418
DefaultPodSecurityContext: &corev1.PodSecurityContext{
358419
RunAsUser: &runAsUser,
359420
},
@@ -378,17 +439,31 @@ func TestBuildResourceSpark(t *testing.T) {
378439
},
379440
Searches: []string{"ns1.svc.cluster-domain.example", "my.dns.search.suffix"},
380441
},
381-
InterruptibleNodeSelector: map[string]string{
382-
"x/interruptible": "true",
442+
DefaultTolerations: []corev1.Toleration{
443+
{
444+
Key: "x/flyte",
445+
Value: "default",
446+
Operator: "Equal",
447+
Effect: "NoSchedule",
448+
},
383449
},
450+
DefaultNodeSelector: defaultNodeSelector,
451+
InterruptibleNodeSelector: interruptibleNodeSelector,
384452
InterruptibleTolerations: []corev1.Toleration{
385453
{
386454
Key: "x/flyte",
387455
Value: "interruptible",
388456
Operator: "Equal",
389457
Effect: "NoSchedule",
390458
},
391-
}}),
459+
},
460+
InterruptibleNodeSelectorRequirement: interruptibleNodeSelectorRequirement,
461+
NonInterruptibleNodeSelectorRequirement: nonInterruptibleNodeSelectorRequirement,
462+
SchedulerName: schedulerName,
463+
EnableHostNetworkingPod: &defaultPodHostNetwork,
464+
DefaultEnvVars: defaultEnvVars,
465+
DefaultEnvVarsFromEnv: defaultEnvVarsFromEnv,
466+
}),
392467
)
393468
resource, err := sparkResourceHandler.BuildResource(context.TODO(), dummySparkTaskContext(taskTemplate, true))
394469
assert.Nil(t, err)
@@ -438,19 +513,40 @@ func TestBuildResourceSpark(t *testing.T) {
438513
assert.Equal(t, dummySparkConf["spark.driver.memory"], *sparkApp.Spec.Driver.Memory)
439514
assert.Equal(t, dummySparkConf["spark.executor.memory"], *sparkApp.Spec.Executor.Memory)
440515
assert.Equal(t, dummySparkConf["spark.batchScheduler"], *sparkApp.Spec.BatchScheduler)
441-
442-
// Validate Interruptible Toleration and NodeSelector set for Executor but not Driver.
443-
assert.Equal(t, 0, len(sparkApp.Spec.Driver.Tolerations))
444-
assert.Equal(t, 0, len(sparkApp.Spec.Driver.NodeSelector))
445-
446-
assert.Equal(t, 1, len(sparkApp.Spec.Executor.Tolerations))
516+
assert.Equal(t, schedulerName, *sparkApp.Spec.Executor.SchedulerName)
517+
assert.Equal(t, schedulerName, *sparkApp.Spec.Driver.SchedulerName)
518+
assert.Equal(t, defaultPodHostNetwork, *sparkApp.Spec.Executor.HostNetwork)
519+
assert.Equal(t, defaultPodHostNetwork, *sparkApp.Spec.Driver.HostNetwork)
520+
521+
// Validate
522+
// * Interruptible Toleration and NodeSelector set for Executor but not Driver.
523+
// * Validate Default NodeSelector set for Driver but overwritten with Interruptible NodeSelector for Executor.
524+
// * Default Tolerations set for both Driver and Executor.
525+
// * Interruptible/Non-Interruptible NodeSelectorRequirements set for Executor Affinity but not Driver Affinity.
526+
assert.Equal(t, 1, len(sparkApp.Spec.Driver.Tolerations))
527+
assert.Equal(t, 1, len(sparkApp.Spec.Driver.NodeSelector))
528+
assert.Equal(t, defaultNodeSelector, sparkApp.Spec.Driver.NodeSelector)
529+
tolDriverDefault := sparkApp.Spec.Driver.Tolerations[0]
530+
assert.Equal(t, tolDriverDefault.Key, "x/flyte")
531+
assert.Equal(t, tolDriverDefault.Value, "default")
532+
assert.Equal(t, tolDriverDefault.Operator, corev1.TolerationOperator("Equal"))
533+
assert.Equal(t, tolDriverDefault.Effect, corev1.TaintEffect("NoSchedule"))
534+
535+
assert.Equal(t, 2, len(sparkApp.Spec.Executor.Tolerations))
447536
assert.Equal(t, 1, len(sparkApp.Spec.Executor.NodeSelector))
448-
449-
tol := sparkApp.Spec.Executor.Tolerations[0]
450-
assert.Equal(t, tol.Key, "x/flyte")
451-
assert.Equal(t, tol.Value, "interruptible")
452-
assert.Equal(t, tol.Operator, corev1.TolerationOperator("Equal"))
453-
assert.Equal(t, tol.Effect, corev1.TaintEffect("NoSchedule"))
537+
assert.Equal(t, interruptibleNodeSelector, sparkApp.Spec.Executor.NodeSelector)
538+
539+
tolExecDefault := sparkApp.Spec.Executor.Tolerations[0]
540+
assert.Equal(t, tolExecDefault.Key, "x/flyte")
541+
assert.Equal(t, tolExecDefault.Value, "default")
542+
assert.Equal(t, tolExecDefault.Operator, corev1.TolerationOperator("Equal"))
543+
assert.Equal(t, tolExecDefault.Effect, corev1.TaintEffect("NoSchedule"))
544+
545+
tolExecInterrupt := sparkApp.Spec.Executor.Tolerations[1]
546+
assert.Equal(t, tolExecInterrupt.Key, "x/flyte")
547+
assert.Equal(t, tolExecInterrupt.Value, "interruptible")
548+
assert.Equal(t, tolExecInterrupt.Operator, corev1.TolerationOperator("Equal"))
549+
assert.Equal(t, tolExecInterrupt.Effect, corev1.TaintEffect("NoSchedule"))
454550
assert.Equal(t, "true", sparkApp.Spec.Executor.NodeSelector["x/interruptible"])
455551

456552
for confKey, confVal := range dummySparkConf {
@@ -485,6 +581,22 @@ func TestBuildResourceSpark(t *testing.T) {
485581
assert.Equal(t, dummySparkConf["spark.flyteorg.feature3.enabled"], sparkApp.Spec.SparkConf["spark.flyteorg.feature3.enabled"])
486582

487583
assert.Equal(t, len(sparkApp.Spec.Driver.EnvVars["FLYTE_MAX_ATTEMPTS"]), 1)
584+
assert.Equal(t, sparkApp.Spec.Driver.EnvVars["foo"], defaultEnvVars["foo"])
585+
assert.Equal(t, sparkApp.Spec.Executor.EnvVars["foo"], defaultEnvVars["foo"])
586+
assert.Equal(t, sparkApp.Spec.Driver.EnvVars["fooEnv"], targetValueFromEnv)
587+
assert.Equal(t, sparkApp.Spec.Executor.EnvVars["fooEnv"], targetValueFromEnv)
588+
assert.Equal(t, sparkApp.Spec.Driver.Affinity, defaultAffinity)
589+
590+
assert.Equal(
591+
t,
592+
sparkApp.Spec.Executor.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions[0],
593+
defaultAffinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions[0],
594+
)
595+
assert.Equal(
596+
t,
597+
sparkApp.Spec.Executor.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions[1],
598+
*interruptibleNodeSelectorRequirement,
599+
)
488600

489601
// Case 2: Driver/Executor request cores set.
490602
dummyConfWithRequest := make(map[string]string)
@@ -514,10 +626,30 @@ func TestBuildResourceSpark(t *testing.T) {
514626
assert.True(t, ok)
515627

516628
// Validate Interruptible Toleration and NodeSelector not set for both Driver and Executors.
517-
assert.Equal(t, 0, len(sparkApp.Spec.Driver.Tolerations))
518-
assert.Equal(t, 0, len(sparkApp.Spec.Driver.NodeSelector))
519-
assert.Equal(t, 0, len(sparkApp.Spec.Executor.Tolerations))
520-
assert.Equal(t, 0, len(sparkApp.Spec.Executor.NodeSelector))
629+
// Validate that the default Toleration and NodeSelector are set for both Driver and Executors.
630+
assert.Equal(t, 1, len(sparkApp.Spec.Driver.Tolerations))
631+
assert.Equal(t, 1, len(sparkApp.Spec.Driver.NodeSelector))
632+
assert.Equal(t, defaultNodeSelector, sparkApp.Spec.Driver.NodeSelector)
633+
assert.Equal(t, 1, len(sparkApp.Spec.Executor.Tolerations))
634+
assert.Equal(t, 1, len(sparkApp.Spec.Executor.NodeSelector))
635+
assert.Equal(t, defaultNodeSelector, sparkApp.Spec.Executor.NodeSelector)
636+
assert.Equal(t, sparkApp.Spec.Executor.Tolerations[0].Key, "x/flyte")
637+
assert.Equal(t, sparkApp.Spec.Executor.Tolerations[0].Value, "default")
638+
assert.Equal(t, sparkApp.Spec.Driver.Tolerations[0].Key, "x/flyte")
639+
assert.Equal(t, sparkApp.Spec.Driver.Tolerations[0].Value, "default")
640+
641+
// Validate correct affinity and nodeselector requirements are set for both Driver and Executors.
642+
assert.Equal(t, sparkApp.Spec.Driver.Affinity, defaultAffinity)
643+
assert.Equal(
644+
t,
645+
sparkApp.Spec.Executor.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions[0],
646+
defaultAffinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions[0],
647+
)
648+
assert.Equal(
649+
t,
650+
sparkApp.Spec.Executor.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions[1],
651+
*nonInterruptibleNodeSelectorRequirement,
652+
)
521653

522654
// Case 4: Invalid Spark Task-Template
523655
taskTemplate.Custom = nil

0 commit comments

Comments
 (0)