make sure top-level timer is closed before writing

ml-agents/mlagents/trainers/learn.py

@@ -30,7 +30,7 @@
 from mlagents_envs.side_channel.side_channel import SideChannel
 from mlagents_envs.side_channel.engine_configuration_channel import EngineConfig
 from mlagents_envs.exception import UnityEnvironmentException
-from mlagents_envs.timers import hierarchical_timer
+from mlagents_envs.timers import hierarchical_timer, get_timer_tree
 from mlagents.logging_util import create_logger
 
 
@@ -329,6 +329,18 @@ def run_training(run_seed: int, options: RunOptions) -> None:
         tc.start_learning(env_manager)
     finally:
         env_manager.close()
+        write_timing_tree(summaries_dir, options.run_id)
+
+
+def write_timing_tree(summaries_dir: str, run_id: str) -> None:
+    timing_path = f"{summaries_dir}/{run_id}_timers.json"
+    try:
+        with open(timing_path, "w") as f:
+            json.dump(get_timer_tree(), f, indent=4)
+    except FileNotFoundError:
+        logging.warning(
+            f"Unable to save to {timing_path}. Make sure the directory exists"
+        )
 
 
 def create_sampler_manager(sampler_config, run_seed=None):

ml-agents/mlagents/trainers/trainer_controller.py

@@ -4,7 +4,6 @@
 
 import os
 import sys
-import json
 import logging
 from typing import Dict, Optional, Set
 from collections import defaultdict
@@ -18,7 +17,7 @@
     UnityCommunicationException,
 )
 from mlagents.trainers.sampler_class import SamplerManager
-from mlagents_envs.timers import hierarchical_timer, get_timer_tree, timed
+from mlagents_envs.timers import hierarchical_timer, timed
 from mlagents.trainers.trainer import Trainer
 from mlagents.trainers.meta_curriculum import MetaCurriculum
 from mlagents.trainers.trainer_util import TrainerFactory
@@ -106,16 +105,6 @@ def _save_model_when_interrupted(self):
         )
         self._save_model()
 
-    def _write_timing_tree(self) -> None:
-        timing_path = f"{self.summaries_dir}/{self.run_id}_timers.json"
-        try:
-            with open(timing_path, "w") as f:
-                json.dump(get_timer_tree(), f, indent=4)
-        except FileNotFoundError:
-            self.logger.warning(
-                f"Unable to save to {timing_path}. Make sure the directory exists"
-            )
-
     def _export_graph(self):
         """
         Exports latest saved models to .nn format for Unity embedding.
@@ -231,7 +220,6 @@ def start_learning(self, env_manager: EnvManager) -> None:
             pass
         if self.train_model:
             self._export_graph()
-        self._write_timing_tree()
 
     def end_trainer_episodes(
         self, env: EnvManager, lessons_incremented: Dict[str, bool]