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 3cb708b05e [CELEBORN-2479] Show shuffle compression ratio in Spark UI
3cb708b05e is described below
commit 3cb708b05e4a21b3d336b2628c461b4990e9f2aa
Author: yew1eb <[email protected]>
AuthorDate: Fri Oct 23 14:00:58 2026 +0800
[CELEBORN-2479] Show shuffle compression ratio in Spark UI
### What changes were proposed in this pull request?
Show the shuffle compression ratio (uncompressed/compressed) on the
Celeborn tab of the Spark UI.
- **common**: new extensible `PbMapperMetrics` message carries both
uncompressed and compressed bytes in `PbMapperEnd` field 10; `MapperEnd` gains
a defaulted `MapperMetrics` parameter, so future mapper-side metrics go there
instead of new top-level fields.
- **client**: `ShuffleClient` keeps the existing abstract 5-argument
`mapperEnd` and adds a 6-argument overload whose default implementation
discards the metrics and delegates to the existing method, preserving
compatibility for external subclasses. `LifecycleManager` invokes the
registered `MapperEndCallback` only for an attempt accepted by the commit
handler.
- **spark**: `HashBasedShuffleWriter` / `SortBasedShuffleWriter` report
both the bytes handed to `ShuffleClient` before compression (including the
giant-record direct-push path) and the post-compression/encryption byte counts
returned by the client. When `CelebornPlugin` is enabled, `SparkShuffleManager`
registers the callback and posts a `CelebornMapperMetricsEvent` to the driver
listener bus. `CelebornListener` aggregates both counters for the live UI and
History Server replay; the [...]
- **docker**: register `CelebornPlugin` in the docker `spark-defaults.conf`
so the docker-compose cluster exercises the UI end-to-end.
The compression ratio derives both counters from the same `MapperMetrics`
payload for each accepted mapper attempt. Retried MapperEnd RPCs and rejected
speculative/rerun attempts do not emit another event: `finishMapperAttempt`
already accepts at most one attempt per mapId, so no additional deduplication
set is needed. This accounting is based on Celeborn commit acceptance rather
than Spark task-success filtering, keeping both counters aligned even if a task
fails after its writer completes.
### Why are the changes needed?
The Spark UI shows the shuffle bytes written to Celeborn
(post-compression), but there is no way to see how much the compression
actually saves. Exposing the ratio helps users evaluate codec choices and
estimate real network/disk traffic during shuffle.
CELEBORN-2479
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
The Celeborn tab of the Spark UI shows a new summary line: "Compression
Ratio (uncompressed/compressed)". Existing shuffle write bytes/time, read
metrics, and task-duration aggregates continue to include all task attempts,
including failed or killed attempts.
### How was this patch tested?
- `UtilsSuite` covers the MapperEnd ser/de roundtrip with both
`MapperMetrics` counters.
- `CelebornListenerSuite` covers aggregation, event JSON roundtrip,
all-attempt task metrics, and consistent ratio accounting when a task fails
after writer completion. All 12 tests passed locally with Spark 3.5 / JDK 17.
- Hash and sort writer suites check uncompressed and compressed byte counts
separately using distinct simulated compressed lengths, including
all-giant-record and mixed paths. Exact-head CI passed 12 hash writer tests, 13
sort writer tests, and 1 pusher test.
- `LifecycleManagerSuite` uses real RPCs to verify that only the accepted
attempt is reported, while retries of both accepted and rejected attempts are
ignored; this regression passed in CI.
- End-to-end on the docker-compose stack (Spark 3.5.9, hash writer, lz4):
confirmed the Celeborn tab shows the compression ratio with a real shuffling
job.
<img width="2832" height="1232" alt="image"
src="https://github.com/user-attachments/assets/f9fcb5a7-04c1-47cf-9463-4b35c3475583";
/>
Closes #3872 from yew1eb/spark-ui-compression-ratio.
Authored-by: yew1eb <[email protected]>
Signed-off-by: zhengtao <[email protected]>
Co-Authored-By: Codex <[email protected]>
AI-Model: gpt-6
AI-Contributed/Feature: 0/376
AI-Contributed/UT: 0/198
---
.../spark/shuffle/celeborn/SortBasedPusher.java | 10 +++
.../spark/shuffle/celeborn/SparkCommonUtils.java | 10 +++
.../spark/shuffle/celeborn/CelebornListener.scala | 45 ++++++----
.../shuffle/celeborn/CelebornStatusStore.scala | 8 +-
.../shuffle/celeborn/ui/CelebornShufflePage.scala | 13 +++
.../shuffle/celeborn/CelebornListenerSuite.scala | 100 ++++++++++++++++++++-
.../celeborn/CelebornMapperMetricsEvent.java | 75 ++++++++++++++++
.../shuffle/celeborn/HashBasedShuffleWriter.java | 16 +++-
.../shuffle/celeborn/SortBasedShuffleWriter.java | 15 +++-
.../shuffle/celeborn/SparkShuffleManager.java | 17 ++++
.../apache/spark/shuffle/celeborn/SparkUtils.java | 20 +++++
.../celeborn/CelebornShuffleWriterSuiteBase.java | 37 +++++++-
.../apache/celeborn/client/DummyShuffleClient.java | 26 +++++-
.../org/apache/celeborn/client/ShuffleClient.java | 17 ++++
.../apache/celeborn/client/ShuffleClientImpl.java | 22 ++++-
.../client/listener/MapperEndCallback.java | 29 ++++++
.../apache/celeborn/client/LifecycleManager.scala | 21 ++++-
common/src/main/proto/TransportMessages.proto | 10 +++
.../common/protocol/message/ControlMessages.scala | 19 +++-
.../apache/celeborn/common/util/UtilsSuite.scala | 6 +-
docker/spark/conf/spark-defaults.conf | 3 +
.../tests/client/LifecycleManagerSuite.scala | 55 ++++++++++++
22 files changed, 533 insertions(+), 41 deletions(-)
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
index 8297aefb60..6e5a7d48a4 100644
---
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
@@ -126,6 +126,9 @@ public class SortBasedPusher extends MemoryConsumer {
private int[] inversedShuffledPartitions = null;
private final SendBufferPool sendBufferPool;
+ // Bytes pushed before compression, for the shuffle compression ratio.
+ private long uncompressedBytesWritten = 0;
+
final MemoryThresholdManager memoryThresholdManager;
private final boolean useAdaptiveThreshold;
@@ -234,6 +237,7 @@ public class SortBasedPusher extends MemoryConsumer {
numPartitions);
mapStatusLengths[currentPartition].add(bytesWritten);
afterPush.accept(bytesWritten);
+ uncompressedBytesWritten += offSet;
memoryThresholdManager.updateStats(offSet, offSet ==
pushBufferMaxSize);
currentPartition = partition;
offSet = 0;
@@ -247,6 +251,7 @@ public class SortBasedPusher extends MemoryConsumer {
if (offSet + recordSize > dataBuf.length) {
try {
dataPusher.addTask(partition, dataBuf, offSet);
+ uncompressedBytesWritten += offSet;
memoryThresholdManager.updateStats(offSet, true);
} catch (InterruptedException e) {
TaskInterruptedHelper.throwTaskKillException();
@@ -262,6 +267,7 @@ public class SortBasedPusher extends MemoryConsumer {
if (offSet > 0) {
try {
dataPusher.addTask(currentPartition, dataBuf, offSet);
+ uncompressedBytesWritten += offSet;
memoryThresholdManager.updateStats(offSet, offSet ==
pushBufferMaxSize);
} catch (InterruptedException e) {
TaskInterruptedHelper.throwTaskKillException();
@@ -475,6 +481,10 @@ public class SortBasedPusher extends MemoryConsumer {
return peakMemoryUsedBytes;
}
+ public long getUncompressedBytesWritten() {
+ return uncompressedBytesWritten;
+ }
+
private long freeMemory() {
updatePeakMemoryUsed();
long memoryFreed = 0;
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
index 08e217390c..dce78460ae 100644
---
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
@@ -20,6 +20,7 @@ package org.apache.spark.shuffle.celeborn;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;
+import java.util.concurrent.atomic.LongAdder;
import scala.Option;
@@ -74,6 +75,15 @@ public class SparkCommonUtils {
return (context.stageAttemptNumber() << 16) | context.attemptNumber();
}
+ /** Sums per-partition pushed (post-compression/encryption) bytes. */
+ public static long sumLongAdders(LongAdder[] adders) {
+ long sum = 0;
+ for (LongAdder adder : adders) {
+ sum += adder.sum();
+ }
+ return sum;
+ }
+
public static void throwSparkOutOfMemoryError() {
try { // for Spark 3.5 and earlier
throw DynConstructors.builder()
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 d0cd03e3c3..a230cff08b 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
@@ -48,6 +48,11 @@ private[celeborn] class CelebornListener(
private val totalReadBytes = new AtomicLong(0L)
private val totalFetchWaitTimeMs = new AtomicLong(0L)
private val totalTaskDurationMs = new AtomicLong(0L)
+ // Total bytes pushed before compression, for the compression ratio.
+ private val totalUncompressedWriteBytes = new AtomicLong(0L)
+ // Total bytes pushed after compression, from the same MapperEnd reports, so
the
+ // compression ratio uses a single attempt set.
+ private val totalCompressedWriteBytes = new AtomicLong(0L)
private val lastUpdateTimestamp = new AtomicLong(-1L)
private val updateIntervalMillis = 5000L
@@ -62,6 +67,7 @@ private[celeborn] class CelebornListener(
return
}
Option(taskEnd.taskMetrics).foreach { metrics =>
+ // Count every attempt regardless of its end reason, like
AppStatusListener.
totalWriteBytes.addAndGet(metrics.shuffleWriteMetrics.bytesWritten)
// writeTime is in nanoseconds; normalize to ms.
totalWriteTimeMs.addAndGet(metrics.shuffleWriteMetrics.writeTime /
1000000L)
@@ -72,6 +78,26 @@ private[celeborn] class CelebornListener(
mayUpdate()
}
+ override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
+ case e: CelebornBuildInfoEvent =>
+ // Receiving the event also proves the application loaded the plugin.
+ kvstore.write(new CelebornBuildInfoUIData(Seq(
+ "Version" -> e.version,
+ "Commit Id" -> e.commitId,
+ "Revision" -> e.revision)))
+ if (!pluginEnabled) {
+ pluginEnabled = true
+ kvstore.write(new CelebornExtensionEnabledUIData())
+ }
+ case e: CelebornMapperMetricsEvent =>
+ if (pluginEnabled) {
+ totalUncompressedWriteBytes.addAndGet(e.getUncompressedBytesWritten)
+ totalCompressedWriteBytes.addAndGet(e.getCompressedBytesWritten)
+ mayUpdate()
+ }
+ case _ =>
+ }
+
override def onJobEnd(jobEnd: SparkListenerJobEnd): Unit = {
// Flush per-job so that tasks finishing within the throttle interval of
the
// previous flush are not lost when the application stays alive but idle.
@@ -85,6 +111,7 @@ private[celeborn] class CelebornListener(
val pluginClass = classOf[CelebornPlugin].getName
// Spark loads plugins from both keys: `spark.plugins.defaultList`
allows a
// default plugin list in the config file that `spark.plugins` does not
overwrite.
+ // Keep in sync with SparkUtils.isCelebornPluginEnabled in the spark-3
module.
val optedIn = sparkProps.exists { case (k, v) =>
(k == "spark.plugins" || k == "spark.plugins.defaultList") &&
v.split(",").exists(_.trim == pluginClass)
@@ -105,20 +132,6 @@ private[celeborn] class CelebornListener(
}
}
- override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
- case e: CelebornBuildInfoEvent =>
- // Receiving the event also proves the application loaded the plugin.
- kvstore.write(new CelebornBuildInfoUIData(Seq(
- "Version" -> e.version,
- "Commit Id" -> e.commitId,
- "Revision" -> e.revision)))
- if (!pluginEnabled) {
- pluginEnabled = true
- kvstore.write(new CelebornExtensionEnabledUIData())
- }
- case _ =>
- }
-
override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd):
Unit = {
mayUpdate(force = true)
logInfo("CelebornListener: application ended, final flush completed")
@@ -151,7 +164,9 @@ private[celeborn] class CelebornListener(
totalWriteTimeMs.get(),
totalReadBytes.get(),
totalFetchWaitTimeMs.get(),
- totalTaskDurationMs.get()))
+ totalTaskDurationMs.get(),
+ totalUncompressedWriteBytes.get(),
+ totalCompressedWriteBytes.get()))
} catch {
case e: Exception =>
logWarning("Failed to flush CelebornListener aggregations", e)
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 cfb255b026..05076cf439 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
@@ -25,7 +25,11 @@ private[celeborn] case class AggregatedTaskInfoUIData(
shuffleWriteTimeMs: Long,
shuffleReadBytes: Long,
shuffleFetchWaitTimeMs: Long,
- taskDurationMs: Long) {
+ taskDurationMs: Long,
+ // Bytes pushed before compression; 0 for records written by older
versions.
+ uncompressedShuffleBytes: Long = 0L,
+ // Bytes pushed after compression; 0 for records written by older versions.
+ compressedShuffleBytes: Long = 0L) {
@JsonIgnore
@KVIndex
@@ -69,7 +73,7 @@ private[celeborn] class CelebornStatusStore(store: KVStore) {
try {
store.read(kClass, kClass.getName)
} catch {
- case _: NoSuchElementException => AggregatedTaskInfoUIData(0L, 0L, 0L,
0L, 0L)
+ case _: NoSuchElementException => AggregatedTaskInfoUIData(0L, 0L, 0L,
0L, 0L, 0L)
}
}
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 33bc4027a7..41edad1130 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
@@ -60,11 +60,20 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
val writeMs = taskInfo.shuffleWriteTimeMs
val readMs = taskInfo.shuffleFetchWaitTimeMs
val durationMs = taskInfo.taskDurationMs
+ val uncompressedBytes = taskInfo.uncompressedShuffleBytes
+ val compressedBytes = taskInfo.compressedShuffleBytes
def mbps(bytes: Long, ms: Long): String =
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%%"
+ // 0 when not reported, e.g. by old clients.
+ def compressionRatio: String =
+ if (uncompressedBytes <= 0 || compressedBytes <= 0) "N/A"
+ else {
+
s"${Utils.bytesToString(uncompressedBytes)}/${Utils.bytesToString(compressedBytes)}
= " +
+ f"${uncompressedBytes.toDouble / compressedBytes}%.2f"
+ }
val summary =
<div>
@@ -83,6 +92,10 @@ private[celeborn] class CelebornShufflePage(parent:
CelebornUITab)
readMs)} | Speed: ${mbps(readBytes, readMs)} MB/s"
}
</li>
+ <li>
+ <strong>Compression Ratio (uncompressed/compressed): </strong>
+ {compressionRatio}
+ </li>
<li>
<strong>Shuffle Duration (write+read) / Task Duration: </strong>
{
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 21af11c0cc..e2314ed4f8 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
@@ -17,10 +17,11 @@
package org.apache.spark.shuffle.celeborn
-import org.apache.spark.{SparkConf, Success, TaskState}
+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.scheduler.{JobSucceeded,
SparkListenerEnvironmentUpdate, SparkListenerJobEnd, SparkListenerTaskEnd,
TaskInfo, TaskLocality}
+import org.apache.spark.scheduler.{JobSucceeded,
SparkListenerEnvironmentUpdate, SparkListenerEvent, SparkListenerJobEnd,
SparkListenerTaskEnd, TaskInfo, TaskLocality}
import org.apache.spark.status.ElementTrackingStore
import org.apache.spark.util.Utils
import org.apache.spark.util.kvstore.InMemoryStore
@@ -38,7 +39,8 @@ class CelebornListenerSuite {
writeBytes: Long,
writeTimeMs: Long,
fetchWaitMs: Long,
- durationMs: Long): SparkListenerTaskEnd = {
+ durationMs: Long,
+ reason: TaskEndReason = Success): SparkListenerTaskEnd = {
val metrics = TaskMetrics.empty
metrics.shuffleWriteMetrics.incBytesWritten(writeBytes)
metrics.shuffleWriteMetrics.incWriteTime(writeTimeMs * 1000000L)
@@ -54,13 +56,103 @@ class CelebornListenerSuite {
TaskLocality.PROCESS_LOCAL,
false)
info.markFinished(TaskState.FINISHED, launchTime + durationMs)
- SparkListenerTaskEnd(0, 0, "ShuffleMapTask", Success, info, null, metrics)
+ SparkListenerTaskEnd(0, 0, "ShuffleMapTask", reason, info, null, metrics)
}
private def envUpdate(sparkProps: (String, String)*):
SparkListenerEnvironmentUpdate = {
SparkListenerEnvironmentUpdate(Map("Spark Properties" -> sparkProps.toSeq))
}
+ @Test
+ def aggregateMapperMetricsEvents(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 100L, 40L))
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 250L, 94L))
+ // Unrelated events are ignored.
+ listener.onOtherEvent(new SparkListenerEvent {})
+ listener.flush()
+
+ assertEquals(350L,
statusStore.aggregatedTaskInfo().uncompressedShuffleBytes)
+ assertEquals(134L, statusStore.aggregatedTaskInfo().compressedShuffleBytes)
+ }
+
+ @Test
+ def shuffleWriteMetricsCountEveryAttempt(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ listener.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L))
+ listener.onTaskEnd(newTaskEnd(200L, 20L, 2L, 6L))
+ listener.onTaskEnd(newTaskEnd(300L, 30L, 3L, 7L, TaskKilled("killed by
test")))
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 100L, 40L))
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 200L, 80L))
+ listener.flush()
+
+ val agg = statusStore.aggregatedTaskInfo()
+ assertEquals(600L, agg.shuffleWriteBytes)
+ assertEquals(60L, agg.shuffleWriteTimeMs)
+ assertEquals(6L, agg.shuffleFetchWaitTimeMs)
+ assertEquals(18L, agg.taskDurationMs)
+ assertEquals(300L, agg.uncompressedShuffleBytes)
+ assertEquals(120L, agg.compressedShuffleBytes)
+ }
+
+ @Test
+ def compressionRatioSidesComeFromTheSameMapperEndAttemptSet(): Unit = {
+ // Regression: a task failing after its writer completed already reported
+ // MapperEnd. Both ratio sides derive from MapperEnd reports, so such
attempts
+ // are counted consistently instead of skewing the ratio.
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ // Partition 0 attempt 0: writer completed, then the task failed.
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 20L, 94L))
+ // Partition 0 attempt 1 (retry) and partition 1 attempt 0: succeeded.
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 20L, 94L))
+ listener.onOtherEvent(new CelebornMapperMetricsEvent(0, 20L, 94L))
+ listener.onTaskEnd(newTaskEnd(94L, 10L, 1L, 5L))
+ listener.onTaskEnd(newTaskEnd(94L, 10L, 1L, 5L))
+ listener.onTaskEnd(newTaskEnd(94L, 10L, 1L, 5L, TaskKilled("failed after
write")))
+ listener.flush()
+
+ val agg = statusStore.aggregatedTaskInfo()
+ assertEquals(282L, agg.shuffleWriteBytes)
+ assertEquals(60L, agg.uncompressedShuffleBytes)
+ assertEquals(282L, agg.compressedShuffleBytes)
+ }
+
+ @Test
+ def mapperMetricsEventJsonRoundtrip(): Unit = {
+ // SparkListenerEvent carries @JsonTypeInfo(property = "Event"), so the
event log
+ // records the class name and History Server replay can restore the event
(the
+ // same mechanism JsonProtocol uses for custom SparkListenerEvents).
+ val mapper = new ObjectMapper()
+ val event = new CelebornMapperMetricsEvent(1, 1024L, 512L)
+ val json = mapper.writeValueAsString(event)
+ assertTrue(json.contains(
+ s""""Event":"${classOf[CelebornMapperMetricsEvent].getName}"""))
+ val restored = mapper
+ .readValue(json, classOf[SparkListenerEvent])
+ .asInstanceOf[CelebornMapperMetricsEvent]
+ assertEquals(1, restored.getShuffleId)
+ assertEquals(1024L, restored.getUncompressedBytesWritten)
+ assertEquals(512L, restored.getCompressedBytesWritten)
+
+ // Event logs written before compressedBytesWritten existed still
deserialize.
+ val legacy = mapper.readValue(
+ s"""{"Event":"${classOf[CelebornMapperMetricsEvent].getName}",""" +
+ """"shuffleId":1,"uncompressedBytesWritten":1024}""",
+ classOf[SparkListenerEvent])
+ .asInstanceOf[CelebornMapperMetricsEvent]
+ assertEquals(1024L, legacy.getUncompressedBytesWritten)
+ assertEquals(0L, legacy.getCompressedBytesWritten)
+ }
+
@Test
def flushOnJobEndWithinThrottleInterval(): Unit = {
val store = new InMemoryStore()
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/CelebornMapperMetricsEvent.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/CelebornMapperMetricsEvent.java
new file mode 100644
index 0000000000..6cde6c3abe
--- /dev/null
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/CelebornMapperMetricsEvent.java
@@ -0,0 +1,75 @@
+/*
+ * 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;
+
+/**
+ * Spark listener event carrying per-mapper metrics reported through
MapperEnd, aggregated by the
+ * Celeborn UI extension. Persisted to the event log and replayed by the
History Server: keep the
+ * no-arg constructor and bean getters/setters for Jackson, and add new
metrics as fields so old
+ * event logs and old readers stay compatible.
+ */
+public class CelebornMapperMetricsEvent implements SparkListenerEvent {
+
+ private int shuffleId;
+
+ // Bytes pushed before compression/encryption; 0 means not reported.
+ private long uncompressedBytesWritten;
+
+ // Bytes pushed after compression/encryption; 0 means not reported.
+ private long compressedBytesWritten;
+
+ public CelebornMapperMetricsEvent() {}
+
+ public CelebornMapperMetricsEvent(
+ int shuffleId, long uncompressedBytesWritten, long
compressedBytesWritten) {
+ this.shuffleId = shuffleId;
+ this.uncompressedBytesWritten = uncompressedBytesWritten;
+ this.compressedBytesWritten = compressedBytesWritten;
+ }
+
+ public int getShuffleId() {
+ return shuffleId;
+ }
+
+ public void setShuffleId(int shuffleId) {
+ this.shuffleId = shuffleId;
+ }
+
+ public long getUncompressedBytesWritten() {
+ return uncompressedBytesWritten;
+ }
+
+ public void setUncompressedBytesWritten(long uncompressedBytesWritten) {
+ this.uncompressedBytesWritten = uncompressedBytesWritten;
+ }
+
+ public long getCompressedBytesWritten() {
+ return compressedBytesWritten;
+ }
+
+ public void setCompressedBytesWritten(long compressedBytesWritten) {
+ this.compressedBytesWritten = compressedBytesWritten;
+ }
+
+ @Override
+ public boolean logEvent() {
+ return true;
+ }
+}
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index a36a0932a5..d1e47e66f0 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -51,6 +51,7 @@ import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.write.DataPusher;
import org.apache.celeborn.client.write.PushTask;
import org.apache.celeborn.common.CelebornConf;
+import
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics;
import org.apache.celeborn.common.util.Utils;
@Private
@@ -99,6 +100,9 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private final boolean unsafeRowFastWrite;
+ // Bytes written before compression, for the shuffle compression ratio.
+ private long uncompressedBytesWritten = 0;
+
// In order to facilitate the writing of unit test code, ShuffleClient needs
to be passed in as
// parameters. By the way, simplify the passed parameters.
public HashBasedShuffleWriter(
@@ -303,6 +307,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
mapStatusLengths[partitionId].add(bytesWritten);
writeMetrics.incBytesWritten(bytesWritten);
writeMetrics.incWriteTime(delta);
+ uncompressedBytesWritten += numBytes;
}
private int getOrUpdateOffset(int partitionId, int serializedRecordSize)
@@ -333,6 +338,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
if (logger.isDebugEnabled()) logger.debug("Flush buffer, size {}.",
Utils.bytesToString(size));
dataPusher.addTask(partitionId, buffer, size);
writeMetrics.incWriteTime(System.nanoTime() - start);
+ uncompressedBytesWritten += size;
}
protected void closeWrite() throws IOException {
@@ -365,6 +371,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
numPartitions);
mapStatusLengths[partitionId].add(bytesWritten);
writeMetrics.incBytesWritten(bytesWritten);
+ uncompressedBytesWritten += length;
}
private void cleanupPusher() throws IOException {
@@ -393,7 +400,14 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
long waitStartTime = System.nanoTime();
dataPusher.waitOnTermination();
sendBufferPool.returnPushTaskQueue(dataPusher.getAndResetIdleQueue());
- shuffleClient.mapperEnd(shuffleId, mapId, encodedAttemptId, numMappers,
numPartitions);
+ shuffleClient.mapperEnd(
+ shuffleId,
+ mapId,
+ encodedAttemptId,
+ numMappers,
+ numPartitions,
+ new MapperMetrics(
+ uncompressedBytesWritten,
SparkCommonUtils.sumLongAdders(mapStatusLengths)));
writeMetrics.incWriteTime(System.nanoTime() - waitStartTime);
BlockManagerId bmId = SparkEnv.get().blockManager().shuffleServerId();
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index d08b36d194..2440a23b01 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -47,6 +47,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
+import
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics;
import org.apache.celeborn.common.util.Utils;
@Private
@@ -79,6 +80,9 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private final LongAdder[] mapStatusLengths;
private long tmpRecordsWritten = 0;
+ // Giant-record bytes pushed directly, bypassing the pusher's uncompressed
counter.
+ private long giantRecordBytesWritten = 0;
+
/**
* Are we in the process of stopping? Because map tasks can call stop() with
success = true and
* then call stop() with success = false if they get an exception, we want
to make sure we don't
@@ -370,6 +374,7 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
mapStatusLengths[partitionId].add(bytesWritten);
writeMetrics.incBytesWritten(bytesWritten);
writeMetrics.incWriteTime(delta);
+ giantRecordBytesWritten += numBytes;
}
private void cleanupPusher() throws IOException {
@@ -396,7 +401,15 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
writeMetrics.incRecordsWritten(tmpRecordsWritten);
long waitStartTime = System.nanoTime();
- shuffleClient.mapperEnd(shuffleId, mapId, encodedAttemptId, numMappers,
numPartitions);
+ shuffleClient.mapperEnd(
+ shuffleId,
+ mapId,
+ encodedAttemptId,
+ numMappers,
+ numPartitions,
+ new MapperMetrics(
+ pusher.getUncompressedBytesWritten() + giantRecordBytesWritten,
+ SparkCommonUtils.sumLongAdders(mapStatusLengths)));
writeMetrics.incWriteTime(System.nanoTime() - waitStartTime);
}
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 8d4749945b..7b5a2fbd65 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
@@ -23,6 +23,7 @@ import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.spark.*;
+import org.apache.spark.SparkContext$;
import org.apache.spark.internal.config.ConfigEntry;
import org.apache.spark.internal.config.package$;
import org.apache.spark.launcher.SparkLauncher;
@@ -191,6 +192,22 @@ public class SparkShuffleManager implements ShuffleManager
{
lifecycleManager = new LifecycleManager(appUniqueId, celebornConf);
lifecycleManager.applicationCount().increment();
lifecycleManager.registerCancelShuffleCallback(SparkUtils::cancelShuffle);
+ // 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)) {
+ lifecycleManager.registerMapperEndCallback(
+ (shuffleId, mapId, attemptId, metrics) -> {
+ SparkContext sc =
SparkContext$.MODULE$.getActive().getOrElse(() -> null);
+ if (sc != null) {
+ sc.listenerBus()
+ .post(
+ new CelebornMapperMetricsEvent(
+ shuffleId,
+ metrics.uncompressedBytesWritten(),
+ metrics.compressedBytesWritten()));
+ }
+ });
+ }
if (celebornConf.clientStageRerunEnabled()) {
MapOutputTrackerMaster mapOutputTracker =
(MapOutputTrackerMaster) SparkEnv.get().mapOutputTracker();
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
index 89cb7c3792..0ad8301b16 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
@@ -570,6 +570,26 @@ public class SparkUtils {
}
}
+ /** Lives in the spark-3-ui module, referenced here by class name. */
+ private static final String CELEBORN_PLUGIN_CLASS_NAME =
+ "org.apache.spark.shuffle.celeborn.CelebornPlugin";
+
+ /** Whether the app opted in to the Celeborn UI extension via its plugin
configuration. */
+ public static boolean isCelebornPluginEnabled(SparkConf conf) {
+ String[] keys = {"spark.plugins", "spark.plugins.defaultList"};
+ for (String key : keys) {
+ String value = conf.get(key, null);
+ if (value != null) {
+ for (String plugin : value.split(",")) {
+ if (plugin.trim().equals(CELEBORN_PLUGIN_CLASS_NAME)) {
+ return true;
+ }
+ }
+ }
+ }
+ return false;
+ }
+
/**
* Lazy holder for isCelebornSkewedShuffle method. The method is initialized
only when this class
* is first accessed, ensuring lazy loading without explicit synchronization.
diff --git
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
index b5bacf489d..fd9fb044f6 100644
---
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
+++
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
@@ -83,6 +83,7 @@ import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.CommitMetadata;
import org.apache.celeborn.common.identity.UserIdentifier;
+import
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.PushState;
@@ -288,6 +289,23 @@ public abstract class CelebornShuffleWriterSuiteBase {
int[] crcPerPartition = pushState.getCRC32PerPartition(true,
numPartitions);
long[] bytesPerPartition = pushState.getBytesWrittenPerPartition(true,
numPartitions);
+ MapperMetrics mapperMetrics = client.getMapperEndMetrics(shuffleId, 0, 0);
+ assertNotNull(mapperMetrics);
+ long totalBytesPushed = 0;
+ for (long bytes : bytesPerPartition) {
+ totalBytesPushed += bytes;
+ }
+ long expectedCompressedBytes = 0;
+ for (List<byte[]> batches : client.getCrcDataByPartition(mapKey).values())
{
+ for (byte[] batch : batches) {
+ expectedCompressedBytes +=
DummyShuffleClient.simulatedCompressedLength(batch.length);
+ }
+ }
+ assertEquals(totalBytesPushed, mapperMetrics.uncompressedBytesWritten());
+ // Distinct expectations so a swap of the two fields would fail.
+ assertEquals(expectedCompressedBytes,
mapperMetrics.compressedBytesWritten());
+ assertTrue(expectedCompressedBytes <= totalBytesPushed);
+
Map<Integer, List<byte[]>> crcData = client.getCrcDataByPartition(mapKey);
Map<Integer, List<byte[]>> pushData =
client.getPushDataByPartition(mapKey);
@@ -353,8 +371,8 @@ public abstract class CelebornShuffleWriterSuiteBase {
final CelebornShuffleHandle<Integer, String, String> handle =
new CelebornShuffleHandle<>(
appId, host, port, userIdentifier, shuffleId, false, numMaps,
dependency);
- final ShuffleClient client = new DummyShuffleClient(conf, tempFile);
- ((DummyShuffleClient) client).initReducePartitionMap(shuffleId,
numPartitions, 1);
+ final DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(shuffleId, numPartitions, 1);
final ShuffleWriter<Integer, String> writer =
createShuffleWriter(handle, taskContext, conf, client,
metrics.shuffleWriteMetrics());
@@ -383,8 +401,21 @@ public abstract class CelebornShuffleWriterSuiteBase {
ShuffleWriteMetrics metrics =
taskContext.taskMetrics().shuffleWriteMetrics();
assertEquals(metrics.recordsWritten(), total.intValue());
- assertEquals(metrics.bytesWritten(), tempFile.length());
+ // Write metrics total the per-batch compressed lengths.
+ String mapKey = Utils.makeMapKey(shuffleId, 0, 0);
+ long expectedCompressedBytes = 0;
+ for (List<byte[]> batches :
client.getPushDataByPartition(mapKey).values()) {
+ for (byte[] batch : batches) {
+ expectedCompressedBytes +=
DummyShuffleClient.simulatedCompressedLength(batch.length);
+ }
+ }
+ assertEquals(expectedCompressedBytes, metrics.bytesWritten());
+ assertTrue(metrics.bytesWritten() <= tempFile.length());
assertTrue(metrics.writeTime() > 0);
+ MapperMetrics mapperMetrics = client.getMapperEndMetrics(shuffleId, 0, 0);
+ assertNotNull(mapperMetrics);
+ assertEquals(tempFile.length(), mapperMetrics.uncompressedBytesWritten());
+ assertEquals(metrics.bytesWritten(),
mapperMetrics.compressedBytesWritten());
try (FileInputStream fis = new FileInputStream(tempFile)) {
Iterator it =
newSerializerInstance(serializer).deserializeStream(fis).asKeyValueIterator();
diff --git
a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
index 86e6a2d191..04643e64b9 100644
--- a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -48,6 +48,7 @@ import
org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.client.TransportClientFactory;
import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.protocol.PbStreamHandler;
+import
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.util.ExceptionMaker;
import org.apache.celeborn.common.util.JavaUtils;
@@ -72,6 +73,7 @@ public class DummyShuffleClient extends ShuffleClient {
JavaUtils.newConcurrentHashMap();
private final Map<String, Map<Integer, List<byte[]>>> pushDataByMapKey =
JavaUtils.newConcurrentHashMap();
+ private final Map<String, MapperMetrics> mapperEndMetrics =
JavaUtils.newConcurrentHashMap();
public AtomicInteger fetchFailureCount = new AtomicInteger();
@@ -113,7 +115,12 @@ public class DummyShuffleClient extends ShuffleClient {
.add(Arrays.copyOfRange(data, offset, offset + length));
os.write(data, offset, length);
- return length;
+ return simulatedCompressedLength(length);
+ }
+
+ /** Deterministic fake compression so tests can tell the two byte counts
apart. */
+ public static int simulatedCompressedLength(int length) {
+ return length * 3 / 4;
}
@Override
@@ -159,7 +166,7 @@ public class DummyShuffleClient extends ShuffleClient {
.add(Arrays.copyOfRange(data, offset, offset + length));
os.write(data, offset, length);
- return length;
+ return simulatedCompressedLength(length);
}
@Override
@@ -169,6 +176,17 @@ public class DummyShuffleClient extends ShuffleClient {
public void mapperEnd(
int shuffleId, int mapId, int attemptId, int numMappers, int
numPartitions) {}
+ @Override
+ public void mapperEnd(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int numMappers,
+ int numPartitions,
+ MapperMetrics metrics) {
+ mapperEndMetrics.put(Utils.makeMapKey(shuffleId, mapId, attemptId),
metrics);
+ }
+
@Override
public void readReducerPartitionEnd(
int shuffleId, int partitionId, int startMapIndex, int endMapIndex, int
crc32, long bytes)
@@ -253,6 +271,10 @@ public class DummyShuffleClient extends ShuffleClient {
return pushDataByMapKey.getOrDefault(mapKey, Collections.emptyMap());
}
+ public MapperMetrics getMapperEndMetrics(int shuffleId, int mapId, int
attemptId) {
+ return mapperEndMetrics.get(Utils.makeMapKey(shuffleId, mapId, attemptId));
+ }
+
@Override
public Tuple2<Integer, Boolean> getShuffleId(
int appShuffleId, String appShuffleIdentifier, boolean isWriter, boolean
isBarrierStage) {
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
index 9fb15fdf7a..dd39d5bc61 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -45,6 +45,7 @@ import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.protocol.PbStreamHandler;
import org.apache.celeborn.common.protocol.StorageInfo;
import org.apache.celeborn.common.protocol.message.ControlMessages;
+import
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.util.CelebornHadoopUtils;
import org.apache.celeborn.common.util.ExceptionMaker;
@@ -375,6 +376,22 @@ public abstract class ShuffleClient {
int shuffleId, int mapId, int attemptId, int numMappers, int
numPartitions)
throws IOException;
+ /**
+ * Same as the 5-arg overload, additionally reporting per-mapper write
metrics. The default
+ * implementation discards the metrics, so external subclasses only need the
abstract 5-arg method
+ * to stay compatible.
+ */
+ public void mapperEnd(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int numMappers,
+ int numPartitions,
+ MapperMetrics metrics)
+ throws IOException {
+ mapperEnd(shuffleId, mapId, attemptId, numMappers, numPartitions);
+ }
+
public abstract void readReducerPartitionEnd(
int shuffleId, int partitionId, int startMapIndex, int endMapIndex, int
crc32, long bytes)
throws IOException;
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index f321924737..6b99c5c82b 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -1854,14 +1854,26 @@ public class ShuffleClientImpl extends ShuffleClient {
@Override
public void mapperEnd(int shuffleId, int mapId, int attemptId, int
numMappers, int numPartitions)
throws IOException {
- mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions, -1);
+ mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions, -1,
null);
+ }
+
+ @Override
+ public void mapperEnd(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int numMappers,
+ int numPartitions,
+ MapperMetrics metrics)
+ throws IOException {
+ mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions, -1,
metrics);
}
@Override
public void mapPartitionMapperEnd(
int shuffleId, int mapId, int attemptId, int numMappers, int
numPartitions, int partitionId)
throws IOException {
- mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions,
partitionId);
+ mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions,
partitionId, null);
}
private void mapEndInternal(
@@ -1870,7 +1882,8 @@ public class ShuffleClientImpl extends ShuffleClient {
int attemptId,
int numMappers,
int numPartitions,
- Integer partitionId)
+ Integer partitionId,
+ MapperMetrics metrics)
throws IOException {
final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
PushState pushState = getPushState(mapKey);
@@ -1896,7 +1909,8 @@ public class ShuffleClientImpl extends ShuffleClient {
numPartitions,
crc32PerPartition,
bytesPerPartition,
- SerdeVersion.V1),
+ SerdeVersion.V1,
+ metrics == null ? new MapperMetrics(0L, 0L) : metrics),
rpcMaxRetries,
rpcRetryWait,
ClassTag$.MODULE$.apply(MapperEndResponse.class));
diff --git
a/client/src/main/java/org/apache/celeborn/client/listener/MapperEndCallback.java
b/client/src/main/java/org/apache/celeborn/client/listener/MapperEndCallback.java
new file mode 100644
index 0000000000..7c3d5cd49d
--- /dev/null
+++
b/client/src/main/java/org/apache/celeborn/client/listener/MapperEndCallback.java
@@ -0,0 +1,29 @@
+/*
+ * 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
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics;
+
+/**
+ * Invoked once per accepted mapper attempt; rejected speculative attempts and
retried RPCs of the
+ * same attempt are not reported.
+ */
+public interface MapperEndCallback {
+
+ void onMapperEnd(int shuffleId, int mapId, int attemptId, MapperMetrics
metrics);
+}
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 facf767b01..77ff0715ae 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.WorkerStatusListener
+import org.apache.celeborn.client.listener.{MapperEndCallback,
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}
@@ -468,7 +468,8 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
numPartitions,
crc32PerPartition,
bytesWrittenPerPartition,
- serdeVersion) =>
+ serdeVersion,
+ metrics) =>
logTrace(s"Received MapperEnd TaskEnd request, " +
s"${Utils.makeMapKey(shuffleId, mapId, attemptId)}")
val partitionType = getPartitionType(shuffleId)
@@ -484,7 +485,8 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
numPartitions,
crc32PerPartition,
bytesWrittenPerPartition,
- serdeVersion)
+ serdeVersion,
+ metrics)
case PartitionType.MAP =>
handleMapPartitionEnd(
context,
@@ -968,7 +970,8 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
numPartitions: Int,
crc32PerPartition: Array[Int],
bytesWrittenPerPartition: Array[Long],
- serdeVersion: SerdeVersion): Unit = {
+ serdeVersion: SerdeVersion,
+ metrics: MapperMetrics): Unit = {
val (mapperAttemptFinishedSuccess, allMapperFinished) =
commitManager.finishMapperAttempt(
@@ -980,6 +983,11 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
numPartitions = numPartitions,
crc32PerPartition = crc32PerPartition,
bytesWrittenPerPartition = bytesWrittenPerPartition)
+ // finishMapperAttempt accepts at most one attempt per mapId, so the
callback fires
+ // exactly once per accepted attempt and no extra dedup is needed.
+ if (mapperAttemptFinishedSuccess) {
+ mapperEndCallback.foreach(_.onMapperEnd(shuffleId, mapId, attemptId,
metrics))
+ }
if (mapperAttemptFinishedSuccess && allMapperFinished) {
// last mapper finished. call mapper end
logInfo(s"Last MapperEnd, call StageEnd with shuffleKey:" +
@@ -2178,6 +2186,11 @@ class LifecycleManager(val appUniqueId: String, val
conf: CelebornConf) extends
unregisterShuffleCallback = Some(callback)
}
+ @volatile private var mapperEndCallback: Option[MapperEndCallback] = None
+ def registerMapperEndCallback(callback: MapperEndCallback): Unit = {
+ mapperEndCallback = Some(callback)
+ }
+
def registerAppShuffleDeterminate(appShuffleId: Int, determinate: Boolean):
Unit = {
appShuffleDeterminateMap.put(appShuffleId, determinate)
}
diff --git a/common/src/main/proto/TransportMessages.proto
b/common/src/main/proto/TransportMessages.proto
index e136d3c0c7..dd671914a0 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -384,6 +384,15 @@ message PbPartitionSplit {
PbPartitionLocation oldPartition = 4;
}
+// Extensible per-mapper metrics carried by PbMapperEnd. New mapper-side
metrics
+// must be added as fields here instead of top-level fields of PbMapperEnd.
+message PbMapperMetrics {
+ // Bytes pushed before compression/encryption; 0 means not reported.
+ int64 uncompressedBytesWritten = 1;
+ // Bytes pushed after compression/encryption; 0 means not reported.
+ int64 compressedBytesWritten = 2;
+}
+
message PbMapperEnd {
int32 shuffleId = 1;
int32 mapId = 2;
@@ -394,6 +403,7 @@ message PbMapperEnd {
int32 numPartitions = 7;
repeated int32 crc32PerPartition = 8;
repeated int64 bytesWrittenPerPartition = 9;
+ PbMapperMetrics metrics = 10;
}
message PbLocationPushFailedBatches {
diff --git
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
index 580291d34b..795dbaaa8f 100644
---
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -209,6 +209,11 @@ object ControlMessages extends Logging {
newLocs: util.Map[Integer, (StatusCode, java.lang.Boolean,
PartitionLocation)],
serdeVersion: SerdeVersion) extends MasterMessage
+ // Extensible per-mapper metrics carried by MapperEnd, mirrors
PbMapperMetrics.
+ case class MapperMetrics(
+ uncompressedBytesWritten: Long = 0L,
+ compressedBytesWritten: Long = 0L)
+
case class MapperEnd(
shuffleId: Int,
mapId: Int,
@@ -219,7 +224,8 @@ object ControlMessages extends Logging {
numPartitions: Int,
crc32PerPartition: Array[Int],
bytesWrittenPerPartition: Array[Long],
- serdeVersion: SerdeVersion)
+ serdeVersion: SerdeVersion,
+ metrics: MapperMetrics = MapperMetrics())
extends MasterMessage
case class ReadReducerPartitionEnd(
@@ -748,7 +754,8 @@ object ControlMessages extends Logging {
numPartitions,
crc32PerPartition,
bytesWrittenPerPartition,
- serdeVersion) =>
+ serdeVersion,
+ metrics) =>
val pushFailedMap = pushFailedBatch.asScala.map { case (k, v) =>
val resultValue = PbSerDeUtils.toPbLocationPushFailedBatches(v)
(k, resultValue)
@@ -764,6 +771,9 @@ object ControlMessages extends Logging {
.addAllCrc32PerPartition(crc32PerPartition.map(Integer.valueOf).toSeq.asJava)
.addAllBytesWrittenPerPartition(bytesWrittenPerPartition.map(
java.lang.Long.valueOf).toSeq.asJava)
+ .setMetrics(PbMapperMetrics.newBuilder()
+ .setUncompressedBytesWritten(metrics.uncompressedBytesWritten)
+ .setCompressedBytesWritten(metrics.compressedBytesWritten))
.build().toByteArray
new TransportMessage(MessageType.MAPPER_END, payload, serdeVersion)
@@ -1274,7 +1284,10 @@ object ControlMessages extends Logging {
pbMapperEnd.getNumPartitions,
crc32Array,
bytesWrittenPerPartitionArray,
- message.getSerdeVersion)
+ message.getSerdeVersion,
+ MapperMetrics(
+ pbMapperEnd.getMetrics.getUncompressedBytesWritten,
+ pbMapperEnd.getMetrics.getCompressedBytesWritten))
case READ_REDUCER_PARTITION_END_VALUE =>
val pbReadReducerPartitionEnd =
PbReadReducerPartitionEnd.parseFrom(message.getPayload)
diff --git
a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
index 7f1725383a..488ebfbef0 100644
--- a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
@@ -35,7 +35,7 @@ import org.apache.celeborn.common.meta.WorkerInfo
import org.apache.celeborn.common.metrics.{ClientMetric, MetricType}
import org.apache.celeborn.common.network.protocol.SerdeVersion
import org.apache.celeborn.common.protocol.{PartitionLocation,
PbRequestWorkers, PbRequestWorkersResponse, PbReviseLostShuffles,
PbReviseLostShufflesResponse, StorageInfo, TransportModuleConstants}
-import
org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFilesResponse,
GetReducerFileGroupResponse, HeartbeatFromApplication, MapperEnd,
ReviseLostShuffles, ReviseLostShufflesResponse}
+import
org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFilesResponse,
GetReducerFileGroupResponse, HeartbeatFromApplication, MapperEnd,
MapperMetrics, ReviseLostShuffles, ReviseLostShufflesResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
class UtilsSuite extends CelebornFunSuite {
@@ -186,7 +186,8 @@ class UtilsSuite extends CelebornFunSuite {
1,
Array.emptyIntArray,
Array.emptyLongArray,
- SerdeVersion.V1)
+ SerdeVersion.V1,
+ MapperMetrics(123456L, 789L))
val mapperEndTrans =
Utils.fromTransportMessage(Utils.toTransportMessage(mapperEnd)).asInstanceOf[MapperEnd]
assert(mapperEnd.shuffleId == mapperEndTrans.shuffleId)
@@ -196,6 +197,7 @@ class UtilsSuite extends CelebornFunSuite {
assert(mapperEnd.partitionId == mapperEndTrans.partitionId)
assert(mapperEnd.failedBatchSet == mapperEndTrans.failedBatchSet)
assert(mapperEnd.numPartitions == mapperEndTrans.numPartitions)
+ assert(mapperEnd.metrics == mapperEndTrans.metrics)
mapperEnd.crc32PerPartition.array should contain theSameElementsInOrderAs
mapperEndTrans.crc32PerPartition
mapperEnd.bytesWrittenPerPartition.array should contain
theSameElementsInOrderAs mapperEndTrans.bytesWrittenPerPartition
}
diff --git a/docker/spark/conf/spark-defaults.conf
b/docker/spark/conf/spark-defaults.conf
index 9692f5ab6f..139e699961 100644
--- a/docker/spark/conf/spark-defaults.conf
+++ b/docker/spark/conf/spark-defaults.conf
@@ -43,6 +43,9 @@ spark.shuffle.service.enabled false
spark.celeborn.client.spark.shuffle.writer hash
spark.celeborn.client.push.replicate.enabled true
+# Register the Celeborn UI extension (Celeborn tab on the Spark application
UI).
+spark.plugins
org.apache.spark.shuffle.celeborn.CelebornPlugin
+
# AQE local shuffle reader is incompatible with remote (Celeborn) shuffle.
spark.sql.adaptive.enabled true
spark.sql.adaptive.localShuffleReader.enabled false
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 b4453c83f9..ade825596f 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
@@ -19,6 +19,8 @@ package org.apache.celeborn.tests.client
import java.util
+import scala.collection.JavaConverters._
+
import org.scalatest.concurrent.Eventually.eventually
import org.scalatest.concurrent.Futures.{interval, timeout}
import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
@@ -26,6 +28,7 @@ import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
import org.apache.celeborn.client.{LifecycleManager, WithShuffleClientSuite}
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.identity.UserIdentifier
+import
org.apache.celeborn.common.protocol.message.ControlMessages.MapperMetrics
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.service.deploy.MiniClusterFeature
@@ -128,6 +131,58 @@ class LifecycleManagerSuite extends WithShuffleClientSuite
with MiniClusterFeatu
lifecycleManager.stop()
}
+ test("CELEBORN-2479: mapper end callback reports only the accepted attempt")
{
+ prepareService()
+ val numMappers = 8
+ val numPartitions = 9
+ val shuffleId = nextShuffleId
+ val mapId = 1
+ val attemptId = 0
+
+ val reported = new util.concurrent.ConcurrentLinkedQueue[(Long, Long)]()
+ lifecycleManager.registerMapperEndCallback((_, _, _, metrics) =>
+ reported.add((metrics.uncompressedBytesWritten,
metrics.compressedBytesWritten)))
+
+ // Register a REDUCE-partition shuffle.
+ shuffleClient.getPartitionLocation(shuffleId, numMappers, numPartitions)
+
+ // Accepted attempt.
+ shuffleClient.mapperEnd(
+ shuffleId,
+ mapId,
+ attemptId,
+ numMappers,
+ numPartitions,
+ new MapperMetrics(100L, 40L))
+ // Retried RPC of the same attempt: finishMapperAttempt is idempotent, not
reported again.
+ shuffleClient.mapperEnd(
+ shuffleId,
+ mapId,
+ attemptId,
+ numMappers,
+ numPartitions,
+ new MapperMetrics(100L, 40L))
+ // Duplicate speculative attempt: rejected by the commit handler, not
reported.
+ shuffleClient.mapperEnd(
+ shuffleId,
+ mapId,
+ attemptId + 1,
+ numMappers,
+ numPartitions,
+ new MapperMetrics(250L, 94L))
+ // Retried RPC of the rejected attempt: not reported again.
+ shuffleClient.mapperEnd(
+ shuffleId,
+ mapId,
+ attemptId + 1,
+ numMappers,
+ numPartitions,
+ new MapperMetrics(250L, 94L))
+
+ // The callback runs synchronously before the RPC reply.
+ assert(reported.asScala.toList == List((100L, 40L)))
+ }
+
override def afterAll(): Unit = {
logInfo("all test complete , stop celeborn mini cluster")
shutdownMiniCluster()
