allow inference test timeouts (and pass from commandline) (#4932)

Author: Chris Elion

Project/Assets/ML-Agents/Examples/SharedAssets/Scripts/ModelOverrider.cs

@@ -28,6 +28,7 @@ public class ModelOverrider : MonoBehaviour
         const string k_CommandLineModelOverrideDirectoryFlag = "--mlagents-override-model-directory";
         const string k_CommandLineModelOverrideExtensionFlag = "--mlagents-override-model-extension";
         const string k_CommandLineQuitAfterEpisodesFlag = "--mlagents-quit-after-episodes";
+        const string k_CommandLineQuitAfterSeconds = "--mlagents-quit-after-seconds";
         const string k_CommandLineQuitOnLoadFailure = "--mlagents-quit-on-load-failure";
 
         // The attached Agent
@@ -45,6 +46,9 @@ public class ModelOverrider : MonoBehaviour
         // Will default to 1 if override models are specified, otherwise 0.
         int m_MaxEpisodes;
 
+        // Deadline - exit if the time exceeds this
+        DateTime m_Deadline = DateTime.MaxValue;
+
         int m_NumSteps;
         int m_PreviousNumSteps;
         int m_PreviousAgentCompletedEpisodes;
@@ -89,6 +93,8 @@ public static string GetOverrideBehaviorName(string originalBehaviorName)
         void GetAssetPathFromCommandLine()
         {
             var maxEpisodes = 0;
+            var timeoutSeconds = 0;
+
             string[] commandLineArgsOverride = null;
             if (!string.IsNullOrEmpty(debugCommandLineOverride) && Application.isEditor)
             {
@@ -120,6 +126,10 @@ void GetAssetPathFromCommandLine()
                 {
                     Int32.TryParse(args[i + 1], out maxEpisodes);
                 }
+                else if (args[i] == k_CommandLineQuitAfterSeconds && i < args.Length - 1)
+                {
+                    Int32.TryParse(args[i + 1], out timeoutSeconds);
+                }
                 else if (args[i] == k_CommandLineQuitOnLoadFailure)
                 {
                     m_QuitOnLoadFailure = true;
@@ -132,6 +142,13 @@ void GetAssetPathFromCommandLine()
                 m_MaxEpisodes = maxEpisodes > 0 ? maxEpisodes : 1;
                 Debug.Log($"setting m_MaxEpisodes to {maxEpisodes}");
             }
+
+            if (timeoutSeconds > 0)
+            {
+                m_Deadline = DateTime.Now + TimeSpan.FromSeconds(timeoutSeconds);
+                Debug.Log($"setting deadline to {timeoutSeconds} from now.");
+
+            }
         }
 
         void OnEnable()
@@ -172,9 +189,21 @@ void FixedUpdate()
                     Application.Quit(0);
 #if UNITY_EDITOR
                     EditorApplication.isPlaying = false;
+#endif
+                }
+                else if (DateTime.Now >= m_Deadline)
+                {
+                    Debug.Log(
+                        $"Deadline exceeded. " +
+                        $"{TotalCompletedEpisodes}/{m_MaxEpisodes} episodes and " +
+                        $"{TotalNumSteps}/{m_MaxEpisodes * m_Agent.MaxStep} steps completed. Exiting.");
+                    Application.Quit(0);
+#if UNITY_EDITOR
+                    EditorApplication.isPlaying = false;
 #endif
                 }
             }
+
             m_NumSteps++;
         }
 

ml-agents/tests/yamato/training_int_tests.py

@@ -133,6 +133,11 @@ def run_inference(env_path: str, output_path: str, model_extension: str) -> bool
 
     log_output_path = f"{get_base_output_path()}/inference.{model_extension}.txt"
 
+    # 10 minutes for inference is more than enough
+    process_timeout = 10 * 60
+    # Try to gracefully exit a few seconds before that.
+    model_override_timeout = process_timeout - 15
+
     exe_path = exes[0]
     args = [
         exe_path,
@@ -147,10 +152,11 @@ def run_inference(env_path: str, output_path: str, model_extension: str) -> bool
         "1",
         "--mlagents-override-model-extension",
         model_extension,
+        "--mlagents-quit-after-seconds",
+        str(model_override_timeout),
     ]
     print(f"Starting inference with args {' '.join(args)}")
-    timeout = 15 * 60  # 15 minutes for inference is more than enough
-    res = subprocess.run(args, timeout=timeout)
+    res = subprocess.run(args, timeout=process_timeout)
     end_time = time.time()
     if res.returncode != 0:
         print("Error running inference!")
@@ -166,7 +172,7 @@ def run_inference(env_path: str, output_path: str, model_extension: str) -> bool
         timer_data = json.load(f)
 
     gauges = timer_data.get("gauges", {})
-    rewards = gauges.get("Override_3DBall.CumulativeReward")
+    rewards = gauges.get("Override_3DBall.CumulativeReward", {})
     max_reward = rewards.get("max")
     if max_reward is None:
         print(