[train][jax_trainer] add jax.distributed.shutdown() for JaxBackend

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

@@ -36,6 +36,20 @@ def _setup_jax_tpu_environment(
         jax.distributed.initialize(master_addr_with_port, num_workers, index)
 
 
+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:
@@ -57,3 +71,23 @@ def on_start(self, worker_group: WorkerGroup, backend_config: JaxConfig):
                 )
             )
         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:
+            return
+
+        # Shutdown JAX distributed on all workers
+        shutdown_futures = worker_group.execute_async(_shutdown_jax_distributed)
+
+        timeout_s = 30
+        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}")