diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark.go index 6bb746c576f..86c66655ef6 100644 --- a/flyteplugins/go/tasks/plugins/k8s/spark/spark.go +++ b/flyteplugins/go/tasks/plugins/k8s/spark/spark.go @@ -76,11 +76,11 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo } sparkConfig := getSparkConfig(taskCtx, &sparkJob) - driverSpec, err := createDriverSpec(ctx, taskCtx, sparkConfig, &sparkJob) + driverSpec, err := createDriverSpec(ctx, taskCtx, &sparkJob) if err != nil { return nil, err } - executorSpec, err := createExecutorSpec(ctx, taskCtx, sparkConfig, &sparkJob) + executorSpec, err := createExecutorSpec(ctx, taskCtx, &sparkJob) if err != nil { return nil, err } @@ -142,6 +142,20 @@ func serviceAccountName(metadata pluginsCore.TaskExecutionMetadata) string { return name } +// getPrimaryContainer returns a pointer to the primary container within the given pod spec, so +// that in-place modifications to the container are reflected in the pod spec. +func getPrimaryContainer(podSpec *v1.PodSpec, name string) (*v1.Container, error) { + for i := range podSpec.Containers { + if podSpec.Containers[i].Name == name { + return &podSpec.Containers[i], nil + } + } + return nil, errors.Errorf(errors.BadTaskSpecification, "invalid TaskSpecification, container [%s] not defined", name) +} + +// createSparkPodSpec builds a SparkPodSpec using the pod `Template` attribute so that the full +// pod spec is passed to the spark-operator as-is, rather than mapping individual fields onto the +// various SparkPodSpec config attributes. func createSparkPodSpec( taskCtx pluginsCore.TaskExecutionContext, podSpec *v1.PodSpec, @@ -164,35 +178,31 @@ func createSparkPodSpec( labels = pluginsUtils.UnionMaps(labels, k8sPod.GetMetadata().GetLabels()) } - sparkEnv := make([]v1.EnvVar, 0) - for _, envVar := range container.Env { - sparkEnv = append(sparkEnv, *envVar.DeepCopy()) - } - sparkEnv = append(sparkEnv, v1.EnvVar{Name: "FLYTE_MAX_ATTEMPTS", Value: strconv.Itoa(int(taskCtx.TaskExecutionMetadata().GetMaxAttempts()))}) - - spec := sparkOp.SparkPodSpec{ - Affinity: podSpec.Affinity, - Annotations: annotations, - Labels: labels, - Env: sparkEnv, - Image: &container.Image, - PodSecurityContext: podSpec.SecurityContext, - SecurityContext: container.SecurityContext, - DNSConfig: podSpec.DNSConfig.DeepCopy(), - Tolerations: podSpec.Tolerations, - SchedulerName: &podSpec.SchedulerName, - NodeSelector: podSpec.NodeSelector, - HostNetwork: &podSpec.HostNetwork, - ServiceAccount: strPtr(serviceAccountName(taskCtx.TaskExecutionMetadata())), - } - return &spec + container.Env = append(container.Env, v1.EnvVar{Name: "FLYTE_MAX_ATTEMPTS", Value: strconv.Itoa(int(taskCtx.TaskExecutionMetadata().GetMaxAttempts()))}) + + serviceAccount := serviceAccountName(taskCtx.TaskExecutionMetadata()) + podSpec.ServiceAccountName = serviceAccount + + return &sparkOp.SparkPodSpec{ + Template: &v1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: annotations, + Labels: labels, + }, + Spec: *podSpec, + }, + // Spark overwrites the pod template's serviceAccountName with the value of + // `spark.kubernetes.authenticate.driver.serviceAccountName` (which the spark-operator + // derives from this field), so it must be set explicitly here as well. + ServiceAccount: strPtr(serviceAccount), + } } type driverSpec struct { sparkSpec *sparkOp.DriverSpec } -func createDriverSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, sparkConfig map[string]string, sparkJob *plugins.SparkJob) (*driverSpec, error) { +func createDriverSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, sparkJob *plugins.SparkJob) (*driverSpec, error) { // Spark driver pods should always run as non-interruptible nonInterruptibleTaskCtx := flytek8s.NewPluginTaskExecutionContext(taskCtx, flytek8s.WithInterruptible(false)) podSpec, _, primaryContainerName, err := flytek8s.ToK8sPodSpec(ctx, nonInterruptibleTaskCtx) @@ -222,25 +232,19 @@ func createDriverSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionCont } } - primaryContainer, err := flytek8s.GetContainer(podSpec, primaryContainerName) + primaryContainer, err := getPrimaryContainer(podSpec, primaryContainerName) if err != nil { return nil, err } + // Rename the primary container to the name Spark expects to find in the driver pod template. + primaryContainer.Name = sparkOpConfig.SparkDriverContainerName sparkPodSpec := createSparkPodSpec(nonInterruptibleTaskCtx, podSpec, primaryContainer, driverPod) - spec := driverSpec{ + return &driverSpec{ &sparkOp.DriverSpec{ SparkPodSpec: *sparkPodSpec, }, - } - spec.sparkSpec.ServiceAccount = strPtr(serviceAccountName(nonInterruptibleTaskCtx.TaskExecutionMetadata())) - spec.sparkSpec.PriorityClassName = strPtr(podSpec.PriorityClassName) - - if cores, err := strconv.ParseInt(sparkConfig["spark.driver.cores"], 10, 32); err == nil { - spec.sparkSpec.Cores = intPtr(int32(cores)) - } - spec.sparkSpec.Memory = strPtr(sparkConfig["spark.driver.memory"]) - return &spec, nil + }, nil } type executorSpec struct { @@ -249,7 +253,7 @@ type executorSpec struct { serviceAccountName string } -func createExecutorSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, sparkConfig map[string]string, sparkJob *plugins.SparkJob) (*executorSpec, error) { +func createExecutorSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionContext, sparkJob *plugins.SparkJob) (*executorSpec, error) { podSpec, _, primaryContainerName, err := flytek8s.ToK8sPodSpec(ctx, taskCtx) if err != nil { return nil, err @@ -276,28 +280,20 @@ func createExecutorSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionCo } } - primaryContainer, err := flytek8s.GetContainer(podSpec, primaryContainerName) + primaryContainer, err := getPrimaryContainer(podSpec, primaryContainerName) if err != nil { return nil, err } + // Rename the primary container to the name Spark expects to find in the executor pod template. + primaryContainer.Name = sparkOpConfig.Spark3DefaultExecutorContainerName sparkPodSpec := createSparkPodSpec(taskCtx, podSpec, primaryContainer, sparkJob.GetExecutorPod()) - serviceAccountName := serviceAccountName(taskCtx.TaskExecutionMetadata()) - spec := executorSpec{ + return &executorSpec{ primaryContainer, &sparkOp.ExecutorSpec{ SparkPodSpec: *sparkPodSpec, }, - serviceAccountName, - } - spec.sparkSpec.PriorityClassName = strPtr(podSpec.PriorityClassName) - if execCores, err := strconv.ParseInt(sparkConfig["spark.executor.cores"], 10, 32); err == nil { - spec.sparkSpec.Cores = intPtr(int32(execCores)) - } - if execCount, err := strconv.ParseInt(sparkConfig["spark.executor.instances"], 10, 32); err == nil { - spec.sparkSpec.Instances = intPtr(int32(execCount)) - } - spec.sparkSpec.Memory = strPtr(sparkConfig["spark.executor.memory"]) - return &spec, nil + serviceAccountName(taskCtx.TaskExecutionMetadata()), + }, nil } func createSparkApplication(sparkJob *plugins.SparkJob, sparkConfig map[string]string, driverSpec *driverSpec, diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go index 886396be099..03a7a1292da 100644 --- a/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go @@ -10,6 +10,7 @@ import ( "time" sparkOp "github.com/kubeflow/spark-operator/api/v1beta2" + sparkOpConfig "github.com/kubeflow/spark-operator/pkg/common" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "google.golang.org/protobuf/types/known/structpb" @@ -704,6 +705,15 @@ func findEnvVarByName(envVars []corev1.EnvVar, name string) *corev1.EnvVar { return nil } +func findContainerByName(containers []corev1.Container, name string) *corev1.Container { + for i := range containers { + if containers[i].Name == name { + return &containers[i] + } + } + return nil +} + func TestBuildResourceContainer(t *testing.T) { sparkResourceHandler := sparkResourceHandler{} @@ -737,65 +747,73 @@ func TestBuildResourceContainer(t *testing.T) { assert.Equal(t, sparkOp.SparkApplicationTypePython, sparkApp.Spec.Type) assert.Equal(t, testArgs, sparkApp.Spec.Arguments) assert.Equal(t, testImage, *sparkApp.Spec.Image) - assert.NotNil(t, sparkApp.Spec.Driver.SparkPodSpec.PodSecurityContext) - assert.Equal(t, *sparkApp.Spec.Driver.SparkPodSpec.PodSecurityContext.RunAsUser, *defaultConfig.DefaultPodSecurityContext.RunAsUser) - assert.NotNil(t, sparkApp.Spec.Driver.DNSConfig) - assert.Equal(t, []string{"8.8.8.8", "8.8.4.4"}, sparkApp.Spec.Driver.DNSConfig.Nameservers) - assert.ElementsMatch(t, defaultConfig.DefaultPodDNSConfig.Options, sparkApp.Spec.Driver.DNSConfig.Options) - assert.Equal(t, []string{"ns1.svc.cluster-domain.example", "my.dns.search.suffix"}, sparkApp.Spec.Driver.DNSConfig.Searches) - assert.NotNil(t, sparkApp.Spec.Executor.SparkPodSpec.PodSecurityContext) - assert.Equal(t, *sparkApp.Spec.Executor.SparkPodSpec.PodSecurityContext.RunAsUser, *defaultConfig.DefaultPodSecurityContext.RunAsUser) - assert.NotNil(t, sparkApp.Spec.Executor.DNSConfig) - assert.NotNil(t, sparkApp.Spec.Executor.DNSConfig) - assert.ElementsMatch(t, defaultConfig.DefaultPodDNSConfig.Options, sparkApp.Spec.Executor.DNSConfig.Options) - assert.Equal(t, []string{"ns1.svc.cluster-domain.example", "my.dns.search.suffix"}, sparkApp.Spec.Executor.DNSConfig.Searches) + assert.NotNil(t, sparkApp.Spec.Driver.Template) + assert.NotNil(t, sparkApp.Spec.Executor.Template) + driverPodSpec := sparkApp.Spec.Driver.Template.Spec + executorPodSpec := sparkApp.Spec.Executor.Template.Spec + driverContainer := findContainerByName(driverPodSpec.Containers, sparkOpConfig.SparkDriverContainerName) + executorContainer := findContainerByName(executorPodSpec.Containers, sparkOpConfig.Spark3DefaultExecutorContainerName) + assert.NotNil(t, driverContainer) + assert.NotNil(t, executorContainer) + assert.NotNil(t, driverPodSpec.SecurityContext) + assert.Equal(t, *driverPodSpec.SecurityContext.RunAsUser, *defaultConfig.DefaultPodSecurityContext.RunAsUser) + assert.NotNil(t, driverPodSpec.DNSConfig) + assert.Equal(t, []string{"8.8.8.8", "8.8.4.4"}, driverPodSpec.DNSConfig.Nameservers) + assert.ElementsMatch(t, defaultConfig.DefaultPodDNSConfig.Options, driverPodSpec.DNSConfig.Options) + assert.Equal(t, []string{"ns1.svc.cluster-domain.example", "my.dns.search.suffix"}, driverPodSpec.DNSConfig.Searches) + assert.NotNil(t, executorPodSpec.SecurityContext) + assert.Equal(t, *executorPodSpec.SecurityContext.RunAsUser, *defaultConfig.DefaultPodSecurityContext.RunAsUser) + assert.NotNil(t, executorPodSpec.DNSConfig) + assert.ElementsMatch(t, defaultConfig.DefaultPodDNSConfig.Options, executorPodSpec.DNSConfig.Options) + assert.Equal(t, []string{"ns1.svc.cluster-domain.example", "my.dns.search.suffix"}, executorPodSpec.DNSConfig.Searches) //Validate Driver/Executor Spec. - driverCores, _ := strconv.ParseInt(dummySparkConf["spark.driver.cores"], 10, 32) - execCores, _ := strconv.ParseInt(dummySparkConf["spark.executor.cores"], 10, 32) - execInstances, _ := strconv.ParseInt(dummySparkConf["spark.executor.instances"], 10, 32) - assert.Equal(t, "new-val", *sparkApp.Spec.Executor.ServiceAccount) assert.Equal(t, "new-val", *sparkApp.Spec.Driver.ServiceAccount) - assert.Equal(t, int32(driverCores), *sparkApp.Spec.Driver.Cores) - assert.Equal(t, int32(execCores), *sparkApp.Spec.Executor.Cores) - assert.Equal(t, int32(execInstances), *sparkApp.Spec.Executor.Instances) - assert.Equal(t, dummySparkConf["spark.driver.memory"], *sparkApp.Spec.Driver.Memory) - assert.Equal(t, dummySparkConf["spark.executor.memory"], *sparkApp.Spec.Executor.Memory) + assert.Equal(t, "new-val", driverPodSpec.ServiceAccountName) + assert.Equal(t, "new-val", executorPodSpec.ServiceAccountName) + + // Resource related configs are passed exclusively through SparkConf, not via the + // Driver/Executor specs. + assert.Nil(t, sparkApp.Spec.Driver.Cores) + assert.Nil(t, sparkApp.Spec.Executor.Cores) + assert.Nil(t, sparkApp.Spec.Executor.Instances) + assert.Nil(t, sparkApp.Spec.Driver.Memory) + assert.Nil(t, sparkApp.Spec.Executor.Memory) assert.Equal(t, dummySparkConf["spark.batchScheduler"], *sparkApp.Spec.BatchScheduler) - assert.Equal(t, defaultConfig.SchedulerName, *sparkApp.Spec.Executor.SchedulerName) - assert.Equal(t, defaultConfig.SchedulerName, *sparkApp.Spec.Driver.SchedulerName) - assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, *sparkApp.Spec.Executor.HostNetwork) - assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, *sparkApp.Spec.Driver.HostNetwork) + assert.Equal(t, defaultConfig.SchedulerName, executorPodSpec.SchedulerName) + assert.Equal(t, defaultConfig.SchedulerName, driverPodSpec.SchedulerName) + assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, executorPodSpec.HostNetwork) + assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, driverPodSpec.HostNetwork) // Validate // * Default tolerations set for both Driver and Executor. // * Interruptible tolerations and node selector set for Executor but not Driver. // * Default node selector set for both Driver and Executor. // * Interruptible node selector requirements set for Executor Affinity, non-interruptiblefir Driver Affinity. - assert.Equal(t, 1, len(sparkApp.Spec.Driver.Tolerations)) - assert.Equal(t, 1, len(sparkApp.Spec.Driver.NodeSelector)) - assert.Equal(t, defaultConfig.DefaultNodeSelector, sparkApp.Spec.Driver.NodeSelector) - tolDriverDefault := sparkApp.Spec.Driver.Tolerations[0] + assert.Equal(t, 1, len(driverPodSpec.Tolerations)) + assert.Equal(t, 1, len(driverPodSpec.NodeSelector)) + assert.Equal(t, defaultConfig.DefaultNodeSelector, driverPodSpec.NodeSelector) + tolDriverDefault := driverPodSpec.Tolerations[0] assert.Equal(t, tolDriverDefault.Key, "x/flyte") assert.Equal(t, tolDriverDefault.Value, "default") assert.Equal(t, tolDriverDefault.Operator, corev1.TolerationOperator("Equal")) assert.Equal(t, tolDriverDefault.Effect, corev1.TaintEffect("NoSchedule")) - assert.Equal(t, 2, len(sparkApp.Spec.Executor.Tolerations)) - assert.Equal(t, 2, len(sparkApp.Spec.Executor.NodeSelector)) + assert.Equal(t, 2, len(executorPodSpec.Tolerations)) + assert.Equal(t, 2, len(executorPodSpec.NodeSelector)) assert.Equal(t, map[string]string{ "x/default": "true", "x/interruptible": "true", - }, sparkApp.Spec.Executor.NodeSelector) + }, executorPodSpec.NodeSelector) - tolExecInterrupt := sparkApp.Spec.Executor.Tolerations[0] + tolExecInterrupt := executorPodSpec.Tolerations[0] assert.Equal(t, tolExecInterrupt.Key, "x/flyte") assert.Equal(t, tolExecInterrupt.Value, "interruptible") assert.Equal(t, tolExecInterrupt.Operator, corev1.TolerationOperator("Equal")) assert.Equal(t, tolExecInterrupt.Effect, corev1.TaintEffect("NoSchedule")) - tolExecDefault := sparkApp.Spec.Executor.Tolerations[1] + tolExecDefault := executorPodSpec.Tolerations[1] assert.Equal(t, tolExecDefault.Key, "x/flyte") assert.Equal(t, tolExecDefault.Value, "default") assert.Equal(t, tolExecDefault.Operator, corev1.TolerationOperator("Equal")) @@ -832,11 +850,11 @@ func TestBuildResourceContainer(t *testing.T) { assert.Greater(t, len(sparkApp.Spec.SparkConf["spark.kubernetes.driverEnv.FLYTE_START_TIME"]), 1) assert.Equal(t, dummySparkConf["spark.flyteorg.feature3.enabled"], sparkApp.Spec.SparkConf["spark.flyteorg.feature3.enabled"]) - assert.Equal(t, len(findEnvVarByName(sparkApp.Spec.Driver.Env, "FLYTE_MAX_ATTEMPTS").Value), 1) - assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(sparkApp.Spec.Driver.Env, "foo").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(sparkApp.Spec.Executor.Env, "foo").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(sparkApp.Spec.Driver.Env, "fooEnv").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(sparkApp.Spec.Executor.Env, "fooEnv").Value) + assert.Equal(t, len(findEnvVarByName(driverContainer.Env, "FLYTE_MAX_ATTEMPTS").Value), 1) + assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(driverContainer.Env, "foo").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(executorContainer.Env, "foo").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(driverContainer.Env, "fooEnv").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(executorContainer.Env, "fooEnv").Value) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ @@ -849,7 +867,7 @@ func TestBuildResourceContainer(t *testing.T) { }, }, }, - }, sparkApp.Spec.Driver.Affinity.NodeAffinity) + }, driverPodSpec.Affinity.NodeAffinity) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ @@ -862,7 +880,7 @@ func TestBuildResourceContainer(t *testing.T) { }, }, }, - }, sparkApp.Spec.Executor.Affinity.NodeAffinity) + }, executorPodSpec.Affinity.NodeAffinity) // Case 2: Driver/Executor request cores set. dummyConfWithRequest := make(map[string]string) @@ -893,16 +911,18 @@ func TestBuildResourceContainer(t *testing.T) { // Validate Interruptible Toleration and NodeSelector not set for both Driver and Executors. // Validate that the default Toleration and NodeSelector are set for both Driver and Executors. - assert.Equal(t, 1, len(sparkApp.Spec.Driver.Tolerations)) - assert.Equal(t, 1, len(sparkApp.Spec.Driver.NodeSelector)) - assert.Equal(t, defaultConfig.DefaultNodeSelector, sparkApp.Spec.Driver.NodeSelector) - assert.Equal(t, 1, len(sparkApp.Spec.Executor.Tolerations)) - assert.Equal(t, 1, len(sparkApp.Spec.Executor.NodeSelector)) - assert.Equal(t, defaultConfig.DefaultNodeSelector, sparkApp.Spec.Executor.NodeSelector) - assert.Equal(t, sparkApp.Spec.Executor.Tolerations[0].Key, "x/flyte") - assert.Equal(t, sparkApp.Spec.Executor.Tolerations[0].Value, "default") - assert.Equal(t, sparkApp.Spec.Driver.Tolerations[0].Key, "x/flyte") - assert.Equal(t, sparkApp.Spec.Driver.Tolerations[0].Value, "default") + driverPodSpec = sparkApp.Spec.Driver.Template.Spec + executorPodSpec = sparkApp.Spec.Executor.Template.Spec + assert.Equal(t, 1, len(driverPodSpec.Tolerations)) + assert.Equal(t, 1, len(driverPodSpec.NodeSelector)) + assert.Equal(t, defaultConfig.DefaultNodeSelector, driverPodSpec.NodeSelector) + assert.Equal(t, 1, len(executorPodSpec.Tolerations)) + assert.Equal(t, 1, len(executorPodSpec.NodeSelector)) + assert.Equal(t, defaultConfig.DefaultNodeSelector, executorPodSpec.NodeSelector) + assert.Equal(t, executorPodSpec.Tolerations[0].Key, "x/flyte") + assert.Equal(t, executorPodSpec.Tolerations[0].Value, "default") + assert.Equal(t, driverPodSpec.Tolerations[0].Key, "x/flyte") + assert.Equal(t, driverPodSpec.Tolerations[0].Value, "default") // Validate correct affinity and nodeselector requirements are set for both Driver and Executors. assert.Equal(t, &corev1.NodeAffinity{ @@ -916,7 +936,7 @@ func TestBuildResourceContainer(t *testing.T) { }, }, }, - }, sparkApp.Spec.Driver.Affinity.NodeAffinity) + }, driverPodSpec.Affinity.NodeAffinity) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ @@ -929,7 +949,7 @@ func TestBuildResourceContainer(t *testing.T) { }, }, }, - }, sparkApp.Spec.Executor.Affinity.NodeAffinity) + }, executorPodSpec.Affinity.NodeAffinity) // Case 4: Invalid Spark Task-Template taskTemplate.Custom = nil @@ -980,27 +1000,32 @@ func TestBuildResourcePodTemplate(t *testing.T) { assert.Equal(t, sparkApplicationFile, *sparkApp.Spec.MainApplicationFile) // Driver - assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultAnnotations, map[string]string{"annotation-1": "val1"}), sparkApp.Spec.Driver.Annotations) - assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultLabels, map[string]string{"label-1": "val1"}), sparkApp.Spec.Driver.Labels) - assert.Equal(t, len(findEnvVarByName(sparkApp.Spec.Driver.Env, "FLYTE_MAX_ATTEMPTS").Value), 1) - assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(sparkApp.Spec.Driver.Env, "foo").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(sparkApp.Spec.Driver.Env, "fooEnv").Value) - assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(sparkApp.Spec.Driver.Env, "SECRET")) - assert.Equal(t, 9, len(sparkApp.Spec.Driver.Env)) - assert.Equal(t, testImage, *sparkApp.Spec.Driver.Image) + assert.NotNil(t, sparkApp.Spec.Driver.Template) + driverPodSpec := sparkApp.Spec.Driver.Template.Spec + driverContainer := findContainerByName(driverPodSpec.Containers, sparkOpConfig.SparkDriverContainerName) + assert.NotNil(t, driverContainer) + assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultAnnotations, map[string]string{"annotation-1": "val1"}), sparkApp.Spec.Driver.Template.Annotations) + assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultLabels, map[string]string{"label-1": "val1"}), sparkApp.Spec.Driver.Template.Labels) + assert.Equal(t, len(findEnvVarByName(driverContainer.Env, "FLYTE_MAX_ATTEMPTS").Value), 1) + assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(driverContainer.Env, "foo").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(driverContainer.Env, "fooEnv").Value) + assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(driverContainer.Env, "SECRET")) + assert.Equal(t, 9, len(driverContainer.Env)) + assert.Equal(t, testImage, driverContainer.Image) assert.Equal(t, flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()), *sparkApp.Spec.Driver.ServiceAccount) - assert.Equal(t, defaultConfig.DefaultPodSecurityContext, sparkApp.Spec.Driver.PodSecurityContext) - assert.Equal(t, defaultConfig.DefaultPodDNSConfig, sparkApp.Spec.Driver.DNSConfig) - assert.Equal(t, defaultConfig.EnableHostNetworkingPod, sparkApp.Spec.Driver.HostNetwork) - assert.Equal(t, defaultConfig.SchedulerName, *sparkApp.Spec.Driver.SchedulerName) + assert.Equal(t, flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()), driverPodSpec.ServiceAccountName) + assert.Equal(t, defaultConfig.DefaultPodSecurityContext, driverPodSpec.SecurityContext) + assert.Equal(t, defaultConfig.DefaultPodDNSConfig, driverPodSpec.DNSConfig) + assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, driverPodSpec.HostNetwork) + assert.Equal(t, defaultConfig.SchedulerName, driverPodSpec.SchedulerName) assert.Equal(t, []corev1.Toleration{ defaultConfig.DefaultTolerations[0], extraToleration, - }, sparkApp.Spec.Driver.Tolerations) + }, driverPodSpec.Tolerations) assert.Equal(t, map[string]string{ "x/default": "true", "x/custom": "foo", - }, sparkApp.Spec.Driver.NodeSelector) + }, driverPodSpec.NodeSelector) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ NodeSelectorTerms: []corev1.NodeSelectorTerm{ @@ -1012,33 +1037,38 @@ func TestBuildResourcePodTemplate(t *testing.T) { }, }, }, - }, sparkApp.Spec.Driver.Affinity.NodeAffinity) - cores, _ := strconv.ParseInt(dummySparkConf["spark.driver.cores"], 10, 32) - assert.Equal(t, intPtr(int32(cores)), sparkApp.Spec.Driver.Cores) - assert.Equal(t, dummySparkConf["spark.driver.memory"], *sparkApp.Spec.Driver.Memory) + }, driverPodSpec.Affinity.NodeAffinity) + assert.Nil(t, sparkApp.Spec.Driver.Cores) + assert.Nil(t, sparkApp.Spec.Driver.Memory) + assert.Equal(t, dummySparkConf["spark.driver.cores"], sparkApp.Spec.SparkConf["spark.driver.cores"]) + assert.Equal(t, dummySparkConf["spark.driver.memory"], sparkApp.Spec.SparkConf["spark.driver.memory"]) // Executor - assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultAnnotations, map[string]string{"annotation-1": "val1"}), sparkApp.Spec.Executor.Annotations) - assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultLabels, map[string]string{"label-1": "val1"}), sparkApp.Spec.Executor.Labels) - assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(sparkApp.Spec.Executor.Env, "foo").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(sparkApp.Spec.Executor.Env, "fooEnv").Value) - assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(sparkApp.Spec.Executor.Env, "SECRET")) - assert.Equal(t, 9, len(sparkApp.Spec.Executor.Env)) - assert.Equal(t, testImage, *sparkApp.Spec.Executor.Image) - assert.Equal(t, defaultConfig.DefaultPodSecurityContext, sparkApp.Spec.Executor.PodSecurityContext) - assert.Equal(t, defaultConfig.DefaultPodDNSConfig, sparkApp.Spec.Executor.DNSConfig) - assert.Equal(t, defaultConfig.EnableHostNetworkingPod, sparkApp.Spec.Executor.HostNetwork) - assert.Equal(t, defaultConfig.SchedulerName, *sparkApp.Spec.Executor.SchedulerName) + assert.NotNil(t, sparkApp.Spec.Executor.Template) + executorPodSpec := sparkApp.Spec.Executor.Template.Spec + executorContainer := findContainerByName(executorPodSpec.Containers, sparkOpConfig.Spark3DefaultExecutorContainerName) + assert.NotNil(t, executorContainer) + assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultAnnotations, map[string]string{"annotation-1": "val1"}), sparkApp.Spec.Executor.Template.Annotations) + assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultLabels, map[string]string{"label-1": "val1"}), sparkApp.Spec.Executor.Template.Labels) + assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(executorContainer.Env, "foo").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(executorContainer.Env, "fooEnv").Value) + assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(executorContainer.Env, "SECRET")) + assert.Equal(t, 9, len(executorContainer.Env)) + assert.Equal(t, testImage, executorContainer.Image) + assert.Equal(t, defaultConfig.DefaultPodSecurityContext, executorPodSpec.SecurityContext) + assert.Equal(t, defaultConfig.DefaultPodDNSConfig, executorPodSpec.DNSConfig) + assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, executorPodSpec.HostNetwork) + assert.Equal(t, defaultConfig.SchedulerName, executorPodSpec.SchedulerName) assert.ElementsMatch(t, []corev1.Toleration{ defaultConfig.DefaultTolerations[0], extraToleration, defaultConfig.InterruptibleTolerations[0], - }, sparkApp.Spec.Executor.Tolerations) + }, executorPodSpec.Tolerations) assert.Equal(t, map[string]string{ "x/default": "true", "x/custom": "foo", "x/interruptible": "true", - }, sparkApp.Spec.Executor.NodeSelector) + }, executorPodSpec.NodeSelector) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ NodeSelectorTerms: []corev1.NodeSelectorTerm{ @@ -1050,12 +1080,13 @@ func TestBuildResourcePodTemplate(t *testing.T) { }, }, }, - }, sparkApp.Spec.Executor.Affinity.NodeAffinity) - cores, _ = strconv.ParseInt(dummySparkConf["spark.executor.cores"], 10, 32) - instances, _ := strconv.ParseInt(dummySparkConf["spark.executor.instances"], 10, 32) - assert.Equal(t, intPtr(int32(instances)), sparkApp.Spec.Executor.Instances) - assert.Equal(t, intPtr(int32(cores)), sparkApp.Spec.Executor.Cores) - assert.Equal(t, dummySparkConf["spark.executor.memory"], *sparkApp.Spec.Executor.Memory) + }, executorPodSpec.Affinity.NodeAffinity) + assert.Nil(t, sparkApp.Spec.Executor.Instances) + assert.Nil(t, sparkApp.Spec.Executor.Cores) + assert.Nil(t, sparkApp.Spec.Executor.Memory) + assert.Equal(t, dummySparkConf["spark.executor.cores"], sparkApp.Spec.SparkConf["spark.executor.cores"]) + assert.Equal(t, dummySparkConf["spark.executor.instances"], sparkApp.Spec.SparkConf["spark.executor.instances"]) + assert.Equal(t, dummySparkConf["spark.executor.memory"], sparkApp.Spec.SparkConf["spark.executor.memory"]) } func TestBuildResourcePriorityClassName(t *testing.T) { @@ -1076,10 +1107,14 @@ func TestBuildResourcePriorityClassName(t *testing.T) { sparkApp, ok := resource.(*sparkOp.SparkApplication) assert.True(t, ok) - assert.NotNil(t, sparkApp.Spec.Driver.PriorityClassName) - assert.Equal(t, priorityClassName, *sparkApp.Spec.Driver.PriorityClassName) - assert.NotNil(t, sparkApp.Spec.Executor.PriorityClassName) - assert.Equal(t, priorityClassName, *sparkApp.Spec.Executor.PriorityClassName) + // The priority class is carried through the pod template rather than the + // Driver/Executor spec attributes. + assert.Nil(t, sparkApp.Spec.Driver.PriorityClassName) + assert.Nil(t, sparkApp.Spec.Executor.PriorityClassName) + assert.NotNil(t, sparkApp.Spec.Driver.Template) + assert.Equal(t, priorityClassName, sparkApp.Spec.Driver.Template.Spec.PriorityClassName) + assert.NotNil(t, sparkApp.Spec.Executor.Template) + assert.Equal(t, priorityClassName, sparkApp.Spec.Executor.Template.Spec.PriorityClassName) } func TestBuildResourceNoPriorityClassName(t *testing.T) { @@ -1096,9 +1131,13 @@ func TestBuildResourceNoPriorityClassName(t *testing.T) { sparkApp, ok := resource.(*sparkOp.SparkApplication) assert.True(t, ok) - // When no priority class is set, the field should be left unset (nil) rather than an empty string. + // When no priority class is set, the fields should be left unset. assert.Nil(t, sparkApp.Spec.Driver.PriorityClassName) assert.Nil(t, sparkApp.Spec.Executor.PriorityClassName) + assert.NotNil(t, sparkApp.Spec.Driver.Template) + assert.Empty(t, sparkApp.Spec.Driver.Template.Spec.PriorityClassName) + assert.NotNil(t, sparkApp.Spec.Executor.Template) + assert.Empty(t, sparkApp.Spec.Executor.Template.Spec.PriorityClassName) } func TestGetPropertiesSpark(t *testing.T) { @@ -1179,35 +1218,40 @@ func TestBuildResourceCustomK8SPod(t *testing.T) { assert.Equal(t, sparkApplicationFile, *sparkApp.Spec.MainApplicationFile) // Driver + assert.NotNil(t, sparkApp.Spec.Driver.Template) + driverTplPodSpec := sparkApp.Spec.Driver.Template.Spec + driverContainer := findContainerByName(driverTplPodSpec.Containers, sparkOpConfig.SparkDriverContainerName) + assert.NotNil(t, driverContainer) assert.Equal(t, utils.UnionMaps( defaultConfig.DefaultAnnotations, map[string]string{ "annotation-1": "val1", "annotation-driver": "val-driver", }, - ), sparkApp.Spec.Driver.Annotations) + ), sparkApp.Spec.Driver.Template.Annotations) assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultLabels, map[string]string{ "label-1": "val1", "label-driver": "val-driver", - }), sparkApp.Spec.Driver.Labels) - assert.Equal(t, len(findEnvVarByName(sparkApp.Spec.Driver.Env, "FLYTE_MAX_ATTEMPTS").Value), 1) - assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(sparkApp.Spec.Driver.Env, "foo").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(sparkApp.Spec.Driver.Env, "fooEnv").Value) - assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(sparkApp.Spec.Driver.Env, "SECRET")) - assert.Equal(t, 11, len(sparkApp.Spec.Driver.Env)) - assert.Equal(t, testImage, *sparkApp.Spec.Driver.Image) + }), sparkApp.Spec.Driver.Template.Labels) + assert.Equal(t, len(findEnvVarByName(driverContainer.Env, "FLYTE_MAX_ATTEMPTS").Value), 1) + assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(driverContainer.Env, "foo").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(driverContainer.Env, "fooEnv").Value) + assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(driverContainer.Env, "SECRET")) + assert.Equal(t, 11, len(driverContainer.Env)) + assert.Equal(t, testImage, driverContainer.Image) assert.Equal(t, flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()), *sparkApp.Spec.Driver.ServiceAccount) - assert.Equal(t, defaultConfig.DefaultPodSecurityContext, sparkApp.Spec.Driver.PodSecurityContext) - assert.Equal(t, defaultConfig.DefaultPodDNSConfig, sparkApp.Spec.Driver.DNSConfig) - assert.Equal(t, defaultConfig.EnableHostNetworkingPod, sparkApp.Spec.Driver.HostNetwork) - assert.Equal(t, defaultConfig.SchedulerName, *sparkApp.Spec.Driver.SchedulerName) + assert.Equal(t, flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()), driverTplPodSpec.ServiceAccountName) + assert.Equal(t, defaultConfig.DefaultPodSecurityContext, driverTplPodSpec.SecurityContext) + assert.Equal(t, defaultConfig.DefaultPodDNSConfig, driverTplPodSpec.DNSConfig) + assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, driverTplPodSpec.HostNetwork) + assert.Equal(t, defaultConfig.SchedulerName, driverTplPodSpec.SchedulerName) assert.Equal(t, []corev1.Toleration{ defaultConfig.DefaultTolerations[0], driverExtraToleration, - }, sparkApp.Spec.Driver.Tolerations) + }, driverTplPodSpec.Tolerations) assert.Equal(t, map[string]string{ "x/default": "true", "x/custom": "foo-driver", - }, sparkApp.Spec.Driver.NodeSelector) + }, driverTplPodSpec.NodeSelector) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ NodeSelectorTerms: []corev1.NodeSelectorTerm{ @@ -1219,39 +1263,44 @@ func TestBuildResourceCustomK8SPod(t *testing.T) { }, }, }, - }, sparkApp.Spec.Driver.Affinity.NodeAffinity) - cores, _ := strconv.ParseInt(dummySparkConf["spark.driver.cores"], 10, 32) - assert.Equal(t, intPtr(int32(cores)), sparkApp.Spec.Driver.Cores) - assert.Equal(t, dummySparkConf["spark.driver.memory"], *sparkApp.Spec.Driver.Memory) + }, driverTplPodSpec.Affinity.NodeAffinity) + assert.Nil(t, sparkApp.Spec.Driver.Cores) + assert.Nil(t, sparkApp.Spec.Driver.Memory) + assert.Equal(t, dummySparkConf["spark.driver.cores"], sparkApp.Spec.SparkConf["spark.driver.cores"]) + assert.Equal(t, dummySparkConf["spark.driver.memory"], sparkApp.Spec.SparkConf["spark.driver.memory"]) // // Executor + assert.NotNil(t, sparkApp.Spec.Executor.Template) + executorTplPodSpec := sparkApp.Spec.Executor.Template.Spec + executorContainer := findContainerByName(executorTplPodSpec.Containers, sparkOpConfig.Spark3DefaultExecutorContainerName) + assert.NotNil(t, executorContainer) assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultAnnotations, map[string]string{ "annotation-1": "val1", "annotation-executor": "val-executor", - }), sparkApp.Spec.Executor.Annotations) + }), sparkApp.Spec.Executor.Template.Annotations) assert.Equal(t, utils.UnionMaps(defaultConfig.DefaultLabels, map[string]string{ "label-1": "val1", "label-executor": "val-executor", - }), sparkApp.Spec.Executor.Labels) - assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(sparkApp.Spec.Executor.Env, "foo").Value) - assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(sparkApp.Spec.Executor.Env, "fooEnv").Value) - assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(sparkApp.Spec.Executor.Env, "SECRET")) - assert.Equal(t, 11, len(sparkApp.Spec.Executor.Env)) - assert.Equal(t, testImage, *sparkApp.Spec.Executor.Image) - assert.Equal(t, defaultConfig.DefaultPodSecurityContext, sparkApp.Spec.Executor.PodSecurityContext) - assert.Equal(t, defaultConfig.DefaultPodDNSConfig, sparkApp.Spec.Executor.DNSConfig) - assert.Equal(t, defaultConfig.EnableHostNetworkingPod, sparkApp.Spec.Executor.HostNetwork) - assert.Equal(t, defaultConfig.SchedulerName, *sparkApp.Spec.Executor.SchedulerName) + }), sparkApp.Spec.Executor.Template.Labels) + assert.Equal(t, defaultConfig.DefaultEnvVars["foo"], findEnvVarByName(executorContainer.Env, "foo").Value) + assert.Equal(t, defaultConfig.DefaultEnvVars["fooEnv"], findEnvVarByName(executorContainer.Env, "fooEnv").Value) + assert.Equal(t, findEnvVarByName(dummyEnvVarsWithSecretRef, "SECRET"), findEnvVarByName(executorContainer.Env, "SECRET")) + assert.Equal(t, 11, len(executorContainer.Env)) + assert.Equal(t, testImage, executorContainer.Image) + assert.Equal(t, defaultConfig.DefaultPodSecurityContext, executorTplPodSpec.SecurityContext) + assert.Equal(t, defaultConfig.DefaultPodDNSConfig, executorTplPodSpec.DNSConfig) + assert.Equal(t, *defaultConfig.EnableHostNetworkingPod, executorTplPodSpec.HostNetwork) + assert.Equal(t, defaultConfig.SchedulerName, executorTplPodSpec.SchedulerName) assert.ElementsMatch(t, []corev1.Toleration{ defaultConfig.DefaultTolerations[0], executorExtraToleration, defaultConfig.InterruptibleTolerations[0], - }, sparkApp.Spec.Executor.Tolerations) + }, executorTplPodSpec.Tolerations) assert.Equal(t, map[string]string{ "x/default": "true", "x/custom": "foo-executor", "x/interruptible": "true", - }, sparkApp.Spec.Executor.NodeSelector) + }, executorTplPodSpec.NodeSelector) assert.Equal(t, &corev1.NodeAffinity{ RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ NodeSelectorTerms: []corev1.NodeSelectorTerm{ @@ -1263,12 +1312,13 @@ func TestBuildResourceCustomK8SPod(t *testing.T) { }, }, }, - }, sparkApp.Spec.Executor.Affinity.NodeAffinity) - cores, _ = strconv.ParseInt(dummySparkConf["spark.executor.cores"], 10, 32) - instances, _ := strconv.ParseInt(dummySparkConf["spark.executor.instances"], 10, 32) - assert.Equal(t, intPtr(int32(instances)), sparkApp.Spec.Executor.Instances) - assert.Equal(t, intPtr(int32(cores)), sparkApp.Spec.Executor.Cores) - assert.Equal(t, dummySparkConf["spark.executor.memory"], *sparkApp.Spec.Executor.Memory) + }, executorTplPodSpec.Affinity.NodeAffinity) + assert.Nil(t, sparkApp.Spec.Executor.Instances) + assert.Nil(t, sparkApp.Spec.Executor.Cores) + assert.Nil(t, sparkApp.Spec.Executor.Memory) + assert.Equal(t, dummySparkConf["spark.executor.cores"], sparkApp.Spec.SparkConf["spark.executor.cores"]) + assert.Equal(t, dummySparkConf["spark.executor.instances"], sparkApp.Spec.SparkConf["spark.executor.instances"]) + assert.Equal(t, dummySparkConf["spark.executor.memory"], sparkApp.Spec.SparkConf["spark.executor.memory"]) } func TestGetEventInfo_DynamicLogLinks(t *testing.T) {