[Train] Add TPU multi-slice support to JaxTrainer

doc/source/ray-core/api/utility.rst

@@ -16,6 +16,7 @@ Utility
    ray.util.tpu.get_current_pod_name
    ray.util.tpu.get_num_tpu_chips_on_node
    ray.util.tpu.get_tpu_coordinator_env_vars
+   ray.util.tpu.get_tpu_num_slices_for_workers
    ray.util.tpu.get_tpu_version_from_type
    ray.util.tpu.get_tpu_worker_resources
 

python/ray/tests/test_tpu.py

@@ -5,7 +5,6 @@
 
 import ray
 from ray._private.accelerators import TPUAcceleratorManager, tpu
-from ray.tests.conftest import _ray_start_cluster
 from ray.util.tpu import SlicePlacementGroup
 
 
@@ -170,77 +169,77 @@ def ray_start_cpu():
 
 
 @pytest.fixture
-def ray_tpu_cluster():
+def ray_tpu_cluster(ray_start_cluster):
     """
     Simulates a Ray cluster with two multi-host TPU v4-16 slices.
     """
     pod_type = "v4-16"
     topology = "2x2x2"
 
-    with _ray_start_cluster() as cluster:
-        slice_0_env_common = {
-            "TPU_NAME": "test-slice-0",
-            "TPU_ACCELERATOR_TYPE": pod_type,
-            "TPU_TOPOLOGY": topology,
-        }
-        slice_0_head_labels = {
-            "ray.io/tpu-slice-name": "test-slice-0",
-            "ray.io/tpu-worker-id": "0",
-            "ray.io/tpu-pod-type": pod_type,
-            "ray.io/tpu-topology": topology,
-        }
-        slice_0_worker_labels = {
-            "ray.io/tpu-slice-name": "test-slice-0",
-            "ray.io/tpu-worker-id": "1",
-            "ray.io/tpu-pod-type": pod_type,
-            "ray.io/tpu-topology": topology,
-        }
-        cluster.add_node(
-            num_cpus=2,
-            resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
-            env_vars={**slice_0_env_common, "TPU_WORKER_ID": "0"},
-            labels=slice_0_head_labels,
-        )
-        cluster.add_node(
-            num_cpus=2,
-            resources={"TPU": 4},
-            env_vars={**slice_0_env_common, "TPU_WORKER_ID": "1"},
-            labels=slice_0_worker_labels,
-        )
-
-        slice_1_env_common = {
-            "TPU_NAME": "test-slice-1",
-            "TPU_ACCELERATOR_TYPE": pod_type,
-            "TPU_TOPOLOGY": topology,
-        }
-        slice_1_head_labels = {
-            "ray.io/tpu-slice-name": "test-slice-1",
-            "ray.io/tpu-worker-id": "0",
-            "ray.io/tpu-pod-type": pod_type,
-            "ray.io/tpu-topology": topology,
-        }
-        slice_1_worker_labels = {
-            "ray.io/tpu-slice-name": "test-slice-1",
-            "ray.io/tpu-worker-id": "1",
-            "ray.io/tpu-pod-type": pod_type,
-            "ray.io/tpu-topology": topology,
-        }
-        cluster.add_node(
-            num_cpus=2,
-            resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
-            env_vars={**slice_1_env_common, "TPU_WORKER_ID": "0"},
-            labels=slice_1_head_labels,
-        )
-        cluster.add_node(
-            num_cpus=2,
-            resources={"TPU": 4},
-            env_vars={**slice_1_env_common, "TPU_WORKER_ID": "1"},
-            labels=slice_1_worker_labels,
-        )
-
-        ray.init(address=cluster.address)
-        yield cluster
-        ray.shutdown()
+    cluster = ray_start_cluster
+    slice_0_env_common = {
+        "TPU_NAME": "test-slice-0",
+        "TPU_ACCELERATOR_TYPE": pod_type,
+        "TPU_TOPOLOGY": topology,
+    }
+    slice_0_head_labels = {
+        "ray.io/tpu-slice-name": "test-slice-0",
+        "ray.io/tpu-worker-id": "0",
+        "ray.io/tpu-pod-type": pod_type,
+        "ray.io/tpu-topology": topology,
+    }
+    slice_0_worker_labels = {
+        "ray.io/tpu-slice-name": "test-slice-0",
+        "ray.io/tpu-worker-id": "1",
+        "ray.io/tpu-pod-type": pod_type,
+        "ray.io/tpu-topology": topology,
+    }
+    cluster.add_node(
+        num_cpus=2,
+        resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
+        env_vars={**slice_0_env_common, "TPU_WORKER_ID": "0"},
+        labels=slice_0_head_labels,
+    )
+    cluster.add_node(
+        num_cpus=2,
+        resources={"TPU": 4},
+        env_vars={**slice_0_env_common, "TPU_WORKER_ID": "1"},
+        labels=slice_0_worker_labels,
+    )
+
+    slice_1_env_common = {
+        "TPU_NAME": "test-slice-1",
+        "TPU_ACCELERATOR_TYPE": pod_type,
+        "TPU_TOPOLOGY": topology,
+    }
+    slice_1_head_labels = {
+        "ray.io/tpu-slice-name": "test-slice-1",
+        "ray.io/tpu-worker-id": "0",
+        "ray.io/tpu-pod-type": pod_type,
+        "ray.io/tpu-topology": topology,
+    }
+    slice_1_worker_labels = {
+        "ray.io/tpu-slice-name": "test-slice-1",
+        "ray.io/tpu-worker-id": "1",
+        "ray.io/tpu-pod-type": pod_type,
+        "ray.io/tpu-topology": topology,
+    }
+    cluster.add_node(
+        num_cpus=2,
+        resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
+        env_vars={**slice_1_env_common, "TPU_WORKER_ID": "0"},
+        labels=slice_1_head_labels,
+    )
+    cluster.add_node(
+        num_cpus=2,
+        resources={"TPU": 4},
+        env_vars={**slice_1_env_common, "TPU_WORKER_ID": "1"},
+        labels=slice_1_worker_labels,
+    )
+
+    ray.init(address=cluster.address)
+    yield cluster
+    ray.shutdown()
 
 
 def test_fetch_tpu_slice_name_from_pg(ray_tpu_cluster):
@@ -367,5 +366,37 @@ def test_get_tpu_version_invalid(invalid_type):
         ray.util.tpu.get_tpu_version_from_type(invalid_type)
 
 
+@pytest.mark.parametrize(
+    "topology, accelerator_type, num_workers, resources_per_worker, expected_slices",
+    [
+        # "2x2x1" has 4 chips, for 4 workers with TPU: 1 each we expect num_slices=1.
+        ("2x2x1", "TPU-V4", 4, {"TPU": 1}, 1),
+        # "2x2x1" has 4 chips, for 8 workers with TPU: 1 each we expect num_slices=2.
+        ("2x2x1", "v4", 8, {"TPU": 1}, 2),
+        # "2x2x2" has 8 chips and 2 hosts, defaulting to 1 TPU worker per host
+        # and requesting 4 workers, we expect num_slices=2.
+        ("2x2x2", "TPU-V4", 4, None, 2),
+        # "2x2x4" has 16 chips and 4 hosts, defaulting to 1 TPU worker per host
+        # and requesting 4 workers, we expect num_slices=1.
+        ("2x2x4", "TPU-V4", 4, None, 1),
+        # 0 workers requested -> fallback to 1 slice.
+        ("2x2x1", "v4", 0, None, 1),
+        # Invalid topology -> fallback to 1 slice.
+        ("", "v4", 4, {"TPU": 1}, 1),
+        ("2x2x1", "", 4, {"TPU": 1}, 1),
+    ],
+)
+def test_get_tpu_num_slices_for_workers(
+    topology, accelerator_type, num_workers, resources_per_worker, expected_slices
+):
+    num_slices = ray.util.tpu.get_tpu_num_slices_for_workers(
+        topology=topology,
+        accelerator_type=accelerator_type,
+        num_workers=num_workers,
+        resources_per_worker=resources_per_worker,
+    )
+    assert num_slices == expected_slices
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

python/ray/train/v2/_internal/callbacks/__init__.py

@@ -2,15 +2,13 @@
 from .backend_setup import BackendSetupCallback
 from .datasets import DatasetsCallback
 from .state_manager import StateManagerCallback
-from .tpu_reservation_callback import TPUReservationCallback
 from .working_dir_setup import WorkingDirectorySetupCallback
 
 __all__ = [
     "AcceleratorSetupCallback",
     "BackendSetupCallback",
     "DatasetsCallback",
     "StateManagerCallback",
-    "TPUReservationCallback",
     "WorkingDirectorySetupCallback",
 ]
 

python/ray/train/v2/_internal/execution/controller/controller.py

@@ -75,6 +75,7 @@
 if TYPE_CHECKING:
     from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
 
+from ray.util.tpu import get_tpu_num_slices_for_workers
 
 logger = logging.getLogger(__name__)
 
@@ -327,13 +328,24 @@ def _start_worker_group(
         except Exception as e:
             return ControllerError(e)
 
+        # Calculate num_slices for the worker group if using TPU.
+        num_slices = 1
+        if scaling_config.use_tpu:
+            num_slices = get_tpu_num_slices_for_workers(
+                topology=scaling_config.topology,
+                accelerator_type=scaling_config.accelerator_type,
+                num_workers=num_workers,
+                resources_per_worker=resources_per_worker,
+            )
+
         worker_group_context = WorkerGroupContext(
             run_attempt_id=self._get_run_attempt_id(),
             train_fn_ref=self._train_fn_ref,
             num_workers=num_workers,
             resources_per_worker=resources_per_worker,
             placement_strategy=placement_strategy,
             label_selector=label_selector,
+            num_slices=num_slices,
         )
         try:
             self._worker_group = self.worker_group_cls.create(

python/ray/train/v2/_internal/execution/worker_group/state.py

@@ -1,13 +1,14 @@
 import logging
 from dataclasses import dataclass
-from typing import List
+from typing import List, Optional
 
 import ray
 from ray.actor import ActorHandle
 from ray.train.v2._internal.execution.checkpoint.sync_actor import SynchronizationActor
 from ray.train.v2._internal.execution.worker_group.worker import Worker
 from ray.train.v2._internal.util import time_monotonic
 from ray.util.placement_group import PlacementGroup, remove_placement_group
+from ray.util.tpu import SlicePlacementGroup
 
 logger = logging.getLogger(__name__)
 
@@ -28,16 +29,21 @@ class WorkerGroupState:
     placement_group: PlacementGroup
     workers: List[Worker]
     sync_actor: ActorHandle
+    slice_placement_group: Optional[SlicePlacementGroup] = None
 
     @property
     def num_workers(self) -> int:
         return len(self.workers)
 
     def shutdown(self):
         _shutdown_workers(self.workers)
-        _shutdown_placement_group(self.placement_group)
         _shutdown_sync_actor(self.sync_actor)
 
+        if self.slice_placement_group:
+            self.slice_placement_group.shutdown()
+        else:
+            _shutdown_placement_group(self.placement_group)
+
 
 class WorkerGroupStateBuilder:
     """Builder for WorkerGroupState.
@@ -58,13 +64,20 @@ def __init__(self):
         self.placement_group = None
         self.workers = None
         self.sync_actor = None
+        self.slice_placement_group = None
 
     def with_placement_group(
         self, placement_group: PlacementGroup
     ) -> "WorkerGroupStateBuilder":
         self.placement_group = placement_group
         return self
 
+    def with_slice_placement_group(
+        self, slice_placement_group: SlicePlacementGroup
+    ) -> "WorkerGroupStateBuilder":
+        self.slice_placement_group = slice_placement_group
+        return self
+
     def with_workers(self, workers: List[Worker]) -> "WorkerGroupStateBuilder":
         self.workers = workers
         return self
@@ -91,19 +104,26 @@ def build(self) -> WorkerGroupState:
             placement_group=self.placement_group,
             workers=self.workers,
             sync_actor=self.sync_actor,
+            slice_placement_group=self.slice_placement_group,
         )
 
     def shutdown(self):
         if self.workers:
             _shutdown_workers(self.workers)
             self.workers = None
-        if self.placement_group:
-            _shutdown_placement_group(self.placement_group)
-            self.placement_group = None
+
         if self.sync_actor:
             _shutdown_sync_actor(self.sync_actor)
             self.sync_actor = None
 
+        if self.slice_placement_group:
+            self.slice_placement_group.shutdown()
+            self.slice_placement_group = None
+            self.placement_group = None
+        elif self.placement_group:
+            _shutdown_placement_group(self.placement_group)
+            self.placement_group = None
+
 
 def _shutdown_workers(workers: List[Worker], patience_s: float = 5):
     # Run the worker shutdown logic on each of the workers. This should