You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by ni...@apache.org on 2026/08/12 07:16:29 UTC
(celeborn) branch branch-0.7 updated: [CELEBORN-2032][FOLLOWUP] Disable replica preference for skewed partitions without map range
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 <li...@meituan.com>
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 <li...@meituan.com>
Signed-off-by: 子懿 <pr...@163.com>
(cherry picked from commit 87aae3fa7a5ef8a2e31ebeb673c18a3a61e4a964)
Signed-off-by: 子懿 <pr...@163.com>
---
.../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<>();