Client exit when am reached startup-timeout

tony-core/src/main/java/com/linkedin/tony/TonyClient.java

@@ -144,6 +144,7 @@ public class TonyClient implements AutoCloseable {
   private String hdfsClasspath = null;
   private String executes;
   private long appTimeout;
+  private long amStartupTimeout;
   private boolean secureMode;
   private Map<String, String> containerEnv = new HashMap<>();
   private String hadoopFrameworkLocation = null;
@@ -544,6 +545,9 @@ public boolean init(String[] args) throws ParseException, IOException {
     appTimeout = tonyConf.getInt(TonyConfigurationKeys.APPLICATION_TIMEOUT,
         TonyConfigurationKeys.DEFAULT_APPLICATION_TIMEOUT);
 
+    amStartupTimeout = tonyConf.getInt(TonyConfigurationKeys.AM_STARTUP_TIMEOUT,
+        TonyConfigurationKeys.DEFAULT_AM_STARTUP_TIMEOUT);
+
     List<String> executionEnvPair = new ArrayList<>();
     if (tonyConf.get(TonyConfigurationKeys.EXECUTION_ENV) != null) {
       String[] envs = tonyConf.getStrings(TonyConfigurationKeys.EXECUTION_ENV);
@@ -1111,6 +1115,15 @@ public boolean monitorApplication() throws YarnException, IOException, Interrupt
         break;
       }
 
+      if (amStartupTimeout > 0 && YarnApplicationState.ACCEPTED == appState
+              && System.currentTimeMillis() - clientStartTime > amStartupTimeout) {
+        LOG.info("Reached AM startup timeout for application. Killing application, ApplicationId: " + appId.getId());
+        forceKillApplication();
+        result = false;
+        break;
+      }
+
+
       if (appTimeout > 0) {
         if (System.currentTimeMillis() > (clientStartTime + appTimeout)) {
           LOG.info("Reached client specified timeout for application. Killing application"

tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java

@@ -190,6 +190,9 @@ private TonyConfigurationKeys() {
   public static final String AM_WAIT_CLIENT_STOP_TIMEOUT = AM_PREFIX + "wait-client-signal-stop-timeout-sec";
   public static final int DEFAULT_AM_WAIT_CLIENT_STOP_TIMEOUT = 15;
 
+  public static final String AM_STARTUP_TIMEOUT = AM_PREFIX + "startup-timeout";
+  public static final int DEFAULT_AM_STARTUP_TIMEOUT = 0;
+
   // Keys/default values for configurable TensorFlow job names
   public static final String INSTANCES_REGEX = "tony\\.([a-z]+)\\.instances";
   public static final String MAX_TOTAL_RESOURCES_REGEX = TONY_TASK_PREFIX + "max-total-([a-z]+)";

tony-core/src/main/java/com/linkedin/tony/util/Utils.java

@@ -452,9 +452,11 @@ private static void ensureStagedTasksIntegrity(List<String> prepareStageTasks, L
   }
 
   public static Set<String> getAllJobTypes(Configuration conf) {
-    return conf.getValByRegex(TonyConfigurationKeys.INSTANCES_REGEX).keySet().stream()
-        .map(Utils::getTaskType)
-        .collect(Collectors.toSet());
+    synchronized (Utils.class) {
+      return conf.getValByRegex(TonyConfigurationKeys.INSTANCES_REGEX).keySet().stream()
+              .map(Utils::getTaskType)
+              .collect(Collectors.toSet());
+    }
   }
 
   public static int getNumTotalTasks(Configuration conf) {

tony-core/src/main/resources/tony-default.xml

@@ -53,6 +53,12 @@
     <value>0</value>
   </property>
 
+  <property>
+    <description>Max waiting startup time of the AM before killing it, in milliseconds.</description>
+    <name>tony.am.startup-timeout</name>
+    <value>0</value>
+  </property>
+
   <property>
     <description>The machine learning framework that will be used for this job - tensorflow or pytorch.</description>
     <name>tony.application.framework</name>

tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java

@@ -11,7 +11,10 @@
 import com.linkedin.tony.client.TaskUpdateListener;
 import com.linkedin.tony.rpc.TaskInfo;
 import com.linkedin.tony.rpc.impl.TaskStatus;
+
+import java.util.ArrayList;
 import java.util.HashSet;
+
 import org.apache.commons.cli.ParseException;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.CommonConfigurationKeys;
@@ -28,7 +31,9 @@
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.nio.file.Files;
+import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 
 import static com.linkedin.tony.TonyConfigurationKeys.TASK_HEARTBEAT_INTERVAL_MS;
 import static com.linkedin.tony.TonyConfigurationKeys.TASK_MAX_MISSED_HEARTBEATS;
@@ -603,6 +608,60 @@ public void testTonyAllocationTimeoutShouldFail() throws ParseException, IOExcep
     Assert.assertEquals(exitCode, -1);
   }
 
+  @Test
+  public void testTonyAMStartupTimeoutShouldFail() throws ParseException, IOException {
+    List<CompletableFuture<Integer>> tasks = new ArrayList<>();
+    for (int i = 0; i < 10; i++) {
+      tasks.add(mockedTask());
+    }
+
+    TonyClient client = new TonyClient(conf);
+    client.init(new String[]{
+            "--src_dir", "tony-core/src/test/resources/scripts",
+            "--executes", "python check_env_and_venv.py",
+            "--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.worker.instances=1",
+            "--conf", "tony.am.memory=2g",
+            "--conf", "tony.am.startup-timeout=10000"
+    });
+    int exitCode = client.start();
+    Assert.assertEquals(exitCode, -1);
+
+    for (CompletableFuture<Integer> task : tasks) {
+      try {
+        task.cancel(true);
+      } catch (Exception e) {
+        // ignore
+      }
+    }
+  }
+
+  private CompletableFuture<Integer> mockedTask() {
+    return CompletableFuture.supplyAsync(() -> {
+      TonyClient client = new TonyClient(conf);
+      try {
+        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=4",
+                "--conf", "tony.ps.command=python sleep_30.py",
+                "--conf", "tony.worker.command=python check_env_and_venv.py"
+        });
+      } catch (Exception e) {
+        e.printStackTrace();
+      }
+      int exitCode = client.start();
+      return exitCode;
+    });
+  }
+
   /**
    * When enable the sidecar tensorboard, it will start the sidecar executor(tensorboard role).
    */