Add option to return SUCCEED when training is completed with some fai…

…led job tasks
tony-framework · Jan 17, 2020 · be28773 · be28773
1 parent 9dcae46
commit be28773
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java b/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java
@@ -246,6 +246,11 @@ public static String getContainerDockerKey() {
   // Job types that we will short circuit when it failed
   public static final String STOP_ON_FAILURE_JOBTYPES = TONY_APPLICATION_PREFIX + "stop.on.failure.jobtypes";
 
+  // Whether to return FAILED when training is completed with some failed job tasks
+  public static final String FAILED_ON_COMPLETED_WITH_FAILURE_ENABLED = TONY_APPLICATION_PREFIX
+      + "failed_on_complete_with_failure.enabled";
+  public static final boolean DEFAULT_FAILED_ON_COMPLETED_WITH_FAILURE_ENABLED = true;
+
   // Training chief configuration
   public static final String CHIEF_PREFIX = TONY_PREFIX + "chief.";
 

diff --git a/tony-core/src/main/java/com/linkedin/tony/tensorflow/TonySession.java b/tony-core/src/main/java/com/linkedin/tony/tensorflow/TonySession.java
@@ -6,6 +6,7 @@
 
 import com.google.common.base.Preconditions;
 import com.linkedin.tony.Constants;
+import com.linkedin.tony.TonyConfigurationKeys;
 import com.linkedin.tony.rpc.TaskInfo;
 import com.linkedin.tony.rpc.impl.TaskStatus;
 import com.linkedin.tony.util.Utils;
@@ -262,7 +263,12 @@ public void onTaskCompleted(String jobName, String jobIndex, int exitCode) {
       if (isChief(jobName, jobIndex) || shouldStopOnFailure(jobName)) {
         trainingFinished = true;
       }
-      setFinalStatus(FinalApplicationStatus.FAILED, "Exit status: " + exitCode);
+      boolean failedOnCompletedWithFailure = tonyConf.getBoolean(
+          TonyConfigurationKeys.FAILED_ON_COMPLETED_WITH_FAILURE_ENABLED,
+          TonyConfigurationKeys.DEFAULT_FAILED_ON_COMPLETED_WITH_FAILURE_ENABLED);
+      if (failedOnCompletedWithFailure) {
+        setFinalStatus(FinalApplicationStatus.FAILED, "Exit status: " + exitCode);
+      }
     }
   }
 
@@ -305,7 +311,10 @@ public void updateSessionStatus() {
       }
     }
 
-    if (failureCount > 0) {
+    boolean failedOnCompletedWithFailure = tonyConf.getBoolean(
+        TonyConfigurationKeys.FAILED_ON_COMPLETED_WITH_FAILURE_ENABLED,
+        TonyConfigurationKeys.DEFAULT_FAILED_ON_COMPLETED_WITH_FAILURE_ENABLED);
+    if (failedOnCompletedWithFailure && failureCount > 0) {
       setFinalStatus(FinalApplicationStatus.FAILED,
           "At least one job task exited with non-zero status, failedCnt=" + failureCount);
     } else {

diff --git a/tony-core/src/main/resources/tony-default.xml b/tony-core/src/main/resources/tony-default.xml
@@ -191,6 +191,13 @@
     <value>ps</value>
   </property>
 
+  <!-- App returned status configurations -->
+  <property>
+    <description>Whether to return FAILED when training is completed with some failed job tasks</description>
+    <name>tony.application.failed_on_complete_with_failure.enabled</name>
+    <value>true</value>
+  </property>
+
   <!-- Docker configurations -->
   <property>
     <description>Whether we use docker container to launch the tasks</description>