Introduce tony.application.x.untracked.timeout to solve partial jobs unfinished

zuston · zuston · commit a8550c399d59 · 2021-10-13T17:12:11.000+08:00
diff --git a/tony-core/src/main/java/com/linkedin/tony/ApplicationMaster.java b/tony-core/src/main/java/com/linkedin/tony/ApplicationMaster.java
@@ -192,6 +192,8 @@ public class ApplicationMaster {
   private String frameworkType;
   private Framework.ApplicationMasterAdapter amRuntimeAdapter;
 
+  private long allTrackedTaskFinishedTime;
+
   private ApplicationMaster() {
     hdfsConf = new Configuration(false);
     yarnConf = new Configuration(false);
@@ -620,6 +622,8 @@ private void reset() {
     session = sessionBuilder.build();
     applicationRpcServer.reset();
     session.sessionId += 1;
+
+    allTrackedTaskFinishedTime = 0L;
   }
 
   /**
@@ -687,18 +691,10 @@ private boolean monitor() {
         break;
       }
 
-      int numTotalTrackedTasks = session.getTotalTrackedTasks();
-      if (numTotalTrackedTasks > 0) {
-        int numCompletedTrackedTasks = session.getNumCompletedTrackedTasks();
-        if (numCompletedTrackedTasks == numTotalTrackedTasks) {
-          Utils.printCompletedTrackedTasks(numCompletedTrackedTasks, numTotalTrackedTasks);
-          break;
-        }
-
-        // Reduce logging frequency to every 100s.
-        if (counter % 20 == 1) {
-          Utils.printCompletedTrackedTasks(numCompletedTrackedTasks, numTotalTrackedTasks);
-        }
+      // Handle job exit when all tracked tasks finished
+      if (allTrackedTaskFinished(counter)) {
+        LOG.info("Application finished due to all tracked executors finished.");
+        break;
       }
 
       // Pause before refresh job status
@@ -718,6 +714,61 @@ private boolean monitor() {
     return status == FinalApplicationStatus.SUCCEEDED;
   }
 
+  private boolean allTrackedTaskFinished(int counter) {
+    int numTotalTrackedTasks = session.getTotalTrackedTasks();
+    if (numTotalTrackedTasks > 0) {
+      int numCompletedTrackedTasks = session.getNumCompletedTrackedTasks();
+      if (numCompletedTrackedTasks == numTotalTrackedTasks) {
+        if (allTrackedTaskFinishedTime == 0L) {
+          allTrackedTaskFinishedTime = System.currentTimeMillis();
+        }
+
+        if (!waitUntrackedTaskWithTimeout()) {
+          Utils.printCompletedTrackedTasks(numCompletedTrackedTasks, numTotalTrackedTasks);
+          return true;
+        }
+      }
+
+      // Reduce logging frequency to every 100s.
+      if (counter % 20 == 1) {
+        Utils.printCompletedTrackedTasks(numCompletedTrackedTasks, numTotalTrackedTasks);
+      }
+    }
+
+    return false;
+  }
+
+  private boolean waitUntrackedTaskWithTimeout() {
+    Map<String, Long> untrackedTaskTimeoutIndex = Utils.getUntrackedTaskWithTimeouts(tonyConf);
+    if (untrackedTaskTimeoutIndex == null || untrackedTaskTimeoutIndex.size() <= 0) {
+      return false;
+    }
+
+    int finishedJobtypeSize = 0;
+    for (Map.Entry<String, Long> untrackedTask : untrackedTaskTimeoutIndex.entrySet()) {
+      String jobtype = untrackedTask.getKey();
+      Long timeout = untrackedTask.getValue();
+
+      // When untracked task finished, ignore timeout
+      if (session.isAllTasksCompletedWithJobtype(jobtype)) {
+        LOG.info("Untracked job type: [" + jobtype + "] has finished.");
+        finishedJobtypeSize += 1;
+        continue;
+      }
+
+      // If waited time > timeout, it will return false directly.
+      if (System.currentTimeMillis() - allTrackedTaskFinishedTime > timeout) {
+        LOG.info("Untracked job type: [" + jobtype + "] exceed timeout of " + timeout + " and app will finish.");
+        return false;
+      }
+    }
+
+    if (finishedJobtypeSize == untrackedTaskTimeoutIndex.size()) {
+      return false;
+    }
+    return true;
+  }
+
   /**
    * Returns the tasks whose containers have launched but not called {@link ApplicationRpc#registerWorkerSpec} yet.
    */
diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java b/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java
@@ -270,6 +270,10 @@ public static String getContainerDockerMountKey() {
   public static final String UNTRACKED_JOBTYPES = TONY_APPLICATION_PREFIX + "untracked.jobtypes";
   public static final String UNTRACKED_JOBTYPES_DEFAULT = "ps";
 
+  // Specified untracked job type timeout when all tracked tasks finished
+  public static final String UNTRACKED_JOBTYPE_TIMEOUT_REGEX = TONY_APPLICATION_PREFIX + "([a-z]+)\\.untracked.timeout";
+  public static final long UNTRACKED_JOBTYPE_TIMEOUT_DEFAULT = 0;
+
   // Job types that we don't wait to finish and ignore its failure.
   public static final String SIDECAR_JOBTYPES = TONY_APPLICATION_PREFIX + "sidecar.jobtypes";
   public static final String DEFAULT_SIDECAR_JOBTYPES = SIDECAR_TB_ROLE_NAME;
diff --git a/tony-core/src/main/java/com/linkedin/tony/TonySession.java b/tony-core/src/main/java/com/linkedin/tony/TonySession.java
@@ -20,6 +20,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
@@ -192,6 +193,13 @@ public int getNumCompletedTrackedTasks() {
         .flatMap(entry -> Arrays.stream(entry.getValue())).filter(task -> task != null && task.isCompleted()).count();
   }
 
+  public boolean isAllTasksCompletedWithJobtype(String jobType) {
+    long taskSize = jobTasks.entrySet().stream().filter(entry -> entry.getKey().equals(jobType)).count();
+    long finishedSize = jobTasks.entrySet().stream().filter(entry -> entry.getKey().equals(jobType))
+            .flatMap(entry -> Arrays.stream(entry.getValue())).filter(task -> task != null && task.isCompleted()).count();
+    return  finishedSize == taskSize;
+  }
+
   public int getNumFailedTasks() {
     return (int) jobTasks.values().stream().flatMap(arr -> Arrays.stream(arr)).filter(task -> task != null && task.isFailed()).count();
   }
diff --git a/tony-core/src/main/java/com/linkedin/tony/util/Utils.java b/tony-core/src/main/java/com/linkedin/tony/util/Utils.java
@@ -459,6 +459,19 @@ public static int getNumTotalTasks(Configuration conf) {
         .sum();
   }
 
+  public static Map<String, Long> getUntrackedTaskWithTimeouts(Configuration conf) {
+    Map<String, String> map = conf.getValByRegex(TonyConfigurationKeys.UNTRACKED_JOBTYPE_TIMEOUT_REGEX);
+    return map.entrySet().stream()
+            .collect(Collectors.toMap(entry -> {
+              Matcher matcher = Pattern.compile(TonyConfigurationKeys.UNTRACKED_JOBTYPE_TIMEOUT_REGEX).matcher(entry.getKey());
+              if (matcher.matches()) {
+                return matcher.group(1);
+              } else {
+                return null;
+              }
+            }, entry -> Long.valueOf(entry.getValue())));
+  }
+
   /**
    * Extracts TensorFlow job name from configuration key of the form "tony.*.instances".
    * @param confKey Name of the configuration key
diff --git a/tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java b/tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java
@@ -603,6 +603,77 @@ public void testTonyAllocationTimeoutShouldFail() throws ParseException, IOExcep
     Assert.assertEquals(exitCode, -1);
   }
 
+  /**
+   * Test untracked tasks specified timeout of 60s.
+   * The untracked task will block all time and it will finished until tracked task finished 60s
+   */
+  @Test
+  public void testUntrackedTaskTimeout60sShouldPass() throws ParseException, IOException {
+    client.init(new String[]{
+            "--src_dir", "tony-core/src/test/resources/scripts",
+            "--hdfs_classpath", libPath,
+            "--shell_env", "ENV_CHECK=ENV_CHECK",
+            "--container_env", Constants.SKIP_HADOOP_PATH + "=true",
+            "--python_venv", "tony-core/src/test/resources/test.zip",
+            "--conf", "tony.ps.instances=1",
+            "--conf", "tony.worker.instances=1",
+            "--conf", "tony.ps.command=python sleep_all.py",
+            "--conf", "tony.worker.command=python sleep_30.py",
+            "--conf", "tony.application.untracked.jobtypes=ps",
+            "--conf", "tony.application.ps.untracked.timeout=60000"
+    });
+
+    int exitCode = client.start();
+    Assert.assertEquals(exitCode, 0);
+  }
+
+
+  /**
+   * Test untracked tasks specified timeout of 60s.
+   * But the untracked task will exit 0 quickly
+   * This case to verify the untracked-timeout-mechanism no block job finish
+   */
+  @Test
+  public void testUntrackedTaskFinishedShouldPass() throws ParseException, IOException {
+    client.init(new String[]{
+            "--src_dir", "tony-core/src/test/resources/scripts",
+            "--hdfs_classpath", libPath,
+            "--shell_env", "ENV_CHECK=ENV_CHECK",
+            "--container_env", Constants.SKIP_HADOOP_PATH + "=true",
+            "--python_venv", "tony-core/src/test/resources/test.zip",
+            "--conf", "tony.ps.instances=1",
+            "--conf", "tony.worker.instances=1",
+            "--conf", "tony.ps.command=python exit_0.py",
+            "--conf", "tony.worker.command=python sleep_30.py",
+            "--conf", "tony.application.untracked.jobtypes=ps",
+            "--conf", "tony.application.ps.untracked.timeout=60000"
+    });
+
+    int exitCode = client.start();
+    Assert.assertEquals(exitCode, 0);
+  }
+
+  @Test
+  public void testUntrackedTaskFailedShouldFail() throws ParseException, IOException {
+    client.init(new String[]{
+            "--src_dir", "tony-core/src/test/resources/scripts",
+            "--hdfs_classpath", libPath,
+            "--shell_env", "ENV_CHECK=ENV_CHECK",
+            "--container_env", Constants.SKIP_HADOOP_PATH + "=true",
+            "--python_venv", "tony-core/src/test/resources/test.zip",
+            "--conf", "tony.chief.instances=1",
+            "--conf", "tony.ps.instances=1",
+            "--conf", "tony.worker.instances=1",
+            "--conf", "tony.ps.command=python sleep_all.py",
+            "--conf", "tony.worker.command=python exit_1.py",
+            "--conf", "tony.chief.command=python sleep_30.py",
+            "--conf", "tony.application.untracked.jobtypes=ps,worker"
+    });
+
+    int exitCode = client.start();
+    Assert.assertEquals(exitCode, -1);
+  }
+
   /**
    * When enable the sidecar tensorboard, it will start the sidecar executor(tensorboard role).
    */
diff --git a/tony-core/src/test/java/com/linkedin/tony/util/TestUtils.java b/tony-core/src/test/java/com/linkedin/tony/util/TestUtils.java
@@ -259,6 +259,17 @@ public void testGetAllJobTypes() {
             new HashSet(Arrays.asList("worker", "evaluator")));
   }
 
+  @Test
+  public void testGetUntrackedTaskTimeout() {
+    Configuration conf = new Configuration();
+    conf.addResource("tony-default.xml");
+    conf.set("tony.application.untracked.jobtypes", "ps,evaluator,worker");
+    conf.setLong("tony.application.worker.untracked.timeout", 1000);
+
+    long timeout = Utils.getUntrackedTaskWithTimeouts(conf).get("worker");
+    assertEquals(timeout, 1000);
+  }
+
   @Test
   public void testGetNumTotalTasks() {
     Configuration conf = new Configuration();
diff --git a/tony-core/src/test/resources/scripts/sleep_all.py b/tony-core/src/test/resources/scripts/sleep_all.py
@@ -0,0 +1,8 @@
+#
+# Copyright 2021 LinkedIn Corporation. All rights reserved. Licensed under the BSD-2 Clause license.
+# See LICENSE in the project root for license information.
+#
+import time
+
+while True:
+    time.sleep(1)
