(celeborn) branch main updated: [CELEBORN-2489] Fix unknown worker lookup in shuffle data loss detection

[ni...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

SteNicholas 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 55474d41bd [CELEBORN-2489] Fix unknown worker lookup in shuffle data loss detection
55474d41bd is described below

commit 55474d41bd92fe818c0a67b463dc86da3f548a07
Author: zhengtao <sh...@alibaba-inc.com>
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 <sh...@alibaba-inc.com>
    Signed-off-by: Nicholas Jiang <pr...@163.com>
---
 .../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)
+}