Use the public API for is_error_due_to_slice_down from pathwaysutils

Author: lukebaumann

axlearn/common/launch_trainer.py

@@ -147,7 +147,7 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
                 f,
             )
 
-    if FLAGS.jax_backend == "proxy":
+    if False and FLAGS.jax_backend == "proxy":
         # pylint: disable-next=import-error,import-outside-toplevel
         from pathwaysutils.elastic import manager
         elastic_manager = manager.Manager()
@@ -158,7 +158,7 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
                 output = trainer.run(prng_key)
                 break
             except jax.errors.JaxRuntimeError as error:
-                if not elastic_manager._is_error_due_to_slice_down(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)