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 c5e6a803c8 [CELEBORN-2491] Fix mapper completion racing with shuffle
integrity metadata aggregation
c5e6a803c8 is described below
commit c5e6a803c89fc839f3184788de1ba9e2992a7a3f
Author: zhengtao <[email protected]>
AuthorDate: Thu Oct 22 20:39:02 2026 +0800
[CELEBORN-2491] Fix mapper completion racing with shuffle integrity
metadata aggregation
### What changes were proposed in this pull request?
Aggregate each accepted mapper's checksums and byte counts under the
existing per-shuffle lock, before publishing its accepted attempt and
completed-mapper count.
Add six deterministic regression tests covering completion visibility, RPC
retries, speculative attempts, the last mapper arriving during an earlier
aggregation, progress across different shuffles, and disabled integrity checks.
### Why are the changes needed?
`finishMapperAttempt` currently marks a mapper complete before aggregating
its integrity metadata outside the lock. During that window, another mapper can
report stage completion, and a repeated request for the pending mapper can
return. `LifecycleManager` replies `SUCCESS` even when that repeated request
did not win the attempt.
A retry or speculative task can therefore be acknowledged before the
accepted request finishes recording its metadata. If another mapper has already
triggered `StageEnd`, downstream validation can observe incomplete expected
checksums/byte counts and report a false mismatch.
The patch makes aggregation and completion publication one critical
section. Duplicate attempts are still counted once, and the lock remains scoped
to a single shuffle. With integrity checks enabled, aggregation within that
shuffle becomes serialized; different shuffles retain independent locks.
The aggregation path only updates in-memory byte counters and checksums; it
performs no RPC or disk I/O. It extends the existing per-shuffle critical
section by O(number of partitions) when integrity checking is enabled.
If aggregation throws before completion publication, the mapper is no
longer marked accepted prematurely. This does not provide rollback: metadata
already accumulated for earlier partitions remains, so malformed-input
validation and failure atomicity still require a separate follow-up. This PR
does not claim that every aggregation exception can be retried safely.
### Does this PR resolve a correctness bug?
- [x] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
Fixes premature completion visibility when shuffle integrity checking is
enabled. The feature still defaults to disabled; configuration and wire formats
are unchanged.
### How was this patch tested?
- On unmodified production sources at
`fca9da23687c4d484f449abce9f87ed66eaf2ae1`, four concurrency regressions failed
and the two compatibility/isolation checks passed. All six pass with this fix.
The final test harness additionally checks the monitor owner to distinguish
shuffle-lock blocking from unrelated contention.
- All 37 Scala tests in `org.apache.celeborn.client.commit` and all 17
tests in `ShuffleClientSuiteJ` pass. After review, the test fixture uses a
package-scoped `VisibleForTesting` accessor instead of reflection into a
private field; all 54 tests also pass after that cleanup.
- The full repository formatter, `dev/reformat`, completed successfully for
all 16 profiles. Common/client compilation, the Maven lifecycle Spotless
checks, and `git diff --check` pass.
- Java 17, Scala 2.12.18, Maven 3.9.16. The linked-worktree git-commit-id
metadata plugin was skipped with `-Dmaven.gitcommitid.skip=true`.
- Validation is at the handler/client unit-test level; no full Spark E2E
workload or throughput benchmark was run.
```sh
dev/reformat
build/mvn -B -ntp -pl client -am test \
-Dtest=ShuffleClientSuiteJ \
-DwildcardSuites=org.apache.celeborn.client.commit \
-Dmaven.gitcommitid.skip=true
```
Closes #3883 from zaynt4606/fix-mapper-end-integrity-race.
Authored-by: zhengtao <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
(cherry picked from commit e21b72656cecaa9c83eba3d4c9893371dbc31bb9)
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../commit/ReducePartitionCommitHandler.scala | 37 ++-
.../commit/ReducePartitionCommitHandlerTest.scala | 297 +++++++++++++++++++++
2 files changed, 320 insertions(+), 14 deletions(-)
diff --git
a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
index 68aec0c977..62c03a3877 100644
---
a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
@@ -25,6 +25,7 @@ import java.util.function
import scala.collection.JavaConverters._
import scala.collection.mutable
+import com.google.common.annotations.VisibleForTesting
import com.google.common.base.Preconditions.checkState
import com.google.common.cache.{Cache, CacheBuilder}
import com.google.common.collect.Sets
@@ -304,6 +305,11 @@ class ReducePartitionCommitHandler(
shuffleMapperAttempts.get(shuffleId)
}
+ @VisibleForTesting
+ private[commit] def getReducerCommitMetadata(shuffleId: Int):
Array[CommitMetadata] = {
+ commitMetadataForReducer.get(shuffleId)
+ }
+
override def areAllMapperAttemptsFinished(shuffleId: Int): Boolean = {
val attempts = shuffleMapperAttempts.get(shuffleId)
if (null != attempts) {
@@ -339,6 +345,23 @@ class ReducePartitionCommitHandler(
val attempts = shuffleMapperAttempts.get(shuffleId)
if (attempts(mapId) < 0) {
+ // Aggregate metadata under the shuffle lock before publishing mapper
completion.
+ // Otherwise another MapperEnd can finish the stage or acknowledge a
duplicate
+ // while this mapper's integrity metadata is still incomplete.
+ if (shuffleIntegrityCheckEnabled) {
+ val commitMetadataArray = commitMetadataForReducer.get(shuffleId)
+ checkState(
+ commitMetadataArray != null,
+ "commitMetadataArray can only be null if shuffleId %s is not
registered!",
+ shuffleId)
+ for (i <- 0 until numPartitions) {
+ if (bytesWrittenPerPartition(i) != 0) {
+ commitMetadataArray(i).addCommitData(
+ crc32PerPartition(i),
+ bytesWrittenPerPartition(i))
+ }
+ }
+ }
attempts(mapId) = attemptId
// increment completed mappers
val completedMappers =
@@ -362,20 +385,6 @@ class ReducePartitionCommitHandler(
(false, false)
}
}
- if (shuffleIntegrityCheckEnabled && mapperAttemptFinishedSuccess) {
- val commitMetadataArray = commitMetadataForReducer.get(shuffleId)
- checkState(
- commitMetadataArray != null,
- "commitMetadataArray can only be null if shuffleId %s is not
registered!",
- shuffleId)
- for (i <- 0 until numPartitions) {
- if (bytesWrittenPerPartition(i) != 0) {
- commitMetadataArray(i).addCommitData(
- crc32PerPartition(i),
- bytesWrittenPerPartition(i))
- }
- }
- }
if (!mapperAttemptFinishedSuccess) {
val another = shuffleMapperAttempts.get(shuffleId)(mapId)
logInfo(
diff --git
a/client/src/test/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandlerTest.scala
b/client/src/test/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandlerTest.scala
new file mode 100644
index 0000000000..3e84c444c3
--- /dev/null
+++
b/client/src/test/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandlerTest.scala
@@ -0,0 +1,297 @@
+/*
+ * 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.commit
+
+import java.lang.management.ManagementFactory
+import java.util.concurrent.{Callable, CountDownLatch, FutureTask,
ScheduledExecutorService, ThreadPoolExecutor, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
+
+import org.scalatest.concurrent.Eventually
+import org.scalatest.time.{Millis, Seconds, Span}
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.client.CommitManager.CommittedPartitionInfo
+import org.apache.celeborn.client.LifecycleManager.ShuffleAllocatedWorkers
+import org.apache.celeborn.client.WorkerStatusTracker
+import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
+import org.apache.celeborn.common.util.ThreadUtils
+
+class ReducePartitionCommitHandlerTest extends CelebornFunSuite with
Eventually {
+ implicit override val patienceConfig: PatienceConfig =
+ PatienceConfig(timeout = Span(10, Seconds), interval = Span(10, Millis))
+
+ private var rpcPool: ThreadPoolExecutor = _
+ private var commitScheduler: ScheduledExecutorService = _
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ rpcPool = ThreadUtils.newDaemonCachedThreadPool("test-reduce-commit-rpc")
+ commitScheduler =
+
ThreadUtils.newDaemonSingleThreadScheduledExecutor("test-reduce-commit-scheduler")
+ }
+
+ override def afterAll(): Unit = {
+ try {
+ if (rpcPool != null) {
+ ThreadUtils.shutdown(rpcPool)
+ }
+ if (commitScheduler != null) {
+ ThreadUtils.shutdown(commitScheduler)
+ }
+ } finally {
+ super.afterAll()
+ }
+ }
+
+ private def newHandler(integrityCheckEnabled: Boolean = true):
ReducePartitionCommitHandler = {
+ val conf = new CelebornConf()
+ .set(CelebornConf.CLIENT_SHUFFLE_INTEGRITY_CHECK_ENABLED,
integrityCheckEnabled)
+ new ReducePartitionCommitHandler(
+ "test-app",
+ conf,
+ new ShuffleAllocatedWorkers(),
+ new CommittedPartitionInfo(),
+ new WorkerStatusTracker(conf, null),
+ rpcPool,
+ commitScheduler,
+ null)
+ }
+
+ private def finishMapper(
+ handler: ReducePartitionCommitHandler,
+ shuffleId: Int = 1,
+ mapId: Int = 0,
+ attemptId: Int = 0,
+ numMappers: Int = 1,
+ checksums: Array[Int] = Array(10, 20),
+ bytes: Array[Long] = Array(100L, 200L)): (Boolean, Boolean) = {
+ handler.finishMapperAttempt(
+ shuffleId,
+ mapId,
+ attemptId,
+ numMappers,
+ partitionId = -1,
+ pushFailedBatches = null,
+ recordWorkerFailure = _ => (),
+ numPartitions = 2,
+ checksums,
+ bytes)
+ }
+
+ private class PausedCommitMetadata extends CommitMetadata {
+ val entered = new CountDownLatch(1)
+ val resume = new CountDownLatch(1)
+ private val pauseNextUpdate = new AtomicBoolean(true)
+
+ override def addCommitData(checksum: Int, numBytes: Long): Unit = {
+ if (pauseNextUpdate.getAndSet(false)) {
+ entered.countDown()
+ assert(resume.await(30, TimeUnit.SECONDS), "metadata aggregation was
not resumed")
+ }
+ super.addCommitData(checksum, numBytes)
+ }
+ }
+
+ private def pauseMetadata(
+ handler: ReducePartitionCommitHandler,
+ numMappers: Int = 1): PausedCommitMetadata = {
+ handler.registerShuffle(1, numMappers, isSegmentGranularityVisible =
false, numPartitions = 2)
+ // Pause the last partition so the first partition has already been
aggregated.
+ val paused = new PausedCommitMetadata()
+ handler.getReducerCommitMetadata(1)(1) = paused
+ paused
+ }
+
+ private def startMapper(
+ body: => (Boolean, Boolean)): (Thread, FutureTask[(Boolean, Boolean)]) =
{
+ val task = new FutureTask[(Boolean, Boolean)](new Callable[(Boolean,
Boolean)] {
+ override def call(): (Boolean, Boolean) = body
+ })
+ val thread = new Thread(task, "test-mapper-end")
+ thread.setDaemon(true)
+ thread.start()
+ (thread, task)
+ }
+
+ private def join(thread: Thread): Unit = {
+ thread.join(10000)
+ assert(!thread.isAlive, "MapperEnd thread did not exit")
+ }
+
+ private def awaitBlockedOrFinished(
+ thread: Thread,
+ task: FutureTask[(Boolean, Boolean)],
+ lockOwner: Thread): Unit = {
+ eventually {
+ val info = ManagementFactory.getThreadMXBean.getThreadInfo(thread.getId)
+ assert(task.isDone || (info != null &&
+ info.getThreadState == Thread.State.BLOCKED && info.getLockOwnerId ==
lockOwner.getId))
+ }
+ }
+
+ private def checkMetadata(
+ handler: ReducePartitionCommitHandler,
+ checksums: Array[Int],
+ bytes: Array[Long]): Unit = {
+ checksums.indices.foreach { partitionId =>
+ val (valid, message) = handler.finishPartition(
+ 1,
+ partitionId,
+ 0,
+ Int.MaxValue,
+ new CommitMetadata(checksums(partitionId), bytes(partitionId)))
+ assert(valid, message)
+ }
+ }
+
+ test("mapper completion is published only after all integrity metadata is
aggregated") {
+ val handler = newHandler()
+ val paused = pauseMetadata(handler)
+ val (thread, task) = startMapper(finishMapper(handler))
+ try {
+ assert(paused.entered.await(10, TimeUnit.SECONDS))
+ assert(!handler.isMapperEnded(1, 0))
+ assert(!handler.areAllMapperAttemptsFinished(1))
+ assert(handler.getMapperAttempts(1).sameElements(Array(-1)))
+
+ paused.resume.countDown()
+ assert(task.get(10, TimeUnit.SECONDS) == ((true, true)))
+ assert(handler.isMapperEnded(1, 0))
+ assert(handler.areAllMapperAttemptsFinished(1))
+ checkMetadata(handler, Array(10, 20), Array(100L, 200L))
+ } finally {
+ paused.resume.countDown()
+ join(thread)
+ }
+ }
+
+ Seq((0, "retried"), (1, "speculative")).foreach { case (attemptId,
description) =>
+ test(s"a $description MapperEnd cannot return before the accepted attempt
records metadata") {
+ val handler = newHandler()
+ val paused = pauseMetadata(handler)
+ val (acceptedThread, acceptedTask) = startMapper(finishMapper(handler))
+ var duplicateThread: Thread = null
+ try {
+ assert(paused.entered.await(10, TimeUnit.SECONDS))
+ val (thread, duplicateTask) = startMapper {
+ // A speculative attempt may produce different bytes; only the
accepted attempt counts.
+ finishMapper(
+ handler,
+ attemptId = attemptId,
+ checksums = if (attemptId == 0) Array(10, 20) else Array(30, 40),
+ bytes = if (attemptId == 0) Array(100L, 200L) else Array(300L,
400L))
+ }
+ duplicateThread = thread
+ awaitBlockedOrFinished(thread, duplicateTask, acceptedThread)
+ // LifecycleManager replies SUCCESS even for a duplicate MapperEnd.
Returning here
+ // would let the mapper finish while the accepted attempt's metadata
is incomplete.
+ assert(!duplicateTask.isDone, "duplicate MapperEnd returned before
metadata was aggregated")
+
+ paused.resume.countDown()
+ assert(acceptedTask.get(10, TimeUnit.SECONDS) == ((true, true)))
+ assert(duplicateTask.get(10, TimeUnit.SECONDS) == ((false, false)))
+ checkMetadata(handler, Array(10, 20), Array(100L, 200L))
+ } finally {
+ paused.resume.countDown()
+ join(acceptedThread)
+ if (duplicateThread != null) {
+ join(duplicateThread)
+ }
+ }
+ }
+ }
+
+ test("another mapper cannot finish the stage before earlier mapper metadata
is aggregated") {
+ val handler = newHandler()
+ val paused = pauseMetadata(handler, numMappers = 2)
+ val (firstThread, firstTask) = startMapper(finishMapper(handler,
numMappers = 2))
+ var lastThread: Thread = null
+ try {
+ assert(paused.entered.await(10, TimeUnit.SECONDS))
+ val (thread, lastTask) = startMapper {
+ finishMapper(
+ handler,
+ mapId = 1,
+ numMappers = 2,
+ checksums = Array(30, 40),
+ bytes = Array(300L, 400L))
+ }
+ lastThread = thread
+ awaitBlockedOrFinished(thread, lastTask, firstThread)
+ assert(!lastTask.isDone, "the last MapperEnd returned before earlier
metadata was aggregated")
+ assert(!handler.areAllMapperAttemptsFinished(1))
+
+ paused.resume.countDown()
+ assert(firstTask.get(10, TimeUnit.SECONDS) == ((true, false)))
+ assert(lastTask.get(10, TimeUnit.SECONDS) == ((true, true)))
+ assert(handler.areAllMapperAttemptsFinished(1))
+ checkMetadata(handler, Array(40, 60), Array(400L, 600L))
+ } finally {
+ paused.resume.countDown()
+ join(firstThread)
+ if (lastThread != null) {
+ join(lastThread)
+ }
+ }
+ }
+
+ test("metadata aggregation does not block MapperEnd for another shuffle") {
+ val handler = newHandler()
+ val paused = pauseMetadata(handler)
+ val (firstThread, firstTask) = startMapper(finishMapper(handler))
+ var otherThread: Thread = null
+ try {
+ assert(paused.entered.await(10, TimeUnit.SECONDS))
+ val (thread, otherTask) = startMapper(finishMapper(handler, shuffleId =
2))
+ otherThread = thread
+ assert(otherTask.get(10, TimeUnit.SECONDS) == ((true, true)))
+ assert(handler.areAllMapperAttemptsFinished(2))
+ assert(!firstTask.isDone)
+ paused.resume.countDown()
+ assert(firstTask.get(10, TimeUnit.SECONDS) == ((true, true)))
+ } finally {
+ paused.resume.countDown()
+ join(firstThread)
+ if (otherThread != null) {
+ join(otherThread)
+ }
+ }
+ }
+
+ test("MapperEnd accepts empty metadata arrays when integrity checking is
disabled") {
+ val handler = newHandler(integrityCheckEnabled = false)
+ assert(finishMapper(
+ handler,
+ numMappers = 2,
+ checksums = Array.emptyIntArray,
+ bytes = Array.emptyLongArray) == ((true, false)))
+ assert(finishMapper(
+ handler,
+ numMappers = 2,
+ attemptId = 1,
+ checksums = Array.emptyIntArray,
+ bytes = Array.emptyLongArray) == ((false, false)))
+ assert(finishMapper(
+ handler,
+ mapId = 1,
+ numMappers = 2,
+ checksums = Array.emptyIntArray,
+ bytes = Array.emptyLongArray) == ((true, true)))
+ assert(handler.areAllMapperAttemptsFinished(1))
+ }
+}
