Improve lifecycle handling of SlicePlacementGroup and support explicit bundle_label_selector

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

Author: ryanaoleary

python/ray/tests/test_tpu.py

@@ -839,5 +839,114 @@ def test_slice_placement_group_chips_per_vm_override(ray_v6e_tpu_cluster):
     assert override_pg.bundle_resources["TPU"] == 4
 
 
+def test_user_bundle_label_selector_merged(ray_tpu_cluster):
+    """Verifies that user-passed bundle_label_selector is merged with dynamic TPU labels."""
+    user_selectors = [{"env": "prod"}, {"env": "test"}]
+
+    # 2x2x2 v4 = 2 hosts = 2 bundles
+    slice_pg = SlicePlacementGroup(
+        topology="2x2x2", accelerator_version="v4", bundle_label_selector=user_selectors
+    )
+
+    assert len(slice_pg._bundle_label_selector) == 2
+
+    # Verify slice 0
+    assert slice_pg._bundle_label_selector[0]["env"] == "prod"
+    assert ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY in slice_pg._bundle_label_selector[0]
+
+    # Verify slice 1
+    assert slice_pg._bundle_label_selector[1]["env"] == "test"
+    assert ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY in slice_pg._bundle_label_selector[1]
+
+
+def test_user_bundle_label_selector_collision_dynamic_wins(ray_v6e_tpu_cluster):
+    """Verifies that dynamic TPU labels take precedence on collision."""
+    user_selectors = [{ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY: "user-requested-slice"}]
+
+    # v6e-8 is single host (1 bundle)
+    slice_pg = SlicePlacementGroup(
+        topology="2x4", accelerator_version="v6e", bundle_label_selector=user_selectors
+    )
+
+    assert len(slice_pg._bundle_label_selector) == 1
+    # The dynamic value should win (it generates test-v6e-slice-N)
+    actual_val = slice_pg._bundle_label_selector[0][
+        ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY
+    ]
+    assert actual_val != "user-requested-slice"
+    assert "test-v6e-slice-" in actual_val
+
+
+def test_user_bundle_label_selector_length_mismatch_raises():
+    """Verifies that providing wrong length of selector list raises ValueError."""
+    user_selectors = [{"env": "prod"}]  # Only 1 provided but 2x2x2 v4 has 2 hosts
+
+    with pytest.raises(ValueError, match="bundle_label_selector length"):
+        SlicePlacementGroup(
+            topology="2x2x2",
+            accelerator_version="v4",
+            bundle_label_selector=user_selectors,
+        )
+
+
+def test_release_head_pgs_idempotent(ray_tpu_cluster):
+    """Verifies that release_head_pgs() is idempotent."""
+    slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
+
+    assert len(slice_pg.head_placement_groups) == 1
+
+    slice_pg.release_head_pgs()
+    assert len(slice_pg.head_placement_groups) == 0
+
+    # Call again, should not raise
+    slice_pg.release_head_pgs()
+    assert len(slice_pg.head_placement_groups) == 0
+
+
+def test_shutdown_idempotent(ray_tpu_cluster):
+    """Verifies that shutdown() is idempotent."""
+    slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
+
+    slice_pg.shutdown()
+    assert slice_pg.placement_group is None
+    assert len(slice_pg.head_placement_groups) == 0
+
+    # Call again, should not raise
+    slice_pg.shutdown()
+
+
+def test_shutdown_safe_after_construction_failure():
+    """Verifies that shutdown() is safe to call on a partially-constructed instance."""
+    with patch(
+        "ray.util.tpu.SlicePlacementGroup._reserve_slice",
+        side_effect=RuntimeError("Test failure"),
+    ):
+        with pytest.raises(RuntimeError, match="Test failure"):
+            SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
+
+    # If the above didn't crash or leak resources, we are good.
+    # We can also manually construct a partial instance and call shutdown.
+    partial_pg = SlicePlacementGroup.__new__(SlicePlacementGroup)
+    partial_pg._head_pgs = []
+    partial_pg._placement_group = None
+
+    # Should not raise even though it's missing attributes
+    partial_pg.shutdown()
+
+
+def test_release_head_pgs_after_ready_then_shutdown(ray_tpu_cluster):
+    """Validates Slice PG lifecycle: wait until ready, release head PGs, then shutdown."""
+    slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
+
+    # Wait for ready
+    ray.get(slice_pg.placement_group.ready())
+
+    slice_pg.release_head_pgs()
+    assert len(slice_pg.head_placement_groups) == 0
+
+    slice_pg.shutdown()
+    assert slice_pg.placement_group is None
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

python/ray/util/tpu.py

@@ -160,8 +160,10 @@ def get_tpu_worker_resources(
     """
     accelerator_version = get_tpu_version_from_type(accelerator_type)
 
-    resolved_chips_per_vm = chips_per_vm or get_chips_per_host(
-        topology, accelerator_version
+    resolved_chips_per_vm = (
+        chips_per_vm
+        if chips_per_vm is not None
+        else get_chips_per_host(topology, accelerator_version)
     )
     total_chips_per_slice = get_num_chips_from_topology(topology)
 
@@ -447,6 +449,8 @@ class SlicePlacementGroup:
             TPU head placement group to become ready. Defaults to
             ``DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S``. Pass ``None`` to wait
             indefinitely.
+        bundle_label_selector: Optional list of label selectors to apply per bundle. These label
+            selectors are applied in addition to dynamic TPU slice name labels, which take precedence.
 
     Examples:
 
@@ -490,7 +494,13 @@ def __init__(
         head_reservation_timeout_s: Optional[float] = (
             DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S
         ),
+        bundle_label_selector: Optional[List[Dict[str, str]]] = None,
     ):
+        self._head_pgs: List[PlacementGroup] = []
+        self._bundle_label_selector: List[Dict[str, str]] = []
+        self._placement_group: Optional[PlacementGroup] = None
+        self._user_bundle_label_selector = bundle_label_selector or []
+
         self._topology = topology.strip().lower()
         self._accelerator_version = get_tpu_version_from_type(
             accelerator_version.strip()
@@ -508,8 +518,10 @@ def __init__(
             chips_per_vm=chips_per_vm,
         )
 
-        self._chips_per_host = chips_per_vm or get_chips_per_host(
-            self._topology, self._accelerator_version
+        self._chips_per_host = (
+            chips_per_vm
+            if chips_per_vm is not None
+            else get_chips_per_host(self._topology, self._accelerator_version)
         )
 
         # Within Ray, a "host" corresponds to a user-visible compute VM.
@@ -518,10 +530,7 @@ def __init__(
         hosts_per_slice = max(1, total_chips // self._chips_per_host)
         self._num_hosts = hosts_per_slice * self._num_slices
 
-        self._head_pgs: List[PlacementGroup] = []
-        self._bundle_label_selector: List[Dict[str, str]] = []
         self._validate_tpu_config()
-        self._placement_group = None
 
         # Reserve a TPU slice of the provided accelerator version and topology.
         self._placement_group = self._reserve_slice(
@@ -549,6 +558,15 @@ def _reserve_slice(
         lifetime: Optional[str] = None,
     ) -> PlacementGroup:
         """Performs the two-step scheduling to reserve a TPU slice."""
+        if (
+            self._user_bundle_label_selector
+            and len(self._user_bundle_label_selector) != self._num_bundles
+        ):
+            raise ValueError(
+                f"bundle_label_selector length ({len(self._user_bundle_label_selector)}) must "
+                f"match the number of bundles ({self._num_bundles})."
+            )
+
         self._bundle_label_selector = []
         bundles = []
         bundles_per_slice = self._num_bundles // self._num_slices
@@ -557,7 +575,7 @@ def _reserve_slice(
         accelerator_type = "TPU-" + self.accelerator_version.upper()
 
         try:
-            for _ in range(self.num_slices):
+            for slice_idx in range(self.num_slices):
                 reservation = reserve_tpu_slice(
                     self._topology,
                     accelerator_type,
@@ -575,10 +593,20 @@ def _reserve_slice(
                 slice_name, head_pg = reservation
                 self._head_pgs.append(head_pg)
 
-                # Reserving a slice is done through constructing num_hosts bundles, each with a label selector for
-                # the unique name of an available TPU slice.
-                selector = {ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY: slice_name}
-                self._bundle_label_selector.extend([selector] * bundles_per_slice)
+                dynamic_labels = {ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY: slice_name}
+
+                for bundle_idx in range(bundles_per_slice):
+                    global_bundle_idx = slice_idx * bundles_per_slice + bundle_idx
+
+                    user_labels = (
+                        self._user_bundle_label_selector[global_bundle_idx]
+                        if global_bundle_idx < len(self._user_bundle_label_selector)
+                        else {}
+                    )
+                    # Dynamic TPU slice labels take precedence; user labels fill in the rest.
+                    merged_labels = {**user_labels, **dynamic_labels}
+                    self._bundle_label_selector.append(merged_labels)
+
                 bundles += [
                     self._bundle_resources.copy() for _ in range(bundles_per_slice)
                 ]
@@ -647,14 +675,47 @@ def bundle_resources(self) -> Dict[str, float]:
         """The resources that are assigned to each bundle."""
         return self._bundle_resources
 
+    @DeveloperAPI(stability="alpha")
+    def release_head_pgs(self) -> None:
+        """Remove all internal head placement groups.
+
+        The head PGs exist only to atomically claim a TPU slice's label during
+        the race window between slice selection and worker-PG construction.
+        Once the worker PG's bundles are scheduled, the worker PG holds the TPU
+        resources on every host in the slice and the head PGs are redundant.
+
+        Callers should invoke this idempotent call after `self.placement_group.ready()`
+        resolves successfully.
+        """
+        head_pgs = getattr(self, "_head_pgs", [])
+        self._head_pgs = []
+        for head_pg in head_pgs:
+            try:
+                remove_placement_group(head_pg)
+            except Exception:
+                logger.exception(
+                    "Failed to remove TPU head placement group %s; the "
+                    "slice reservation marker may leak until the creator "
+                    "process exits.",
+                    getattr(head_pg, "id", head_pg),
+                )
+
     def shutdown(self):
-        """Removes the worker placement group and all internal head PGs."""
-        if self._placement_group:
-            remove_placement_group(self._placement_group)
+        """Remove the worker placement group and all internal head PGs.
+
+        Idempotent. Safe to call on a partially-constructed instance.
+        """
+        worker_pg = getattr(self, "_placement_group", None)
+        if worker_pg is not None:
             self._placement_group = None
-        for head_pg in self._head_pgs:
-            remove_placement_group(head_pg)
-        self._head_pgs = []
+            try:
+                remove_placement_group(worker_pg)
+            except Exception:
+                logger.exception(
+                    "Failed to remove TPU worker placement group %s.",
+                    getattr(worker_pg, "id", worker_pg),
+                )
+        self.release_head_pgs()
 
 
 @PublicAPI(stability="alpha")