Adding support for Pathways proxy

axlearn/common/launch.py

@@ -113,7 +113,9 @@ def setup():
     logging.info("Devices: %s", devices)
     local_devices = jax.local_devices()
     logging.info("Local Devices: %s", local_devices)
-    if not devices or not all(device.platform == FLAGS.jax_backend for device in devices):
+    if FLAGS.jax_backend != "proxy" and (
+        not devices or not all(device.platform == FLAGS.jax_backend for device in devices)
+    ):
         raise RuntimeError(f"Expected backend {FLAGS.jax_backend}. Got {devices}.")
     if FLAGS.data_dir:
         # TODO(ruoming): Get rid of --data_dir and use only env var DATA_DIR.

axlearn/common/launch_trainer_main.py

@@ -2,6 +2,7 @@
 
 """Main function for launching the trainer."""
 
+import pathwaysutils  # pylint: disable=unused-import
 from absl import app, flags
 
 from axlearn.common import launch, launch_trainer, measurement

axlearn/common/utils_spmd.py

@@ -53,7 +53,7 @@ def setup(
         if initialization_timeout is not None:
             init_kwargs["initialization_timeout"] = initialization_timeout
 
-        if jax_backend == "tpu":
+        if jax_backend in ("tpu", "proxy"):
             if not (
                 distributed_coordinator is None and num_processes is None and process_id is None
             ):
@@ -115,5 +115,6 @@ def setup(
                     f"({initialization_timeout} seconds)."
                 )
         else:
-            jax.distributed.initialize(**init_kwargs)
-            _jax_distributed_initialized = True
+            if jax_backend != "proxy":
+                jax.distributed.initialize(**init_kwargs)
+                _jax_distributed_initialized = True

pyproject.toml

@@ -143,6 +143,11 @@ orbax = [
     "orbax-checkpoint==0.5.23",
 ]
 
+# Pathways utilities.
+pathways = [
+    "pathwaysutils@git+https://github.com/google/[email protected]", # for JAX+Pathways single-controller accelerator coordinator
+]
+
 [tool.flit.module]
 # This defines the import name. https://flit.pypa.io/en/stable/pyproject_toml.html#module-section
 name = "axlearn"