Adds pause-resume elasticity if the proxy backend is used.

Added the changes to the jobset for elastic training to enable elasticity.
Added changes to launch_trainer so that the pause_resume decorator is used.
Set logging.raiseExceptions=True so that DATA_LOSS errors that occur in debug/info/other log calls are raise exceptions immediately.

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -343,6 +343,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
             f"--resource_manager_address=localhost:{_PATHWAYS_RESOURCE_MANAGER_PORT}",
             f"--server_port={_PATHWAYS_PROXY_PORT}",
             f"--gcs_scratch_location={staging_location}",
+            f"--num_elastic_slices={cfg.accelerator.num_replicas}",
         ]
         cmd_args.extend(xla_flags_from_options(self._xla_options).split())
 
@@ -604,14 +605,21 @@ def _build_pathways_worker_job(
         annotations.update(
             {"alpha.jobset.sigs.k8s.io/exclusive-topology": "cloud.google.com/gke-nodepool"}
         )
+        # Default value for suspend and resume.
+        # References:
+        # https://github.com/google/pathways-job/blob/4417de7aa23d3c2316e400a3a327512834374475/internal/controller/pathwaysjob_controller.go#L651
+        backoffLimit = system.vms_per_slice * _PATHWAYS_BACK_OFF_LIMIT
+
+        # This backoffLimit is just for verifying elastic pause-resume
+        backoffLimit *= 1000
 
         spec = dict(
             parallelism=system.vms_per_slice,
             completions=system.vms_per_slice,
             # Default value for suspend and resume.
             # References:
             # https://github.com/google/pathways-job/blob/4417de7aa23d3c2316e400a3a327512834374475/internal/controller/pathwaysjob_controller.go#L651
-            backoffLimit=system.vms_per_slice * _PATHWAYS_BACK_OFF_LIMIT,
+            backoffLimit=backoffLimit,
             template=self._build_pathways_worker_pod(pathways_worker_replicated_job_index),
         )
         worker_job = dict(

axlearn/common/launch_trainer.py

@@ -3,6 +3,7 @@
 """Utilities to launch a trainer."""
 
 import json
+import logging as py_logging
 import os
 from typing import Any, Optional
 
@@ -15,6 +16,8 @@
 from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape
 from axlearn.experiments import TrainerConfigFn, get_named_trainer_config
 
+py_logging.raiseException = True
+
 # Trainer-specific flags.
 flags.DEFINE_string(
     "module",
@@ -158,8 +161,24 @@ def run_trainer(trainer_config: SpmdTrainer.Config) -> Any:
                 f,
             )
 
-    trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
-    prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
-    output = trainer.run(prng_key)
-    measurement.record_event(measurement.Event.END_JOB)
+    def run() -> Any:
+        trainer: SpmdTrainer = trainer_config.instantiate(parent=None)
+        prng_key = jax.random.PRNGKey(seed=FLAGS.trainer_prng_seed)
+        output = trainer.run(prng_key)
+        measurement.record_event(measurement.Event.END_JOB)
+        return output
+
+    if FLAGS.jax_backend == "proxy":
+        # pylint: disable-next=import-error,import-outside-toplevel
+        from pathwaysutils.elastic import manager
+        elastic_manager = manager.Manager()
+        max_retries = 5
+        timeout = 10 * 60  # ten minutes
+        run = elastic_manager.pause_resume(
+            max_retries=max_retries,
+            timeout=timeout,
+        )(run)
+
+    output = run()
+
     return output

pyproject.toml

@@ -108,7 +108,7 @@ gcp = [
 tpu = [
     "axlearn[gcp]",
     "jax[tpu]==0.5.3",  # must be >=0.4.19 for compat with v5p.
-    "pathwaysutils==0.1.1",  # For JAX+Pathways single-controller accelerator coordinator.
+    "pathwaysutils @ git+https://github.com/AI-Hypercomputer/pathways-utils.git@test_796970321",  # For JAX+Pathways single-controller accelerator coordinator.
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
 vertexai_tensorboard = [