set priorityClassName for spark tasks (#7476)

Signed-off-by: madiyar-wayve <madiyar.aitzhanov@wayve.ai>

Author: madiyar-wayve

flyteplugins/go/tasks/plugins/k8s/spark/spark.go

@@ -234,6 +234,7 @@ func createDriverSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionCont
 		},
 	}
 	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))
@@ -288,6 +289,7 @@ func createExecutorSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionCo
 		},
 		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))
 	}

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

@@ -1058,6 +1058,49 @@ func TestBuildResourcePodTemplate(t *testing.T) {
 	assert.Equal(t, dummySparkConf["spark.executor.memory"], *sparkApp.Spec.Executor.Memory)
 }
 
+func TestBuildResourcePriorityClassName(t *testing.T) {
+	defaultConfig := defaultPluginConfig()
+	assert.NoError(t, config.SetK8sPluginConfig(defaultConfig))
+
+	const priorityClassName = "high-priority"
+	podSpec := dummyPodSpec()
+	podSpec.PriorityClassName = priorityClassName
+	taskTemplate := dummySparkTaskTemplatePod("blah-1", dummySparkConf, podSpec)
+
+	sparkResourceHandler := sparkResourceHandler{}
+	taskCtx := dummySparkTaskContext(taskTemplate, true, k8s.PluginState{})
+	resource, err := sparkResourceHandler.BuildResource(context.TODO(), taskCtx)
+
+	assert.Nil(t, err)
+	assert.NotNil(t, resource)
+	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)
+}
+
+func TestBuildResourceNoPriorityClassName(t *testing.T) {
+	defaultConfig := defaultPluginConfig()
+	assert.NoError(t, config.SetK8sPluginConfig(defaultConfig))
+
+	taskTemplate := dummySparkTaskTemplateContainer("blah-1", dummySparkConf)
+	sparkResourceHandler := sparkResourceHandler{}
+	taskCtx := dummySparkTaskContext(taskTemplate, true, k8s.PluginState{})
+	resource, err := sparkResourceHandler.BuildResource(context.TODO(), taskCtx)
+
+	assert.Nil(t, err)
+	assert.NotNil(t, resource)
+	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.
+	assert.Nil(t, sparkApp.Spec.Driver.PriorityClassName)
+	assert.Nil(t, sparkApp.Spec.Executor.PriorityClassName)
+}
+
 func TestGetPropertiesSpark(t *testing.T) {
 	sparkResourceHandler := sparkResourceHandler{}
 	expected := k8s.PluginProperties{}