[CELEBORN-2336] Support decommission shutdown for worker scale-down scenarios

[chenghuichen]

What changes were proposed in this pull request?

Add celeborn.worker.decommission.shutdown.enabled configuration. When set to true, the worker will walk the decommission path (send ReportWorkerDecommission to master, wait for all shuffle data to be consumed or expired) upon receiving SIGTERM, instead of the graceful shutdown path.

When enabled, this overrides celeborn.worker.graceful.shutdown.enabled. The shutdown hook timeout is also extended to celeborn.worker.decommission.forceExitTimeout to match the decommission wait window.

Why are the changes needed?

In scheduled auto-scaling scenarios (e.g., scale up at peak hours, scale down at off-peak), operators want to simply shrink the node pool and let the PaaS layer send SIGTERM without writing custom preStop scripts or manually invoking the decommission REST API.

The existing graceful shutdown (celeborn.worker.graceful.shutdown.enabled=true) is designed for rolling upgrades. For permanent scale-down, the correct semantic is decommission. Previously this could only be triggered via REST API or master-pushed events, requiring additional scripting in the teardown workflow.

With this change, operators only need to set one config and align celeborn.worker.decommission.forceExitTimeout with the pod's terminationGracePeriodSeconds.

Does this PR resolve a correctness bug?

• Yes

Does this PR introduce any user-facing change?

• Yes

How was this patch tested?

• Verified default behavior unchanged (config defaults to false, existing graceful shutdown unaffected)
• Verified that when enabled, WorkerStatusManager.exitEventType is set to Decommission and workerGracefulShutdown returns false (suppressing recovery DB initialization)
• Verified shutdown hook extends timeout to decommission.forceExitTimeout before calling decommissionWorker()

[Copilot]

Pull request overview

Adds an opt-in “decommission-on-SIGTERM” mode for workers to better support scale-down scenarios, by introducing a new config flag that makes SIGTERM follow the decommission flow (report to master and wait for shuffle data consumption/expiration) instead of the graceful-restart flow.

Changes:

• Add celeborn.worker.decommission.shutdown.enabled config and surface it via CelebornConf.
• Make workerGracefulShutdown effectively disabled when decommission-on-shutdown is enabled, and set WorkerStatusManager.exitEventType accordingly.
• Attempt to extend shutdown hook timeout to workerDecommissionForceExitTimeout before decommissioning during shutdown.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala	Initialize exitEventType to Decommission when decommission-on-shutdown is enabled (overriding graceful).
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala	On decommission shutdown path, call ShutdownHookManager.updateTimeout before decommissionWorker().
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala	Add new config entry + getter, and ensure graceful shutdown is suppressed when decommission shutdown is enabled.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[SteNicholas]

@chenghuichen, thanks for first contribution. Could you please add test cases for this support?

[chenghuichen]

@SteNicholas Thanks for the review. Fixed the timeout issue, added tests, and updated worker.md.

[SteNicholas]

@chenghuichen, please use dev/reformat command to format code for failure of style check.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

[SteNicholas]

@chenghuichen, thanks for update. I left some comments for udpates. PTAL.

[RexXiong]

Overall: Clean, well-scoped change. The decommission-on-SIGTERM approach is the right abstraction for scale-down — avoids requiring custom preStop scripts. The Copilot/SteNicholas feedback has been addressed well (explicit timeout registration, WORKER_DECOMMISSION exit kind, assertion order fix, config override test).

Config version should likely be 0.7.0

The new config celeborn.worker.decommission.shutdown.enabled is versioned 0.6.0, but other new configs in recent PRs (e.g., PR #3695) use 0.7.0. If main is targeting 0.7.0, this should be updated to match.

Reviewed with Claude Code

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

[afterincomparableyum]

this is awesome and a known issue with auto-scaling down, thank you for your contribution @chenghuichen. I will review this PR soon.

[afterincomparableyum]

Overall lgtm. Could you pls address my comments and the other comments.

[chenghuichen]

@afterincomparableyum Thanks for the review! Updated: added assertions, fixed version to 0.7.0, clarified docs.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

[SteNicholas]

@chenghuichen, thanks for updates. I left some comments for the updates. PTAL.

[afterincomparableyum]

overall lgtm, thanks for contributing. Could you address the last comments by CoPilot please and @SteNicholas.

[SteNicholas]

@chenghuichen, any update?

[SteNicholas]

Reviewed the decommission-shutdown path. The state-machine wiring and the new WORKER_DECOMMISSION exit-kind plumbing through stop()/close() look correct, and the config version + docs row check out. One correctness concern (the shutdown-hook timeout budget doesn't account for stop()), one behavior change worth confirming, and a few cleanup/altitude nits inline.

[chenghuichen]

@SteNicholas Thanks for the review! All fixed.

[SteNicholas]

@chenghuichen, re-review of the current revision. This version addresses my earlier feedback well — the effectiveWorkerGracefulShutdown accessor de-duplicates the override, drainBeforeExit consolidates the exit-kind checks, and the waitTime + interval <= timeout bound makes the budget self-consistent in the common case. close(WORKER_DECOMMISSION) is db-null-safe and the master stops allocating slots (heartbeat reports InDecommission).

Remaining items are mostly about the shutdown timing budget in the worst case and a couple of intent/behavior inconsistencies — inline below. Two low-severity notes not inlined: (a) a spurious transition status ... is not allowed warning is logged on every runtime REST decommission because exit() already moved the state to InDecommission before decommissionWorker() re-transitions it; (b) exitEventType is sticky — recommissionWorker() resets state to Normal but not exitEventType, so a Recommission can't restore non-decommission exit semantics (likely intended for a decommission-shutdown worker, worth confirming).

[SteNicholas]

Worst case can truncate stop() mid-cleanup. This comment claims the wait loop leaves budget for stop(), but the loop below can consume the entire forceExitTimeout (floor(6h/30s)·30s = 6h with defaults), leaving only ~checkInterval (30s) of the forceExitTimeout + checkInterval hook budget for stop(WORKER_DECOMMISSION). That path runs StorageManager.cleanupExpiredShuffleKey(shuffleKeySet(), false) — unbounded synchronous DFS delete() of all unreleased shuffle (exactly the slow case, since shuffle didn't drain) — plus flusher + 3 netty server shutdowns. If it exceeds 30s, ShutdownHookManager.executeShutdown does future.cancel(true), interrupting teardown → orphaned DFS objects and transitionState(Exit) skipped. The sendWorkerDecommissionToMaster() askSync before the loop is also unbudgeted (retries on a slow master). Consider bounding the wait to leave a teardown headroom proportional to the real stop() cost, not a fixed checkInterval.

[SteNicholas]

forceExitTimeout < checkInterval → no wait at all → data loss. With waitTimes=0 the guard is 0 + interval <= timeout, i.e. checkInterval <= forceExitTimeout. If an operator sets forceExitTimeout below checkInterval, the loop body never executes — decommissionWorker() returns immediately and stop() deletes shuffle data consumers still need → FetchFailure / SHUFFLE_DATA_LOST. The previous waitTime < timeout always waited at least once. There's no checkValue guarding forceExitTimeout >= checkInterval; consider adding one, or special-casing interval > timeout to still wait (bounded by timeout).

[SteNicholas]

Behavior change for runtime REST decommission. This branch is reached not only by the new SIGTERM feature but also by the existing exit("DECOMMISSION") REST API (via the exit-thread → System.exit → hook). Pre-PR that path ran stop(EXIT_IMMEDIATELY); now it runs stop(WORKER_DECOMMISSION), which via drainBeforeExit makes the heartbeat/fast-fail tasks cancel(false) and the thread pools shutdown() (drain) instead of cancel(true)/shutdownNow(). So every cluster using the REST decommission API gets changed teardown semantics even without enabling celeborn.worker.decommission.shutdown.enabled. Worth calling out explicitly as intended.

[chenghuichen]

Not a regression — the existing REST exit("DECOMMISSION") path already swallows this, and the master falls back to heartbeat-timeout eviction.

[SteNicholas]

Drain intent isn't carried through to the partition sorter. This comment/drainBeforeExit declare that decommission "has drained data, so in-flight tasks are allowed to finish" — and that holds for the thread pools here. But stop() later calls partitionsSorter.close(WORKER_DECOMMISSION), and PartitionFilesSorter.close only takes the await-termination branch for exitKind == WORKER_GRACEFUL_SHUTDOWN; WORKER_DECOMMISSION falls into the else and calls fileSorterExecutors.shutdownNow() (PartitionFilesSorter.java:415-417), force-interrupting an in-flight on-demand sort → truncated sorted file/index → fetch failure for a reader mid-sort. Not a regression vs the old immediate path, but inconsistent with the new drain semantics — either group WORKER_DECOMMISSION with graceful in the sorter, or soften this comment.

[SteNicholas]

Decommission report failure has no retry/fallback on the SIGTERM path. decommissionWorker() calls sendWorkerDecommissionToMaster(), which swallows any exception and proceeds. If the master is briefly unreachable when SIGTERM fires, no decommission record reaches it and it must wait for heartbeat-timeout eviction. The other two exit paths are more robust — exitImmediately() uses WorkerLost and shutdownGracefully() uses ReportWorkerUnavailable. Consider a single retry or a ReportWorkerUnavailable fallback to match them.

[SteNicholas]

Altitude: updateTimeout is process-wide. ShutdownHookManager.updateTimeout does hooks.forEach(setTimeout), so this raises the timeout of every registered hook (not just the worker hook) to forceExitTimeout + checkInterval (6h+ by default). An unrelated stuck hook could then hang JVM exit for hours where the default would force-terminate it. Prefer registering the worker hook with the 4-arg addShutdownHook(hook, priority, timeout, unit) so only that hook gets the extended budget — which also removes the register-then-update sequencing.

[SteNicholas]

Duplicated hook-timeout expression. conf.workerDecommissionForceExitTimeout + conf.workerDecommissionCheckInterval appears here (REST exit handler) and again in the construction block (~line 1107). Extract a single decommissionHookTimeoutMs helper so the SIGTERM and REST decommission paths can't drift to different timeouts.

[SteNicholas]

exitEventType is read across threads without synchronization. It's written here and in exit() under this.synchronized, but the shutdown-hook thread reads workerStatusManager.exitEventType (Worker.scala:1087) without the lock, and the field (declared line 38) isn't @volatile. A runtime REST exit(...) that updates it shortly before an independent SIGTERM-triggered hook read has no happens-before guarantee → the hook could take a stale branch (e.g. exitImmediately instead of decommissionWorker). Mark exitEventType @volatile.

[SteNicholas]

terminationGracePeriodSeconds guidance understates the budget. This tells operators to set it to forceExitTimeout + checkInterval + small buffer, but stop() teardown runs inside that same forceExitTimeout + checkInterval hook timeout, and its cost is unbounded (DFS deletes of all unreleased shuffle). When the wait loop runs close to forceExitTimeout, only ~checkInterval is left for stop(), so a "small buffer" can't cover it and the pod is SIGKILLed mid-cleanup. The in-code comments acknowledge stop() needs headroom; the operator-facing doc should say so too (and the buffer should scale with the worst-case teardown, not be "small").

[SteNicholas]

Non-defensive test setup. conf1 above explicitly sets both keys with the comment "so the assertion does not depend on system properties leaked from other tests." But conf2 here sets only graceful.shutdown.enabled=true (not decommission.shutdown.enabled=false). Since new CelebornConf() loads system properties, a leaked celeborn.worker.decommission.shutdown.enabled=true would make exitEventType=Decommission and this assertEquals(Graceful, mgr2.exitEventType) fail. Apply conf1's both-keys defensiveness to conf2/conf3.

[chenghuichen]

The real bugs are fixed; the remaining points are intentional tradeoffs for an opt-in flag.

One ask for future rounds: please reconcile suggestions against the full review history. The hook-timeout registration, for instance, has now round-tripped 4-arg → updateTimeout → 4-arg across reviews — the current explicit 4-arg form is a deliberate decision and I won't flip it again.