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 1272744e67 [CELEBORN-2489] Fix unknown worker lookup in shuffle data
loss detection
1272744e67 is described below
commit 1272744e6726b895c7b2d095d802dc4517a71189
Author: zhengtao <[email protected]>
AuthorDate: Wed Oct 21 10:47:27 2026 +0800
[CELEBORN-2489] Fix unknown worker lookup in shuffle data loss detection
### What changes were proposed in this pull request?
Use `allocatedWorkers.containsKey(workerId.toUniqueId)` when checking
whether an unknown worker belongs to a shuffle.
Add focused tests for an allocated unknown worker, an unrelated unknown
worker, disabled detection, enabled replication, a different failure status,
and a shuffle without allocated workers.
JIRA: https://issues.apache.org/jira/browse/CELEBORN-2489
### Why are the changes needed?
`allocatedWorkers` is a `ConcurrentHashMap[String,
ShufflePartitionLocationInfo]`. Its legacy `contains` method checks values,
while the lookup argument is a `WorkerInfo`. The condition therefore remains
false even when the lost worker is allocated to the shuffle.
With `celeborn.client.shuffleDataLostOnUnknownWorker.enabled=true` and
replication disabled, the client should report shuffle data loss as soon as an
allocated worker becomes `WORKER_UNKNOWN`. The current check misses this and
waits for a later failure path. This fix restores the existing early failure
behavior and preserves the configuration and replication guards.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
Restores the configured early failure behavior for a lost allocated worker.
### How was this patch tested?
- Before the fix, the new suite runs six tests: the allocated
unknown-worker case fails and the other five pass.
- After the fix, 19 tests pass across `ReducePartitionCommitHandlerSuite`,
`WorkerStatusTrackerSuite`, and `MapPartitionCommitHandlerTest` with Java 17 /
Scala 2.12.
- `spotless:check` passes for the client and its upstream modules.
Closes #3881 from zaynt4606/CELEBORN-2489.
Authored-by: zhengtao <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
(cherry picked from commit 55474d41bd92fe818c0a67b463dc86da3f548a07)
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../commit/ReducePartitionCommitHandler.scala | 10 +--
.../commit/ReducePartitionCommitHandlerSuite.scala | 96 ++++++++++++++++++++++
2 files changed, 101 insertions(+), 5 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 febee90f8d..7477cb9dd3 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
@@ -152,11 +152,11 @@ class ReducePartitionCommitHandler(
if (conf.clientShuffleDataLostOnUnknownWorkerEnabled &&
!conf.clientPushReplicateEnabled) {
val allocatedWorkers = shuffleAllocatedWorkers.get(shuffleId)
if (allocatedWorkers != null) {
- return workerStatusTracker.excludedWorkers.asScala.collect {
- case (workerId, (status, _))
- if status == StatusCode.WORKER_UNKNOWN &&
allocatedWorkers.contains(workerId) =>
- workerId
- }.nonEmpty
+ return workerStatusTracker.excludedWorkers.asScala.exists {
+ case (workerId, (status, _)) =>
+ status == StatusCode.WORKER_UNKNOWN &&
+ allocatedWorkers.containsKey(workerId.toUniqueId)
+ }
}
}
false
diff --git
a/client/src/test/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandlerSuite.scala
b/client/src/test/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandlerSuite.scala
new file mode 100644
index 0000000000..bca906f341
--- /dev/null
+++
b/client/src/test/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandlerSuite.scala
@@ -0,0 +1,96 @@
+/*
+ * 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.util.concurrent.{ScheduledExecutorService, ThreadPoolExecutor}
+
+import org.mockito.Mockito.mock
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.client.{LifecycleManager, WorkerStatusTracker}
+import org.apache.celeborn.client.CommitManager.CommittedPartitionInfo
+import org.apache.celeborn.client.LifecycleManager.ShuffleAllocatedWorkers
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo,
WorkerInfo}
+import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.common.util.JavaUtils
+
+class ReducePartitionCommitHandlerSuite extends CelebornFunSuite {
+ private val shuffleId = 1
+
+ test("mark data lost when an allocated worker becomes unknown") {
+ // Heartbeats may report a different WorkerInfo instance for the same
worker identity.
+ val handler = newHandler(worker("allocated"))
+ assert(handler.isStageDataLost(shuffleId))
+ }
+
+ test("do not mark data lost when an unrelated worker becomes unknown") {
+ val handler = newHandler(worker("unrelated"))
+ assert(!handler.isStageDataLost(shuffleId))
+ }
+
+ test("do not mark data lost when unknown-worker detection is disabled") {
+ val handler = newHandler(worker("allocated"), detectionEnabled = false)
+ assert(!handler.isStageDataLost(shuffleId))
+ }
+
+ test("do not mark data lost on an unknown worker when replication is
enabled") {
+ val handler = newHandler(worker("allocated"), replicationEnabled = true)
+ assert(!handler.isStageDataLost(shuffleId))
+ }
+
+ test("do not mark data lost for an allocated worker with another failure
status") {
+ val handler = newHandler(worker("allocated"), status =
StatusCode.WORKER_UNRESPONSIVE)
+ assert(!handler.isStageDataLost(shuffleId))
+ }
+
+ test("do not mark data lost for a shuffle without allocated workers") {
+ val handler = newHandler(worker("allocated"))
+ assert(!handler.isStageDataLost(shuffleId + 1))
+ }
+
+ private def newHandler(
+ excludedWorker: WorkerInfo,
+ detectionEnabled: Boolean = true,
+ replicationEnabled: Boolean = false,
+ status: StatusCode = StatusCode.WORKER_UNKNOWN):
ReducePartitionCommitHandler = {
+ val conf = new CelebornConf()
+ .set(CelebornConf.CLIENT_SHUFFLE_DATA_LOST_ON_UNKNOWN_WORKER_ENABLED,
detectionEnabled)
+ .set(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED, replicationEnabled)
+ val allocatedWorker = worker("allocated")
+ val workers = JavaUtils.newConcurrentHashMap[String,
ShufflePartitionLocationInfo]()
+ workers.put(allocatedWorker.toUniqueId, new
ShufflePartitionLocationInfo(allocatedWorker))
+ val allocatedWorkers = new ShuffleAllocatedWorkers()
+ allocatedWorkers.put(shuffleId, workers)
+ val statusTracker = new WorkerStatusTracker(conf, null)
+ statusTracker.excludedWorkers.put(excludedWorker, (status,
System.currentTimeMillis()))
+
+ new ReducePartitionCommitHandler(
+ "test-app",
+ conf,
+ allocatedWorkers,
+ new CommittedPartitionInfo(),
+ statusTracker,
+ mock(classOf[ThreadPoolExecutor]),
+ mock(classOf[ScheduledExecutorService]),
+ mock(classOf[LifecycleManager]))
+ }
+
+ private def worker(host: String): WorkerInfo =
+ new WorkerInfo(host, 1001, 1002, 1003, 1004)
+}
