[Serve][2/N] Implement AcceleratorConfig to enable custom scheduling logic for accelerators with Serve deployments

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
 
 from starlette.types import Scope
 
@@ -18,6 +18,9 @@
 from ray.util.annotations import PublicAPI
 from ray.util.placement_group import PlacementGroup
 
+if TYPE_CHECKING:
+    from ray.serve.config import AcceleratorConfig
+
 REPLICA_ID_FULL_ID_STR_PREFIX = "SERVE_REPLICA::"
 GANG_PG_NAME_PREFIX = "SERVE_GANG::"
 
@@ -893,13 +896,25 @@ class ReplicaQueueLengthInfo:
 
 @dataclass(frozen=True)
 class CreatePlacementGroupRequest:
-    bundles: List[Dict[str, float]]
-    strategy: str
-    target_node_id: str
-    name: str
+    """Internal request for creating a per-replica placement group.
+
+    Either ``bundles`` or ``accelerator_config`` must be provided:
+    - For plain CPU/GPU deployments, the caller provides ``bundles`` and the
+      default path creates a standard PlacementGroup.
+    - For accelerator deployments (e.g. TPU), the caller provides
+      ``accelerator_config`` and the dispatch derives bundles from the
+      structured config (e.g. TPU topology -> per-host bundles).
+    """
+
+    bundles: Optional[List[Dict[str, float]]] = None
+    strategy: str = "PACK"
+    target_node_id: Optional[str] = None
+    name: str = ""
     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
+    lifetime: Optional[str] = "detached"
 
 
 @dataclass

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,
@@ -190,6 +191,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
@@ -322,6 +327,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"])
@@ -429,6 +436,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,8 +44,9 @@
     inside_ray_client_context,
     resolve_deployment_response,
 )
-from ray.serve.config import ControllerOptions
-from ray.util.placement_group import PlacementGroup
+from ray.serve.config import ControllerOptions, TPUAcceleratorConfig
+from ray.util.placement_group import PlacementGroup, remove_placement_group
+from ray.util.tpu import SlicePlacementGroup, slice_placement_group
 
 # NOTE: Please read carefully before changing!
 #
@@ -52,11 +55,112 @@
 # 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: Optional[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.
+
+    Dispatches on ``request.accelerator_config``:
+    - TPUAcceleratorConfig: derive bundles from topology via
+      slice_placement_group; ``request.bundles`` is ignored.
+    - None: use ``request.bundles`` to create a standard PlacementGroup.
+
+    Raises ValueError if neither bundles nor a recognized accelerator
+    config is provided - this catches users setting an unrecognized
+    accelerator_config type without explicit bundles, which would
+    otherwise schedule with no PG at all.
+    """
+    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=request.lifetime,
+            bundle_label_selector=request.bundle_label_selector,
+        )
+        return ReplicaPlacementGroup(
+            placement_group=slice_pg.placement_group,
+            _slice_pg=slice_pg,
+        )
+
+    if request.bundles is None:
+        raise ValueError(
+            "CreatePlacementGroupRequest requires either non-None bundles "
+            "or a recognized accelerator_config. Got accelerator_config="
+            f"{type(accelerator_config).__name__ if accelerator_config else None}, "
+            "bundles=None."
+        )
+
+    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,
+        resources_per_bundle=tpu_config.resources_per_bundle,
+        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(
@@ -67,7 +171,7 @@ def _default_create_placement_group(
         request.strategy,
         _soft_target_node_id=request.target_node_id,
         name=request.name,
-        lifetime="detached",
+        lifetime=request.lifetime,
         bundle_label_selector=request.bundle_label_selector,
     )
 
@@ -82,7 +186,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,12 +638,16 @@ def _schedule_replica(
         replica_id = scheduling_request.replica_id
         deployment_id = replica_id.deployment_id
         placement_group = None
+        replica_pg = None
 
         scheduling_strategy = default_scheduling_strategy
 
         if scheduling_request.gang_placement_group is not None:
-            # Gang scheduling -- use the reserved gang placement group
+            # Gang scheduling -- use the reserved gang placement group.
+            # Gang PGs are always bare PlacementGroup objects; accelerator
+            # deployments bypass gang scheduling entirely (see deployment_state).
             placement_group = scheduling_request.gang_placement_group
+
             scheduling_strategy = PlacementGroupSchedulingStrategy(
                 placement_group=placement_group,
                 placement_group_bundle_index=scheduling_request.gang_pg_index,
@@ -650,22 +656,40 @@ def _schedule_replica(
             # TODO (jeffreywang): Add support for target labels and node affinity
             target_labels = None
             target_node_id = None
-        elif scheduling_request.placement_group_bundles is not None:
+        elif (
+            scheduling_request.placement_group_bundles is not None
+            or scheduling_request.accelerator_config is not None
+        ):
+            # Per-replica PG path. Entered when either:
+            # - The user provided explicit bundles (CPU/GPU deployments), or
+            # - The user provided an accelerator_config that derives its own
+            #   bundles from structured fields (e.g. TPUAcceleratorConfig
+            #   derives bundles from topology via slice_placement_group).
             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,
+                    ),
                 )
+                # Import ReplicaPlacementGroup inline here to avoid circular dependency with default_impl
+                from ray.serve._private.default_impl import ReplicaPlacementGroup
+
+                if isinstance(pg_result, ReplicaPlacementGroup):
+                    placement_group = pg_result.placement_group
+                    replica_pg = pg_result
+                else:
+                    placement_group = pg_result
+                    replica_pg = None
             except Exception:
                 # We add a defensive exception here, so the controller can
                 # make progress even if the placement group isn't created.
@@ -678,7 +702,7 @@ def _schedule_replica(
                 )
                 return False
             scheduling_strategy = PlacementGroupSchedulingStrategy(
-                placement_group=pg,
+                placement_group=placement_group,
                 placement_group_capture_child_tasks=True,
             )
             target_labels = None
@@ -720,6 +744,14 @@ def _schedule_replica(
             scheduling_request.status = (
                 ReplicaSchedulingRequestStatus.ACTOR_CREATION_FAILED
             )
+
+            # Only clean up single-replica PGs. Gang PGs are managed elsewhere.
+            if scheduling_request.gang_placement_group is None:
+                if replica_pg is not None:
+                    replica_pg.shutdown()
+                elif placement_group is not None:
+                    ray.util.remove_placement_group(placement_group)
+
             return False
 
         del self._pending_replicas[deployment_id][replica_id]
@@ -731,7 +763,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=replica_pg,
+        )
         return True
 
     @abstractmethod
@@ -859,7 +895,7 @@ def _prepare_gangs_for_deployment(
             )
 
             try:
-                pg = self._create_placement_group_fn(
+                pg_result = self._create_placement_group_fn(
                     CreatePlacementGroupRequest(
                         bundles=bundles,
                         strategy=request.gang_placement_strategy,
@@ -869,6 +905,18 @@ def _prepare_gangs_for_deployment(
                         fallback_strategy=fallback_strategy,
                     )
                 )
+
+                # Unwrap the ReplicaPlacementGroup to get the underyling PlacementGroup.
+                # Gang scheduling currently does not support accelerator_config (since it's
+                # handled by the specific accelerator backend), so we don't need the
+                # wrapper. Inline import here is required to avoid circular dependencies.
+                from ray.serve._private.default_impl import ReplicaPlacementGroup
+
+                if isinstance(pg_result, ReplicaPlacementGroup):
+                    pg = pg_result.placement_group
+                else:
+                    pg = pg_result
+
                 gang_pgs.append(pg)
                 gang_ids.append(gang_id)
                 gang_pg_names.append(pg_name)