[SPARK-36699][Core] Reuse compatible executors for stage-level scheduling

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

@@ -110,6 +110,8 @@ private[spark] class ExecutorAllocationManager(
 
   import ExecutorAllocationManager._
 
+  private val reuseExecutors = conf.get(SCHEDULER_REUSE_COMPATIBLE_EXECUTORS)
+
   // Lower and upper bounds on the number of executors.
   private val minNumExecutors = conf.get(DYN_ALLOCATION_MIN_EXECUTORS)
   private val maxNumExecutors = conf.get(DYN_ALLOCATION_MAX_EXECUTORS)
@@ -297,11 +299,13 @@ private[spark] class ExecutorAllocationManager(
     val numRunningOrPendingTasks = pendingTask + pendingSpeculative + running
     val rp = resourceProfileManager.resourceProfileFromId(rpId)
     val tasksPerExecutor = rp.maxTasksPerExecutor(conf)
-    logDebug(s"max needed for rpId: $rpId numpending: $numRunningOrPendingTasks," +
-      s" tasksperexecutor: $tasksPerExecutor")
+
     val maxNeeded = math.ceil(numRunningOrPendingTasks * executorAllocationRatio /
       tasksPerExecutor).toInt
 
+    logDebug(s"max needed for rpId: $rpId numpending: $numRunningOrPendingTasks," +
+      s" tasksperexecutor: $tasksPerExecutor = $maxNeeded")
+
     val maxNeededWithSpeculationLocalityOffset =
       if (tasksPerExecutor > 1 && maxNeeded == 1 && pendingSpeculative > 0) {
       // If we have pending speculative tasks and only need a single executor, allocate one more
@@ -490,6 +494,11 @@ private[spark] class ExecutorAllocationManager(
     numExecutorsTargetPerResourceProfileId(rpId) - oldNumExecutorsTarget
   }
 
+  private def numExecutorsTargetsCompatibleProfiles(rpId: Int): Int = {
+    val compatibleProfileIds = resourceProfileManager.getCompatibleProfileIds(rpId)
+    compatibleProfileIds.map { numExecutorsTargetPerResourceProfileId(_) }.sum
+  }
+
   /**
    * Update the target number of executors and figure out how many to add.
    * If the cap on the number of executors is reached, give up and reset the
@@ -502,14 +511,27 @@ private[spark] class ExecutorAllocationManager(
    */
   private def addExecutors(maxNumExecutorsNeeded: Int, rpId: Int): Int = {
     val oldNumExecutorsTarget = numExecutorsTargetPerResourceProfileId(rpId)
-    // Do not request more executors if it would put our target over the upper bound
-    // this is doing a max check per ResourceProfile
-    if (oldNumExecutorsTarget >= maxNumExecutors) {
-      logDebug("Not adding executors because our current target total " +
-        s"is already ${oldNumExecutorsTarget} (limit $maxNumExecutors)")
-      numExecutorsToAddPerResourceProfileId(rpId) = 1
-      return 0
+
+    if (!reuseExecutors) {
+      // Do not request more executors if it would put our target over the upper bound
+      // this is doing a max check per ResourceProfile
+      if (oldNumExecutorsTarget >= maxNumExecutors) {
+        logDebug("Not adding executors because our current target total " +
+          s"is already ${oldNumExecutorsTarget} (limit $maxNumExecutors)")
+        numExecutorsToAddPerResourceProfileId(rpId) = 1
+        return 0
+      }
+    } else {
+      val numCompatibleExecutors = numExecutorsTargetsCompatibleProfiles(rpId)
+      if (oldNumExecutorsTarget +  numCompatibleExecutors >= maxNumExecutors) {
+        logDebug("Not adding executors because our current target total is already " +
+          s"${oldNumExecutorsTarget} (limit: max $maxNumExecutors - " +
+          s"reused $numCompatibleExecutors)")
+        numExecutorsToAddPerResourceProfileId(rpId) = 1
+        return 0
+      }
     }
+
     // There's no point in wasting time ramping up to the number of executors we already have, so
     // make sure our target is at least as much as our current allocation:
     var numExecutorsTarget = math.max(numExecutorsTargetPerResourceProfileId(rpId),
@@ -518,11 +540,25 @@ private[spark] class ExecutorAllocationManager(
     numExecutorsTarget += numExecutorsToAddPerResourceProfileId(rpId)
     // Ensure that our target doesn't exceed what we need at the present moment:
     numExecutorsTarget = math.min(numExecutorsTarget, maxNumExecutorsNeeded)
-    // Ensure that our target fits within configured bounds:
-    numExecutorsTarget = math.max(math.min(numExecutorsTarget, maxNumExecutors), minNumExecutors)
+    numExecutorsTarget = if (!reuseExecutors) {
+      // Ensure that our target fits within configured bounds:
+      math.max(math.min(numExecutorsTarget, maxNumExecutors), minNumExecutors)
+    } else {
+      // Ensure that our target fits within adjusted bounds:
+      val numCompatibleExecutors = numExecutorsTargetsCompatibleProfiles(rpId)
+      val adjustedMinNumExecutors = math.max(0, minNumExecutors - numCompatibleExecutors)
+      val adjustedMaxNumExecutors = math.max(1, maxNumExecutors - numCompatibleExecutors)
+
+      math.max(math.min(numExecutorsTarget, adjustedMaxNumExecutors),
+        adjustedMinNumExecutors)
+    }
+
     val delta = numExecutorsTarget - oldNumExecutorsTarget
     numExecutorsTargetPerResourceProfileId(rpId) = numExecutorsTarget
 
+    logDebug("new numExecutorsTargetPerResourceProfileId: " +
+      numExecutorsTargetPerResourceProfileId.mkString(", "))
+
     // If our target has not changed, do not send a message
     // to the cluster manager and reset our exponential growth
     if (delta == 0) {
@@ -712,7 +748,13 @@ private[spark] class ExecutorAllocationManager(
 
         if (!numExecutorsTargetPerResourceProfileId.contains(profId)) {
           numExecutorsTargetPerResourceProfileId.put(profId, initialNumExecutors)
-          if (initialNumExecutors > 0) {
+          val adjustedinitialNumExecutors = if (reuseExecutors) {
+            val numCompatibleExecutors = numExecutorsTargetsCompatibleProfiles(profId)
+            math.max(0, initialNumExecutors - numCompatibleExecutors)
+          } else {
+            initialNumExecutors
+          }
+          if (adjustedinitialNumExecutors > 0) {
             logDebug(s"requesting executors, rpId: $profId, initial number is $initialNumExecutors")
             // we need to trigger a schedule since we add an initial number here.
             client.requestTotalExecutors(

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -1958,6 +1958,12 @@ package object config {
       .timeConf(TimeUnit.MILLISECONDS)
       .createOptional
 
+  private[spark] val SCHEDULER_REUSE_COMPATIBLE_EXECUTORS =
+    ConfigBuilder("spark.scheduler.reuseCompatibleExecutors")
+      .version("3.3.0")
+      .booleanConf
+      .createWithDefault(false)
+
   private[spark] val SPECULATION_ENABLED =
     ConfigBuilder("spark.speculation")
       .version("0.6.0")

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -45,6 +45,7 @@ import org.apache.spark.partial.CountEvaluator
 import org.apache.spark.partial.GroupedCountEvaluator
 import org.apache.spark.partial.PartialResult
 import org.apache.spark.resource.ResourceProfile
+import org.apache.spark.resource.ResourceProfileCompatiblePolicy.ResourceProfileCompatiblePolicy
 import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 import org.apache.spark.util.{BoundedPriorityQueue, Utils}
 import org.apache.spark.util.collection.{ExternalAppendOnlyMap, OpenHashMap,
@@ -1809,6 +1810,22 @@ abstract class RDD[T: ClassTag](
     this
   }
 
+  /**
+   * Specify a ResourceProfile and reuse existing compatible executors to use when calculating
+   * this RDD.
+   * @param reuseResourceNames specify what resource should be checked when reusing executors
+   * @param reusePolicy specify executor reuse policy
+   */
+  @Experimental
+  @Since("3.3.0")
+  def withResources(rp: ResourceProfile,
+    reuseResourceNames: Set[String], reusePolicy: ResourceProfileCompatiblePolicy): this.type = {
+    resourceProfile = Option(rp)
+    sc.resourceProfileManager.updateResourceReusePolicy(reuseResourceNames, reusePolicy)
+    sc.resourceProfileManager.addResourceProfile(resourceProfile.get)
+    this
+  }
+
   /**
    * Get the ResourceProfile specified with this RDD or null if it wasn't specified.
    * @return the user specified ResourceProfile or null (for Java compatibility) if

core/src/main/scala/org/apache/spark/resource/ResourceProfile.scala

@@ -241,6 +241,13 @@ class ResourceProfile(
     rp.taskResources == taskResources && rp.executorResources == executorResources
   }
 
+  private[spark] def resourcesCompatible(rp: ResourceProfile,
+    resourceProfileCompatibleFunc: (ResourceProfile, ResourceProfile) => Boolean
+    = ResourceProfileCompatiblePolicy.resourceProfileCompatibleWithEqualCores)
+  : Boolean = {
+    resourceProfileCompatibleFunc(this, rp)
+  }
+
   override def hashCode(): Int = Seq(taskResources, executorResources).hashCode()
 
   override def toString(): String = {
@@ -249,6 +256,67 @@ class ResourceProfile(
   }
 }
 
+/**
+ * Resource profile compatibility based on user policy
+ *
+ * EQUAL_RESOURCES:
+ *     only reuse executors with same resources (including cores, memory and all 3rd party
+ *     resources)
+ * MORE_RESOURCES:
+ *     reuse executors with equal or more resources (including cores, memory and all 3rd party
+ *     resources)
+ * Notes:
+ *     if EQUAL_RESOURCES is specified, there is no resource waste but less
+ *     user flexibility.
+ *     if MORE_RESOURCES is specified, users should know there are some resources
+ *     wasted. They need to tradeoff between reusing executors and creating new ones.
+ */
+object ResourceProfileCompatiblePolicy extends Enumeration {
+
+  type ResourceProfileCompatiblePolicy = Value
+
+  val EQUAL_RESOURCES, MORE_RESOURCES = Value
+
+  def resourceProfileCompatibleWithPolicy(
+      prevRP: ResourceProfile, curRP: ResourceProfile,
+      resourceNames: Set[String], policy: ResourceProfileCompatiblePolicy): Boolean = {
+    resourceNames.forall { resourceName: String =>
+      policy match {
+        case EQUAL_RESOURCES =>
+          ! prevRP.executorResources.get(resourceName).isEmpty &&
+          ! curRP.executorResources.get(resourceName).isEmpty &&
+          prevRP.executorResources(resourceName).amount ==
+            curRP.executorResources(resourceName).amount
+        case MORE_RESOURCES =>
+          ! prevRP.executorResources.get(resourceName).isEmpty &&
+          ! curRP.executorResources.get(resourceName).isEmpty &&
+          prevRP.executorResources(resourceName).amount >=
+            curRP.executorResources(resourceName).amount
+      }
+    }
+  }
+
+  def resourceProfileCompatibleWithEqualCores(prevRP: ResourceProfile,
+                                              curRP: ResourceProfile): Boolean = {
+    resourceProfileCompatibleWithPolicy(prevRP, curRP, Set("cores"), EQUAL_RESOURCES)
+  }
+
+  def resourceProfileCompatibleWithMoreCores(prevRP: ResourceProfile,
+                                             curRP: ResourceProfile): Boolean = {
+    resourceProfileCompatibleWithPolicy(prevRP, curRP, Set("cores"), MORE_RESOURCES)
+  }
+
+  def resourceProfileCompatibleWithEqualResources(
+    prevRP: ResourceProfile, curRP: ResourceProfile, resourceNames: Set[String]): Boolean = {
+    resourceProfileCompatibleWithPolicy(prevRP, curRP, resourceNames, EQUAL_RESOURCES)
+  }
+
+  def resourceProfileCompatibleWithMoreResources(
+      prevRP: ResourceProfile, curRP: ResourceProfile, resourceNames: Set[String]): Boolean = {
+    resourceProfileCompatibleWithPolicy(prevRP, curRP, resourceNames, MORE_RESOURCES)
+  }
+}
+
 object ResourceProfile extends Logging {
   // task resources
   /**
@@ -373,6 +441,11 @@ object ResourceProfile extends Logging {
     }
   }
 
+  private[spark] def hasCustomExecutorResources(rp: ResourceProfile): Boolean = {
+    rp.executorResources.keys.exists(
+      k => !ResourceProfile.allSupportedExecutorResources.contains(k))
+  }
+
   private[spark] def getCustomTaskResources(
       rp: ResourceProfile): Map[String, TaskResourceRequest] = {
     rp.taskResources.filterKeys(k => !k.equals(ResourceProfile.CPUS)).toMap

core/src/main/scala/org/apache/spark/resource/ResourceProfileManager.scala

@@ -24,7 +24,9 @@ import scala.collection.mutable.HashMap
 import org.apache.spark.{SparkConf, SparkException}
 import org.apache.spark.annotation.Evolving
 import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config.SCHEDULER_REUSE_COMPATIBLE_EXECUTORS
 import org.apache.spark.internal.config.Tests._
+import org.apache.spark.resource.ResourceProfileCompatiblePolicy.{EQUAL_RESOURCES, ResourceProfileCompatiblePolicy}
 import org.apache.spark.scheduler.{LiveListenerBus, SparkListenerResourceProfileAdded}
 import org.apache.spark.util.Utils
 import org.apache.spark.util.Utils.isTesting
@@ -40,6 +42,9 @@ private[spark] class ResourceProfileManager(sparkConf: SparkConf,
     listenerBus: LiveListenerBus) extends Logging {
   private val resourceProfileIdToResourceProfile = new HashMap[Int, ResourceProfile]()
 
+  private val reuseExecutors = sparkConf.get(SCHEDULER_REUSE_COMPATIBLE_EXECUTORS)
+  private val resourceProfileIdToCompatibleResourceProfileIds = new HashMap[Int, Set[Int]]()
+
   private val (readLock, writeLock) = {
     val lock = new ReentrantReadWriteLock()
     (lock.readLock(), lock.writeLock())
@@ -50,6 +55,9 @@ private[spark] class ResourceProfileManager(sparkConf: SparkConf,
 
   def defaultResourceProfile: ResourceProfile = defaultProfile
 
+  var reuseResourceNames: Set[String] = Set("cores")
+  var reusePolicy = EQUAL_RESOURCES
+
   private val dynamicEnabled = Utils.isDynamicAllocationEnabled(sparkConf)
   private val master = sparkConf.getOption("spark.master")
   private val isYarn = master.isDefined && master.get.equals("yarn")
@@ -77,13 +85,27 @@ private[spark] class ResourceProfileManager(sparkConf: SparkConf,
     true
   }
 
+  def updateResourceReusePolicy(resourceNames: Set[String],
+    policy: ResourceProfileCompatiblePolicy): Unit = {
+    writeLock.lock()
+    try {
+      reuseResourceNames = resourceNames
+      reusePolicy = policy
+    } finally {
+      writeLock.unlock()
+    }
+  }
+
   def addResourceProfile(rp: ResourceProfile): Unit = {
     isSupported(rp)
     var putNewProfile = false
     writeLock.lock()
     try {
       if (!resourceProfileIdToResourceProfile.contains(rp.id)) {
         val prev = resourceProfileIdToResourceProfile.put(rp.id, rp)
+        if (reuseExecutors) {
+          calculateCompatibleProfileIds()
+        }
         if (prev.isEmpty) putNewProfile = true
       }
     } finally {
@@ -94,7 +116,8 @@ private[spark] class ResourceProfileManager(sparkConf: SparkConf,
       // force the computation of maxTasks and limitingResource now so we don't have cost later
       rp.limitingResource(sparkConf)
       logInfo(s"Added ResourceProfile id: ${rp.id}")
-      listenerBus.post(SparkListenerResourceProfileAdded(rp))
+      listenerBus.post(SparkListenerResourceProfileAdded(rp,
+        resourceProfileIdToCompatibleResourceProfileIds.get(rp.id)))
     }
   }
 
@@ -127,4 +150,48 @@ private[spark] class ResourceProfileManager(sparkConf: SparkConf,
       readLock.unlock()
     }
   }
+
+  private [spark] def getCompatibleProfileIds(rpId: Int): Set[Int] = {
+    readLock.lock()
+    try {
+      resourceProfileIdToCompatibleResourceProfileIds.getOrElse(rpId, Set())
+    } finally {
+      readLock.unlock()
+    }
+  }
+
+  /*
+   * Calculate and cache all compatible profile IDs excluding itself
+   * This operation is done when new ResourceProfile added
+   */
+  private def calculateCompatibleProfileIds(): Unit = {
+    resourceProfileIdToResourceProfile.foreach {
+      case (id: Int, rpEntry: ResourceProfile) =>
+        val resourceProfile = resourceProfileFromId(id)
+        val compatibleIdSet = resourceProfileIdToResourceProfile.filter {
+          case (otherId: Int, _: ResourceProfile) =>
+            id != otherId && rpEntry.resourcesCompatible(resourceProfile)
+        }.map(_._1).toSet
+        resourceProfileIdToCompatibleResourceProfileIds.put(id, compatibleIdSet)
+    }
+  }
+
+  private [spark] def dumpResourceProfiles(): Unit = {
+    readLock.lock()
+    try {
+      logDebug("dump resource profiles:")
+      resourceProfileIdToResourceProfile.foreach { case (id: Int, rpEntry: ResourceProfile) =>
+        logDebug("-- rpId " + id + ": " + rpEntry)
+      }
+      if (reuseExecutors) {
+        logDebug("dump compatible resource profile ids:")
+        resourceProfileIdToCompatibleResourceProfileIds.foreach {
+          case (id: Int, compatibleIds: Set[Int]) =>
+            logDebug("-- rpId " + id + ": " + compatibleIds)
+        }
+      }
+    } finally {
+      readLock.unlock()
+    }
+  }
 }

core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala

@@ -283,8 +283,9 @@ case class SparkListenerApplicationEnd(time: Long) extends SparkListenerEvent
 case class SparkListenerLogStart(sparkVersion: String) extends SparkListenerEvent
 
 @DeveloperApi
-@Since("3.1.0")
-case class SparkListenerResourceProfileAdded(resourceProfile: ResourceProfile)
+@Since("3.3.0")
+case class SparkListenerResourceProfileAdded(resourceProfile: ResourceProfile,
+  compatibleResourceProfileIds: Option[Set[Int]])
   extends SparkListenerEvent
 
 /**