[SPARK-51250][K8S] Add Support for K8s PriorityClass Configuration fo…

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala

@@ -749,6 +749,13 @@ private[spark] object Config extends Logging {
       .checkValue(value => value > 0, "Gracefully shutdown period must be a positive time value")
       .createWithDefaultString("20s")
 
+  val KUBERNETES_PRIORITY_CLASS_NAME =
+    ConfigBuilder("spark.kubernetes.priorityclass.name")
+      .doc("Priority class name of driver and executor pods")
+      .version("4.1.0")
+      .stringConf
+      .createOptional
+
   val KUBERNETES_DRIVER_LABEL_PREFIX = "spark.kubernetes.driver.label."
   val KUBERNETES_DRIVER_ANNOTATION_PREFIX = "spark.kubernetes.driver.annotation."
   val KUBERNETES_DRIVER_SERVICE_LABEL_PREFIX = "spark.kubernetes.driver.service.label."

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/PriorityClassFeatureStep.scala

@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.deploy.k8s.features
+
+import io.fabric8.kubernetes.api.model.PodBuilder
+
+import org.apache.spark.deploy.k8s.{KubernetesConf, SparkPod}
+import org.apache.spark.deploy.k8s.Config.KUBERNETES_PRIORITY_CLASS_NAME
+
+private[spark] class PriorityClassFeatureStep(conf: KubernetesConf)
+  extends KubernetesFeatureConfigStep {
+
+  private val priorityClassNameConfig = conf.get(KUBERNETES_PRIORITY_CLASS_NAME)
+
+  def configurePod(pod: SparkPod): SparkPod = {
+    priorityClassNameConfig match {
+      case Some(priorityClassName) =>
+        val updatedPod = new PodBuilder(pod.pod)
+          .editSpec()
+          .withPriorityClassName(priorityClassName)
+          .endSpec()
+          .build()
+        SparkPod(updatedPod, pod.container)
+      case None => pod
+    }
+  }
+}

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/KubernetesDriverBuilder.scala

@@ -83,7 +83,8 @@ class KubernetesDriverBuilder {
       new HadoopConfDriverFeatureStep(conf),
       new KerberosConfDriverFeatureStep(conf),
       new PodTemplateConfigMapStep(conf),
-      new LocalDirsFeatureStep(conf)) ++ userFeatures
+      new LocalDirsFeatureStep(conf),
+      new PriorityClassFeatureStep(conf)) ++ userFeatures
 
     val spec = KubernetesDriverSpec(
       initialPod,

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBuilder.scala

@@ -72,7 +72,8 @@ private[spark] class KubernetesExecutorBuilder {
       new EnvSecretsFeatureStep(conf),
       new MountVolumesFeatureStep(conf),
       new HadoopConfExecutorFeatureStep(conf),
-      new LocalDirsFeatureStep(conf)) ++ userFeatures
+      new LocalDirsFeatureStep(conf),
+      new PriorityClassFeatureStep(conf)) ++ userFeatures
 
     val spec = KubernetesExecutorSpec(
       initialPod,

resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/PriorityClassFeatureStepSuite.scala

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.deploy.k8s.features
+
+import org.apache.spark.{SparkConf, SparkFunSuite}
+import org.apache.spark.deploy.k8s._
+import org.apache.spark.deploy.k8s.Config._
+
+class PriorityClassFeatureStepSuite extends SparkFunSuite {
+
+  test("Test driver priority class name is set") {
+    val sparkConf = new SparkConf(false).set(KUBERNETES_PRIORITY_CLASS_NAME, "high-criticality-job")
+    val kubernetesDriverConf = KubernetesTestConf.createDriverConf(sparkConf = sparkConf)
+    val priorityClassFeatureStep = new PriorityClassFeatureStep(kubernetesDriverConf)
+    val driverPod = priorityClassFeatureStep.configurePod(SparkPod.initialPod())
+    assert(driverPod.pod.getSpec.getPriorityClassName === "high-criticality-job")
+  }
+
+  test("Test executor priority class name is set") {
+    val sparkConf = new SparkConf(false).set(KUBERNETES_PRIORITY_CLASS_NAME, "high-criticality-job")
+    val kubernetesExecutorConf = KubernetesTestConf.createExecutorConf(sparkConf = sparkConf)
+    val priorityClassFeatureStep = new PriorityClassFeatureStep(kubernetesExecutorConf)
+    val executorPod = priorityClassFeatureStep.configurePod(SparkPod.initialPod())
+    assert(executorPod.pod.getSpec.getPriorityClassName === "high-criticality-job")
+  }
+
+  test("Test pod spec remain unchanged when no priority class name is set") {
+    val sparkConf = new SparkConf(false)
+    val kubernetesExecutorConf = KubernetesTestConf.createExecutorConf(sparkConf = sparkConf)
+    val executorStep = new PriorityClassFeatureStep(kubernetesExecutorConf)
+    val executorPod = executorStep.configurePod(SparkPod.initialPod())
+    val podWithoutStep = SparkPod.initialPod()
+    assert(executorPod.pod.getSpec.getPriorityClassName === null)
+    assert(executorPod.pod.getSpec === podWithoutStep.pod.getSpec)
+  }
+}