Moved elastic logic from launch_trainer_main to launch_trainer.

Added guards to only use fast-resume if the proxy backend is used.

Author: lukebaumann

axlearn/common/launch_trainer.py

@@ -149,6 +149,21 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
 
     trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
     prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
-    output = trainer.run(prng_key)
+
+    if FLAGS.jax_backend == "proxy":
+        # pylint: disable-next=import-error,import-outside-toplevel
+        from pathwaysutils.elastic import manager
+        elastic_manager = manager.Manager()
+        while True:
+            try:
+                output = trainer.run(prng_key)
+            except jax.errors.JaxRuntimeError as error:
+                if not elastic_manager._is_error_due_to_slice_down(error):
+                    raise
+                ten_minutes = 10 * 60
+                elastic_manager.wait_for_slices(timeout=ten_minutes)
+    else:
+        output = trainer.run(prng_key)
+
     measurement.record_event(measurement.Event.END_JOB)
     return output

axlearn/common/launch_trainer_main.py

@@ -6,25 +6,15 @@
 
 from axlearn.common import launch, launch_trainer, measurement
 from axlearn.common.config import config_for_function
-from pathwaysutils.elastic import manager
 
 
 def main(_):
-    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
-          ten_minutes = 10 * 60
-          elastic_manager.wait_for_slices(timeout=ten_minutes)
+    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)
 
 
 if __name__ == "__main__":