CR feedback, move adjustExecutors to a common utility function

Author: holdenk

core/src/main/scala/org/apache/spark/ExecutorAllocationClient.scala

@@ -93,7 +93,7 @@ private[spark] trait ExecutorAllocationClient {
    * @return the ids of the executors acknowledged by the cluster manager to be removed.
    */
   def decommissionExecutors(
-    executorsAndDecomInfo: Seq[(String, ExecutorDecommissionInfo)],
+    executorsAndDecomInfo: Array[(String, ExecutorDecommissionInfo)],
     adjustTargetNumExecutors: Boolean): Seq[String] = {
     killExecutors(executorsAndDecomInfo.map(_._1),
       adjustTargetNumExecutors,
@@ -103,19 +103,18 @@ private[spark] trait ExecutorAllocationClient {
 
   /**
    * Request that the cluster manager decommission the specified executor.
-   * Default implementation delegates to decommissionExecutors, scheduler can override
-   * if it supports graceful decommissioning.
+   * Delegates to decommissionExecutors.
    *
    * @param executorId identifiers of executor to decommission
    * @param decommissionInfo information about the decommission (reason, host loss)
    * @param adjustTargetNumExecutors if we should adjust the target number of executors.
    * @return whether the request is acknowledged by the cluster manager.
    */
-  def decommissionExecutor(executorId: String,
+  final def decommissionExecutor(executorId: String,
       decommissionInfo: ExecutorDecommissionInfo,
       adjustTargetNumExecutors: Boolean): Boolean = {
     val decommissionedExecutors = decommissionExecutors(
-      Seq((executorId, decommissionInfo)),
+      Array((executorId, decommissionInfo)),
       adjustTargetNumExecutors = adjustTargetNumExecutors)
     decommissionedExecutors.nonEmpty && decommissionedExecutors(0).equals(executorId)
   }

core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala

@@ -575,7 +575,7 @@ private[spark] class ExecutorAllocationManager(
       // when the task backlog decreased.
       if (conf.get(WORKER_DECOMMISSION_ENABLED)) {
         val executorIdsWithoutHostLoss = executorIdsToBeRemoved.toSeq.map(
-          id => (id, ExecutorDecommissionInfo("spark scale down", false)))
+          id => (id, ExecutorDecommissionInfo("spark scale down", false))).toArray
         client.decommissionExecutors(executorIdsWithoutHostLoss, adjustTargetNumExecutors = false)
       } else {
         client.killExecutors(executorIdsToBeRemoved.toSeq, adjustTargetNumExecutors = false,

core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala

@@ -513,7 +513,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
    * @return the ids of the executors acknowledged by the cluster manager to be removed.
    */
   override def decommissionExecutors(
-      executorsAndDecomInfo: Seq[(String, ExecutorDecommissionInfo)],
+      executorsAndDecomInfo: Array[(String, ExecutorDecommissionInfo)],
       adjustTargetNumExecutors: Boolean): Seq[String] = {
 
     val executorsToDecommission = executorsAndDecomInfo.filter { case (executorId, _) =>
@@ -530,21 +530,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
 
     // If we don't want to replace the executors we are decommissioning
     if (adjustTargetNumExecutors) {
-      executorsToDecommission.foreach { case (exec, _) =>
-        val rpId = withLock {
-          executorDataMap(exec).resourceProfileId
-        }
-        val rp = scheduler.sc.resourceProfileManager.resourceProfileFromId(rpId)
-        if (requestedTotalExecutorsPerResourceProfile.isEmpty) {
-          // Assume that we are killing an executor that was started by default and
-          // not through the request api
-          requestedTotalExecutorsPerResourceProfile(rp) = 0
-        } else {
-          val requestedTotalForRp = requestedTotalExecutorsPerResourceProfile(rp)
-          requestedTotalExecutorsPerResourceProfile(rp) = math.max(requestedTotalForRp - 1, 0)
-        }
-      }
-      doRequestTotalExecutors(requestedTotalExecutorsPerResourceProfile.toMap)
+      adjustExecutors(executorsToDecommission.map(_._1))
     }
 
     val decommissioned = executorsToDecommission.filter { case (executorId, decomInfo) =>
@@ -846,6 +832,27 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
       resourceProfileToTotalExecs: Map[ResourceProfile, Int]): Future[Boolean] =
     Future.successful(false)
 
+  /**
+   * Adjust the number of executors being requested to no longer include the provided executors.
+   */
+  private def adjustExecutors(executorIds: Seq[String]) = {
+      executorIds.foreach { exec =>
+        withLock {
+          val rpId = executorDataMap(exec).resourceProfileId
+          val rp = scheduler.sc.resourceProfileManager.resourceProfileFromId(rpId)
+          if (requestedTotalExecutorsPerResourceProfile.isEmpty) {
+            // Assume that we are killing an executor that was started by default and
+            // not through the request api
+            requestedTotalExecutorsPerResourceProfile(rp) = 0
+          } else {
+            val requestedTotalForRp = requestedTotalExecutorsPerResourceProfile(rp)
+            requestedTotalExecutorsPerResourceProfile(rp) = math.max(requestedTotalForRp - 1, 0)
+          }
+        }
+      }
+      doRequestTotalExecutors(requestedTotalExecutorsPerResourceProfile.toMap)
+  }
+
   /**
    * Request that the cluster manager kill the specified executors.
    *
@@ -884,19 +891,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
       // take into account executors that are pending to be added or removed.
       val adjustTotalExecutors =
         if (adjustTargetNumExecutors) {
-          executorsToKill.foreach { exec =>
-            val rpId = executorDataMap(exec).resourceProfileId
-            val rp = scheduler.sc.resourceProfileManager.resourceProfileFromId(rpId)
-            if (requestedTotalExecutorsPerResourceProfile.isEmpty) {
-              // Assume that we are killing an executor that was started by default and
-              // not through the request api
-              requestedTotalExecutorsPerResourceProfile(rp) = 0
-            } else {
-              val requestedTotalForRp = requestedTotalExecutorsPerResourceProfile(rp)
-              requestedTotalExecutorsPerResourceProfile(rp) = math.max(requestedTotalForRp - 1, 0)
-            }
-          }
-          doRequestTotalExecutors(requestedTotalExecutorsPerResourceProfile.toMap)
+          adjustExecutors(executorsToKill)
         } else {
           Future.successful(true)
         }

core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala

@@ -177,7 +177,7 @@ private[spark] class StandaloneSchedulerBackend(
   override def executorDecommissioned(fullId: String, decommissionInfo: ExecutorDecommissionInfo) {
     logInfo("Asked to decommission executor")
     val execId = fullId.split("/")(1)
-    decommissionExecutors(Seq((execId, decommissionInfo)), adjustTargetNumExecutors = false)
+    decommissionExecutors(Array((execId, decommissionInfo)), adjustTargetNumExecutors = false)
     logInfo("Executor %s decommissioned: %s".format(fullId, decommissionInfo))
   }
 

core/src/test/scala/org/apache/spark/scheduler/WorkerDecommissionSuite.scala

@@ -77,7 +77,7 @@ class WorkerDecommissionSuite extends SparkFunSuite with LocalSparkContext {
     val sched = sc.schedulerBackend.asInstanceOf[StandaloneSchedulerBackend]
     val execs = sched.getExecutorIds()
     // Make the executors decommission, finish, exit, and not be replaced.
-    val execsAndDecomInfo = execs.map((_, ExecutorDecommissionInfo("", false)))
+    val execsAndDecomInfo = execs.map((_, ExecutorDecommissionInfo("", false))).toArray
     sched.decommissionExecutors(execsAndDecomInfo, adjustTargetNumExecutors = true)
     val asyncCountResult = ThreadUtils.awaitResult(asyncCount, 20.seconds)
     assert(asyncCountResult === 10)