This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/branch-0.7 by this push:
new 00889ddf71 [CELEBORN-2032][FOLLOWUP] Disable replica preference for
skewed partitions without map range
00889ddf71 is described below
commit 00889ddf71ce53b85030c72941f1de8730b15414
Author: lijianfu03 <[email protected]>
AuthorDate: Wed Aug 12 15:15:49 2026 +0800
[CELEBORN-2032][FOLLOWUP] Disable replica preference for skewed partitions
without map range
### What changes were proposed in this pull request?
CELEBORN-2032 introduced attempt-based primary/replica switching
(`preferReplicaRead = context.attemptNumber % 2 == 1`) in
`CelebornShuffleReader` so that odd-numbered task attempts prefer reading the
replica `PartitionLocation` instead of the primary, improving fault tolerance
across retries/speculative execution.
This PR disables that replica preference specifically when
`celeborn.client.adaptive.optimizeSkewedPartitionRead.enabled` is on and the
partition is being read as a skewed partition without map range
(`splitSkewPartitionWithoutMapRange`). In that mode, all attempts for the same
skewed partition will now consistently read the primary (or
previously-resolved) locations instead of alternating between primary and
replica.
A unit test
(`CelebornPartitionUtilSuiteJ#testSkewPartitionSplitDiffersBetweenPrimaryAndReplicaChunkOffsets`)
is added to demonstrate the root cause directly against
`CelebornPartitionUtil#splitSkewedPartitionLocations`.
### Why are the changes needed?
When `celeborn.client.adaptive.optimizeSkewedPartitionRead.enabled=true`, a
skewed reduce partition is not read by map-id range. Instead, Celeborn treats
all `PartitionLocation`s of that partition as one logical byte stream and
splits it into `subPartitionSize` sub-partitions purely by byte offset
(`CelebornPartitionUtil#splitSkewedPartitionLocations`). For a given
`subPartitionIndex`, this method computes a `chunkRange` (physical chunk index
interval) by walking the `chunkOffsets` of [...]
The primary and its replica are flushed independently by two different
Workers. Even though they hold logically identical data and share the same
`uniqueId`, their physical `chunkOffsets` (the byte positions at which each
flush produced a new chunk) are **not guaranteed to be identical**.
Because of CELEBORN-2032, whether a task attempt reads the primary or the
replica depends on `attemptNumber % 2`. So:
- Attempt 0 (first run) reads the primary and resolves `chunkRange` from
the primary's chunk offsets.
- Attempt 1 (retry / speculative execution) reads the replica and resolves
`chunkRange` from the replica's (possibly different) chunk offsets.
For the exact same logical `subPartitionIndex`, this can produce two
**different physical byte ranges**, e.g. primary resolves to chunk range `[2,
3]` while replica resolves to `[3, 3]` (dropping chunk 2 entirely). The two
attempts then read different bytes for what should be the identical logical
sub-partition, so their computed byte-count/CRC diverge and fail
`SkewHandlingWithoutMapRangeValidator`, surfacing
as:org.apache.celeborn.common.exception.CelebornIOException: AQE Partition
[...]
<img width="1711" height="605" alt="image"
src="https://github.com/user-attachments/assets/b4bff79e-7287-4bb5-8ba1-f2bf7cf21f20";
/>
### Does this PR resolve a correctness bug?
- [x] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
- Added
`CelebornPartitionUtilSuiteJ#testSkewPartitionSplitDiffersBetweenPrimaryAndReplicaChunkOffsets`,
which constructs a primary and a replica `PartitionLocation` sharing the same
`uniqueId` but with different `chunkOffsets` (simulating independent flush
behavior), and asserts that
`CelebornPartitionUtil#splitSkewedPartitionLocations` resolves different chunk
ranges (`[2, 3]` vs `[3, 3]`) for the identical `subPartitionIndex`, proving
the root cause of the non-idempotent read.
- Ran the full `CelebornPartitionUtilSuiteJ` (5 tests) and
`CelebornShuffleReaderSuite` (10 tests) in `client-spark/spark-3` — all passed.
Closes #3798 from buska88/celeborn-2032-fix.
Authored-by: lijianfu03 <[email protected]>
Signed-off-by: 子懿 <[email protected]>
(cherry picked from commit 87aae3fa7a5ef8a2e31ebeb673c18a3a61e4a964)
Signed-off-by: 子懿 <[email protected]>
---
.../shuffle/celeborn/CelebornShuffleReader.scala | 10 +++-
.../celeborn/CelebornPartitionUtilSuiteJ.java | 64 ++++++++++++++++++++++
2 files changed, 73 insertions(+), 1 deletion(-)
diff --git
a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
index 87bf9ffd6e..c0c2b20ec8 100644
---
a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
+++
b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
@@ -291,11 +291,19 @@ class CelebornShuffleReader[K, C](
if (fileGroups.partitionGroups.containsKey(partitionId)) {
// CELEBORN-2032. For the first time of open stream and
// attemptNumber % 2 = 1, we should read the replica data first.
+ // Note: when splitSkewPartitionWithoutMapRange is enabled, the chunk
range is
+ // computed against a single physical location (primary or replica)
based on the
+ // byte offsets recorded by that location's own flush. Since primary
and replica
+ // don't guarantee identical chunk offsets, switching which replica is
read across
+ // task attempts (retries/speculation) can make
splitSkewedPartitionLocations resolve
+ // a different physical chunk range for the same logical
sub-partition, causing
+ // CRC/byte-count mismatches on retry. So we skip the replica
preference in that case
+ // and always stick with the primary (or previously chosen) location.
val originLocations = fileGroups.partitionGroups.get(partitionId)
val hasReplicate = pushReplicateEnabled &&
originLocations.asScala.exists(p => p != null && p.hasPeer)
var locations =
- if (preferReplicaRead && hasReplicate) {
+ if (preferReplicaRead && hasReplicate &&
!splitSkewPartitionWithoutMapRange) {
originLocations.asScala.map { p =>
if (p != null && p.hasPeer) p.getPeer else p
}.asJava
diff --git
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornPartitionUtilSuiteJ.java
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornPartitionUtilSuiteJ.java
index 06c431bf83..98ad1fd24d 100644
---
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornPartitionUtilSuiteJ.java
+++
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornPartitionUtilSuiteJ.java
@@ -115,6 +115,70 @@ public class CelebornPartitionUtilSuiteJ {
}
}
+ /**
+ * CELEBORN-2032 combined with skewed-partition reading
(readSkewPartitionWithoutMapRange):
+ * splitSkewedPartitionLocations resolves a logical sub-partition's chunk
range purely from the
+ * chunk offsets of whichever physical PartitionLocation instance is passed
in. The primary and
+ * its replica are flushed independently by two different Workers, so their
chunk offsets are not
+ * guaranteed to match even though they hold the same data and share the
same uniqueId.
+ *
+ * <p>This test simulates a first attempt (reads primary) and a
retry/speculative attempt (reads
+ * replica, per CELEBORN-2032's odd-attemptNumber-prefers-replica policy)
for the exact same
+ * logical sub-partition index, and shows that the resolved chunk range
differs between the two:
+ * reading the primary resolves to chunk range [2, 3], while reading the
replica for the very same
+ * sub-partition resolves to [3, 3] and drops chunk 2 entirely. If the
caller were to switch
+ * between primary/replica across attempts (as happened before this fix),
the bytes/CRC actually
+ * read would differ across attempts and fail the AQE skew validation
+ * (SkewHandlingWithoutMapRangeValidator) on retry.
+ */
+ @Test
+ public void
testSkewPartitionSplitDiffersBetweenPrimaryAndReplicaChunkOffsets() {
+ // Primary and replica hold the same logical data (same total file size
900) but flush
+ // independently, resulting in different physical chunk boundaries around
the sub-partition
+ // split point (step = 900 / 3 = 300).
+ PartitionLocation primary =
+ genPartitionLocation(0, new Long[] {0L, 100L, 200L, 340L, 600L, 900L});
+ PartitionLocation replica =
+ genPartitionLocation(0, new Long[] {0L, 100L, 200L, 260L, 600L, 900L});
+ Assert.assertEquals(
+ "primary and replica must represent the same logical partition",
+ primary.getUniqueId(),
+ replica.getUniqueId());
+
+ int subPartitionSize = 3;
+ int subPartitionIndex = 1; // e.g. the sub-partition assigned to this
reduce task
+
+ Map<String, Pair<Integer, Integer>> primaryResult =
+ CelebornPartitionUtil.splitSkewedPartitionLocations(
+ new ArrayList<>(Collections.singletonList(primary)),
+ subPartitionSize,
+ subPartitionIndex);
+ Map<String, Pair<Integer, Integer>> replicaResult =
+ CelebornPartitionUtil.splitSkewedPartitionLocations(
+ new ArrayList<>(Collections.singletonList(replica)),
+ subPartitionSize,
+ subPartitionIndex);
+
+ Map<String, Pair<Integer, Integer>> expectedPrimaryRange =
+ genRanges(new Object[][] {{"0-0", 2, 3}});
+ Map<String, Pair<Integer, Integer>> expectedReplicaRange =
+ genRanges(new Object[][] {{"0-0", 3, 3}});
+ Assert.assertEquals(expectedPrimaryRange, primaryResult);
+ Assert.assertEquals(expectedReplicaRange, replicaResult);
+
+ // The core bug: for the identical (uniqueId, subPartitionIndex), the
chunk range resolved
+ // depends on which replica's chunk offsets were used, so it is NOT
idempotent across
+ // attempts unless the same replica is consistently read every time. Here
the replica-based
+ // range [3, 3] is a strict subset of the primary-based range [2, 3] and
drops chunk 2
+ // entirely, which is exactly the kind of mismatch that fails
+ // SkewHandlingWithoutMapRangeValidator when different attempts read
different replicas.
+ Assert.assertNotEquals(
+ "chunk range must not depend on which replica is read, otherwise
retries/speculative "
+ + "attempts will read a different byte range for the same logical
sub-partition",
+ primaryResult.get(primary.getUniqueId()),
+ replicaResult.get(replica.getUniqueId()));
+ }
+
@Test
public void testSplitStable() {
ArrayList<PartitionLocation> locations = new ArrayList<>();
