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 89fc9bd549 [CELEBORN-2516] Worker decommission does not wait for
shuffles held only in memory files
89fc9bd549 is described below
commit 89fc9bd5491e93552c6903f48582273545c2d7c9
Author: strelok89 <[email protected]>
AuthorDate: Fri Oct 30 23:17:02 2026 +0800
[CELEBORN-2516] Worker decommission does not wait for shuffles held only in
memory files
### What changes were proposed in this pull request?
`StorageManager.shuffleKeySet()` now includes `memoryFileInfos.keySet()`
along with `diskFileInfos.keySet()`.
**Chart:** also keeps a draining worker resolvable. A terminating pod is
dropped from its headless service immediately, so executors fetching from a
draining worker hit `UnknownHostException` about a second into the drain, even
with the wait fixed. With `worker.drain.enabled` the worker service now sets
`publishNotReadyAddresses: true`. Note: #3854 sets this flag unconditionally on
the worker service; if it merges first, this hunk becomes redundant.
### Why are the changes needed?
With the MEMORY storage tier enabled
(`celeborn.storage.availableTypes=MEMORY,SSD`), a shuffle whose partitions are
still in memory files is invisible to `shuffleKeySet()`.
`Worker.decommissionWorker()` waits until that set is empty, so a
decommissioning worker that holds only memory-file shuffles exits immediately
(`Waiting for all shuffle expired cost 0ms`) with slots still active. Clients
reading those shuffles then fail with fetch failures and resubmit stages. We
hit this on 0.7.0 [...]
Every other caller also wants "shuffles this worker holds", whatever the
storage tier:
- the graceful-shutdown check in `WorkerStatusManager`
- the `UnreleasedShuffleCount` gauge
- the active shuffle keys reported in the heartbeat
- the shuffle list on the HTTP endpoint
- the expired-shuffle cleanup on shutdown, which already handles memory
keys in `cleanupExpiredShuffleKey`
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes. With the MEMORY tier enabled, decommission and graceful shutdown
now wait for memory-file shuffles, the `UnreleasedShuffleCount` gauge counts
them, and `/api/v1/shuffles` lists them.
### How was this patch tested?
New `StorageManagerSuite` case: a shuffle held only in a memory file is
returned by `shuffleKeySet()`. Existing `StorageManagerSuite` passes. Plus
helm-unittest: two new worker-service cases (absent by default, set when the
drain is enabled); 136/136 pass.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes #3916 from strelok89/fix-decommission-memory-shuffle.
Authored-by: strelok89 <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
charts/celeborn/README.md | 5 +++++
charts/celeborn/templates/worker/service.yaml | 4 ++++
charts/celeborn/tests/worker/service_test.yaml | 15 +++++++++++++++
.../service/deploy/worker/storage/StorageManager.scala | 2 ++
.../deploy/worker/storage/StorageManagerSuite.scala | 16 +++++++++++++++-
5 files changed, 41 insertions(+), 1 deletion(-)
diff --git a/charts/celeborn/README.md b/charts/celeborn/README.md
index c679f5c364..c2886828ce 100644
--- a/charts/celeborn/README.md
+++ b/charts/celeborn/README.md
@@ -171,6 +171,11 @@ drain, so the chart refuses to render `worker.drain`
together with either settin
`celeborn`. It can only see what is set there: enabling TLS or authentication
some other way,
such as an extra config file, still bypasses the drain.
+A terminating pod is dropped from its headless service straight away, so a
draining worker
+would stop resolving the moment the drain starts, and clients still fetching
from it would fail
+with `UnknownHostException`. With the drain enabled the worker service
therefore sets
+`publishNotReadyAddresses: true`, keeping the worker's DNS record until the
pod is gone.
+
It is off by default. Turn it on alongside `worker.autoscaling`, or leave it
off and treat
scale-in as destructive, only letting it happen when the fleet is idle.
diff --git a/charts/celeborn/templates/worker/service.yaml
b/charts/celeborn/templates/worker/service.yaml
index 90b5940275..c4fb93f618 100644
--- a/charts/celeborn/templates/worker/service.yaml
+++ b/charts/celeborn/templates/worker/service.yaml
@@ -29,4 +29,8 @@ spec:
{{- include "celeborn.worker.selectorLabels" . | nindent 4 }}
type: {{ .Values.service.type }}
clusterIP: None
+ {{- if .Values.worker.drain.enabled }}
+ # Keep a draining worker resolvable until it exits. See README.md.
+ publishNotReadyAddresses: true
+ {{- end }}
{{- include "celeborn.worker.service.port" . | nindent 2 }}
diff --git a/charts/celeborn/tests/worker/service_test.yaml
b/charts/celeborn/tests/worker/service_test.yaml
index 27a50e9949..291f5cce86 100644
--- a/charts/celeborn/tests/worker/service_test.yaml
+++ b/charts/celeborn/tests/worker/service_test.yaml
@@ -47,3 +47,18 @@ tests:
- equal:
path: spec.type
value: NodePort
+
+ - it: Should not publish not-ready addresses by default
+ asserts:
+ - notExists:
+ path: spec.publishNotReadyAddresses
+
+ - it: Should publish not-ready addresses when the worker drain is enabled
+ set:
+ worker:
+ drain:
+ enabled: true
+ asserts:
+ - equal:
+ path: spec.publishNotReadyAddresses
+ value: true
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 80aee90a77..50f2b5f888 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -637,6 +637,8 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
def shuffleKeySet(): util.HashSet[String] = {
val hashSet = new util.HashSet[String]()
hashSet.addAll(diskFileInfos.keySet())
+ // A shuffle can live only in memory files; decommission and graceful
shutdown wait on it too.
+ hashSet.addAll(memoryFileInfos.keySet())
hashSet
}
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
index 57505e115c..c6a1e72024 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManagerSuite.scala
@@ -36,7 +36,7 @@ import org.apache.celeborn.CelebornFunSuite
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.CelebornConf.{WORKER_DISK_RESERVE_SIZE,
WORKER_GRACEFUL_SHUTDOWN_ENABLED, WORKER_GRACEFUL_SHUTDOWN_RECOVER_PATH,
WORKER_GRACEFUL_SHUTDOWN_SAVE_COMMITTED_FILEINFO_INTERVAL, WORKER_STORAGE_DIRS}
import org.apache.celeborn.common.identity.UserIdentifier
-import org.apache.celeborn.common.meta.{DiskFileInfo, DiskInfo, DiskStatus}
+import org.apache.celeborn.common.meta.{DiskFileInfo, DiskInfo, DiskStatus,
MemoryFileInfo, ReduceFileMeta}
import org.apache.celeborn.common.protocol.{PartitionLocation,
PartitionSplitMode, PartitionType, StorageInfo}
import org.apache.celeborn.common.util.{CelebornExitKind, JavaUtils, Utils}
import org.apache.celeborn.service.deploy.worker.WorkerSource
@@ -57,6 +57,20 @@ class StorageManagerSuite extends CelebornFunSuite with
MockitoHelper {
storageManager.saveAllCommittedFileInfosToDB()
}
+ test("shuffleKeySet includes shuffles held only in memory files") {
+ val conf = new CelebornConf()
+ val storageManager = new StorageManager(conf, new WorkerSource(conf))
+ val fileInfo =
+ new MemoryFileInfo(
+ UserIdentifier("tenant", "user"),
+ false,
+ new ReduceFileMeta(8 * 1024 * 1024))
+ storageManager.memoryFileInfos
+ .computeIfAbsent("app-1-0", _ => JavaUtils.newConcurrentHashMap[String,
MemoryFileInfo]())
+ .put("0-0-0", fileInfo)
+ assert(storageManager.shuffleKeySet().asScala == Set("app-1-0"))
+ }
+
test("[CELEBORN-2495] committedFileInfos should not leak when graceful
shutdown is disabled") {
withDiskFile(gracefulShutdown = false) { (storageManager, fileInfo, _) =>
storageManager.notifyFileInfoCommitted("app-1-0", "0-0-0", fileInfo)
