feat: support driver and executor pod use different priority (kubeflo…

…w#2146) * feat: support driver and executor pod use different priority Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> * feat: if *app.Spec.Driver.PriorityClassName and *app.Spec.Executor.PriorityClassName specifically defined, then can precedence over spec.batchSchedulerOptions.priorityClassName Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> * feat: merge the logic of setPodPriorityClassName into addPriorityClassName Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> * feat: support driver and executor pod use different priority Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> * feat: if *app.Spec.Driver.PriorityClassName and *app.Spec.Executor.PriorityClassName specifically defined, then can precedence over spec.batchSchedulerOptions.priorityClassName Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> * feat: merge the logic of setPodPriorityClassName into addPriorityClassName Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> * feat: add adjust pointer if is nil Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> * feat: remove spec.batchSchedulerOptions.priorityClassName define , split driver and executor pod priorityClass Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> * feat: remove spec.batchSchedulerOptions.priorityClassName define , split driver and executor pod priorityClass Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> * feat: Optimize code to avoid null pointer exceptions Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> * fix: remove backup crd files Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> * fix: remove BatchSchedulerOptions.PriorityClassName test code Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> * fix: add driver and executor pod priorityClassName test code Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> --------- Signed-off-by: Kevin Wu <kevin.wu@momenta.ai> Signed-off-by: Kevin.Wu <kevin.wu@momenta.ai> Co-authored-by: Kevin Wu <kevin.wu@momenta.ai>
jbhalodia-slack · Sep 10, 2024 · 6ae1b2f · 6ae1b2f
1 parent 95d202e
commit 6ae1b2f
Show file tree

Hide file tree

Showing 9 changed files with 84 additions and 13 deletions.
diff --git a/api/v1beta2/sparkapplication_types.go b/api/v1beta2/sparkapplication_types.go
@@ -536,6 +536,9 @@ type DriverSpec struct {
 	// Ports settings for the pods, following the Kubernetes specifications.
 	// +optional
 	Ports []Port `json:"ports,omitempty"`
+	// PriorityClassName is the name of the PriorityClass for the driver pod.
+	// +optional
+	PriorityClassName *string `json:"priorityClassName,omitempty"`
 }
 
 // ExecutorSpec is specification of the executor.
@@ -563,6 +566,9 @@ type ExecutorSpec struct {
 	// Ports settings for the pods, following the Kubernetes specifications.
 	// +optional
 	Ports []Port `json:"ports,omitempty"`
+	// PriorityClassName is the name of the PriorityClass for the executor pod.
+	// +optional
+	PriorityClassName *string `json:"priorityClassName,omitempty"`
 }
 
 // NamePath is a pair of a name and a path to which the named objects should be mounted to.

diff --git a/api/v1beta2/zz_generated.deepcopy.go b/api/v1beta2/zz_generated.deepcopy.go
diff --git a/charts/spark-operator-chart/crds/sparkoperator.k8s.io_scheduledsparkapplications.yaml b/charts/spark-operator-chart/crds/sparkoperator.k8s.io_scheduledsparkapplications.yaml
@@ -3179,6 +3179,10 @@ spec:
                           - protocol
                           type: object
                         type: array
+                      priorityClassName:
+                        description: PriorityClassName is the name of the PriorityClass
+                          for the driver pod.
+                        type: string
                       schedulerName:
                         description: SchedulerName specifies the scheduler that will
                           be used for scheduling
@@ -7946,6 +7950,10 @@ spec:
                           - protocol
                           type: object
                         type: array
+                      priorityClassName:
+                        description: PriorityClassName is the name of the PriorityClass
+                          for the executor pod.
+                        type: string
                       schedulerName:
                         description: SchedulerName specifies the scheduler that will
                           be used for scheduling

diff --git a/charts/spark-operator-chart/crds/sparkoperator.k8s.io_sparkapplications.yaml b/charts/spark-operator-chart/crds/sparkoperator.k8s.io_sparkapplications.yaml
@@ -3127,6 +3127,10 @@ spec:
                       - protocol
                       type: object
                     type: array
+                  priorityClassName:
+                    description: PriorityClassName is the name of the PriorityClass
+                      for the driver pod.
+                    type: string
                   schedulerName:
                     description: SchedulerName specifies the scheduler that will be
                       used for scheduling
@@ -7864,6 +7868,10 @@ spec:
                       - protocol
                       type: object
                     type: array
+                  priorityClassName:
+                    description: PriorityClassName is the name of the PriorityClass
+                      for the executor pod.
+                    type: string
                   schedulerName:
                     description: SchedulerName specifies the scheduler that will be
                       used for scheduling

diff --git a/config/crd/bases/sparkoperator.k8s.io_scheduledsparkapplications.yaml b/config/crd/bases/sparkoperator.k8s.io_scheduledsparkapplications.yaml
@@ -3179,6 +3179,10 @@ spec:
                           - protocol
                           type: object
                         type: array
+                      priorityClassName:
+                        description: PriorityClassName is the name of the PriorityClass
+                          for the driver pod.
+                        type: string
                       schedulerName:
                         description: SchedulerName specifies the scheduler that will
                           be used for scheduling
@@ -7946,6 +7950,10 @@ spec:
                           - protocol
                           type: object
                         type: array
+                      priorityClassName:
+                        description: PriorityClassName is the name of the PriorityClass
+                          for the executor pod.
+                        type: string
                       schedulerName:
                         description: SchedulerName specifies the scheduler that will
                           be used for scheduling

diff --git a/config/crd/bases/sparkoperator.k8s.io_sparkapplications.yaml b/config/crd/bases/sparkoperator.k8s.io_sparkapplications.yaml
@@ -3127,6 +3127,10 @@ spec:
                       - protocol
                       type: object
                     type: array
+                  priorityClassName:
+                    description: PriorityClassName is the name of the PriorityClass
+                      for the driver pod.
+                    type: string
                   schedulerName:
                     description: SchedulerName specifies the scheduler that will be
                       used for scheduling
@@ -7864,6 +7868,10 @@ spec:
                       - protocol
                       type: object
                     type: array
+                  priorityClassName:
+                    description: PriorityClassName is the name of the PriorityClass
+                      for the executor pod.
+                    type: string
                   schedulerName:
                     description: SchedulerName specifies the scheduler that will be
                       used for scheduling

diff --git a/docs/api-docs.md b/docs/api-docs.md
@@ -625,6 +625,18 @@ executors to connect to the driver.</p>
 <p>Ports settings for the pods, following the Kubernetes specifications.</p>
 </td>
 </tr>
+<tr>
+<td>
+<code>priorityClassName</code><br/>
+<em>
+string
+</em>
+</td>
+<td>
+<em>(Optional)</em>
+<p>PriorityClassName is the name of the PriorityClass for the driver pod.</p>
+</td>
+</tr>
 </tbody>
 </table>
 <h3 id="sparkoperator.k8s.io/v1beta2.DriverState">DriverState
@@ -840,6 +852,18 @@ Maps to <code>spark.kubernetes.executor.deleteOnTermination</code> that is avail
 <p>Ports settings for the pods, following the Kubernetes specifications.</p>
 </td>
 </tr>
+<tr>
+<td>
+<code>priorityClassName</code><br/>
+<em>
+string
+</em>
+</td>
+<td>
+<em>(Optional)</em>
+<p>PriorityClassName is the name of the PriorityClass for the executor pod.</p>
+</td>
+</tr>
 </tbody>
 </table>
 <h3 id="sparkoperator.k8s.io/v1beta2.ExecutorState">ExecutorState

diff --git a/internal/webhook/sparkpod_defaulter.go b/internal/webhook/sparkpod_defaulter.go
@@ -477,19 +477,19 @@ func addSchedulerName(pod *corev1.Pod, app *v1beta2.SparkApplication) error {
 
 func addPriorityClassName(pod *corev1.Pod, app *v1beta2.SparkApplication) error {
 	var priorityClassName *string
-	if app.Spec.BatchSchedulerOptions != nil {
-		priorityClassName = app.Spec.BatchSchedulerOptions.PriorityClassName
+
+	if util.IsDriverPod(pod) {
+		priorityClassName = app.Spec.Driver.PriorityClassName
+	} else if util.IsExecutorPod(pod) {
+		priorityClassName = app.Spec.Executor.PriorityClassName
 	}
 
 	if priorityClassName != nil && *priorityClassName != "" {
 		pod.Spec.PriorityClassName = *priorityClassName
-		if pod.Spec.Priority != nil {
-			pod.Spec.Priority = nil
-		}
-		if pod.Spec.PreemptionPolicy != nil {
-			pod.Spec.PreemptionPolicy = nil
-		}
+		pod.Spec.Priority = nil
+		pod.Spec.PreemptionPolicy = nil
 	}
+
 	return nil
 }
 

diff --git a/internal/webhook/sparkpod_defaulter_test.go b/internal/webhook/sparkpod_defaulter_test.go
@@ -781,14 +781,13 @@ func TestPatchSparkPod_PriorityClassName(t *testing.T) {
 			UID:  "spark-test-1",
 		},
 		Spec: v1beta2.SparkApplicationSpec{
-			BatchSchedulerOptions: &v1beta2.BatchSchedulerConfiguration{
-				PriorityClassName: &priorityClassName,
-			},
 			Driver: v1beta2.DriverSpec{
-				SparkPodSpec: v1beta2.SparkPodSpec{},
+				SparkPodSpec:      v1beta2.SparkPodSpec{},
+				PriorityClassName: &priorityClassName,
 			},
 			Executor: v1beta2.ExecutorSpec{
-				SparkPodSpec: v1beta2.SparkPodSpec{},
+				SparkPodSpec:      v1beta2.SparkPodSpec{},
+				PriorityClassName: &priorityClassName,
 			},
 		},
 	}