[CELEBORN-2479] Show shuffle compression ratio in Spark UI by yew1eb · Pull Request #3872 · apache/celeborn

yew1eb · 2026-10-13T15:43:10Z

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 page shows N/A when the required metrics are unavailable, such as with old clients.
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?

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.

zaynt4606

Thanks for adding the compression-ratio metric. I found two accounting gaps that can make the displayed ratio incorrect. Please address the inline comments and add regression coverage for these paths.

zaynt4606 · 2026-10-14T02:48:00Z

+        encodedAttemptId,
+        numMappers,
+        numPartitions,
+        new MapperMetrics(pusher.getUncompressedBytesWritten()));


Could we include bytes sent by pushGiantRecord() in this metric? That path pushes directly through shuffleClient.pushDataWithCRC() and increments writeMetrics.bytesWritten, but never updates SortBasedPusher's uncompressed-byte counter, which is the sole source used here. If all records exceed pushBufferMaxSize, this reports zero and the UI shows N/A despite a successful write. Please account for the direct-push path and cover it with a giant-record test.

Good catch. pushGiantRecord now increments a writer-level giant-bytes counter that is summed with the pusher's counter in the MapperMetrics reported at mapperEnd, mirroring HashBasedShuffleWriter. The writer suites now assert the reported uncompressed bytes equal the bytes that reached the shuffle client — the all-giant-record tests (testGiantRecord, testGiantRecordWithFastWrite, which force every record above pushBufferMaxSize) fail without the fix, and the mixed-path checkWithIntegrity covers giant + buffered records together.

zaynt4606 · 2026-10-14T02:48:00Z

+    event match {
+      case e: CelebornMapperMetricsEvent =>
+        if (pluginEnabled) {
+          totalUncompressedWriteBytes.addAndGet(e.getUncompressedBytesWritten)


Could the numerator and denominator use the same set of mapper attempts? onTaskEnd adds shuffleWriteMetrics.bytesWritten for every task end with metrics, while LifecycleManager.handleMapperEnd emits this event only when finishMapperAttempt accepts that attempt. If a failed or speculative attempt writes data and a later attempt succeeds, its compressed bytes remain in the denominator but its uncompressed bytes are absent from the numerator, understating the ratio. Please align the accounting and test a retry/duplicate-attempt case.

Aligned. Both sides now use the same set of mapper attempts:

The listener counts shuffleWriteMetrics only for successful task ends (failed/killed attempts push bytes whose MapperEnd is never reported).

LifecycleManager reports metrics for every successfully closed writer attempt — including speculative/rerun attempts whose MapperEnd the commit handler rejects — deduplicated by attemptId so a retried MapperEnd RPC of the same attempt is not counted twice.

Covered by CelebornListenerSuite.shuffleWriteMetricsCountOnlySuccessfulAttempts (success + killed attempts + mapper events) and a LifecycleManagerSuite test driving the accepted / retried / duplicate-attempt cases over real RPC.

Thanks for the update. There is still a gap in the assumption that failed attempts never report MapperEnd: the writer can finish and report MapperEnd before Spark finishes the task. A task-completion listener can then throw, causing that attempt to fail and be retried.

I reproduced this on ea7a4fa with Spark 3.5.8, local[2,2], two map partitions, and an injected task-completion-listener failure on partition 0 / attempt 0 after the shuffle writer completes. The retry succeeds and the job result is correct, but the metrics are:

MapperMetrics events (uncompressed bytes): 20, 20, 20, including the failed attempt.

Successful map attempts (compressed bytes): 94, 94.

Failed map attempt (compressed bytes): 94, excluded by onTaskEnd.

UI totals: 60 / 188, rather than 40 / 188 for the successful attempts, overstating the ratio by 50%.

The onTaskEnd and mapper-event accumulation code, both writers, and the event payload are unchanged in 96f8cc8. The event still lacks map/attempt identity to reconcile it with the final task outcome. Could we derive both counters from the same set of attempts (with enough identity to reconcile them if using task-success filtering) and add a regression where the task fails after writer completion? This is distinct from deduplicating MapperEnd RPC retries.

Addressed in cd1636e by deriving both counters from the same attempt set: writers now also report post-compression bytes via a new additive PbMapperMetrics.compressedBytesWritten field, and the listener accumulates both ratio sides from MapperEnd reports. In your repro the page would now show 60/282 — the same ratio as the successful-only 40/188 — instead of the skewed 60/188. Added a listener regression for the fail-after-write scenario, plus JSON round-trip coverage for event logs written before the new field existed.

Track the bytes each mapper hands to ShuffleClient before compression and show the aggregated compression ratio (uncompressed/compressed) on the Celeborn tab of the Spark UI. - Writers report the bytes at mapperEnd via a new extensible MapperMetrics message carried by PbMapperEnd (field 10), keeping MapperEnd source compatible through a default value. - LifecycleManager forwards the metrics to a registered MapperEndCallback; SparkShuffleManager posts CelebornMapperMetricsEvent on the driver listener bus so the value is available in the live UI and replayable by the History Server from the event log. - CelebornListener aggregates the events and CelebornShufflePage renders the ratio; unreported values degrade to N/A. - Register CelebornPlugin in the docker spark-defaults.conf so the local docker-compose cluster exercises the UI end-to-end.

yew1eb · 2026-10-20T09:33:09Z

@zaynt4606

zaynt4606

Thanks for syncing with main. I checked 96f8cc8: the listener currently fails to compile, and the two accounting paths discussed below still need changes. The runtime reproductions were run on ea7a4fa; the relevant handleMapperEnd, task-end accounting, mapper-event handler, and writer code are unchanged in this update. I verified the duplicate-method compilation error on 96f8cc8; I could not rerun the runtime probes on this head until that compilation error is fixed.

zaynt4606 · 2026-10-21T03:32:19Z

    mayUpdate()
  }

+  override def onOtherEvent(event: SparkListenerEvent): Unit = {


[P1] Combine the two onOtherEvent overrides

After merging main, this new override coexists with the override at line 127 that handles CelebornBuildInfoEvent from #3870. Compiling the current listener at 96f8cc8 fails with method onOtherEvent is defined twice (lines 83 and 127). Please put the CelebornMapperMetricsEvent and CelebornBuildInfoEvent cases into a single onOtherEvent override and retain coverage for both event types.

Fixed in cd1636e — the two overrides are merged into a single onOtherEvent that handles both CelebornBuildInfoEvent and CelebornMapperMetricsEvent; the existing tests cover both paths.

zaynt4606 · 2026-10-21T03:32:19Z

+    if (mapperAttemptFinishedSuccess ||
+      attemptId != commitManager.getMapperAttempts(shuffleId)(mapId)) {
+      mapperEndCallback.foreach(_.onMapperEnd(shuffleId, mapId, attemptId, metrics))


[P2] Deduplicate retried rejected attempts as well

getMapperAttempts(shuffleId)(mapId) retains the accepted attempt ID. If attempt 0 is accepted and speculative attempt 1 is rejected, every retry of attempt 1 still satisfies 1 != 0, so its metrics are emitted repeatedly.

A real-RPC probe on ea7a4fa sent MapperEnd for attempts 0, 0, 1, 1 with metrics 100, 100, 250, 250. The callbacks were 100, 250, 250 (600 total), rather than 100, 250 (350 total). This entire handleMapperEnd method is unchanged in 96f8cc8. The existing test retries the accepted attempt, so it misses this case.

Could we deduplicate every reported (shuffleId, mapId, attemptId), including rejected attempts, clean up the bookkeeping with the shuffle, and add a regression that retries the rejected attempt?

Addressed in cd1636e. LifecycleManager now tracks every reported (shuffleId, mapId, attemptId) — including rejected attempts — and invokes the callback exactly once per finished writer attempt; retried MapperEnd RPCs of rejected attempts are suppressed, and the bookkeeping is removed with the shuffle in removeExpiredShuffle. The regression test now replays your probe (attempts 0,0,1,1 with metrics 100,100,250,250) and asserts exactly [100, 250].

Now that compressed bytes also come through MapperEnd, do we still need to report rejected attempts?

If we only report when mapperAttemptFinishedSuccess is true, both sides of the ratio still use the same attempt set, and finishMapperAttempt already dedups by mapId. Then this set can go away.

It's a boxed Long per mapper on the driver, and it's populated even when no callback is registered.

Agreed, adopted in 90d35c1. Since both ratio sides travel in the same MapperMetrics payload, reporting only the accepted attempt keeps them aligned, finishMapperAttempt already dedups by mapId (same-attempt retries and rejected attempts all return false), and the dedup set is gone along with its expiry cleanup. The regression test now replays your probe (attempts 0, 0, 1, 1 / 100, 100, 250, 250) and asserts exactly [100].

…ted mapper attempts - Merge the two onOtherEvent overrides in CelebornListener (compile error after syncing with main): handle CelebornBuildInfoEvent and CelebornMapperMetricsEvent in a single override. - LifecycleManager tracks every reported (shuffleId, mapId, attemptId) so each finished writer attempt is reported exactly once, including retried MapperEnd RPCs of rejected speculative attempts; bookkeeping is cleaned up with the shuffle. - Writers also report post-compression bytes in MapperMetrics (new PbMapperMetrics.compressedBytesWritten, additive), so both sides of the compression ratio derive from the same MapperEnd attempt set. A task that fails after its writer completed (e.g. completion listener failure) no longer skews the ratio.

yew1eb · 2026-10-21T07:08:22Z

Hi @zaynt4606, thanks for the detailed review and the runtime reproductions! All three comments are addressed in cd1636e — please see the inline replies for details.

PTAL, thanks!

ForVic · 2026-10-22T00:49:58Z

          lifecycleManager.applicationCount().increment();
          lifecycleManager.registerCancelShuffleCallback(SparkUtils::cancelShuffle);
+          // Forward mapper metrics to the Celeborn UI extension.
+          lifecycleManager.registerMapperEndCallback(


This is registered unconditionally, so every app logs one event per map task even without CelebornPlugin enabled. EventLoggingListener also flushes the event log for each custom event.

Could we gate this on the plugin? Or, building on the LifecycleManager comment, aggregate per shuffle and post once at StageEnd.

Gated in da2330f — the callback is now registered only when SparkUtils.isCelebornPluginEnabled(conf) finds the plugin class in spark.plugins or spark.plugins.defaultList (the same keys the listener checks on replay), so apps that never opted in no longer post one event per map task or flush the event log per event.

Aggregating per shuffle and posting once at StageEnd would additionally cut the event volume for opted-in apps; I have left it as a follow-up to keep this change small.

ForVic · 2026-10-22T00:50:58Z

-      totalWriteTimeMs.addAndGet(metrics.shuffleWriteMetrics.writeTime / 1000000L)
+      taskEnd.reason match {
+        case Success =>
+          // Write-side metrics only for successful attempts. The compression ratio


I think this is left over from the earlier fix. The ratio comes entirely from MapperEnd now.

It also makes "Shuffle Duration / Task Duration" inconsistent: write time only counts successful attempts, but task duration still counts all of them. And Spark's AppStatusListener counts metrics from failed/killed attempts too, so this doesn't match Spark's stage accounting either. Can we revert it?

You are right, it was left over from the earlier split-source ratio — reverted in 5643bb7. Write bytes/time now count every attempt like the read side and task duration, matching AppStatusListener's stage accounting. shuffleWriteMetricsCountEveryAttempt covers a killed attempt being counted.

ForVic · 2026-10-22T00:51:59Z

  // Report partition locations written by the completed map task of ReducePartition Shuffle Type.
  public abstract void mapperEnd(
-      int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions)
+      int shuffleId,


Making the new overload abstract breaks external ShuffleClient subclasses, and it's the opposite of what the description says.

Could we keep the 5-arg one abstract and make this one a default that ignores the metrics?

Fixed in 45ee042 — the 5-arg mapperEnd is abstract again, and the new 6-arg overload is a default method that discards the metrics and delegates to the 5-arg one, so external subclasses keep compiling with just the 5-arg implementation. ShuffleClientImpl overrides the 6-arg one to pass the metrics through.

ForVic · 2026-10-22T00:52:20Z

+    for (long bytes : bytesPerPartition) {
+      totalBytesPushed += bytes;
+    }
+    assertEquals(totalBytesPushed, mapperMetrics.uncompressedBytesWritten());


DummyShuffleClient returns the input length, so uncompressed == compressed here and a swap in the writer would still pass.

Could the dummy client return a different byte count so these assertions actually distinguish the two?

Good catch — and it was slightly worse than that: the SortBased async path pushes through DataPusher's bare pushData (CRC already recorded at enqueue time), which bypassed the WithCRC overloads entirely. Fixed in 45ee042 by making the dummy's pushData/mergeData themselves return length * 3 / 4 — like the real client, every push return value is a post-compression byte count — so all paths agree. The suites now assert uncompressed == bytes pushed and compressed == Σ 3/4·batch separately; I verified a temporary swap of the two MapperMetrics fields turns 6 of 13 SortBased tests red.

…ics tests distinguish compression sides Keep the 5-arg mapperEnd abstract as before and make the new 6-arg overload a default method that discards the metrics and delegates to the 5-arg one, so external ShuffleClient subclasses stay source- and binary-compatible; ShuffleClientImpl implements both and DummyShuffleClient records the reported metrics. The dummy client now simulates a deterministic compression (length * 3/4) in its pushData/mergeData return values, covering the async DataPusher bare-pushData path as well, so the writer suites can assert the uncompressed and compressed byte counts separately; swapping the MapperMetrics fields fails the tests. All push call sites return post-compression bytes like the real client. Verified: HashBasedShuffleWriterSuiteJ 12/12, SortBasedShuffleWriterSuiteJ 13/13, SortBasedPusherSuiteJ 1/1; a temporary swap of the MapperMetrics fields turned 6 of 13 SortBased tests red, restored green after revert.

…y the commit handler Since both sides of the compression ratio now travel in the same MapperMetrics payload, only the accepted attempt needs to be reported. finishMapperAttempt already accepts at most one attempt per mapId (retried and duplicate attempts return false), so the callback fires exactly once per accepted attempt and the reportedMapperAttempts dedup set — one boxed Long per mapper, populated even with no callback registered — is removed along with its expiry cleanup. The regression test now replays attempts 0, 0, 1, 1 with metrics 100, 100, 250, 250 and asserts the callback observes exactly [100].

…gin is enabled The callback posts one CelebornMapperMetricsEvent per map task to the listener bus even when no CelebornPlugin is configured, and EventLoggingListener serializes and flushes each custom event to the event log. Gate the registration on SparkUtils.isCelebornPluginEnabled, which checks spark.plugins and spark.plugins.defaultList for the plugin class — the same keys the CelebornListener checks on replay. Aggregating per shuffle and posting once at StageEnd is left as a follow-up.

The Success-only filter was needed when the compressed side of the ratio came from onTaskEnd; the ratio now derives both sides from MapperEnd reports, so the filter only made Shuffle Write inconsistent with the other aggregates (read metrics and task duration count every attempt) and with Spark's own stage accounting in AppStatusListener, which accumulates task metrics regardless of the end reason. Revert to counting every attempt. Verified: CelebornListenerSuite 12/12.

ForVic

one minor optional thing but lgtm:

The MapperEndCallback Javadoc is also out of date. It still says “including duplicate speculative attempts”; it should say the callback runs once per accepted attempt.

zaynt4606

LGTM. The previous accounting and compatibility concerns are addressed on e077148. Both sides of the compression ratio now come from the same accepted MapperEnd attempt, and the targeted UI listener tests pass.

zaynt4606 · 2026-10-23T02:32:06Z

Thanks @yew1eb for addressing the review feedback! I have approved e077148 and updated the PR description to match the current implementation.

Please sync this branch with the latest main and resolve the import conflict in SparkShuffleManager.java, then let CI run on the updated head before merging. The current SBT service failure is the existing Master HA test race in exportedGauge (head of empty list), unrelated to this change.

…n-ratio # Conflicts: # client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java

…ttempt

zaynt4606

LGTM. Rechecked ddf4b46 after syncing main and correcting the MapperEndCallback Javadoc. The import conflict is resolved, the reviewed compression-ratio behavior is preserved, and all 109 checks pass.

zaynt4606 · 2026-10-23T06:04:20Z

Merge Summary:

merged into main 3cb708b0

Posted by merge_pr.py

zaynt4606 · 2026-10-23T06:05:55Z

Thanks @yew1eb for the contribution!
Merged to main(v1.0.0).

Commit: 3cb708b

yew1eb force-pushed the spark-ui-compression-ratio branch 2 times, most recently from b6b7269 to 737d4fb Compare October 13, 2026 16:12

github-actions Bot added module:client module:spark kind:build module:common labels Oct 13, 2026

zaynt4606 requested changes Oct 14, 2026

View reviewed changes

yew1eb force-pushed the spark-ui-compression-ratio branch 5 times, most recently from fcd341e to 22b8ad3 Compare October 14, 2026 09:00

yew1eb force-pushed the spark-ui-compression-ratio branch from 22b8ad3 to ea7a4fa Compare October 14, 2026 09:11

github-actions Bot added the module:tests label Oct 14, 2026

Merge branch 'main' into spark-ui-compression-ratio

96f8cc8

zaynt4606 reviewed Oct 21, 2026

View reviewed changes

yew1eb force-pushed the spark-ui-compression-ratio branch from 6db7a32 to cd1636e Compare October 21, 2026 07:06

ForVic reviewed Oct 22, 2026

View reviewed changes

yew1eb added 5 commits October 22, 2026 19:46

[CELEBORN-2479] Trim comments

e077148

ForVic approved these changes Oct 22, 2026

View reviewed changes

zaynt4606 approved these changes Oct 23, 2026

View reviewed changes

yew1eb added 2 commits October 23, 2026 10:32

Merge remote-tracking branch 'upstream/main' into spark-ui-compressio…

63642fc

…n-ratio # Conflicts: # client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java

[CELEBORN-2479] Update MapperEndCallback Javadoc: once per accepted a…

ddf4b46

…ttempt

zaynt4606 approved these changes Oct 23, 2026

View reviewed changes

zaynt4606 closed this in 3cb708b Oct 23, 2026

Conversation

yew1eb commented Oct 13, 2026 • edited by zaynt4606 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR resolve a correctness bug?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

zaynt4606 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yew1eb commented Oct 20, 2026

Uh oh!

zaynt4606 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yew1eb commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ForVic left a comment

Choose a reason for hiding this comment

Uh oh!

zaynt4606 left a comment

Choose a reason for hiding this comment

Uh oh!

zaynt4606 commented Oct 23, 2026

Uh oh!

zaynt4606 left a comment

Choose a reason for hiding this comment

Uh oh!

zaynt4606 commented Oct 23, 2026

Uh oh!

zaynt4606 commented Oct 23, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

yew1eb commented Oct 13, 2026 •

edited by zaynt4606

Loading

yew1eb commented Oct 21, 2026 •

edited

Loading