This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new d6fa6f7f2a [CELEBORN-2499] Show shuffle slot assignments in the Spark
UI Celeborn tab
d6fa6f7f2a is described below
commit d6fa6f7f2a4c765dd0a165536f5069e69c8529af
Author: ForVic <[email protected]>
AuthorDate: Mon Oct 26 10:00:14 2026 +0800
[CELEBORN-2499] Show shuffle slot assignments in the Spark UI Celeborn tab
### What changes were proposed in this pull request?
JIRA: https://issues.apache.org/jira/browse/CELEBORN-2499
Show the slots reserved on each worker for every Celeborn shuffle on the
Celeborn tab of the Spark UI and the History Server. This is the "Shuffle
Assignments" follow-up planned in CELEBORN-2415. The driver-side
LifecycleManager already knows the assignment, so no executor instrumentation
is needed.
- **client**: once a shuffle is registered, `LifecycleManager` reports the
slots reserved on each worker to a `ShuffleAssignmentListener`, together with
the Celeborn shuffle id and the application (Spark) shuffle id. With stage
rerun enabled, the two ids differ.
`SparkShuffleManager.registerShuffleAssignmentListener` hands the listener to
its lazily created LifecycleManager.
- **spark-ui**:
- `CelebornPlugin` registers the listener and posts a
`CelebornShuffleAssignmentEvent` per shuffle, so applications without the
plugin post nothing. Registration happens in `registerMetrics` because Spark 4
initializes the shuffle manager after the driver plugins' `init`.
- `CelebornListener` persists two things, for both the live UI and
History Server replay:
- a compact per-shuffle summary, retaining up to
`spark.ui.retainedStages` shuffles;
- the shuffles and slots of each worker over all shuffles, which stay
exact after old shuffles are dropped.
- The tab gains a *Shuffle Assignments* summary line and two collapsed
tables:
- *Shuffle Assignments*: Spark and Celeborn shuffle ids, workers,
reserved slots, and slots per worker as min / avg / max.
- *Worker Slots*: the shuffles and reserved slots of each worker.
### Why are the changes needed?
The Celeborn tab shows nothing about which workers serve an application's
shuffles or how evenly their slots are spread, which is needed to understand
worker hotspots for a specific application.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
The Celeborn tab gains the *Shuffle Assignments* summary line and two
collapsed tables. Nothing changes for applications that don't enable
`CelebornPlugin`. Some cases are not covered:
- Slots that partition splits or revives reallocate later are not included.
- Shuffle managers that create their own LifecycleManager instead of using
Celeborn's `SparkShuffleManager` (e.g. Gluten's) report no assignments.
Same workload and cluster (Spark 3.5.9, 3 workers, replication enabled).
Before (main):
After, with the collapsed sections expanded:
### How was this patch tested?
- `CelebornListenerSuite`:
- per-shuffle summaries and per-worker totals are persisted from the
event;
- per-worker totals stay exact beyond `spark.ui.retainedStages`;
- the History Server KVStore serializer round trip works.
- `CelebornPluginSuite`: a local SparkContext with `CelebornPlugin` and
`SparkShuffleManager` has the listener registered, and reported assignments
reach the UI store.
- `LifecycleManagerSuite`: against a mini cluster, the listener receives
both shuffle ids and primary plus replica slots for every partition.
- `./build/mvn -P<profile> -pl client-spark/spark-3-ui test` passes for
`spark-3.0`, `spark-3.5` and `spark-4.0`.
- Manual: the screenshots above, from a local cluster with 3 Celeborn
workers, Spark 3.5.9 standalone and the History Server, which shows the same
data on replay.
Closes #3891 from ForVic/ForVic/spark-ui-shuffle-assignments.
Lead-authored-by: ForVic <[email protected]>
Co-authored-by: zhengtao <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/303
AI-Contributed/UT: 0/128
---
client-spark/spark-3-ui/pom.xml | 5 ++
.../spark/shuffle/celeborn/CelebornListener.scala | 44 +++++++++++
.../spark/shuffle/celeborn/CelebornPlugin.scala | 16 +++-
.../celeborn/CelebornShuffleAssignmentEvent.scala | 31 ++++++++
.../shuffle/celeborn/CelebornStatusStore.scala | 42 +++++++++++
.../shuffle/celeborn/ui/CelebornShufflePage.scala | 87 ++++++++++++++++++++++
.../shuffle/celeborn/CelebornListenerSuite.scala | 50 ++++++++++++-
.../shuffle/celeborn/CelebornPluginSuite.scala | 53 +++++++++++++
.../shuffle/celeborn/SparkShuffleManager.java | 15 ++++
.../client/listener/ShuffleAssignmentListener.java | 33 ++++++++
.../apache/celeborn/client/LifecycleManager.scala | 29 +++++++-
project/CelebornBuild.scala | 1 +
.../tests/client/LifecycleManagerSuite.scala | 25 +++++++
13 files changed, 428 insertions(+), 3 deletions(-)
diff --git a/client-spark/spark-3-ui/pom.xml b/client-spark/spark-3-ui/pom.xml
index d11f0a29a6..eb31e3315f 100644
--- a/client-spark/spark-3-ui/pom.xml
+++ b/client-spark/spark-3-ui/pom.xml
@@ -45,6 +45,11 @@
</exclusion>
</exclusions>
</dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>javax.servlet</groupId>
<artifactId>javax.servlet-api</artifactId>
diff --git
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
index a230cff08b..efcc56841a 100644
---
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
+++
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
@@ -19,8 +19,12 @@ package org.apache.spark.shuffle.celeborn
import java.util.concurrent.atomic.AtomicLong
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
import org.apache.spark.SparkConf
import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config.Status.MAX_RETAINED_STAGES
import org.apache.spark.scheduler._
import org.apache.spark.util.Utils
import org.apache.spark.util.kvstore.KVStore
@@ -57,6 +61,13 @@ private[celeborn] class CelebornListener(
private val lastUpdateTimestamp = new AtomicLong(-1L)
private val updateIntervalMillis = 5000L
+ // Bound the retained shuffles like Spark bounds its retained stages.
+ private val retainedShuffles = conf.get(MAX_RETAINED_STAGES)
+ // Only accessed from the listener bus thread.
+ private var assignedShuffles = 0L
+ private val workerShuffles = mutable.HashMap[String, Int]()
+ private val workerSlots = mutable.HashMap[String, Int]()
+
def register(sc: org.apache.spark.SparkContext): Unit = {
sc.addSparkListener(this)
logInfo("CelebornListener registered successfully")
@@ -95,6 +106,7 @@ private[celeborn] class CelebornListener(
totalCompressedWriteBytes.addAndGet(e.getCompressedBytesWritten)
mayUpdate()
}
+ case e: CelebornShuffleAssignmentEvent => onShuffleAssignment(e)
case _ =>
}
@@ -132,6 +144,38 @@ private[celeborn] class CelebornListener(
}
}
+ private def onShuffleAssignment(event: CelebornShuffleAssignmentEvent): Unit
= {
+ val slots = event.workerSlots.values
+ kvstore.write(ShuffleAssignmentUIData(
+ event.shuffleId,
+ event.appShuffleId,
+ event.workerSlots.size,
+ slots.sum,
+ if (slots.isEmpty) 0 else slots.min,
+ if (slots.isEmpty) 0 else slots.max))
+ dropLowestShuffleAssignments()
+
+ assignedShuffles += 1
+ event.workerSlots.foreach { case (worker, numSlots) =>
+ workerShuffles(worker) = workerShuffles.getOrElse(worker, 0) + 1
+ workerSlots(worker) = workerSlots.getOrElse(worker, 0) + numSlots
+ }
+ kvstore.write(WorkerSlotsUIData(assignedShuffles, workerShuffles.toMap,
workerSlots.toMap))
+ }
+
+ // Drops the lowest Celeborn shuffle ids, which roughly are the oldest
shuffles.
+ private def dropLowestShuffleAssignments(): Unit = {
+ val uiDataClass = classOf[ShuffleAssignmentUIData]
+ val excess = kvstore.count(uiDataClass) - retainedShuffles
+ if (excess > 0) {
+ val dropped = Utils.tryWithResource(
+ kvstore.view(uiDataClass).max(excess).closeableIterator()) { iter =>
+ iter.asScala.map(_.shuffleId).toList
+ }
+ dropped.foreach(shuffleId => kvstore.delete(uiDataClass, shuffleId))
+ }
+ }
+
override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd):
Unit = {
mayUpdate(force = true)
logInfo("CelebornListener: application ended, final flush completed")
diff --git
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala
index a38a419cac..54c85dc2ef 100644
---
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala
+++
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala
@@ -19,7 +19,9 @@ package org.apache.spark.shuffle.celeborn
import java.util
-import org.apache.spark.SparkContext
+import scala.collection.JavaConverters._
+
+import org.apache.spark.{SparkContext, SparkEnv}
import org.apache.spark.api.plugin.{DriverPlugin, PluginContext, SparkPlugin}
import org.apache.spark.internal.Logging
import org.apache.spark.shuffle.celeborn.ui.CelebornUITab
@@ -66,6 +68,18 @@ private class CelebornDriverPlugin extends DriverPlugin with
Logging {
sc.ui.foreach { ui =>
new CelebornUITab(new CelebornStatusStore(ui.store.store), ui)
}
+ // Not in init: Spark 4 initializes the shuffle manager after the driver
plugins.
+ SparkEnv.get.shuffleManager match {
+ case manager: SparkShuffleManager =>
+ manager.registerShuffleAssignmentListener((shuffleId, appShuffleId,
workerSlots) =>
+ sc.listenerBus.post(CelebornShuffleAssignmentEvent(
+ shuffleId,
+ appShuffleId,
+ workerSlots.asScala.map { case (worker, slots) => worker ->
slots.intValue }.toMap)))
+ case other =>
+ logInfo("Shuffle assignments are not reported by shuffle manager " +
+ s"${Option(other).map(_.getClass.getName).orNull}.")
+ }
}
override def shutdown(): Unit = {}
diff --git
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleAssignmentEvent.scala
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleAssignmentEvent.scala
new file mode 100644
index 0000000000..fa0240fb69
--- /dev/null
+++
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleAssignmentEvent.scala
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.celeborn
+
+import org.apache.spark.scheduler.SparkListenerEvent
+
+/**
+ * Slots reserved on each worker when a Celeborn shuffle is registered, keyed
by the
+ * worker's host:rpcPort:pushPort:fetchPort:replicatePort. Recorded in the
event log;
+ * the listener writes it to the KVStore both for the live UI and on History
Server replay.
+ */
+private[celeborn] case class CelebornShuffleAssignmentEvent(
+ shuffleId: Int,
+ appShuffleId: Int,
+ workerSlots: Map[String, Int])
+ extends SparkListenerEvent
diff --git
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala
index 05076cf439..dbc4f64940 100644
---
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala
+++
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala
@@ -17,7 +17,10 @@
package org.apache.spark.shuffle.celeborn
+import scala.collection.JavaConverters._
+
import com.fasterxml.jackson.annotation.JsonIgnore
+import org.apache.spark.util.Utils
import org.apache.spark.util.kvstore.{KVIndex, KVStore}
private[celeborn] case class AggregatedTaskInfoUIData(
@@ -52,6 +55,30 @@ private[celeborn] class CelebornBuildInfoUIData(
def id: String = classOf[CelebornBuildInfoUIData].getName
}
+private[celeborn] case class ShuffleAssignmentUIData(
+ shuffleId: Int,
+ appShuffleId: Int,
+ workers: Int,
+ slots: Int,
+ minWorkerSlots: Int,
+ maxWorkerSlots: Int) {
+
+ @JsonIgnore
+ @KVIndex
+ def id: Int = shuffleId
+}
+
+/** Shuffles and slots of each worker, over all shuffles of the application. */
+private[celeborn] case class WorkerSlotsUIData(
+ shuffles: Long,
+ workerShuffles: Map[String, Int],
+ workerSlots: Map[String, Int]) {
+
+ @JsonIgnore
+ @KVIndex
+ def id: String = classOf[WorkerSlotsUIData].getName
+}
+
/**
* Marker persisted during History Server replay when the application's
recorded
* `spark.plugins` or `spark.plugins.defaultList` contains [[CelebornPlugin]].
@@ -95,6 +122,21 @@ private[celeborn] class CelebornStatusStore(store: KVStore)
{
}
}
+ def shuffleAssignments(): Seq[ShuffleAssignmentUIData] = {
+
Utils.tryWithResource(store.view(classOf[ShuffleAssignmentUIData]).closeableIterator())
{
+ iter => iter.asScala.toList
+ }
+ }
+
+ def workerSlots(): WorkerSlotsUIData = {
+ val kClass = classOf[WorkerSlotsUIData]
+ try {
+ store.read(kClass, kClass.getName)
+ } catch {
+ case _: NoSuchElementException => WorkerSlotsUIData(0L, Map.empty,
Map.empty)
+ }
+ }
+
def extensionEnabled(): Boolean = {
val kClass = classOf[CelebornExtensionEnabledUIData]
try {
diff --git
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
index 41edad1130..911a3831da 100644
---
a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
+++
b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
@@ -20,6 +20,7 @@ package org.apache.spark.shuffle.celeborn.ui
import scala.xml.Node
import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.celeborn.ShuffleAssignmentUIData
import
org.apache.spark.shuffle.celeborn.ui.SparkServletBridge.HttpServletRequest
import org.apache.spark.ui.{UIUtils, WebUIPage}
import org.apache.spark.util.Utils
@@ -54,6 +55,11 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
val taskInfo = store.aggregatedTaskInfo()
val properties = store.celebornProperties()
val buildInfo = store.buildInfo()
+ val assignments = store.shuffleAssignments().sortBy(-_.shuffleId)
+ val workerSlots = store.workerSlots()
+ val workerRows = workerSlots.workerSlots.toSeq.map { case (worker, slots)
=>
+ (worker, workerSlots.workerShuffles.getOrElse(worker, 0), slots)
+ }.sortBy { case (worker, _, slots) => (-slots, worker) }
val writeBytes = taskInfo.shuffleWriteBytes
val readBytes = taskInfo.shuffleReadBytes
@@ -67,6 +73,19 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
if (ms <= 0) "N/A" else f"${bytes.toDouble / 1000.0 / 1000.0 /
(ms.toDouble / 1000.0)}%.2f"
def pct(part: Long, total: Long): String =
if (total <= 0) "N/A" else f"${part.toDouble * 100.0 /
total.toDouble}%.1f%%"
+ def assignmentSummary: String =
+ if (workerSlots.shuffles == 0) "N/A"
+ else {
+ s"Shuffles: ${workerSlots.shuffles} | Workers: ${workerRows.size} | " +
+ s"Reserved Slots: ${workerRows.map(_._3.toLong).sum}"
+ }
+ val assignmentsTitle =
+ if (workerSlots.shuffles > assignments.size) {
+ s"Shuffle Assignments (${workerSlots.shuffles}, only showing
${assignments.size})"
+ } else {
+ "Shuffle Assignments"
+ }
+
// 0 when not reported, e.g. by old clients.
def compressionRatio: String =
if (uncompressedBytes <= 0 || compressedBytes <= 0) "N/A"
@@ -104,6 +123,10 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
durationMs)}, Read ${pct(readMs, durationMs)})"
}
</li>
+ <li>
+ <strong>Shuffle Assignments: </strong>
+ {assignmentSummary}
+ </li>
</ul>
</div>
@@ -114,6 +137,25 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
fixedWidth = true,
headerClasses = headerClasses)
+ val shuffleAssignmentTable = UIUtils.listingTable(
+ Seq(
+ "Shuffle ID",
+ "Celeborn Shuffle ID",
+ "Workers",
+ "Reserved Slots",
+ "Slots per Worker (Min / Avg / Max)"),
+ shuffleAssignmentRow,
+ assignments,
+ fixedWidth = true,
+ headerClasses = Seq.fill(5)("sorttable_numeric"))
+
+ val workerSlotsTable = UIUtils.listingTable(
+ Seq("Worker", "Shuffles", "Reserved Slots"),
+ workerSlotsRow,
+ workerRows,
+ fixedWidth = true,
+ headerClasses = Seq("sorttable_alpha", "sorttable_numeric",
"sorttable_numeric"))
+
val buildInfoTable = UIUtils.listingTable(
propertyHeader,
propertyRow,
@@ -124,6 +166,28 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
val content =
<span>
{summary}
+ <span class="collapse-aggregated-celebornShuffleAssignments
collapse-table"
+
onClick="collapseTable('collapse-aggregated-celebornShuffleAssignments',
+ 'aggregated-celebornShuffleAssignments')">
+ <h4>
+ <span class="collapse-table-arrow arrow-closed"></span>
+ <a>{assignmentsTitle}</a>
+ </h4>
+ </span>
+ <div class="aggregated-celebornShuffleAssignments collapsible-table
collapsed">
+ {shuffleAssignmentTable}
+ </div>
+ <span class="collapse-aggregated-celebornWorkerSlots collapse-table"
+ onClick="collapseTable('collapse-aggregated-celebornWorkerSlots',
+ 'aggregated-celebornWorkerSlots')">
+ <h4>
+ <span class="collapse-table-arrow arrow-closed"></span>
+ <a>Worker Slots</a>
+ </h4>
+ </span>
+ <div class="aggregated-celebornWorkerSlots collapsible-table
collapsed">
+ {workerSlotsTable}
+ </div>
<span class="collapse-aggregated-celebornBuildInfo collapse-table"
onClick="collapseTable('collapse-aggregated-celebornBuildInfo',
'aggregated-celebornBuildInfo')">
@@ -154,6 +218,29 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
private def propertyHeader: Seq[String] = Seq("Name", "Value")
private def headerClasses: Seq[String] = Seq("sorttable_alpha",
"sorttable_alpha")
+ private def shuffleAssignmentRow(assignment: ShuffleAssignmentUIData):
Seq[Node] = {
+ val perWorker =
+ if (assignment.workers == 0) "N/A"
+ else {
+ f"${assignment.minWorkerSlots} / ${assignment.slots.toDouble /
assignment.workers}%.1f / " +
+ s"${assignment.maxWorkerSlots}"
+ }
+ <tr>
+ <td>{assignment.appShuffleId}</td>
+ <td>{assignment.shuffleId}</td>
+ <td>{assignment.workers}</td>
+ <td>{assignment.slots}</td>
+ <td
sorttable_customkey={assignment.maxWorkerSlots.toString}>{perWorker}</td>
+ </tr>
+ }
+
+ private def workerSlotsRow(row: (String, Int, Int)): Seq[Node] =
+ <tr>
+ <td>{row._1}</td>
+ <td>{row._2}</td>
+ <td>{row._3}</td>
+ </tr>
+
private def propertyRow(kv: (String, String)): Seq[Node] =
<tr>
<td>{kv._1}</td>
diff --git
a/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
index e2314ed4f8..f7192039f1 100644
---
a/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
+++
b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
@@ -20,9 +20,10 @@ package org.apache.spark.shuffle.celeborn
import com.fasterxml.jackson.databind.ObjectMapper
import org.apache.spark.{SparkConf, Success, TaskEndReason, TaskKilled,
TaskState}
import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.internal.config.Status.ASYNC_TRACKING_ENABLED
+import org.apache.spark.internal.config.Status.{ASYNC_TRACKING_ENABLED,
MAX_RETAINED_STAGES}
import org.apache.spark.scheduler.{JobSucceeded,
SparkListenerEnvironmentUpdate, SparkListenerEvent, SparkListenerJobEnd,
SparkListenerTaskEnd, TaskInfo, TaskLocality}
import org.apache.spark.status.ElementTrackingStore
+import org.apache.spark.status.KVUtils.KVStoreScalaSerializer
import org.apache.spark.util.Utils
import org.apache.spark.util.kvstore.InMemoryStore
import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
@@ -317,4 +318,51 @@ class CelebornListenerSuite {
assertEquals("abcdef12", info("Commit Id"))
assertEquals("abcdef1234567890", info("Revision"))
}
+
+ @Test
+ def shuffleAssignmentPersistedFromPostedEvent(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ listener.onOtherEvent(CelebornShuffleAssignmentEvent(3, 1, Map("w1" -> 2,
"w2" -> 4)))
+ listener.onOtherEvent(CelebornShuffleAssignmentEvent(4, 2, Map("w2" -> 1)))
+
+ assertEquals(
+ Seq(ShuffleAssignmentUIData(3, 1, 2, 6, 2, 4),
ShuffleAssignmentUIData(4, 2, 1, 1, 1, 1)),
+ statusStore.shuffleAssignments())
+ assertEquals(
+ WorkerSlotsUIData(2, Map("w1" -> 1, "w2" -> 2), Map("w1" -> 2, "w2" ->
5)),
+ statusStore.workerSlots())
+ }
+
+ @Test
+ def workerSlotsCoverShuffleAssignmentsBeyondRetainedStages(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new
SparkConf().set(MAX_RETAINED_STAGES, 2))
+
+ (0 until 3).foreach { id =>
+ listener.onOtherEvent(CelebornShuffleAssignmentEvent(id, id, Map("w1" ->
1)))
+ }
+
+ assertEquals(Seq(1, 2), statusStore.shuffleAssignments().map(_.shuffleId))
+ assertEquals(WorkerSlotsUIData(3, Map("w1" -> 3), Map("w1" -> 3)),
statusStore.workerSlots())
+ }
+
+ @Test
+ def shuffleAssignmentKVStoreRoundTrip(): Unit = {
+ // History Server disk stores persist UI data through this serializer.
+ val serializer = new KVStoreScalaSerializer()
+ val assignment = ShuffleAssignmentUIData(5, 1, 2, 4, 1, 3)
+ val workerSlots = WorkerSlotsUIData(1, Map("w1" -> 1, "w2" -> 1), Map("w1"
-> 3, "w2" -> 1))
+
+ assertEquals(
+ assignment,
+ serializer.deserialize(serializer.serialize(assignment),
classOf[ShuffleAssignmentUIData]))
+ val restored =
+ serializer.deserialize(serializer.serialize(workerSlots),
classOf[WorkerSlotsUIData])
+ assertEquals(workerSlots, restored)
+ assertEquals(4, restored.workerSlots.values.sum)
+ }
}
diff --git
a/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornPluginSuite.scala
b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornPluginSuite.scala
new file mode 100644
index 0000000000..07f20578a1
--- /dev/null
+++
b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornPluginSuite.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to You under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific
+ * language governing permissions and limitations under the License.
+ */
+package org.apache.spark.shuffle.celeborn
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.{SparkConf, SparkContext, SparkEnv}
+import org.junit.Assert.assertEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(classOf[JUnit4])
+class CelebornPluginSuite {
+
+ @Test
+ def shuffleAssignmentsReportedThroughSparkShuffleManager(): Unit = {
+ val conf = new SparkConf()
+ .setMaster("local")
+ .setAppName("CelebornPluginSuite")
+ .set("spark.plugins", classOf[CelebornPlugin].getName)
+ .set("spark.shuffle.manager", classOf[SparkShuffleManager].getName)
+ .set("spark.ui.enabled", "false")
+ val sc = new SparkContext(conf)
+ try {
+ val listener =
+
SparkEnv.get.shuffleManager.asInstanceOf[SparkShuffleManager].getShuffleAssignmentListener
+ listener.onShuffleAssigned(3, 1, Map[String, Integer]("w1" -> 2, "w2" ->
4).asJava)
+ sc.listenerBus.waitUntilEmpty()
+
+ assertEquals(
+ Seq(ShuffleAssignmentUIData(3, 1, 2, 6, 2, 4)),
+ new CelebornStatusStore(sc.statusStore.store).shuffleAssignments())
+ } finally {
+ sc.stop()
+ }
+ }
+}
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index 7b5a2fbd65..b731d4923f 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -36,6 +36,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.listener.ShuffleAssignmentListener;
import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.FallbackPolicy;
@@ -99,6 +100,7 @@ public class SparkShuffleManager implements ShuffleManager {
private String appUniqueId;
private LifecycleManager lifecycleManager;
+ private ShuffleAssignmentListener shuffleAssignmentListener;
private ShuffleClient shuffleClient;
private volatile SortShuffleManager _sortShuffleManager;
private final ConcurrentHashMap.KeySetView<Integer, Boolean> sortShuffleIds =
@@ -192,6 +194,9 @@ public class SparkShuffleManager implements ShuffleManager {
lifecycleManager = new LifecycleManager(appUniqueId, celebornConf);
lifecycleManager.applicationCount().increment();
lifecycleManager.registerCancelShuffleCallback(SparkUtils::cancelShuffle);
+ if (shuffleAssignmentListener != null) {
+
lifecycleManager.registerShuffleAssignmentListener(shuffleAssignmentListener);
+ }
// Register only when the UI plugin is opted in, otherwise every map
task
// posts an event nobody consumes and each one flushes the event log.
if (SparkUtils.isCelebornPluginEnabled(conf)) {
@@ -559,6 +564,16 @@ public class SparkShuffleManager implements ShuffleManager
{
}
}
+ /** Registers a listener for the LifecycleManager, which is created on the
first shuffle. */
+ public synchronized void
registerShuffleAssignmentListener(ShuffleAssignmentListener listener) {
+ shuffleAssignmentListener = listener;
+ }
+
+ // for testing
+ public synchronized ShuffleAssignmentListener getShuffleAssignmentListener()
{
+ return shuffleAssignmentListener;
+ }
+
// for testing
public LifecycleManager getLifecycleManager() {
return this.lifecycleManager;
diff --git
a/client/src/main/java/org/apache/celeborn/client/listener/ShuffleAssignmentListener.java
b/client/src/main/java/org/apache/celeborn/client/listener/ShuffleAssignmentListener.java
new file mode 100644
index 0000000000..41b082eab4
--- /dev/null
+++
b/client/src/main/java/org/apache/celeborn/client/listener/ShuffleAssignmentListener.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.client.listener;
+
+import java.util.Map;
+
+/** The Listener about the slots reserved for newly registered shuffles. */
+public interface ShuffleAssignmentListener {
+
+ /**
+ * Notifies the slots reserved on each worker once a shuffle is registered.
+ *
+ * @param shuffleId the Celeborn shuffle id
+ * @param appShuffleId the application shuffle id, which differs on stage
rerun
+ * @param workerSlots slots reserved on each worker, keyed by its unique id
+ */
+ void onShuffleAssigned(int shuffleId, int appShuffleId, Map<String, Integer>
workerSlots);
+}
diff --git
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 77ff0715ae..7074899fc1 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -39,7 +39,7 @@ import com.google.common.annotations.VisibleForTesting
import com.google.common.cache.{Cache, CacheBuilder}
import org.apache.celeborn.client.LifecycleManager.{ShuffleAllocatedWorkers,
ShuffleFailedWorkers}
-import org.apache.celeborn.client.listener.{MapperEndCallback,
WorkerStatusListener}
+import org.apache.celeborn.client.listener.{MapperEndCallback,
ShuffleAssignmentListener, WorkerStatusListener}
import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
import org.apache.celeborn.common.CelebornConf.ACTIVE_STORAGE_TYPES
import org.apache.celeborn.common.client.{ApplicationInfoProvider,
MasterClient}
@@ -904,6 +904,28 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
StatusCode.SUCCESS,
allPrimaryPartitionLocations,
serdeVersion))
+ reportShuffleAssignment(shuffleId, slots)
+ }
+ }
+
+ private def reportShuffleAssignment(shuffleId: Int, slots: WorkerResource):
Unit = {
+ shuffleAssignmentListener.foreach { listener =>
+ val workerSlots = new util.HashMap[String, Integer]()
+ slots.asScala.foreach { case (workerInfo, (primaryLocations,
replicaLocations)) =>
+ val numSlots = primaryLocations.size + replicaLocations.size
+ if (numSlots > 0) {
+ workerSlots.put(workerInfo.toUniqueId, numSlots)
+ }
+ }
+ try {
+ listener.onShuffleAssigned(
+ shuffleId,
+ celebornShuffleIdToAppShuffleIdMap.getOrDefault(shuffleId,
shuffleId),
+ workerSlots)
+ } catch {
+ case NonFatal(e) =>
+ logWarning(s"Failed to report the slot assignment of shuffle
$shuffleId.", e)
+ }
}
}
@@ -2227,6 +2249,11 @@ class LifecycleManager(val appUniqueId: String, val
conf: CelebornConf) extends
celebornSkewShuffleCheckCallback = Some(callback)
}
+ @volatile private var shuffleAssignmentListener:
Option[ShuffleAssignmentListener] = None
+ def registerShuffleAssignmentListener(listener: ShuffleAssignmentListener):
Unit = {
+ shuffleAssignmentListener = Some(listener)
+ }
+
// Initialize at the end of LifecycleManager construction.
initialize()
diff --git a/project/CelebornBuild.scala b/project/CelebornBuild.scala
index 8595f62f58..bf57c8c7b4 100644
--- a/project/CelebornBuild.scala
+++ b/project/CelebornBuild.scala
@@ -1140,6 +1140,7 @@ trait SparkClientProjects {
commonSettings,
libraryDependencies ++= Seq(
"org.apache.spark" %% "spark-core" % sparkVersion % "provided",
+ "org.apache.spark" %% "spark-sql" % sparkVersion % "test",
Dependencies.javaxServletApi % "provided",
Dependencies.jakartaServletApi % "provided"
) ++ commonUnitTestDependencies,
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/LifecycleManagerSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/LifecycleManagerSuite.scala
index ade825596f..75892d70c2 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/LifecycleManagerSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/LifecycleManagerSuite.scala
@@ -18,6 +18,7 @@
package org.apache.celeborn.tests.client
import java.util
+import java.util.concurrent.ConcurrentLinkedQueue
import scala.collection.JavaConverters._
@@ -131,6 +132,30 @@ class LifecycleManagerSuite extends WithShuffleClientSuite
with MiniClusterFeatu
lifecycleManager.stop()
}
+ test("CELEBORN-2499: report the slots reserved on each worker for a
registered shuffle") {
+ prepareService()
+ val reported = new ConcurrentLinkedQueue[(Int, Int, Map[String, Int])]()
+ lifecycleManager.registerShuffleAssignmentListener { (shuffleId,
appShuffleId, workerSlots) =>
+ val slots = workerSlots.asScala.map { case (worker, n) => worker ->
n.intValue }.toMap
+ reported.add((shuffleId, appShuffleId, slots))
+ }
+ val shuffleId = nextShuffleId
+ val numPartitions = 9
+
+ shuffleClient.getPartitionLocation(shuffleId, 8, numPartitions)
+
+ eventually(timeout(3.seconds), interval(10.milliseconds)) {
+ assert(reported.size == 1)
+ }
+ val (reportedShuffleId, appShuffleId, workerSlots) = reported.peek()
+ assert(reportedShuffleId == shuffleId)
+ assert(appShuffleId == shuffleId)
+ // Replication is enabled: a primary and a replica slot per partition.
+ assert(workerSlots.values.sum == 2 * numPartitions)
+ assert(workerSlots.values.forall(_ > 0))
+
assert(workerSlots.keySet.subsetOf(workerInfos.keySet.map(_.workerInfo.toUniqueId)))
+ }
+
test("CELEBORN-2479: mapper end callback reports only the accepted attempt")
{
prepareService()
val numMappers = 8
