Add AcceleratorConfig to Serve and fix gang scheduling

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

Author: ryanaoleary

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

@@ -323,8 +323,6 @@ def test_tpu_serve_deployment_explicit_per_chip_bundles(ray_tpu_cluster):
     Verifies that a user can explicitly request chip-level bundles (1 TPU per bundle)
     for a full multi-host TPU slice via placement_group_config.
     """
-    from ray.llm._internal.serve.core.configs.accelerators import TPUConfig
-
     llm_config = LLMConfig(
         model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"),
         accelerator_type="TPU-V6E",

python/ray/serve/_private/common.py

@@ -1,7 +1,7 @@
 import json
 from dataclasses import asdict, dataclass, field
 from enum import Enum
-from typing import Any, Awaitable, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Optional, Union
 
 from starlette.types import Scope
 
@@ -18,6 +18,10 @@
 from ray.util.annotations import PublicAPI
 from ray.util.placement_group import PlacementGroup
 
+if TYPE_CHECKING:
+    from ray.serve._private.default_impl import _ReplicaPlacementGroup
+    from ray.serve.config import AcceleratorConfig
+
 REPLICA_ID_FULL_ID_STR_PREFIX = "SERVE_REPLICA::"
 GANG_PG_NAME_PREFIX = "SERVE_GANG::"
 
@@ -897,6 +901,7 @@ class CreatePlacementGroupRequest:
     runtime_env: Optional[str] = None
     bundle_label_selector: Optional[List[Dict[str, str]]] = None
     fallback_strategy: Optional[List[Dict[str, Any]]] = None
+    accelerator_config: Optional["AcceleratorConfig"] = None
 
 
 @dataclass
@@ -922,6 +927,7 @@ class GangPlacementGroupRequest:
     """Label selector for per-replica placement group bundles."""
 
     replica_pg_fallback_strategy: Optional[List[Dict[str, Any]]] = None
+    accelerator_config: Optional["AcceleratorConfig"] = None
     """Fallback strategy for per-replica placement group bundles."""
 
 
@@ -932,7 +938,7 @@ class GangReservationResult:
     success: bool
     """True when all gang PGs were created successfully."""
     error_message: Optional[str] = None
-    gang_pgs: Optional[List[PlacementGroup]] = None
+    gang_pgs: Optional[List[Union[PlacementGroup, "_ReplicaPlacementGroup"]]] = None
     gang_ids: Optional[List[str]] = None
     gang_pg_names: Optional[List[str]] = None
 

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(self.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, Union
 
 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 TPUAcceleratorConfig
+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,11 +57,93 @@
 # 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,
+) -> _ReplicaPlacementGroup:
+    """Internal entry point that supports accelerator-specific dispatch."""
+    accelerator_config = request.accelerator_config
+
+    if isinstance(accelerator_config, TPUAcceleratorConfig):
+        slice_pg = _default_create_tpu_placement_group(
+            tpu_config=accelerator_config,
+            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_config: TPUAcceleratorConfig,
+    strategy: str,
+    name: str,
+    lifetime: Optional[str],
+    bundle_label_selector: Optional[List[Dict[str, str]]] = None,
+) -> SlicePlacementGroup:
+    return slice_placement_group(
+        topology=tpu_config.topology,
+        accelerator_version=tpu_config.accelerator_version,
+        num_slices=tpu_config.num_slices,
+        chips_per_vm=tpu_config.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)
 
 
-CreatePlacementGroupFn = Callable[[CreatePlacementGroupRequest], PlacementGroup]
+CreatePlacementGroupFn = Callable[
+    [CreatePlacementGroupRequest], Union[PlacementGroup, _ReplicaPlacementGroup]
+]
 
 
 def _default_create_placement_group(
@@ -81,7 +169,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

@@ -7,7 +7,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from functools import total_ordering
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import ray
 from ray._raylet import node_labels_match_selector
@@ -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,
@@ -35,6 +36,9 @@
     PlacementGroupSchedulingStrategy,
 )
 
+if TYPE_CHECKING:
+    from ray.serve._private.default_impl import _ReplicaPlacementGroup
+
 logger = logging.getLogger(SERVE_LOGGER_NAME)
 
 
@@ -198,6 +202,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,12 +641,21 @@ def _schedule_replica(
         replica_id = scheduling_request.replica_id
         deployment_id = replica_id.deployment_id
         placement_group = None
+        slice_pg = None
 
         scheduling_strategy = default_scheduling_strategy
 
         if scheduling_request.gang_placement_group is not None:
             # Gang scheduling -- use the reserved gang placement group
-            placement_group = scheduling_request.gang_placement_group
+            pg_wrapper = scheduling_request.gang_placement_group
+            placement_group = (
+                pg_wrapper.placement_group
+                if hasattr(pg_wrapper, "placement_group")
+                else pg_wrapper
+            )
+            # Preserve the wrapper for cleanup of head PGs
+            slice_pg = pg_wrapper if hasattr(pg_wrapper, "placement_group") else None
+
             scheduling_strategy = PlacementGroupSchedulingStrategy(
                 placement_group=placement_group,
                 placement_group_bundle_index=scheduling_request.gang_pg_index,
@@ -651,21 +665,32 @@ def _schedule_replica(
             target_labels = None
             target_node_id = None
         elif scheduling_request.placement_group_bundles is not None:
+            slice_pg = 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
+                    slice_pg = pg_result
+                else:
+                    pg = pg_result
+                    slice_pg = None
             except Exception:
                 # We add a defensive exception here, so the controller can
                 # make progress even if the placement group isn't created.
@@ -720,6 +745,15 @@ def _schedule_replica(
             scheduling_request.status = (
                 ReplicaSchedulingRequestStatus.ACTOR_CREATION_FAILED
             )
+
+            if slice_pg is not None:
+                slice_pg.shutdown()
+            elif (
+                placement_group is not None
+                and scheduling_request.placement_group_bundles is not None
+            ):
+                ray.util.remove_placement_group(placement_group)
+
             return False
 
         del self._pending_replicas[deployment_id][replica_id]
@@ -731,7 +765,11 @@ 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=slice_pg,
+        )
         return True
 
     @abstractmethod
@@ -816,7 +854,7 @@ def _prepare_gangs_for_deployment(
 
         # Flatten per-replica bundles to form a placement group to atomically reserve resources
         # required for each gang
-        gang_pgs: List[PlacementGroup] = []
+        gang_pgs: List[Union[PlacementGroup, "_ReplicaPlacementGroup"]] = []
         gang_ids: List[str] = []
         gang_pg_names: List[str] = []
         for gang_index in range(num_gangs):
@@ -867,6 +905,7 @@ def _prepare_gangs_for_deployment(
                         name=pg_name,
                         bundle_label_selector=label_selector,
                         fallback_strategy=fallback_strategy,
+                        accelerator_config=request.accelerator_config,
                     )
                 )
                 gang_pgs.append(pg)