chore: import upstream snapshot with attribution

python/ray/train/v2/jax/config.py

@@ -0,0 +1,227 @@
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import ray
+from ray._private import ray_constants
+from ray.train._internal.utils import get_address_and_port
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import Backend, BackendConfig
+from ray.train.constants import (
+    DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+    JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+)
+from ray.train.v2._internal.util import TrainingFramework
+from ray.util import PublicAPI
+from ray.util.tpu import get_tpu_coordinator_env_vars, get_tpu_worker_resources
+
+logger = logging.getLogger(__name__)
+
+# JAX multi-slice (megascale) coordination env vars whose values are
+# authoritatively computed by the controller for the current worker group.
+# These must always be overwritten on each worker (even when the underlying
+# TPU node provider has already set them in the pod environment), because
+# stale values from a previous worker group configuration -- e.g. after a
+# slice was preempted and replaced -- would otherwise cause libtpu's
+# megascale topology coordinator to wait for slices that no longer exist
+# and hang TPU initialization.
+_JAX_MULTISLICE_OVERRIDE_KEYS = frozenset(
+    {
+        "MEGASCALE_COORDINATOR_ADDRESS",
+        "MEGASCALE_NUM_SLICES",
+        "MEGASCALE_SLICE_ID",
+    }
+)
+
+
+@PublicAPI(stability="alpha")
+@dataclass
+class JaxConfig(BackendConfig):
+    use_tpu: bool = False
+    use_gpu: bool = False
+
+    @property
+    def backend_cls(self):
+        return _JaxBackend
+
+    @property
+    def framework(self):
+        return TrainingFramework.JAX
+
+    def to_dict(self) -> Dict[str, Any]:
+        config_dict = {
+            "use_tpu": self.use_tpu,
+            "use_gpu": self.use_gpu,
+        }
+        return config_dict
+
+
+def _setup_jax_distributed_environment(
+    master_addr_with_port: str,
+    num_workers: int,
+    index: int,
+    use_tpu: bool,
+    use_gpu: bool,
+    resources_per_worker: dict,
+    jax_env_vars: Optional[dict] = None,
+):
+    """Set up distributed Jax training information.
+
+    This function should be called on each worker. It sets JAX environment
+    variables and initializes JAX distributed training.
+
+    Args:
+        master_addr_with_port: The master address with port for coordination.
+        num_workers: Total number of workers.
+        index: Index of this worker.
+        use_tpu: Whether to configure for TPU. If True and JAX_PLATFORMS is not
+            already set, it will be set to "tpu".
+        use_gpu: Whether to configure for GPU. If True and JAX_PLATFORMS is not
+            already set, it will be set to "cuda".
+        resources_per_worker: The resources per worker.
+        jax_env_vars: The JAX coordinator env vars to inject for multi-slice.
+            Multi-slice coordination keys (``MEGASCALE_*``) always overwrite
+            any pre-existing value in the environment; all other keys are
+            only set if they are not already present.
+    """
+    # Get JAX_PLATFORMS from environment if already set
+    jax_platforms = os.environ.get("JAX_PLATFORMS", "").lower()
+
+    if not jax_platforms and use_tpu:
+        os.environ["JAX_PLATFORMS"] = "tpu"
+        jax_platforms = "tpu"
+
+    if jax_env_vars:
+        for k, v in jax_env_vars.items():
+            # For multi-slice coordination keys, always override -- the
+            # controller's freshly computed value is the source of truth and
+            # may differ from what the TPU node provider baked into the pod
+            # environment (e.g. after a slice replacement following preemption).
+            # For all other keys, respect any pre-existing value.
+            if k in _JAX_MULTISLICE_OVERRIDE_KEYS or k not in os.environ:
+                os.environ[k] = v
+
+    if not jax_platforms and use_gpu:
+        os.environ["JAX_PLATFORMS"] = "cuda"
+        jax_platforms = "cuda"
+
+    if "cuda" in jax_platforms.split(","):
+        num_gpus_per_worker = resources_per_worker.get("GPU", 0)
+        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
+            str(i) for i in range(num_gpus_per_worker)
+        )
+
+    import jax
+
+    jax_platforms_list = jax_platforms.split(",")
+
+    if "tpu" in jax_platforms_list:
+        jax.distributed.initialize(master_addr_with_port, num_workers, index)
+        logger.info("Initialized JAX distributed on TPU.")
+    elif "cuda" in jax_platforms_list:
+        if num_gpus_per_worker > 0:
+            local_device_ids = list(range(num_gpus_per_worker))
+        else:
+            local_device_ids = 0
+        jax.distributed.initialize(
+            master_addr_with_port, num_workers, index, local_device_ids
+        )
+        logger.info("Initialized JAX distributed on CUDA.")
+    elif "cpu" in jax_platforms_list:
+        jax.distributed.initialize(master_addr_with_port, num_workers, index)
+        logger.info("Initialized JAX distributed on CPU.")
+
+
+def _shutdown_jax_distributed():
+    """Shutdown JAX distributed environment.
+
+    This function should be called on each worker during cleanup.
+    If JAX distributed was not initialized, this is a no-op.
+    """
+    try:
+        import jax
+
+        jax.distributed.shutdown()
+    except Exception as e:
+        logger.warning(f"Error during JAX distributed shutdown: {e}")
+
+
+class _JaxBackend(Backend):
+    def on_start(self, worker_group: WorkerGroup, backend_config: JaxConfig):
+        if not backend_config.use_tpu and not backend_config.use_gpu:
+            return
+
+        master_addr, master_port = worker_group.execute_single(0, get_address_and_port)
+        master_addr_with_port = f"{master_addr}:{master_port}"
+
+        if backend_config.use_tpu and hasattr(worker_group, "get_worker_group_context"):
+            num_slices = worker_group.get_worker_group_context().num_slices
+        else:
+            num_slices = 1
+
+        # Calculate the number of workers per slice for multi-slice env setup.
+        if backend_config.use_tpu and num_slices > 1:
+            # Handle the case where a user requests less workers than the total
+            # capacity of the TPU slice.
+            scaling_config = worker_group._train_run_context.scaling_config
+            workers_per_slice, _ = get_tpu_worker_resources(
+                topology=scaling_config.topology,
+                accelerator_type=scaling_config.accelerator_type,
+                resources_per_unit=scaling_config.resources_per_worker,
+                num_slices=1,
+            )
+        else:
+            # Assume even distribution based on the requested number of workers.
+            workers_per_slice = max(1, len(worker_group) // num_slices)
+
+        # Set up JAX distributed environment on all workers
+        num_workers_total = len(worker_group)
+        setup_futures = []
+        for i in range(num_workers_total):
+            env_vars = {}
+            if num_slices > 1:
+                slice_id = min(i // workers_per_slice, num_slices - 1)
+                env_vars = get_tpu_coordinator_env_vars(
+                    coordinator_address=master_addr,
+                    num_slices=num_slices,
+                    slice_id=slice_id,
+                )
+
+            setup_futures.append(
+                worker_group.execute_single_async(
+                    i,
+                    _setup_jax_distributed_environment,
+                    master_addr_with_port=master_addr_with_port,
+                    num_workers=len(worker_group),
+                    index=i,
+                    use_tpu=backend_config.use_tpu,
+                    use_gpu=backend_config.use_gpu,
+                    resources_per_worker=worker_group.get_resources_per_worker(),
+                    jax_env_vars=env_vars,
+                )
+            )
+        ray.get(setup_futures)
+
+    def on_shutdown(self, worker_group: WorkerGroup, backend_config: JaxConfig):
+        """Cleanup JAX distributed resources when shutting down worker group."""
+        if not backend_config.use_tpu and not backend_config.use_gpu:
+            return
+
+        # Shutdown JAX distributed on all workers
+        shutdown_futures = worker_group.execute_async(_shutdown_jax_distributed)
+
+        timeout_s = ray_constants.env_integer(
+            JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+            DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
+        )
+        try:
+            ray.get(shutdown_futures, timeout=timeout_s)
+            logger.debug("JAX distributed shutdown completed")
+        except ray.exceptions.GetTimeoutError:
+            logger.warning(
+                f"JAX distributed shutdown timed out after {timeout_s} seconds. "
+                "This may indicate workers are hung or unresponsive."
+            )
+        except Exception as e:
+            logger.warning(f"Error during JAX distributed shutdown: {e}")