(celeborn) branch main updated: [CELEBORN-2472] Unregister Celeborn application from MR AppMaster on shutdown

[za...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zaynt4606 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 a649f32db5 [CELEBORN-2472] Unregister Celeborn application from MR AppMaster on shutdown
a649f32db5 is described below

commit a649f32db5c6c655e3dbde5414e4e9471ab445ce
Author: senthh <se...@acceldata.io>
AuthorDate: Tue Oct 13 10:05:40 2026 +0800

    [CELEBORN-2472] Unregister Celeborn application from MR AppMaster on shutdown
    
    ### What changes were proposed in this pull request?
    
    Store the `LifecycleManager` as a field in `MRAppMasterWithCeleborn` and stop it from `serviceStop()` (rather than `stop()`), so the application is unregistered from the Celeborn master on MR AppMaster shutdown. `LifecycleManager.stop() -> heartbeater.stop() -> unregisterApplication()` sends `ApplicationLost` (gated by `celeborn.client.application.unregister.enabled`), releasing worker slots and shuffle files immediately.
    
    Only `client-mr/mr/src/main/java/org/apache/celeborn/mapreduce/v2/app/MRAppMasterWithCeleborn.java` is changed.
    
    ### Why are the changes needed?
    
    `MRAppMasterWithCeleborn` creates a `LifecycleManager` and registers the application with the Celeborn master, but never stops it or unregisters on shutdown. The master therefore only frees the application (and the worker's shuffle files / slots) via heartbeat expiry (`appHeartbeatTimeoutMs`, default 120s), leaking worker disk after every MR job. Stopping from `serviceStop()` runs exactly once even though both `MRAppMaster#shutDownJob()` and the JVM shutdown hook can reach `stop()`. T [...]
    
    ### Does this PR resolve a correctness bug?
    
    - [ ] Yes
    
    ### Does this PR introduce _any_ user-facing change?
    
    - [ ] Yes
    
    ### How was this patch tested?
    
    - Compiles clean; the change follows the identical shutdown pattern already used by the other engine clients (Spark, Flink, Tez).
    
    Closes #3864 from senthh/CELEBORN-2472.
    
    Authored-by: senthh <se...@acceldata.io>
    Signed-off-by: zhengtao <sh...@alibaba-inc.com>
    
    Co-Authored-By: Codex <no...@openai.com>
    AI-Model: gpt-6
    AI-Contributed/Feature: 0/28
    AI-Contributed/UT: 0/0
---
 .../mapreduce/v2/app/MRAppMasterWithCeleborn.java  | 28 +++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/client-mr/mr/src/main/java/org/apache/celeborn/mapreduce/v2/app/MRAppMasterWithCeleborn.java b/client-mr/mr/src/main/java/org/apache/celeborn/mapreduce/v2/app/MRAppMasterWithCeleborn.java
index a23eab494f..61cf4dcaad 100644
--- a/client-mr/mr/src/main/java/org/apache/celeborn/mapreduce/v2/app/MRAppMasterWithCeleborn.java
+++ b/client-mr/mr/src/main/java/org/apache/celeborn/mapreduce/v2/app/MRAppMasterWithCeleborn.java
@@ -52,6 +52,8 @@ public class MRAppMasterWithCeleborn extends MRAppMaster {
 
   private static final String MASTER_ENDPOINTS_ENV = "CELEBORN_MASTER_ENDPOINTS";
 
+  private LifecycleManager lifecycleManager;
+
   public MRAppMasterWithCeleborn(
       ApplicationAttemptId applicationAttemptId,
       ContainerId containerId,
@@ -67,7 +69,7 @@ public class MRAppMasterWithCeleborn extends MRAppMaster {
     if (numReducers > 0) {
       CelebornConf conf = HadoopUtils.fromYarnConf(jobConf);
       String appUniqueId = conf.appUniqueIdWithUUIDSuffix(applicationAttemptId.toString());
-      LifecycleManager lifecycleManager = new LifecycleManager(appUniqueId, conf);
+      this.lifecycleManager = new LifecycleManager(appUniqueId, conf);
       String lmHost = lifecycleManager.getHost();
       int lmPort = lifecycleManager.getPort();
       logger.info("MRAppMaster initialized with {} {} {}", lmHost, lmPort, appUniqueId);
@@ -126,6 +128,30 @@ public class MRAppMasterWithCeleborn extends MRAppMaster {
     return value;
   }
 
+  /**
+   * In addition to the standard MR AppMaster teardown, stop the Celeborn LifecycleManager so the
+   * application is explicitly unregistered from the Celeborn master and its shuffle data is
+   * released. Without this, the Celeborn master only frees the application (and the worker's
+   * shuffle files) via heartbeat timeout, leaking disk space after the job finishes.
+   *
+   * <p>This is overridden in {@code serviceStop()} rather than {@code stop()}: the service
+   * transition is state-guarded and runs exactly once, whereas {@code stop()} can be reached twice
+   * on the AM ({@code MRAppMaster#shutDownJob()} then the JVM-shutdown hook), which would stop the
+   * LifecycleManager twice. {@code serviceStop()} mirrors the Tez client ({@code
+   * CelebornDagAppMaster#serviceStop()}).
+   */
+  @Override
+  protected void serviceStop() throws Exception {
+    if (lifecycleManager != null) {
+      try {
+        lifecycleManager.stop();
+      } catch (Throwable t) {
+        logger.warn("Error stopping Celeborn LifecycleManager", t);
+      }
+    }
+    super.serviceStop();
+  }
+
   public static void main(String[] args) {
     JobConf rmAppConf = new JobConf(new YarnConfiguration());
     rmAppConf.addResource(new Path(MRJobConfig.JOB_CONF_FILE));