[CELEBORN-2370] Scope reducer metadata by partition range

[sunchao]

JIRA: CELEBORN-2370

Supersedes #3687.

Why are the changes needed?

Before a Spark reducer can read shuffle data, Celeborn calls GetReducerFileGroup to obtain reducer file locations and related metadata.

Today that response is shuffle-wide. For a shuffle with N reducers, it contains metadata for all N reducers, even though a Spark task normally reads only [startPartition, endPartition).

For example:

• a shuffle has 1,000,000 reducers;
• a task reads only [42, 43); but
• the executor still downloads and materializes metadata for all 1,000,000 reducers.

The response is cached independently on every executor, so metadata transfer and executor memory grow with the total reducer count multiplied by the number of executors, rather than with the reducer ranges those executors actually read. On very large shuffles this can cause:

• large driver RPC responses;
• repeated transfer of the same shuffle-wide metadata;
• executor heap and GC pressure; and
• task threads waiting behind the same shuffle-wide metadata load.

The driver still needs to retain complete shuffle commit metadata. This PR reduces the metadata transferred to and cached by each executor.

What changes were proposed in this pull request?

Add optional partition-range fields to GetReducerFileGroup requests and responses.

When Spark reads a shuffle:

1. The reader sends its actual [startPartition, endPartition) range.
2. The driver returns only the reducer file groups, successful partition IDs, and failed-batch metadata that intersect that range.
3. The executor records which ranges it has loaded and requests only missing ranges.

The executor cache also:

• shares one in-flight RPC between callers requesting the same cold range;
• lets unrelated cached ranges proceed independently;
• fetches mapper-attempt metadata on the first request and omits it from later range requests;
• prevents an in-flight request from repopulating state after shuffle cleanup;
• preserves interruption and failure propagation; and
• treats an unscoped response from an older driver as complete shuffle metadata.

Partition-scoped responses bypass the existing shuffle-wide RPC cache and Spark broadcast path. Legacy full-shuffle requests continue to use the existing cache and broadcast behavior.

This PR changes the Spark client and the Spark driver-side lifecycle/commit endpoint. It does not change the Celeborn Master or Worker.

Does this PR introduce any user-facing change?

No configuration or public API change is required.

When both the Spark client and driver contain this change, reducer metadata transfer and executor cache size become proportional to the partition ranges read by that executor instead of the total reducer count.

The wire protocol remains backward compatible:

• New client with old driver: the old driver ignores the optional request fields and returns full metadata; the new client detects the unscoped response and caches it as complete.
• Old client with new driver: unset range fields select the existing full-response cache and broadcast path.
• Flink clients continue to use the legacy full-shuffle request.

Mixed-version deployments remain correct, but the optimization applies only when the driver supports scoped responses.

How was this patch tested?

• ShuffleClientSuiteJ: 29 tests covering range caching, concurrent loads, cleanup races, interruption, and old-driver fallback
• UtilsSuite: 28 tests, including V1 and V2 protocol round trips
• ReducerFileGroupFilterSuite: 2 range-filtering tests
• ConfigurationSuite: 8 tests
• Spark 4.0 / Scala 2.13 focused integration run: CelebornHashSuite and CelebornSortSuite (4 tests)
• Spark 4.0 / Scala 2.13 package compile
• Flink 1.20 common client package compile
• Spotless apply/check

[afterincomparableyum]

I'll take a look at this soon. It's an interesting PR.

[Copilot]

Pull request overview

This PR introduces partition-range scoping for GetReducerFileGroup so Spark executors request/cache only the reducer metadata needed for their [startPartition, endPartition) read range, reducing RPC payload size and executor-side memory/GC pressure on very large shuffles while keeping wire compatibility with legacy drivers/clients.

Changes:

• Extend GetReducerFileGroup request/response (protobuf + transport serde) with optional partition-range fields and an omitMapAttempts optimization.
• Add executor-side range cache/single-flight loading in ShuffleClientImpl and update Spark reader to request metadata by partition range (bypassing broadcast for scoped requests).
• Add/adjust unit and integration tests, plus documentation updates clarifying broadcast applies to legacy full-shuffle responses only.

Reviewed changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornSortSuite.scala	Updates Spark IT assertion to ensure scoped requests bypass broadcast.
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashSuite.scala	Same as above for hash shuffle IT.
docs/configuration/client.md	Clarifies broadcast configs apply only to legacy shuffle-wide requests.
common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala	Adds round-trip serde tests for new request/response fields and interruption test.
common/src/main/scala/org/apache/celeborn/common/util/Utils.scala	Preserves interruption during retry backoff by restoring interrupt status and throwing InterruptedException.
common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala	Adds transport (de)serialization for new partition-range fields.
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala	Updates config docs to note broadcast is legacy-only.
common/src/main/proto/TransportMessages.proto	Extends protobuf messages for range scoping and omitMapAttempts.
client/src/test/scala/org/apache/celeborn/client/commit/ReducerFileGroupFilterSuite.scala	Adds tests for reducer metadata filtering by partition range.
client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java	Adds extensive concurrency/range-cache tests for ShuffleClientImpl.updateFileGroup behavior.
client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala	Plumbs range fields through the driver endpoint and validates range inputs.
client/src/main/scala/org/apache/celeborn/client/CommitManager.scala	Passes new range parameters through to commit handlers.
client/src/main/scala/org/apache/celeborn/client/commit/ReducerFileGroupFilter.scala	Implements server-side filtering helpers for file groups/partition IDs/failed batches.
client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala	Builds/scopes reducer metadata responses and bypasses broadcast for scoped requests.
client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala	Applies range filtering for mapper-side handler responses and supports omitting map attempts.
client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala	Extends handler API to accept range + omit flags.
client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java	Adds executor-side range cache with in-flight sharing, cleanup safety, and old-driver fallback.
client/src/main/java/org/apache/celeborn/client/ShuffleClient.java	Adds a range overload (default implementation) for updateFileGroup.
client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java	Implements the new range overload as a no-op for dummy client.
client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala	Switches Spark reader to request reducer metadata for the actual partition range.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sunchao]

I'll take a look at this soon. It's an interesting PR.

Thanks @afterincomparableyum and @SteNicholas! We ran into this issue internally. In our case, a Spark job with an extremely large number of reducers (over 1 million) caused the Spark driver to become stuck for hours without making any progress.

[github-actions]

This PR is stale because it has been open 20 days with no activity. Remove stale label or comment or this will be closed in 10 days.

[SteNicholas]

@sunchao, please rebase the latest main branch for stable CI.

[Copilot]

Pull request overview

Copilot reviewed 20 out of 20 changed files in this pull request and generated no new comments.

Suppressed comments (1)

client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala:996

• The driver-side range validation rejects endPartition == startPartition (endPartition <= startPartition), but the client-side range API treats an empty [start, end) interval as valid and returns an empty result without loading metadata. Rejecting empty ranges makes the RPC contract less robust/inconsistent (e.g., if another client forwards an empty range to the driver) and turns a well-defined query into REQUEST_FAILED.

Consider allowing empty ranges by treating only endPartition < startPartition as invalid (still rejecting negative starts). The commit handlers’ range filtering already naturally returns an empty map for an empty interval.

    if (hasPartitionRange && (startPartition < 0 || endPartition <= startPartition)) {
      logWarning(
        s"Reject invalid reducer file group partition range [$startPartition, $endPartition) " +
          s"for shuffle $shuffleId")
      context.reply(

[sunchao]

Thanks @SteNicholas . Just did it.

[Copilot]

Pull request overview

Copilot reviewed 20 out of 20 changed files in this pull request and generated 1 comment.