Clean up the remote executor stuff and create a dedicated heartbeat

[python_utils.git] / ml / model_trainer.py

diff --git a/ml/model_trainer.py b/ml/model_trainer.py
index 79ce7062b5b4a05616cddcf7b3d35d59cfbef007..041f0f805cc5958fb948a48cc9bc160bc8578956 100644 (file)
--- a/ml/model_trainer.py
+++ b/ml/model_trainer.py
@@ -22,6 +22,7 @@ from ansi import bold, reset
 import argparse_utils
 import config
 from decorator_utils import timed
+import executors
 import parallelize as par
 
 logger = logging.getLogger(__file__)
@@ -171,9 +172,11 @@ class TrainingBlueprint(ABC):
                         )
 
         if not self.spec.quiet:
+            executors.DefaultExecutors().shutdown()
             msg = f"Done training; best test set score was: {best_test_score:.1f}%"
             print(msg)
             logger.info(msg)
+
         scaler_filename, model_filename, model_info_filename = (
             self.maybe_persist_scaler_and_model(
                 best_training_score,