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

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

@@ -1357,7 +1357,12 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
   // //////////////////////////////////////////////////////
   //            Graceful Shutdown & Recover              //
   // //////////////////////////////////////////////////////
+  def workerDecommissionShutdown: Boolean = get(WORKER_DECOMMISSION_SHUTDOWN_ENABLED)
   def workerGracefulShutdown: Boolean = get(WORKER_GRACEFUL_SHUTDOWN_ENABLED)
+  // Decommission shutdown overrides graceful shutdown: a decommissioned worker will not
+  // restart, so recovery state (recovery DB, sorter state) should not be persisted.
+  def effectiveWorkerGracefulShutdown: Boolean =
+    workerGracefulShutdown && !workerDecommissionShutdown
   def workerGracefulShutdownTimeoutMs: Long = get(WORKER_GRACEFUL_SHUTDOWN_TIMEOUT)
   def workerGracefulShutdownCheckSlotsFinishedInterval: Long =
     get(WORKER_CHECK_SLOTS_FINISHED_INTERVAL)
@@ -4477,6 +4482,23 @@ object CelebornConf extends Logging {
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("6h")
 
+  val WORKER_DECOMMISSION_SHUTDOWN_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.worker.decommission.shutdown.enabled")
+      .categories("worker")
+      .doc("When true, the worker will decommission on shutdown signal (e.g. SIGTERM), " +
+        "waiting for all shuffle data to be consumed or expired before exiting. " +
+        "This is suitable for permanent scale-down scenarios where the worker will not restart. " +
+        "When enabled, this overrides celeborn.worker.graceful.shutdown.enabled " +
+        "(recovery state will not be saved since the worker is not expected to come back). " +
+        "Operators should set the pod's terminationGracePeriodSeconds to " +
+        "celeborn.worker.decommission.forceExitTimeout + " +
+        "celeborn.worker.decommission.checkInterval plus a small buffer, to ensure " +
+        "the shutdown hook has enough time to complete resource cleanup before " +
+        "being killed.")
+      .version("0.7.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val WORKER_GRACEFUL_SHUTDOWN_ENABLED: ConfigEntry[Boolean] =
     buildConf("celeborn.worker.graceful.shutdown.enabled")
       .categories("worker")

docs/configuration/worker.md

@@ -78,6 +78,7 @@ license: |
 | celeborn.worker.congestionControl.workerProduceSpeed.low.watermark | 9223372036854775807b | false | Stop congestion If worker total produce speed less than this configuration | 0.6.0 |  | 
 | celeborn.worker.decommission.checkInterval | 30s | false | The wait interval of checking whether all the shuffle expired during worker decommission | 0.4.0 |  | 
 | celeborn.worker.decommission.forceExitTimeout | 6h | false | The wait time of waiting for all the shuffle expire during worker decommission. | 0.4.0 |  | 
+| celeborn.worker.decommission.shutdown.enabled | false | false | When true, the worker will decommission on shutdown signal (e.g. SIGTERM), waiting for all shuffle data to be consumed or expired before exiting. This is suitable for permanent scale-down scenarios where the worker will not restart. When enabled, this overrides celeborn.worker.graceful.shutdown.enabled (recovery state will not be saved since the worker is not expected to come back). Operators should set the pod's terminationGracePeriodSeconds to celeborn.worker.decommission.forceExitTimeout + celeborn.worker.decommission.checkInterval plus a small buffer, to ensure the shutdown hook has enough time to complete resource cleanup before being killed. | 0.7.0 |  | 
 | celeborn.worker.directMemoryRatioForMemoryFileStorage | 0.0 | false | Max ratio of direct memory to store shuffle data. This feature is experimental and disabled by default. | 0.5.0 |  | 
 | celeborn.worker.directMemoryRatioForReadBuffer | 0.35 | false | Max ratio of direct memory for read buffer | 0.2.0 |  | 
 | celeborn.worker.directMemoryRatioToMergeBuffer | 0.4 | false | If direct memory usage is above this limit, the worker will merge low utilization push data's body buffer | 0.6.2 |  | 

worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java

@@ -113,7 +113,7 @@ public PartitionFilesSorter(
     long indexCacheMaxWeight = conf.workerPartitionSorterIndexCacheMaxWeight();
     this.source = source;
     this.cleaner = new PartitionFilesCleaner(this);
-    boolean gracefulShutdown = conf.workerGracefulShutdown();
+    boolean gracefulShutdown = conf.effectiveWorkerGracefulShutdown();
     // Assume a chunk won't be larger than 2GB
     // ShuffleClient can fetch shuffle data from a restarted worker only
     // when the worker's fetching port is stable and enables graceful shutdown.

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala

@@ -156,7 +156,7 @@ private[celeborn] class Worker(
 
   private val WORKER_SHUTDOWN_PRIORITY = 100
   val shutdown = new AtomicBoolean(false)
-  private val gracefulShutdown = conf.workerGracefulShutdown
+  private val gracefulShutdown = conf.effectiveWorkerGracefulShutdown
   if (gracefulShutdown) {
     var checkPortMap = Map(
       WORKER_RPC_PORT -> conf.workerRpcPort,
@@ -619,29 +619,34 @@ private[celeborn] class Worker(
     if (!stopped) {
       logInfo("Stopping Worker.")
 
+      // Both graceful shutdown and decommission have drained data, so in-flight
+      // tasks are allowed to finish instead of being force-cancelled.
+      val drainBeforeExit = exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN ||
+        exitKind == CelebornExitKind.WORKER_DECOMMISSION
+
       if (jvmProfiler != null) {
         jvmProfiler.stop()
       }
       if (jvmQuake != null) {
         jvmQuake.stop()
       }
       if (sendHeartbeatTask != null) {
-        if (exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN) {
+        if (drainBeforeExit) {
           sendHeartbeatTask.cancel(false)
         } else {
           sendHeartbeatTask.cancel(true)
         }
         sendHeartbeatTask = null
       }
       if (checkFastFailTask != null) {
-        if (exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN) {
+        if (drainBeforeExit) {
           checkFastFailTask.cancel(false)
         } else {
           checkFastFailTask.cancel(true)
         }
         checkFastFailTask = null
       }
-      if (exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN) {
+      if (drainBeforeExit) {
         forwardMessageScheduler.shutdown()
         replicateThreadPool.shutdown()
         commitThreadPool.shutdown()
@@ -950,7 +955,7 @@ private[celeborn] class Worker(
     exitType.toUpperCase(Locale.ROOT) match {
       case "DECOMMISSION" =>
         ShutdownHookManager.get().updateTimeout(
-          conf.workerDecommissionForceExitTimeout,
+          conf.workerDecommissionForceExitTimeout + conf.workerDecommissionCheckInterval,
           TimeUnit.MILLISECONDS)
         workerStatusManager.doTransition(WorkerEventType.Decommission)
       case "GRACEFUL" =>
@@ -1031,7 +1036,9 @@ private[celeborn] class Worker(
 
     def waitTime: Long = waitTimes * interval
 
-    while (!storageManager.shuffleKeySet().isEmpty && waitTime < timeout) {
+    // Bound the total wait strictly by the timeout so that the remaining shutdown hook
+    // budget is left for stop(WORKER_DECOMMISSION) to clean up resources.
+    while (!storageManager.shuffleKeySet().isEmpty && waitTime + interval <= timeout) {
       Thread.sleep(interval)
       waitTimes += 1
     }
@@ -1080,22 +1087,28 @@ private[celeborn] class Worker(
           workerStatusManager.exitEventType match {
             case WorkerEventType.Graceful =>
               shutdownGracefully()
+              stop(CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN)
             case WorkerEventType.Decommission =>
               decommissionWorker()
+              stop(CelebornExitKind.WORKER_DECOMMISSION)
             case _ =>
               exitImmediately()
-          }
-
-          if (workerStatusManager.exitEventType == WorkerEventType.Graceful) {
-            stop(CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN)
-          } else {
-            stop(CelebornExitKind.EXIT_IMMEDIATELY)
+              stop(CelebornExitKind.EXIT_IMMEDIATELY)
           }
         }
       },
       "worker-shutdown-hook-thread"),
     WORKER_SHUTDOWN_PRIORITY)
 
+  if (conf.workerDecommissionShutdown) {
+    // The wait loop in decommissionWorker() is bounded by forceExitTimeout, so the extra
+    // checkInterval reserves headroom for stop(WORKER_DECOMMISSION) to finish cleanup
+    // before the hook is cancelled.
+    ShutdownHookManager.get().updateTimeout(
+      conf.workerDecommissionForceExitTimeout + conf.workerDecommissionCheckInterval,
+      TimeUnit.MILLISECONDS)
+  }
+
   @VisibleForTesting
   def getPushFetchServerPort: (Int, Int) = (pushPort, fetchPort)
 

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala

@@ -39,8 +39,13 @@ private[celeborn] class WorkerStatusManager(conf: CelebornConf) extends Logging
   private var worker: Worker = _
   private var shutdown: AtomicBoolean = _
   private var storageManager: StorageManager = _
+  private val decommissionShutdown = conf.workerDecommissionShutdown
   private val gracefulShutdown = conf.workerGracefulShutdown
-  if (gracefulShutdown) {
+  if (decommissionShutdown) {
+    exitEventType = WorkerEventType.Decommission
+    logInfo("Decommission shutdown enabled, worker will decommission on SIGTERM" +
+      " (overrides graceful shutdown)")
+  } else if (gracefulShutdown) {
     exitEventType = WorkerEventType.Graceful
   }
 

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala

@@ -295,7 +295,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
     JavaUtils.newConcurrentHashMap[String, ConcurrentHashMap[String, DiskFileInfo]]()
   // ShuffleClient can fetch data from a restarted worker only
   // when the worker's fetching port is stable.
-  val workerGracefulShutdown = conf.workerGracefulShutdown
+  val workerGracefulShutdown = conf.effectiveWorkerGracefulShutdown
   if (workerGracefulShutdown) {
     try {
       val dbBackend = DBBackend.byName(conf.workerGracefulShutdownRecoverDbBackend)

worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManagerSuite.scala

@@ -57,24 +57,55 @@ class WorkerStatusManagerSuite extends AnyFunSuite {
     statusManager.init(worker)
 
     statusManager.doTransition(WorkerEventType.DecommissionThenIdle)
-    Assert.assertEquals(statusManager.getWorkerState(), PbWorkerStatus.State.InDecommissionThenIdle)
+    Assert.assertEquals(PbWorkerStatus.State.InDecommissionThenIdle, statusManager.getWorkerState())
     Assert.assertEquals(
-      worker.workerInfo.getWorkerStatus().getStateValue,
-      PbWorkerStatus.State.InDecommissionThenIdle.getNumber)
+      PbWorkerStatus.State.InDecommissionThenIdle.getNumber,
+      worker.workerInfo.getWorkerStatus().getStateValue)
 
     // Rerun state Transition
     statusManager.doTransition(WorkerEventType.DecommissionThenIdle)
-    Assert.assertEquals(statusManager.getWorkerState(), PbWorkerStatus.State.InDecommissionThenIdle)
+    Assert.assertEquals(PbWorkerStatus.State.InDecommissionThenIdle, statusManager.getWorkerState())
 
     // Reset shuffleKeys
     shuffleKeys.clear()
     statusManager.doTransition(WorkerEventType.DecommissionThenIdle)
-    Assert.assertEquals(statusManager.getWorkerState(), PbWorkerStatus.State.Idle)
+    Assert.assertEquals(PbWorkerStatus.State.Idle, statusManager.getWorkerState())
 
     statusManager.doTransition(WorkerEventType.Recommission)
-    Assert.assertEquals(statusManager.getWorkerState(), PbWorkerStatus.State.Normal)
+    Assert.assertEquals(PbWorkerStatus.State.Normal, statusManager.getWorkerState())
 
     statusManager.doTransition(WorkerEventType.Recommission)
-    Assert.assertEquals(statusManager.getWorkerState(), PbWorkerStatus.State.Normal)
+    Assert.assertEquals(PbWorkerStatus.State.Normal, statusManager.getWorkerState())
+  }
+
+  test("Test exitEventType initialization based on config") {
+    // Neither graceful nor decommission → Immediately. Set both keys explicitly so the
+    // assertion does not depend on system properties leaked from other tests.
+    val conf1 = new CelebornConf()
+    conf1.set("celeborn.worker.graceful.shutdown.enabled", "false")
+    conf1.set("celeborn.worker.decommission.shutdown.enabled", "false")
+    val mgr1 = new WorkerStatusManager(conf1)
+    Assert.assertEquals(WorkerEventType.Immediately, mgr1.exitEventType)
+
+    // Graceful shutdown only → Graceful
+    val conf2 = new CelebornConf()
+    conf2.set("celeborn.worker.graceful.shutdown.enabled", "true")
+    val mgr2 = new WorkerStatusManager(conf2)
+    Assert.assertEquals(WorkerEventType.Graceful, mgr2.exitEventType)
+
+    // Decommission shutdown only → Decommission
+    val conf3 = new CelebornConf()
+    conf3.set("celeborn.worker.decommission.shutdown.enabled", "true")
+    val mgr3 = new WorkerStatusManager(conf3)
+    Assert.assertEquals(WorkerEventType.Decommission, mgr3.exitEventType)
+
+    // Both enabled → Decommission overrides graceful
+    val conf4 = new CelebornConf()
+    conf4.set("celeborn.worker.graceful.shutdown.enabled", "true")
+    conf4.set("celeborn.worker.decommission.shutdown.enabled", "true")
+    val mgr4 = new WorkerStatusManager(conf4)
+    Assert.assertEquals(WorkerEventType.Decommission, mgr4.exitEventType)
+    Assert.assertTrue(conf4.workerGracefulShutdown)
+    Assert.assertTrue(conf4.workerDecommissionShutdown)
   }
 }