This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 71724db239 [CELEBORN-2361] Unblock stage-end when CommitFilesResponse
carries an empty map-id bitmap
71724db239 is described below
commit 71724db23980cc80cb920fee2a057b96f26893f7
Author: strelok89 <[email protected]>
AuthorDate: Mon Oct 12 17:47:47 2026 +0800
[CELEBORN-2361] Unblock stage-end when CommitFilesResponse carries an empty
map-id bitmap
This revives #3736 by afterincomparableyum, which was closed by the stale
bot rather than
rejected. Rebased onto current `main`, credit preserved via
`Co-Authored-By`.
A worker serializes an empty `RoaringBitmap` as an empty protobuf
bytestring, and
`Utils.byteStringToRoaringBitmap` deserializes empty bytes back to `null`.
That `null` was
stored in `CommitFilesResponse.committedMapIdBitMap` (a `HashMap`, which
allows null values)
and then merged into the `ConcurrentHashMap`
`ShuffleCommittedInfo.committedMapIdBitmap` via
`putAll`, which rejects null values and throws NPE in
`CommitHandler.processResponse`.
Since `ReducePartitionCommitHandler.tryFinalCommit` had no `try/finally`,
the shuffle was
left in `inProcessStageEndShuffleSet` and `setStageEnd` was never called,
so every
`GetReducerFileGroup` request parked until the executor RPC timeout,
hanging the Spark
client process.
Fixes:
- `ControlMessages`: skip null bitmaps when deserializing
`CommitFilesResponse` so `null`
never reaches `ConcurrentHashMap.putAll`.
- `ReducePartitionCommitHandler`: wrap the commit in `try/finally` so the
in-process
stage-end marker is always cleared, allowing a retry.
The original revision of this PR also made
`CelebornInputStream.skipLocation` treat a null
mapId bitmap as "do not skip". That has since landed on `main` via
CELEBORN-2442
(`adadbd924`), which moved the logic into
`CelebornInputStream.shouldSkipLocation` with the
same null-safe behavior, so it is dropped here.
The null-means-unknown semantics are kept deliberately. The wire format
collapses null and
empty to `ByteString.EMPTY`, so making `byteStringToRoaringBitmap` return
an empty bitmap
would turn "unknown" into "no map ids here" at the `PbSerDeUtils` call
sites and make
`shouldSkipLocation` skip locations that must be read.
Without this, Spark clients can hang for a long time and emit repeated NPEs.
- [x] Yes
- [ ] Yes
Added a `UtilsSuite` regression test asserting a `CommitFilesResponse` with
an empty bitmap
survives the transport round-trip without null map values. Verified it
fails without the
`ControlMessages` fix and passes with it. `celeborn-client/Test/compile` and
`spotless:check -pl common,client` pass locally; the rest is left to CI.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes #3833 from strelok89/CELEBORN-2361.
Authored-by: strelok89 <[email protected]>
Signed-off-by: zhengtao <[email protected]>
Co-Authored-By: afterincomparableyum
<[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Co-Authored-By: Codex <[email protected]>
AI-Model: gpt-6
AI-Contributed/Feature: 0/44
AI-Contributed/UT: 0/30
---
.../commit/ReducePartitionCommitHandler.scala | 37 ++++++++++++----------
.../common/protocol/message/ControlMessages.scala | 7 +++-
.../apache/celeborn/common/util/UtilsSuite.scala | 30 +++++++++++++++++-
3 files changed, 56 insertions(+), 18 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 b3e7aa90ab..febee90f8d 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
@@ -210,23 +210,28 @@ class ReducePartitionCommitHandler(
}
}
- // ask allLocations workers holding partitions to commit files
- val allocatedWorkers = shuffleAllocatedWorkers.get(shuffleId)
- val (dataLost, commitFailedWorkers) = handleFinalCommitFiles(shuffleId,
allocatedWorkers)
- recordWorkerFailure(commitFailedWorkers)
- // reply
- if (!dataLost) {
- logInfo(s"Succeed to handle stageEnd for $shuffleId.")
- // record in stageEndShuffleSet
- setStageEnd(shuffleId)
- } else {
- logError(s"Failed to handle stageEnd for $shuffleId, lost file!")
- dataLostShuffleSet.add(shuffleId)
- // record in stageEndShuffleSet
- setStageEnd(shuffleId)
+ try {
+ // ask allLocations workers holding partitions to commit files
+ val allocatedWorkers = shuffleAllocatedWorkers.get(shuffleId)
+ val (dataLost, commitFailedWorkers) = handleFinalCommitFiles(shuffleId,
allocatedWorkers)
+ recordWorkerFailure(commitFailedWorkers)
+ // reply
+ if (!dataLost) {
+ logInfo(s"Succeed to handle stageEnd for $shuffleId.")
+ // record in stageEndShuffleSet
+ setStageEnd(shuffleId)
+ } else {
+ logError(s"Failed to handle stageEnd for $shuffleId, lost file!")
+ dataLostShuffleSet.add(shuffleId)
+ // record in stageEndShuffleSet
+ setStageEnd(shuffleId)
+ }
+ true
+ } finally {
+ // Clear the in-process marker even if the commit threw, so stage-end
can be retried
+ // instead of parking every GetReducerFileGroup until the RPC timeout.
+ inProcessStageEndShuffleSet.remove(shuffleId)
}
- inProcessStageEndShuffleSet.remove(shuffleId)
- true
}
private def handleFinalCommitFiles(
diff --git
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
index b305539b65..580c986ae6 100644
---
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -1492,7 +1492,12 @@ object ControlMessages extends Logging {
pbCommitFilesResponse.getCommittedReplicaStorageInfosMap.asScala.foreach(entry
=>
committedReplicaStorageInfos.put(entry._1,
StorageInfo.fromPb(entry._2)))
pbCommitFilesResponse.getMapIdBitmapMap.asScala.foreach { entry =>
- committedBitMap.put(entry._1,
Utils.byteStringToRoaringBitmap(entry._2))
+ // An empty bitmap is serialized as ByteString.EMPTY and
deserializes back to null.
+ // Skip it: null values break the ConcurrentHashMap.putAll in
processResponse.
+ val bitmap = Utils.byteStringToRoaringBitmap(entry._2)
+ if (bitmap != null) {
+ committedBitMap.put(entry._1, bitmap)
+ }
}
CommitFilesResponse(
StatusCode.fromValue(pbCommitFilesResponse.getStatus),
diff --git
a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
index a3e03f7c9b..7b115f044d 100644
--- a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
@@ -22,6 +22,7 @@ import java.util
import java.util.Collections
import java.util.concurrent.atomic.AtomicInteger
+import org.roaringbitmap.RoaringBitmap
import org.scalatest.matchers.must.Matchers.contain
import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
@@ -33,7 +34,7 @@ import
org.apache.celeborn.common.identity.{DefaultIdentityProvider, UserIdentif
import org.apache.celeborn.common.meta.WorkerInfo
import org.apache.celeborn.common.network.protocol.SerdeVersion
import org.apache.celeborn.common.protocol.{PartitionLocation,
PbRequestWorkers, PbRequestWorkersResponse, PbReviseLostShuffles,
PbReviseLostShufflesResponse, StorageInfo, TransportModuleConstants}
-import
org.apache.celeborn.common.protocol.message.ControlMessages.{GetReducerFileGroupResponse,
MapperEnd, ReviseLostShuffles, ReviseLostShufflesResponse}
+import
org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFilesResponse,
GetReducerFileGroupResponse, MapperEnd, ReviseLostShuffles,
ReviseLostShufflesResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
class UtilsSuite extends CelebornFunSuite {
@@ -311,6 +312,33 @@ class UtilsSuite extends CelebornFunSuite {
assert(set.size == 0)
}
+ test("CommitFilesResponse with empty map-id bitmap survives transport
round-trip") {
+ // An empty RoaringBitmap serializes to ByteString.EMPTY and deserializes
back to null,
+ // which used to NPE when merged into the ConcurrentHashMap in
CommitHandler.processResponse.
+ val committedMapIdBitMap = new util.HashMap[String, RoaringBitmap]()
+ committedMapIdBitMap.put("0-0", new RoaringBitmap())
+ val nonEmpty = RoaringBitmap.bitmapOf(1, 2, 3)
+ committedMapIdBitMap.put("0-1", nonEmpty)
+
+ val response = CommitFilesResponse(
+ StatusCode.SUCCESS,
+ Collections.emptyList[String](),
+ Collections.emptyList[String](),
+ Collections.emptyList[String](),
+ Collections.emptyList[String](),
+ committedMapIdBitMap = committedMapIdBitMap)
+ val responseTrans =
Utils.fromTransportMessage(Utils.toTransportMessage(response))
+ .asInstanceOf[CommitFilesResponse]
+
+ assert(!responseTrans.committedMapIdBitMap.containsValue(null))
+ assert(!responseTrans.committedMapIdBitMap.containsKey("0-0"))
+ assert(responseTrans.committedMapIdBitMap.get("0-1") == nonEmpty)
+
+ val concurrent = JavaUtils.newConcurrentHashMap[String, RoaringBitmap]()
+ concurrent.putAll(responseTrans.committedMapIdBitMap)
+ assert(concurrent.size() == 1)
+ }
+
test("validate number of client/server netty threads") {
val celebornConf = new CelebornConf()
celebornConf.set("celeborn.io.maxDefaultNettyThreads", "100")
