[kyuubi] branch master updated: [KYUUBI #4636] Improve eagerly kill redundant executors

[ul...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ulyssesyou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new 97aedf504 [KYUUBI #4636] Improve eagerly kill redundant executors
97aedf504 is described below

commit 97aedf50488b86187fb14ae92c3d380ddc2c3749
Author: ulysses-you <ul...@gmail.com>
AuthorDate: Thu Mar 30 15:39:44 2023 +0800

    [KYUUBI #4636] Improve eagerly kill redundant executors
    
    ### _Why are the changes needed?_
    
    This pr improves the behavoir of kill redundant executors.
    1. support kill executors even if AQE can not optimize shuffle read. e.g., people call `.repartition(2)`
    2. fix a issue that avoid always kill executors which holds shuffle data
    
    ### _How was this patch tested?_
    test manually
    
    Closes #4636 from ulysses-you/kill-executors.
    
    Closes #4636
    
    19ac808d3 [ulysses-you] Improve eagerly kill redundant executors
    
    Authored-by: ulysses-you <ul...@gmail.com>
    Signed-off-by: ulyssesyou <ul...@apache.org>
---
 .../spark/sql/FinalStageResourceManager.scala      | 59 ++++++++++++----------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala
index f568f9df8..2bf7ae6b7 100644
--- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala
+++ b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/FinalStageResourceManager.scala
@@ -64,41 +64,45 @@ case class FinalStageResourceManager(session: SparkSession)
       return plan
     }
 
-    val stage = findFinalRebalanceStage(plan)
-    if (stage.isEmpty) {
+    val stageOpt = findFinalRebalanceStage(plan)
+    if (stageOpt.isEmpty) {
       return plan
     }
 
     // Since we are in `prepareQueryStage`, the AQE shuffle read has not been applied.
     // So we need to apply it by self.
-    val shuffleRead = queryStageOptimizerRules.foldLeft(stage.get.asInstanceOf[SparkPlan]) {
+    val shuffleRead = queryStageOptimizerRules.foldLeft(stageOpt.get.asInstanceOf[SparkPlan]) {
       case (latest, rule) => rule.apply(latest)
     }
-    shuffleRead match {
+    val (targetCores, stage) = shuffleRead match {
       case AQEShuffleReadExec(stage: ShuffleQueryStageExec, partitionSpecs) =>
-        // The condition whether inject custom resource profile:
-        // - target executors < active executors
-        // - active executors - target executors > min executors
-        val numActiveExecutors = sc.getExecutorIds().length
-        val targetCores = partitionSpecs.length
-        val targetExecutors = (math.ceil(targetCores.toFloat / coresPerExecutor) * factor).toInt
-          .max(1)
-        val hasBenefits = targetExecutors < numActiveExecutors &&
-          (numActiveExecutors - targetExecutors) > minExecutors
-        if (hasBenefits) {
-          val shuffleId = stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleDependency.shuffleId
-          val numReduce = stage.plan.asInstanceOf[ShuffleExchangeExec].numPartitions
-          // Now, there is only a final rebalance stage waiting to execute and all tasks of previous
-          // stage are finished. Kill redundant existed executors eagerly so the tasks of final
-          // stage can be centralized scheduled.
-          killExecutors(sc, targetExecutors, shuffleId, numReduce)
-        } else {
-          logInfo(s"Has no benefits to kill executors or inject custom resource profile, " +
-            s"active executors: $numActiveExecutors, min executor: $minExecutors, " +
-            s"target executors: $targetExecutors.")
-        }
-
+        (partitionSpecs.length, stage)
+      case stage: ShuffleQueryStageExec =>
+        // we can still kill executors if no AQE shuffle read, e.g., `.repartition(2)`
+        (stage.shuffle.numPartitions, stage)
       case _ =>
+        // it should never happen in current Spark, but to be safe do nothing if happens
+        logWarning("BUG, Please report to Apache Kyuubi community")
+        return plan
+    }
+    // The condition whether inject custom resource profile:
+    // - target executors < active executors
+    // - active executors - target executors > min executors
+    val numActiveExecutors = sc.getExecutorIds().length
+    val targetExecutors = (math.ceil(targetCores.toFloat / coresPerExecutor) * factor).toInt
+      .max(1)
+    val hasBenefits = targetExecutors < numActiveExecutors &&
+      (numActiveExecutors - targetExecutors) > minExecutors
+    logInfo(s"The snapshot of current executors view, " +
+      s"active executors: $numActiveExecutors, min executor: $minExecutors, " +
+      s"target executors: $targetExecutors, has benefits: $hasBenefits")
+    if (hasBenefits) {
+      val shuffleId = stage.plan.asInstanceOf[ShuffleExchangeExec].shuffleDependency.shuffleId
+      val numReduce = stage.plan.asInstanceOf[ShuffleExchangeExec].numPartitions
+      // Now, there is only a final rebalance stage waiting to execute and all tasks of previous
+      // stage are finished. Kill redundant existed executors eagerly so the tasks of final
+      // stage can be centralized scheduled.
+      killExecutors(sc, targetExecutors, shuffleId, numReduce)
     }
 
     plan
@@ -144,9 +148,8 @@ case class FinalStageResourceManager(session: SparkSession)
     // shuffle status.
     // We should evict executors by their alive time first and retain all of executors which
     // have better locality for shuffle block.
-    val numExecutorToKillWithNoShuffle = expectedNumExecutorToKill - executorToBlockSize.size
     executorsWithRegistrationTs.toSeq.sortBy(_._2).foreach { case (id, _) =>
-      if (executorIdsToKill.length < numExecutorToKillWithNoShuffle &&
+      if (executorIdsToKill.length < expectedNumExecutorToKill &&
         !executorToBlockSize.contains(id)) {
         executorIdsToKill.append(id)
       }