Adding basic elastic training

Author: lukebaumann

axlearn/common/launch_trainer_main.py

@@ -6,15 +6,24 @@
 
 from axlearn.common import launch, launch_trainer, measurement
 from axlearn.common.config import config_for_function
+from pathwaysutils.elastic import manager
 
 
 def main(_):
-    measurement.initialize(flags.FLAGS)
-    launch.setup()
-    trainer_config = launch_trainer.get_trainer_config()
-    trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
-    measurement.start_monitoring()
-    launch_trainer.run_trainer(trainer_config)
+    elastic_manager = manager.Manager()
+    while True:
+        try:
+            measurement.initialize(flags.FLAGS)
+            launch.setup()
+            trainer_config = launch_trainer.get_trainer_config()
+            trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
+            measurement.start_monitoring()
+            launch_trainer.run_trainer(trainer_config)
+            break
+        except jax.errors.JaxRuntimeError as error:
+          if not elastic_manager._is_error_due_to_slice_down(error):
+            raise
+          elastic_manager.wait_for_slices(timeout)
 
 
 if __name__ == "__main__":