Skip to content

[CELEBORN-2489] Fix unknown worker lookup in shuffle data loss detection - #3881

Open
zaynt4606 wants to merge 2 commits into
apache:mainfrom
zaynt4606:CELEBORN-2489
Open

zaynt4606 wants to merge 2 commits into
apache:mainfrom
zaynt4606:CELEBORN-2489

Conversation

@zaynt4606

@zaynt4606 zaynt4606 commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

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?

  • 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.

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) =>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, updated to use exists.

@github-actions github-actions Bot added module:client correctness Correctness bugfix labels Oct 19, 2026
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
@github-actions github-actions Bot removed the correctness Correctness bugfix label Oct 20, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants