Add AcceleratorConfig to Serve

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

Author: ryanaoleary

python/ray/llm/_internal/serve/core/configs/accelerators.py

@@ -199,7 +199,7 @@ def default_bundles(
 
         num_hosts = max(1, num_devices // chips_per_host)
 
-        bundle = {"TPU": float(chips_per_host)}
+        bundle = {"TPU": chips_per_host}
         bundle[format_ray_accelerator_resource(accelerator_type_str)] = 0.001
 
         return [bundle.copy() for _ in range(num_hosts)]
@@ -293,7 +293,3 @@ def shutdown(self):
                 logger.warning(f"Failed to shut down TPU slice PG: {e}")
             finally:
                 self._slice_pg_wrapper = None
-
-    def __del__(self):
-        """Ensure placement groups are cleaned up when this backend is garbage collected."""
-        self.shutdown()

python/ray/llm/_internal/serve/core/server/llm_server.py

@@ -22,6 +22,7 @@
     MODEL_RESPONSE_BATCH_TIMEOUT_MS,
     RAYLLM_VLLM_ENGINE_CLS_ENV,
 )
+from ray.llm._internal.serve.core.configs.accelerators import TPUConfig
 from ray.llm._internal.serve.core.configs.llm_config import (
     DiskMultiplexConfig,
     LLMConfig,
@@ -39,6 +40,8 @@
 from ray.llm._internal.serve.utils.server_utils import (
     get_serve_request_id,
 )
+from ray.serve.config import AcceleratorConfig, TPUSliceSpec
+from ray.util.tpu import get_tpu_version_from_type
 
 if TYPE_CHECKING:
     from ray.llm._internal.serve.core.configs.openai_api_models import (
@@ -737,4 +740,23 @@ def get_deployment_options(cls, llm_config: "LLMConfig"):
         }
         deployment_options["ray_actor_options"] = ray_actor_options
 
+        if llm_config.accelerator_config is not None and isinstance(
+            llm_config.accelerator_config, TPUConfig
+        ):
+            if not llm_config.accelerator_type:
+                raise ValueError(
+                    "llm_config.accelerator_type must be specified when "
+                    "accelerator_config is a TPUConfig."
+                )
+            version = get_tpu_version_from_type(llm_config.accelerator_type)
+
+            deployment_options["accelerator_config"] = AcceleratorConfig(
+                accelerator_type="tpu",
+                tpu=TPUSliceSpec(
+                    topology=llm_config.accelerator_config.topology,
+                    accelerator_version=version,
+                    num_slices=1,
+                ),
+            )
+
         return deployment_options

python/ray/llm/tests/serve/cpu/configs/test_models.py

@@ -420,13 +420,13 @@ def test_requires_deferred_placement_group(self):
     @pytest.mark.parametrize(
         "topology,num_devices,accelerator_type_str,expected_bundles_count,expected_chips_per_host",
         [
-            ("1x1", 1, "v6e", 1, 1.0),
-            ("1x1", 1, "v7x", 1, 4.0),
-            ("4x4", 16, "v6e", 4, 4.0),
-            ("2x2x2", 8, "v5p", 2, 4.0),
-            ("2x2", 4, "v5litepod", 1, 4.0),
-            ("2x2x1", 4, "v4", 1, 4.0),
-            ("2x4", 8, "v6e", 1, 8.0),
+            ("1x1", 1, "v6e", 1, 1),
+            ("1x1", 1, "v7x", 1, 4),
+            ("4x4", 16, "v6e", 4, 4),
+            ("2x2x2", 8, "v5p", 2, 4),
+            ("2x2", 4, "v5litepod", 1, 4),
+            ("2x2x1", 4, "v4", 1, 4),
+            ("2x4", 8, "v6e", 1, 8),
         ],
     )
     def test_default_bundles_topology(
@@ -457,6 +457,40 @@ def test_default_bundles_topology_missing_accelerator_type_raises(self):
         ):
             tpu_accel.default_bundles(num_devices=16, accelerator_type_str=None)
 
+    def test_default_bundles_v6e_4x4(self):
+        """Test that v6e 4x4 topology returns per-host bundles."""
+        tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4"))
+        bundles = tpu_accel.default_bundles(num_devices=16, accelerator_type_str="v6e")
+
+        # 4x4 v6e = 16 chips. 4 chips per host -> 4 hosts.
+        assert len(bundles) == 4
+        for bundle in bundles:
+            assert bundle["TPU"] == 4.0
+            assert "accelerator_type:v6e" in bundle
+
+    def test_default_bundles_v5p_2x2x2(self):
+        """Test that v5p 2x2x2 topology returns per-host bundles."""
+        tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="2x2x2"))
+        bundles = tpu_accel.default_bundles(num_devices=8, accelerator_type_str="v5p")
+
+        # 2x2x2 v5p = 8 chips. 4 chips per host -> 2 hosts.
+        assert len(bundles) == 2
+        for bundle in bundles:
+            assert bundle["TPU"] == 4.0
+            assert "accelerator_type:v5p" in bundle
+
+    def test_default_bundles_single_host_topology(self):
+        """Test that a single-host topology returns a single bundle."""
+        tpu_accel = TPUAccelerator(TPUConfig(kind="tpu", topology="2x2"))
+        bundles = tpu_accel.default_bundles(
+            num_devices=4, accelerator_type_str="v5litepod"
+        )
+
+        # 2x2 v5litepod = 4 chips on 1 host.
+        assert len(bundles) == 1
+        assert bundles[0]["TPU"] == 4.0
+        assert "accelerator_type:v5litepod" in bundles[0]
+
 
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine_tpu.py

@@ -37,11 +37,11 @@ def test_tpu_slice_placement_group_creation_default_resources(ray_tpu_cluster):
     pg_table = placement_group_table(pg)
     assert pg_table["strategy"] == "PACK"
 
-    # 4x4 v6e = 16 chips. We default to 1 TPU chip per bundle.
-    assert len(pg_table["bundles"]) == 16
+    # 4x4 v6e = 16 chips. We default to 4 TPU chips per bundle (per-host).
+    assert len(pg_table["bundles"]) == 4
     for bundle in pg_table["bundles"].values():
         assert "TPU" in bundle
-        assert bundle["TPU"] == 1
+        assert bundle["TPU"] == 4.0
 
     # Let the backend tear down its own resources if it has any
     engine_config.accelerator.shutdown()
@@ -62,7 +62,7 @@ def test_tpu_slice_placement_group_creation_host_resources(ray_tpu_cluster):
         accelerator_config={"kind": "tpu", "topology": "4x4"},
         placement_group_config={
             "strategy": "STRICT_SPREAD",
-            "bundles": [{"TPU": 4}],
+            "bundles": [{"TPU": 4}] * 4,
         },
     )
 
@@ -256,10 +256,10 @@ def test_tpu_serve_deployment_default_chip_level_bundles(ray_tpu_cluster):
     worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0]
 
     assert worker_pg["strategy"] == "PACK"
-    # 4x4 topology = 16 chips. Default is 16 bundles of 1 TPU.
-    assert len(worker_pg["bundles"]) == 16
+    # 4x4 topology = 16 chips. Default is 4 bundles of 4 TPUs (per-host).
+    assert len(worker_pg["bundles"]) == 4
     for bundle in worker_pg["bundles"].values():
-        assert bundle.get("TPU", 0) == 1
+        assert bundle.get("TPU", 0) == 4.0
 
     serve.shutdown()
 

python/ray/serve/_private/config.py

@@ -32,6 +32,7 @@
 )
 from ray.serve._private.utils import DEFAULT, DeploymentOptionUpdateType
 from ray.serve.config import (
+    AcceleratorConfig,
     AggregationFunction,
     AutoscalingConfig,
     DeploymentActorConfig,
@@ -191,6 +192,10 @@ class DeploymentConfig(BaseModel):
         update_type=DeploymentOptionUpdateType.NeedsActorReconfigure,
     )
 
+    accelerator_config: Optional[AcceleratorConfig] = Field(
+        default=None, update_type=DeploymentOptionUpdateType.HeavyWeight
+    )
+
     # This flag is used to let replica know they are deployed from
     # a different language.
     is_cross_language: bool = False
@@ -323,6 +328,8 @@ def needs_pickle(self):
 
     def to_proto(self):
         data = self.model_dump()
+        if data.get("accelerator_config") is not None:
+            data["accelerator_config"] = cloudpickle.dumps(data["accelerator_config"])
         if data.get("user_config") is not None:
             if self.needs_pickle():
                 data["user_config"] = cloudpickle.dumps(data["user_config"])
@@ -430,6 +437,11 @@ def from_proto(cls, proto: DeploymentConfigProto):
             data["is_cross_language"] if "is_cross_language" in data else False
         )
         needs_pickle = _needs_pickle(deployment_language, is_cross_language)
+        if "accelerator_config" in data:
+            if data["accelerator_config"] != b"":
+                data["accelerator_config"] = cloudpickle.loads(proto.accelerator_config)
+            else:
+                data["accelerator_config"] = None
         if "user_config" in data:
             if data["user_config"] != b"":
                 if needs_pickle:

python/ray/serve/_private/default_impl.py

@@ -1,5 +1,7 @@
 import asyncio
-from typing import Callable, Optional, Tuple
+import logging
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Tuple
 
 import ray
 from ray._common.constants import HEAD_NODE_RESOURCE_NAME
@@ -42,7 +44,11 @@
     inside_ray_client_context,
     resolve_deployment_response,
 )
-from ray.util.placement_group import PlacementGroup
+from ray.serve.config import AcceleratorConfig, TPUSliceSpec
+from ray.util.placement_group import PlacementGroup, remove_placement_group
+from ray.util.tpu import SlicePlacementGroup, slice_placement_group
+
+logger = logging.getLogger(__name__)
 
 # NOTE: Please read carefully before changing!
 #
@@ -51,6 +57,86 @@
 # API modified w/o substantial enough justification
 
 
+@dataclass
+class _ReplicaPlacementGroup:
+    """Internal Serve handle for a replica's placement group(s).
+
+    Wraps the worker PG and any accelerator-specific cleanup hooks so the
+    controller doesn't need to know whether the underlying request was a
+    plain CPU/GPU PG or a TPU slice reservation.
+    """
+
+    placement_group: PlacementGroup
+    _slice_pg: Optional[SlicePlacementGroup] = None
+
+    def release_reservation_holders(self) -> None:
+        """Call after ``placement_group.ready()`` resolves successfully.
+
+        Releases any internal reservation-holder PGs (e.g. TPU head PGs)
+        that were only needed to claim resources during scheduling. No-op
+        for non-accelerator deployments.
+        """
+        if self._slice_pg is not None:
+            self._slice_pg.release_head_pgs()
+
+    def shutdown(self) -> None:
+        """Tear down the replica's PG(s). Idempotent."""
+        if self._slice_pg is not None:
+            self._slice_pg.shutdown()
+            self._slice_pg = None
+            self.placement_group = None
+        elif self.placement_group is not None:
+
+            try:
+                remove_placement_group(self.placement_group)
+            except Exception:
+                logger.exception("Failed to remove placement group.")
+            finally:
+                self.placement_group = None
+
+
+def _create_replica_placement_group(
+    request: CreatePlacementGroupRequest,
+    *,
+    accelerator_config: Optional[AcceleratorConfig] = None,
+) -> _ReplicaPlacementGroup:
+    """Internal entry point that supports accelerator-specific dispatch."""
+    if accelerator_config is not None and accelerator_config.accelerator_type == "tpu":
+        slice_pg = _default_create_tpu_placement_group(
+            tpu_spec=accelerator_config.tpu,
+            strategy=request.strategy,
+            name=request.name,
+            lifetime="detached",
+            bundle_label_selector=request.bundle_label_selector,
+        )
+        return _ReplicaPlacementGroup(
+            placement_group=slice_pg.placement_group,
+            _slice_pg=slice_pg,
+        )
+
+    pg = _default_create_placement_group(request)
+    return _ReplicaPlacementGroup(placement_group=pg)
+
+
+def _default_create_tpu_placement_group(
+    tpu_spec: TPUSliceSpec,
+    strategy: str,
+    name: str,
+    lifetime: Optional[str],
+    bundle_label_selector: Optional[List[Dict[str, str]]] = None,
+) -> SlicePlacementGroup:
+    return slice_placement_group(
+        topology=tpu_spec.topology,
+        accelerator_version=tpu_spec.accelerator_version,
+        num_slices=tpu_spec.num_slices,
+        chips_per_vm=tpu_spec.chips_per_vm,
+        strategy=strategy,
+        name=name,
+        lifetime=lifetime,
+        bundle_label_selector=bundle_label_selector,
+    )
+
+
 def create_cluster_node_info_cache(gcs_client: GcsClient) -> ClusterNodeInfoCache:
     return DefaultClusterNodeInfoCache(gcs_client)
 
@@ -81,7 +167,7 @@ def create_deployment_scheduler(
         cluster_node_info_cache,
         head_node_id,
         create_placement_group_fn=create_placement_group_fn_override
-        or _default_create_placement_group,
+        or _create_replica_placement_group,
     )
 
 

python/ray/serve/_private/deployment_scheduler.py

@@ -27,6 +27,7 @@
     RAY_SERVE_USE_PACK_SCHEDULING_STRATEGY,
     SERVE_LOGGER_NAME,
 )
+from ray.serve.config import AcceleratorConfig
 from ray.util.placement_group import PlacementGroup
 from ray.util.scheduling_strategies import (
     LabelMatchExpressionsT,
@@ -198,6 +199,7 @@ class ReplicaSchedulingRequest:
     placement_group_strategy: Optional[str] = None
     placement_group_bundle_label_selector: Optional[List[Dict[str, str]]] = None
     placement_group_fallback_strategy: Optional[List[Dict[str, Any]]] = None
+    accelerator_config: Optional[AcceleratorConfig] = None
     max_replicas_per_node: Optional[int] = None
     # Gang scheduling fields -- if set, replica should be scheduled on
     # the reserved gang placement group at the specified bundle index.
@@ -636,6 +638,7 @@ def _schedule_replica(
         replica_id = scheduling_request.replica_id
         deployment_id = replica_id.deployment_id
         placement_group = None
+        sp = None
 
         scheduling_strategy = default_scheduling_strategy
 
@@ -651,21 +654,32 @@ def _schedule_replica(
             target_labels = None
             target_node_id = None
         elif scheduling_request.placement_group_bundles is not None:
+            sp = None
             placement_group_strategy = (
                 scheduling_request.placement_group_strategy
                 if scheduling_request.placement_group_strategy
                 else "PACK"
             )
             try:
-                pg = self._create_placement_group_fn(
+                pg_result = self._create_placement_group_fn(
                     CreatePlacementGroupRequest(
                         bundles=scheduling_request.placement_group_bundles,
                         strategy=placement_group_strategy,
                         target_node_id=target_node_id,
                         name=scheduling_request.actor_options["name"],
                         bundle_label_selector=scheduling_request.placement_group_bundle_label_selector,
-                    )
+                    ),
+                    accelerator_config=scheduling_request.accelerator_config,
                 )
+
+                from ray.serve._private.default_impl import _ReplicaPlacementGroup
+
+                if isinstance(pg_result, _ReplicaPlacementGroup):
+                    pg = pg_result.placement_group
+                    sp = pg_result
+                else:
+                    pg = pg_result
+                    sp = None
             except Exception:
                 # We add a defensive exception here, so the controller can
                 # make progress even if the placement group isn't created.
@@ -731,7 +745,9 @@ def _schedule_replica(
             placement_group = scheduling_strategy.placement_group
 
         scheduling_request.status = ReplicaSchedulingRequestStatus.SUCCEEDED
-        scheduling_request.on_scheduled(actor_handle, placement_group=placement_group)
+        scheduling_request.on_scheduled(
+            actor_handle, placement_group=placement_group, placement_group_manager=sp
+        )
         return True
 
     @abstractmethod
@@ -869,6 +885,10 @@ def _prepare_gangs_for_deployment(
                         fallback_strategy=fallback_strategy,
                     )
                 )
+                from ray.serve._private.default_impl import _ReplicaPlacementGroup
+
+                if isinstance(pg, _ReplicaPlacementGroup):
+                    pg = pg.placement_group
                 gang_pgs.append(pg)
                 gang_ids.append(gang_id)
                 gang_pg_names.append(pg_name)