This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 001165d0fd57 [SPARK-57467][SCHEDULER] Reuse identical resource
profiles when available
001165d0fd57 is described below
commit 001165d0fd5799384316fd8f65778fd5a8e27f83
Author: PA Savalle <[email protected]>
AuthorDate: Sat Aug 29 09:25:55 2026 -0700
[SPARK-57467][SCHEDULER] Reuse identical resource profiles when available
### What changes were proposed in this pull request?
When registering a resource profile, check if one with the same requests
already exists. If so, return it, so that the same executors can be reused.
`ResourceProfileManager` has a method `getEquivalentProfile`, but it is
only called in `DAGScheduler.mergeResourceProfilesForStage` when
`stageResourceProfiles.size > 1`. Even there, it seems it may be prone to a
race condition since it does not acquire the lock for both steps (checking for
a profile, adding one if needed).
### Why are the changes needed?
Currently, every time a profile is registered, new executors will need to
spin up -- even if the resource requests are the same. This is especially
problematic with Spark Connect, where different isolated sessions may try and
register the same profile, but all end up with isolated executors.
We could also consider putting this behavior behind a configuration flag,
if it is desirable to retain the previous behavior by default.
### Does this PR introduce _any_ user-facing change?
Previously, equivalent resource profiles would get different profile IDs,
and separate executors.
With this change, equivalent resource profiles will get the same ID, and
share executors.
### How was this patch tested?
Added tests
### Was this patch authored or co-authored using generative AI tooling?
Yes, modified initial changes provided by Claude.
`Generated-by: 2.1.159 (Claude Code)`
Closes #56516 from psavalle/pa/resource-profile-get-or-add.
Authored-by: PA Savalle <[email protected]>
Signed-off-by: Chao Sun <[email protected]>
(cherry picked from commit a6e760750bca8dce0a2ee5814437d6574dc2ca36)
Signed-off-by: Chao Sun <[email protected]>
---
core/src/main/scala/org/apache/spark/rdd/RDD.scala | 5 +-
.../spark/resource/ResourceProfileManager.scala | 115 ++++++++++++-----
.../org/apache/spark/scheduler/DAGScheduler.scala | 21 +--
.../resource/ResourceProfileManagerSuite.scala | 143 +++++++++++++++++----
.../apache/spark/scheduler/DAGSchedulerSuite.scala | 21 +++
.../sql/connect/planner/SparkConnectPlanner.scala | 9 +-
6 files changed, 239 insertions(+), 75 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index e4cbbe6302af..8b3ecb943518 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1934,8 +1934,9 @@ abstract class RDD[T: ClassTag](
@Experimental
@Since("3.1.0")
def withResources(rp: ResourceProfile): this.type = {
- resourceProfile = Option(rp)
- sc.resourceProfileManager.addResourceProfile(resourceProfile.get)
+ // Reuse an already-registered profile with equal resources if one exists
so that equivalent
+ // profiles share a single id, allowing executors to be reused instead of
newly allocated.
+ resourceProfile =
Option(sc.resourceProfileManager.getOrAddEquivalentProfile(rp))
this
}
diff --git
a/core/src/main/scala/org/apache/spark/resource/ResourceProfileManager.scala
b/core/src/main/scala/org/apache/spark/resource/ResourceProfileManager.scala
index fed5129bcf9d..b4f1a29d7c14 100644
--- a/core/src/main/scala/org/apache/spark/resource/ResourceProfileManager.scala
+++ b/core/src/main/scala/org/apache/spark/resource/ResourceProfileManager.scala
@@ -125,41 +125,93 @@ private[spark] class ResourceProfileManager(sparkConf:
SparkConf,
taskRpId == executorRpId || (!dynamicEnabled &&
taskRp.isInstanceOf[TaskResourceProfile])
}
+ /**
+ * Register the given ResourceProfile unless its id is already registered.
Unlike
+ * [[getOrAddEquivalentProfile]], this method does not reuse a profile with
equal resources.
+ */
def addResourceProfile(rp: ResourceProfile): Unit = {
isSupported(rp)
- // Validate before inserting so a malformed profile never enters the
registry, where it
- // would stay visible to the whole application. The cpus amount is checked
under the map
- // key -- the identity scheduling uses -- with the same rule as the
request entry points:
- // construction stays lenient (deserialization of persisted data must
accept anything an
- // earlier release wrote), so registration is the enforcement point for
every live path,
- // including raw profile construction and Spark Connect.
- rp.taskResources.get(ResourceProfile.CPUS).foreach { treq =>
- require(!treq.amount.isNaN && !treq.amount.isInfinity &&
-
CpuAmount.isInRange(CpuAmount.normalize(BigDecimal(treq.amount.toString))),
- s"The cpus amount ${treq.amount} must be at least 1e-9 and at most
${Int.MaxValue}.")
+ registerProfileIfAbsent(rp, findProfileById)
+ }
+
+ /**
+ * Get the registered ResourceProfile whose resources are equal to the given
one, registering
+ * the given profile first if no equivalent one exists yet.
+ */
+ def getOrAddEquivalentProfile(rp: ResourceProfile): ResourceProfile = {
+ isSupported(rp)
+ registerProfileIfAbsent(rp, findEquivalentOrSameIdProfile)
+ }
+
+ private def registerProfileIfAbsent(
+ rp: ResourceProfile,
+ findRegistered: ResourceProfile => Option[ResourceProfile]):
ResourceProfile = {
+ val existingProfile = {
+ readLock.lock()
+ try {
+ findRegistered(rp)
+ } finally {
+ readLock.unlock()
+ }
}
- // Force the computation of maxTasks and limitingResource now so we don't
have cost later;
- // doing it before the insert also surfaces any other malformed shape
(e.g. a task
- // resource without a matching executor resource) before registration
instead of after.
- // The result is cached in the profile, so re-adding an existing profile
stays cheap.
- rp.limitingResource(sparkConf)
- var putNewProfile = false
- writeLock.lock()
- try {
- if (!resourceProfileIdToResourceProfile.contains(rp.id)) {
- val prev = resourceProfileIdToResourceProfile.put(rp.id, rp)
- if (prev.isEmpty) putNewProfile = true
+ existingProfile.getOrElse {
+ // Validate before inserting so a malformed profile never enters the
registry, where it
+ // would stay visible to the whole application. The cpus amount is
checked under the map
+ // key -- the identity scheduling uses -- with the same rule as the
request entry points:
+ // construction stays lenient (deserialization of persisted data must
accept anything an
+ // earlier release wrote), so registration is the enforcement point for
every live path,
+ // including raw profile construction and Spark Connect.
+ rp.taskResources.get(ResourceProfile.CPUS).foreach { treq =>
+ require(!treq.amount.isNaN && !treq.amount.isInfinity &&
+
CpuAmount.isInRange(CpuAmount.normalize(BigDecimal(treq.amount.toString))),
+ s"The cpus amount ${treq.amount} must be at least 1e-9 and at most
${Int.MaxValue}.")
}
- } finally {
- writeLock.unlock()
+ // Force the computation of maxTasks and limitingResource now so we
don't have cost later;
+ // doing it before the insert also surfaces any other malformed shape
(e.g. a task
+ // resource without a matching executor resource) before registration
instead of after.
+ rp.limitingResource(sparkConf)
+ var addedProfile: Option[ResourceProfile] = None
+ val resolvedProfile = {
+ writeLock.lock()
+ try {
+ // Another thread may have added this profile after the read-lock
fast path, so check
+ // again while holding the write lock.
+ findRegistered(rp).getOrElse {
+ resourceProfileIdToResourceProfile.put(rp.id, rp)
+ addedProfile = Some(rp)
+ rp
+ }
+ } finally {
+ writeLock.unlock()
+ }
+ }
+ // do this outside the write lock only when we add a new profile
+ addedProfile.foreach(onProfileAdded)
+ resolvedProfile
}
- // do this outside the write lock only when we add a new profile
- if (putNewProfile) {
- logInfo(log"Added ResourceProfile id: ${MDC(LogKeys.RESOURCE_PROFILE_ID,
rp.id)}")
- listenerBus.post(SparkListenerResourceProfileAdded(rp))
+ }
+
+ private def findProfileById(rp: ResourceProfile): Option[ResourceProfile] = {
+ resourceProfileIdToResourceProfile.get(rp.id)
+ }
+
+ private def findEquivalentOrSameIdProfile(rp: ResourceProfile):
Option[ResourceProfile] = {
+
resourceProfileIdToResourceProfile.get(rp.id).orElse(findEquivalentProfile(rp))
+ }
+
+ private def findEquivalentProfile(rp: ResourceProfile):
Option[ResourceProfile] = {
+ resourceProfileIdToResourceProfile.values.find { existing =>
+ val sameDefaultStatus =
+ (existing eq defaultProfile) == (rp eq defaultProfile)
+ sameDefaultStatus && existing.resourcesEqual(rp)
}
}
+ private def onProfileAdded(rp: ResourceProfile): Unit = {
+ logInfo(log"Added ResourceProfile id: ${MDC(LogKeys.RESOURCE_PROFILE_ID,
rp.id)}")
+ listenerBus.post(SparkListenerResourceProfileAdded(rp))
+ }
+
/*
* Gets the ResourceProfile associated with the id, if a profile doesn't
exist
* it returns the default ResourceProfile created from the application level
configs.
@@ -176,15 +228,14 @@ private[spark] class ResourceProfileManager(sparkConf:
SparkConf,
}
/*
- * If the ResourceProfile passed in is equivalent to an existing one, return
the
- * existing one, other return None
+ * If the ResourceProfile passed in is equivalent to an existing one, return
the existing one.
+ * The actual default profile is kept distinct because cluster managers give
its id special
+ * semantics that do not apply to explicit profiles with otherwise equal
resources.
*/
def getEquivalentProfile(rp: ResourceProfile): Option[ResourceProfile] = {
readLock.lock()
try {
- resourceProfileIdToResourceProfile.find { case (_, rpEntry) =>
- rpEntry.resourcesEqual(rp)
- }.map(_._2)
+ findEquivalentProfile(rp)
} finally {
readLock.unlock()
}
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index fa2ab84c64e7..67b7823091e2 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -889,16 +889,17 @@ private[spark] class DAGScheduler(
val startResourceProfile = stageResourceProfiles.head
val mergedProfile = stageResourceProfiles.drop(1)
.foldLeft(startResourceProfile)((a, b) => mergeResourceProfiles(a,
b))
- // compared merged profile with existing ones so we don't add it over
and over again
- // if the user runs the same operation multiple times
- val resProfile =
sc.resourceProfileManager.getEquivalentProfile(mergedProfile)
- resProfile match {
- case Some(existingRp) => existingRp
- case None =>
- // this ResourceProfile could be different if it was merged so we
have to add it to
- // our ResourceProfileManager
- sc.resourceProfileManager.addResourceProfile(mergedProfile)
- mergedProfile
+ val defaultProfile = sc.resourceProfileManager.defaultResourceProfile
+ if (stageResourceProfiles.exists(_ eq defaultProfile) &&
+ defaultProfile.resourcesEqual(mergedProfile)) {
+ // The default profile id has special meaning to cluster managers.
Preserve it when the
+ // actual default was an input and the merge did not add any
requirements.
+ defaultProfile
+ } else {
+ // Compare the merged profile with existing ones so we don't add it
over and over again
+ // if the user runs the same operation multiple times. The merged
ResourceProfile could
+ // be different from any existing one, in which case it is
registered here.
+ sc.resourceProfileManager.getOrAddEquivalentProfile(mergedProfile)
}
} else {
throw new IllegalArgumentException("Multiple ResourceProfiles
specified in the RDDs for " +
diff --git
a/core/src/test/scala/org/apache/spark/resource/ResourceProfileManagerSuite.scala
b/core/src/test/scala/org/apache/spark/resource/ResourceProfileManagerSuite.scala
index ce57f053a3e7..799e04f2c939 100644
---
a/core/src/test/scala/org/apache/spark/resource/ResourceProfileManagerSuite.scala
+++
b/core/src/test/scala/org/apache/spark/resource/ResourceProfileManagerSuite.scala
@@ -17,12 +17,17 @@
package org.apache.spark.resource
+import org.mockito.ArgumentMatchers.isA
+import org.mockito.Mockito.{never, reset, times, verify}
+import org.scalatestplus.mockito.MockitoSugar
+
import org.apache.spark.{SparkConf, SparkException, SparkFunSuite}
import org.apache.spark.internal.config._
import org.apache.spark.internal.config.Tests._
-import org.apache.spark.scheduler.LiveListenerBus
+import org.apache.spark.scheduler.{LiveListenerBus,
SparkListenerResourceProfileAdded}
+import org.apache.spark.util.ThreadUtils
-class ResourceProfileManagerSuite extends SparkFunSuite {
+class ResourceProfileManagerSuite extends SparkFunSuite with MockitoSugar {
override def beforeAll(): Unit = {
try {
@@ -251,36 +256,118 @@ class ResourceProfileManagerSuite extends SparkFunSuite {
" with dynamic allocation"))
}
- test("ResourceProfileManager has equivalent profile") {
+ test("getOrAddEquivalentProfile reuses an equivalent profile") {
val conf = new SparkConf().set(EXECUTOR_CORES, 4)
- val rpmanager = new ResourceProfileManager(conf, listenerBus)
- var rpAlreadyExist: Option[ResourceProfile] = None
- val checkId = 500
- for (i <- 1 to 1000) {
+ val mockListenerBus = mock[LiveListenerBus]
+ val rpmanager = new ResourceProfileManager(conf, mockListenerBus)
+ reset(mockListenerBus)
+
+ def buildProfile(cores: Int): ResourceProfile = {
val rprofBuilder = new ResourceProfileBuilder()
val ereqs = new ExecutorResourceRequests()
- ereqs.cores(i).memory("4g").memoryOverhead("2000m")
+ ereqs.cores(cores).memory("4g").memoryOverhead("2000m")
val treqs = new TaskResourceRequests()
- treqs.cpus(i)
- rprofBuilder.require(ereqs).require(treqs)
- val rprof = rprofBuilder.build()
- rpmanager.addResourceProfile(rprof)
- if (i == checkId) rpAlreadyExist = Some(rprof)
+ treqs.cpus(1)
+ rprofBuilder.require(ereqs).require(treqs).build()
}
- val rpNotMatch = new ResourceProfileBuilder().build()
- assert(rpmanager.getEquivalentProfile(rpNotMatch).isEmpty,
- s"resourceProfile should not have existed")
-
- val rprofBuilder = new ResourceProfileBuilder()
- val ereqs = new ExecutorResourceRequests()
- ereqs.cores(checkId).memory("4g").memoryOverhead("2000m")
- val treqs = new TaskResourceRequests()
- treqs.cpus(checkId)
- rprofBuilder.require(ereqs).require(treqs)
- val rpShouldMatch = rprofBuilder.build()
-
- val equivProf = rpmanager.getEquivalentProfile(rpShouldMatch)
- assert(equivProf.nonEmpty)
- assert(equivProf.get.id == rpAlreadyExist.get.id, s"resourceProfile should
have existed")
+
+ val first = buildProfile(8)
+ val registered = rpmanager.getOrAddEquivalentProfile(first)
+ // A brand-new profile is registered and returned as-is.
+ assert(registered.id == first.id)
+
+ // A distinct profile object with equal resources resolves to the
already-registered one,
+ // so they share a single id and can therefore reuse the same executors.
+ val equivalent = buildProfile(8)
+ assert(equivalent.id != first.id, "the new profile object should have a
different id")
+ val resolved = rpmanager.getOrAddEquivalentProfile(equivalent)
+ assert(resolved.id == first.id, "equivalent profile should resolve to the
existing id")
+
+ // A profile with different resources is registered under its own id.
+ val different = buildProfile(16)
+ val resolvedDifferent = rpmanager.getOrAddEquivalentProfile(different)
+ assert(resolvedDifferent.id == different.id)
+
+ verify(mockListenerBus,
times(2)).post(isA(classOf[SparkListenerResourceProfileAdded]))
+ }
+
+ test("getOrAddEquivalentProfile atomically registers equivalent profiles") {
+ val conf = new SparkConf().set(EXECUTOR_CORES, 4)
+ val mockListenerBus = mock[LiveListenerBus]
+ val rpmanager = new ResourceProfileManager(conf, mockListenerBus)
+ reset(mockListenerBus)
+
+ val profiles = (1 to 16).map { _ =>
+ val ereqs = new ExecutorResourceRequests().cores(8)
+ val treqs = new TaskResourceRequests().cpus(1)
+ new ResourceProfileBuilder().require(ereqs).require(treqs).build()
+ }
+ val resolved = ThreadUtils.parmap(profiles, "register-resource-profiles",
profiles.size) {
+ rpmanager.getOrAddEquivalentProfile
+ }
+
+ assert(resolved.map(_.id).toSet.size === 1)
+ verify(mockListenerBus,
times(1)).post(isA(classOf[SparkListenerResourceProfileAdded]))
+ }
+
+ test("getOrAddEquivalentProfile validates before registering") {
+ val conf = new SparkConf().setMaster("yarn").set(EXECUTOR_CORES, 4)
+ .set(DYN_ALLOCATION_ENABLED, true)
+ val mockListenerBus = mock[LiveListenerBus]
+ val rpmanager = new ResourceProfileManager(conf, mockListenerBus)
+ reset(mockListenerBus)
+
+ def invalidProfile(): ResourceProfile = {
+ val ereqs = new ExecutorResourceRequests().resource("gpu", 1,
"discoveryScript")
+ val treqs = new TaskResourceRequests().resource("gpu", 2)
+ new ResourceProfileBuilder().require(ereqs).require(treqs).build()
+ }
+
+ val first = invalidProfile()
+ val equivalent = invalidProfile()
+ Seq(first, equivalent).foreach { profile =>
+ val error = intercept[SparkException] {
+ rpmanager.getOrAddEquivalentProfile(profile)
+ }
+ assert(error.getMessage.contains("needs to be >= the task resource
request amount"))
+ intercept[SparkException] {
+ rpmanager.resourceProfileFromId(profile.id)
+ }
+ }
+ assert(rpmanager.getEquivalentProfile(first).isEmpty)
+ verify(mockListenerBus,
never()).post(isA(classOf[SparkListenerResourceProfileAdded]))
+ }
+
+ test("getOrAddEquivalentProfile keeps explicit profiles distinct from the
default") {
+ Seq("yarn", "k8s://test").foreach { master =>
+ ResourceProfile.clearDefaultProfile()
+ val conf = new SparkConf().setMaster(master).set(EXECUTOR_CORES, 4)
+ .set(DYN_ALLOCATION_ENABLED, true)
+ val rpmanager = new ResourceProfileManager(conf, listenerBus)
+ val defaultProfile = rpmanager.defaultResourceProfile
+ val explicitProfile = new ResourceProfile(
+ defaultProfile.executorResources, defaultProfile.taskResources)
+
+ val resolved = rpmanager.getOrAddEquivalentProfile(explicitProfile)
+ assert(resolved eq explicitProfile)
+ assert(resolved.id !== ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)
+
assert(rpmanager.getEquivalentProfile(explicitProfile).contains(explicitProfile))
+ }
+ }
+
+ test("getOrAddEquivalentProfile preserves an existing id mapping") {
+ val conf = new SparkConf().set(EXECUTOR_CORES, 4)
+ val rpmanager = new ResourceProfileManager(conf, listenerBus)
+ val defaultProfile = rpmanager.defaultResourceProfile
+ val collidingProfile = new ResourceProfileBuilder()
+ .require(new ExecutorResourceRequests().cores(8))
+ .require(new TaskResourceRequests().cpus(1))
+ .build()
+ collidingProfile.setToDefaultProfile()
+
+ val resolved = rpmanager.getOrAddEquivalentProfile(collidingProfile)
+ assert(resolved eq defaultProfile)
+
assert(rpmanager.resourceProfileFromId(ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)
eq
+ defaultProfile)
}
}
diff --git
a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index a44ed77b6511..518347b1b090 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -4843,6 +4843,27 @@ class DAGSchedulerSuite extends SparkFunSuite with
TempLocalSparkContext with Ti
assert(mergedRp.getExecutorCores.get == 4)
}
+ test("merge of default and no-op task-only resource profiles uses default
profile") {
+ conf.set(config.RESOURCE_PROFILE_MERGE_CONFLICTS.key, "true")
+ conf.set(config.DYN_ALLOCATION_ENABLED, false)
+
+ val defaultProfile = sc.resourceProfileManager.defaultResourceProfile
+ val taskOnlyProfile = new ResourceProfileBuilder()
+ .require(new TaskResourceRequests().cpus(1))
+ .build()
+ val rdd = sc.parallelize(1 to 10, 2)
+ .withResources(defaultProfile)
+ .map(x => (x, x))
+ .withResources(taskOnlyProfile)
+
+ submit(rdd, Array(0, 1))
+
+ assert(taskSets.length === 1)
+ assert(taskSets.head.resourceProfileId ===
ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)
+
assert(sc.resourceProfileManager.resourceProfileFromId(taskSets.head.resourceProfileId)
eq
+ defaultProfile)
+ }
+
test("test multiple resource profiles created from merging use same rp") {
conf.set(config.RESOURCE_PROFILE_MERGE_CONFLICTS.key, "true")
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index 78562e8ebf92..f0ca4154c154 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -3957,13 +3957,16 @@ class SparkConnectPlanner(
name -> new TaskResourceRequest(res.getResourceName, res.getAmount)
}.toMap
- // Create ResourceProfile add add it to ResourceProfileManager
- val profile = if (ereqs.isEmpty) {
+ // Create the ResourceProfile and register it, reusing an
already-registered profile with
+ // equal resources if one exists so that equivalent profiles share a
single id (and thus
+ // can reuse the same executors instead of triggering new allocations).
+ val newProfile = if (ereqs.isEmpty) {
new TaskResourceProfile(treqs)
} else {
new ResourceProfile(ereqs, treqs)
}
- session.sparkContext.resourceProfileManager.addResourceProfile(profile)
+ val profile =
+
session.sparkContext.resourceProfileManager.getOrAddEquivalentProfile(newProfile)
executeHolder.eventsManager.postFinished()
responseObserver.onNext(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
