Repository navigation
Conversation
Look up allocated workers by their unique string IDs so an unknown worker correctly triggers the configured early shuffle data loss detection. Cover allocated and unrelated workers together with the existing guards. Co-Authored-By: Codex <noreply@openai.com> AI-Model: gpt-6 AI-Contributed/Feature: 3/3 AI-Contributed/UT: 96/96
| case (workerId, (status, _)) | ||
| if status == StatusCode.WORKER_UNKNOWN && allocatedWorkers.contains(workerId) => | ||
| if status == StatusCode.WORKER_UNKNOWN && | ||
| allocatedWorkers.containsKey(workerId.toUniqueId) => |
There was a problem hiding this comment.
Verified locally: allocatedWorkers is a ConcurrentHashMap[String, ShufflePartitionLocationInfo] whose keys are written as workerInfo.toUniqueId (see LifecycleManager#handleRegisterShuffleResponse). ConcurrentHashMap.contains(Object) is the legacy Hashtable-style method that checks values, and since the values are ShufflePartitionLocationInfo (never equal to a WorkerInfo), the old guard was always false — the early SHUFFLE_DATA_LOST path introduced by CELEBORN-2166 has been dead code since then.
containsKey(workerId.toUniqueId) matches the key format exactly and is consistent with the lookups in WorkerStatusTracker#excludeWorkerFromPartition and CommitManager#handleCommitFiles. LGTM on the fix itself.
Optional nit: since we're touching this expression anyway, collect { ... }.nonEmpty could be simplified to exists { case (workerId, (status, _)) => ... } to avoid building an intermediate collection for a membership check.
There was a problem hiding this comment.
Thanks, updated to use exists.
Avoid allocating an intermediate collection and stop at the first matching allocated unknown worker. Co-Authored-By: Codex <noreply@openai.com> AI-Model: gpt-6 AI-Contributed/Feature: 11/11 AI-Contributed/UT: 0/0
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?
allocatedWorkersis aConcurrentHashMap[String, ShufflePartitionLocationInfo]. Its legacycontainsmethod checks values, while the lookup argument is aWorkerInfo. The condition therefore remains false even when the lost worker is allocated to the shuffle.With
celeborn.client.shuffleDataLostOnUnknownWorker.enabled=trueand replication disabled, the client should report shuffle data loss as soon as an allocated worker becomesWORKER_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?
Does this PR introduce any user-facing change?
Restores the configured early failure behavior for a lost allocated worker.
How was this patch tested?
ReducePartitionCommitHandlerSuite,WorkerStatusTrackerSuite, andMapPartitionCommitHandlerTestwith Java 17 / Scala 2.12.spotless:checkpasses for the client and its upstream modules.