Keep all resources close when client is killed (#600)

Co-authored-by: Keqiu Hu <[email protected]>
tony-framework · Sep 29, 2021 · f78e171 · f78e171
1 parent d0f3176
commit f78e171
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/tony-cli/src/main/java/com/linkedin/tony/cli/ClusterSubmitter.java b/tony-cli/src/main/java/com/linkedin/tony/cli/ClusterSubmitter.java
@@ -41,6 +41,8 @@
 public class ClusterSubmitter extends TonySubmitter {
   private static final Log LOG = LogFactory.getLog(ClusterSubmitter.class);
   private TonyClient client;
+  private Path cachedLibPath;
+  private Configuration hdfsConf = new Configuration();
 
   public ClusterSubmitter(TonyClient client) {
     this.client = client;
@@ -49,12 +51,10 @@ public ClusterSubmitter(TonyClient client) {
   public int submit(String[] args) throws ParseException, URISyntaxException {
     LOG.info("Starting ClusterSubmitter..");
     String jarLocation = new File(ClusterSubmitter.class.getProtectionDomain().getCodeSource().getLocation().toURI()).getPath();
-    Configuration hdfsConf = new Configuration();
     hdfsConf.addResource(new Path(System.getenv(HADOOP_CONF_DIR) + File.separatorChar + CORE_SITE_CONF));
     hdfsConf.addResource(new Path(System.getenv(HADOOP_CONF_DIR) + File.separatorChar + HDFS_SITE_CONF));
     LOG.info(hdfsConf);
     int exitCode;
-    Path cachedLibPath = null;
     try (FileSystem fs = FileSystem.get(hdfsConf)) {
       cachedLibPath = new Path(fs.getHomeDirectory(), TONY_FOLDER + Path.SEPARATOR + UUID.randomUUID().toString());
       Utils.uploadFileAndSetConfResources(cachedLibPath, new Path(jarLocation), TONY_JAR_NAME, client.getTonyConf(),
@@ -66,13 +66,7 @@ public int submit(String[] args) throws ParseException, URISyntaxException {
         return -1;
       }
       // ensure application is killed when this process terminates
-      Runtime.getRuntime().addShutdownHook(new Thread(() -> {
-        try {
-          client.forceKillApplication();
-        } catch (YarnException | IOException e) {
-          LOG.error("Failed to kill application during shutdown.", e);
-        }
-      }));
+      addShutdownHook();
       exitCode = client.start();
     } catch (IOException e) {
       LOG.fatal("Failed to create FileSystem: ", e);
@@ -83,6 +77,18 @@ public int submit(String[] args) throws ParseException, URISyntaxException {
     return exitCode;
   }
 
+  private void addShutdownHook() {
+    Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+      try {
+        client.forceKillApplication();
+        client.close();
+        Utils.cleanupHDFSPath(hdfsConf, cachedLibPath);
+      } catch (YarnException | IOException e) {
+        LOG.error("Failed to kill application during shutdown.", e);
+      }
+    }));
+  }
+
   public static void main(String[] args) throws ParseException, URISyntaxException {
     int exitCode;
     try (TonyClient tonyClient = new TonyClient(new Configuration())) {