[Serve][1/N] Add Serve accelerator config schemas

python/ray/serve/_private/config.py

@@ -21,6 +21,7 @@
 from ray._common.serialization import pickle_dumps
 from ray._common.utils import resources_from_ray_options
 from ray.serve._private.constants import (
+    ACCELERATOR_KIND_TPU,
     DEFAULT_CONSTRUCTOR_RETRY_COUNT,
     DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_S,
     DEFAULT_GRACEFUL_SHUTDOWN_WAIT_LOOP_S,
@@ -32,6 +33,7 @@
 )
 from ray.serve._private.utils import DEFAULT, DeploymentOptionUpdateType
 from ray.serve.config import (
+    AcceleratorConfig,
     AggregationFunction,
     AutoscalingConfig,
     DeploymentActorConfig,
@@ -41,7 +43,9 @@
     HTTPOptions,
     ProxyLocation,
     RequestRouterConfig,
+    TPUAcceleratorConfig,
 )
+from ray.serve.generated import serve_pb2
 from ray.serve.generated.serve_pb2 import (
     AutoscalingConfig as AutoscalingConfigProto,
     DeploymentActorConfig as DeploymentActorConfigProto,
@@ -87,9 +91,12 @@ def _proto_to_dict(proto: Message) -> Dict:
             # `google.protobuf.internal.containers.RepeatedScalarFieldContainer
             # Explicitly convert to list
             if field.type == FieldDescriptor.TYPE_MESSAGE:
-                data[field.name] = [
-                    _proto_to_dict(v) for v in value
-                ]  # Convert each item
+                if field.message_type.GetOptions().map_entry:
+                    data[field.name] = dict(value)
+                else:
+                    data[field.name] = [
+                        _proto_to_dict(v) for v in value
+                    ]  # Convert each item
             else:
                 data[field.name] = list(value)  # Convert to list directly
         # Recursively call if the field is another protobuf.
@@ -190,6 +197,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 +333,20 @@ def needs_pickle(self):
 
     def to_proto(self):
         data = self.model_dump()
+        if self.accelerator_config is not None:
+            if isinstance(self.accelerator_config, TPUAcceleratorConfig):
+                tpu_proto = serve_pb2.TPUAcceleratorConfig(
+                    topology=self.accelerator_config.topology,
+                    accelerator_version=self.accelerator_config.accelerator_version,
+                    num_slices=self.accelerator_config.num_slices,
+                )
+                if self.accelerator_config.chips_per_vm is not None:
+                    tpu_proto.chips_per_vm = self.accelerator_config.chips_per_vm
+                if self.accelerator_config.resources_per_bundle is not None:
+                    tpu_proto.resources_per_bundle.update(
+                        self.accelerator_config.resources_per_bundle
+                    )
+                data["accelerator_config"] = serve_pb2.AcceleratorConfig(tpu=tpu_proto)
         if data.get("user_config") is not None:
             if self.needs_pickle():
                 data["user_config"] = cloudpickle.dumps(data["user_config"])
@@ -429,6 +454,26 @@ 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 proto.HasField("accelerator_config"):
+            ac_proto = proto.accelerator_config
+            field = ac_proto.WhichOneof("config")
+            if field == ACCELERATOR_KIND_TPU:
+                tpu_proto = ac_proto.tpu
+                data["accelerator_config"] = TPUAcceleratorConfig(
+                    topology=tpu_proto.topology,
+                    accelerator_version=tpu_proto.accelerator_version,
+                    num_slices=tpu_proto.num_slices,
+                    chips_per_vm=tpu_proto.chips_per_vm
+                    if tpu_proto.HasField("chips_per_vm")
+                    else None,
+                    resources_per_bundle=dict(tpu_proto.resources_per_bundle)
+                    if tpu_proto.resources_per_bundle
+                    else None,
+                )
+            else:
+                data["accelerator_config"] = None
+        else:
+            data["accelerator_config"] = None
         if "user_config" in data:
             if data["user_config"] != b"":
                 if needs_pickle:

python/ray/serve/_private/constants.py

@@ -43,6 +43,8 @@
 #: Ray namespace used for all Serve actors
 SERVE_NAMESPACE = "serve"
 
+ACCELERATOR_KIND_TPU = "tpu"
+
 DEFAULT_HTTP_HOST = os.environ.get("RAY_SERVE_DEFAULT_HTTP_HOST")
 
 #: HTTP Port

python/ray/serve/_private/version.py

@@ -79,6 +79,7 @@ def requires_actor_restart(self, new_version):
             or self.max_replicas_per_node != new_version.max_replicas_per_node
             or self.gang_scheduling_config_hash
             != new_version.gang_scheduling_config_hash
+            or self.accelerator_config_hash != new_version.accelerator_config_hash
         )
 
     def requires_actor_reconfigure(self, new_version):
@@ -124,6 +125,12 @@ def compute_hashes(self):
             else {}
         )
         self.gang_scheduling_config_hash = crc32(serialized_gang_scheduling_config)
+        serialized_accelerator_config = (
+            self.deployment_config.accelerator_config.model_dump_json().encode("utf-8")
+            if self.deployment_config.accelerator_config is not None
+            else b""
+        )
+        self.accelerator_config_hash = crc32(serialized_accelerator_config)
         # Include app-level route prefix in the version hashes so changing
         # it triggers an in-place reconfigure of running replicas.
         serialized_route_prefix = _serialize(self.route_prefix)
@@ -152,6 +159,7 @@ def compute_hashes(self):
                 ]
             )
             + serialized_gang_scheduling_config
+            + serialized_accelerator_config
         )
 
     def to_proto(self) -> bytes:

python/ray/serve/api.py

@@ -40,13 +40,15 @@
     wait_for_interrupt,
 )
 from ray.serve.config import (
+    AcceleratorConfig,
     AutoscalingConfig,
     ControllerOptions,
     DeploymentActorConfig,
     GangSchedulingConfig,
     HTTPOptions,
     ProxyLocation,
     RequestRouterConfig,
+    _resolve_accelerator_config,
     gRPCOptions,
 )
 from ray.serve.context import (
@@ -484,6 +486,7 @@ def deployment(
     user_config: Default[Optional[Any]] = DEFAULT.VALUE,
     max_ongoing_requests: Default[int] = DEFAULT.VALUE,
     max_queued_requests: Default[int] = DEFAULT.VALUE,
+    accelerator_config: Default[Union[Dict, AcceleratorConfig, None]] = DEFAULT.VALUE,
     autoscaling_config: Default[Union[Dict, AutoscalingConfig, None]] = DEFAULT.VALUE,
     graceful_shutdown_wait_loop_s: Default[float] = DEFAULT.VALUE,
     graceful_shutdown_timeout_s: Default[float] = DEFAULT.VALUE,
@@ -554,6 +557,9 @@ class MyDeployment:
             Once this limit is reached, subsequent requests will raise a
             BackPressureError (for handles) or return an HTTP 503 status code (for HTTP
             requests). Defaults to -1 (no limit).
+        accelerator_config: Configuration for hardware accelerators, such as TPUs.
+            Can be passed as an unstructured dictionary or a structured `AcceleratorConfig`
+            subclass (e.g. `TPUAcceleratorConfig`). See `AcceleratorConfig` for options.
         autoscaling_config: Parameters to configure autoscaling behavior. If this
             is set, `num_replicas` should be "auto" or not set.
         graceful_shutdown_wait_loop_s: Duration that replicas wait until there is
@@ -654,11 +660,32 @@ class MyDeployment:
     if isinstance(logging_config, LoggingConfig):
         logging_config = logging_config.model_dump()
 
+    if accelerator_config is not DEFAULT.VALUE and accelerator_config is not None:
+        accelerator_config = _resolve_accelerator_config(accelerator_config)
+
+        if (
+            gang_scheduling_config is not DEFAULT.VALUE
+            and gang_scheduling_config is not None
+        ):
+            # TODO(ryanaoleary@): Revisit this mutual exclusivity restriction once
+            # Data Parallel (DP) attention or more complex multi-slice gang
+            # scheduling is supported for TPUs.
+            #
+            # The only supported accelerator_config currently is for TPU, which utilizes
+            # SlicePlacementGroup internally for atomic scheduling of SPMD workers. This
+            # check can be loosened if additional accelerator configs are added in the
+            # future that don't manage their own gang scheduling.
+            raise ValueError(
+                "Cannot specify both `accelerator_config` and `gang_scheduling_config`. "
+                "Accelerator configurations automatically manage their own gang scheduling."
+            )
+
     deployment_config = DeploymentConfig.from_default(
         num_replicas=num_replicas if num_replicas is not None else 1,
         user_config=user_config,
         max_ongoing_requests=max_ongoing_requests,
         max_queued_requests=max_queued_requests,
+        accelerator_config=accelerator_config,
         autoscaling_config=autoscaling_config,
         graceful_shutdown_wait_loop_s=graceful_shutdown_wait_loop_s,
         graceful_shutdown_timeout_s=graceful_shutdown_timeout_s,

python/ray/serve/config.py

@@ -4,7 +4,7 @@
 import warnings
 from enum import Enum
 from functools import cached_property
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 from pydantic import (
     BaseModel,
@@ -27,6 +27,7 @@
 # Import types needed for AutoscalingContext
 from ray.serve._private.common import DeploymentID, ReplicaID, TimeSeries
 from ray.serve._private.constants import (
+    ACCELERATOR_KIND_TPU,
     DEFAULT_AUTOSCALING_POLICY_NAME,
     DEFAULT_GRPC_PORT,
     DEFAULT_HTTP_HOST,
@@ -709,6 +710,77 @@ def get_target_ongoing_requests(self) -> PositiveFloat:
         return self.target_ongoing_requests
 
 
+@PublicAPI(stability="alpha")
+class AcceleratorConfig(BaseModel):
+    """Base class for structured accelerator configurations.
+
+    Use a concrete subclass — e.g. :class:`TPUAcceleratorConfig` — when
+    declaring a deployment's accelerator requirements via
+    ``serve.deployment(accelerator_config=...)``.
+    """
+
+    kind: str = Field(
+        ..., description="Discriminator identifying the accelerator config type."
+    )
+
+    model_config = {"frozen": True, "extra": "forbid"}
+
+    @model_validator(mode="after")
+    def validate_concrete_subclass(self):
+        if type(self) is AcceleratorConfig:
+            raise ValueError(
+                "AcceleratorConfig is an abstract base class. "
+                "Please use a concrete subclass like TPUAcceleratorConfig."
+            )
+        return self
+
+
+@PublicAPI(stability="alpha")
+class TPUAcceleratorConfig(AcceleratorConfig):
+    """TPU slice specification for a Serve deployment.
+
+    Mirrors the parameters of :func:`ray.util.tpu.slice_placement_group`.
+    Ray Serve uses this config to provision a TPU slice placement group
+    per replica and to manage its lifecycle through the controller.
+
+    When set on a deployment, this config drives placement-group creation
+    entirely. The deployment's ``placement_group_bundles`` and
+    ``placement_group_strategy`` fields are ignored - the bundles are
+    derived from ``topology`` (or optionally ``resources_per_bundle``),
+    and the strategy is chosen internally to honor slice gang scheduling.
+
+    Example:
+        >>> from ray.serve.config import TPUAcceleratorConfig
+        >>> config = TPUAcceleratorConfig(topology="4x4", accelerator_version="v6e")
+    """
+
+    kind: Literal["tpu"] = ACCELERATOR_KIND_TPU
+
+    topology: str = Field(
+        ..., description="TPU pod topology, e.g. '2x2', '4x4', '2x2x2'."
+    )
+    accelerator_version: str = Field(
+        ..., description="TPU accelerator version, e.g. 'v4', 'v5p', 'v6e'."
+    )
+    num_slices: int = Field(default=1, ge=1, description="Number of slices to reserve.")
+    chips_per_vm: Optional[int] = Field(
+        default=None,
+        description=(
+            "Override for chips per host. Defaults to the canonical value "
+            "for the given accelerator_version."
+        ),
+    )
+    resources_per_bundle: Optional[Dict[str, float]] = Field(
+        default=None,
+        description=(
+            "Resources to include in every worker bundle. When unspecified, "
+            "SlicePlacementGroup defaults to one bundle per TPU host with "
+            "the bundle resources set to the number of chips on that host. "
+            "See ray.util.tpu.slice_placement_group for details."
+        ),
+    )
+
+
 @PublicAPI(stability="stable")
 class ProxyLocation(str, Enum):
     """Config for where to run proxies to receive ingress traffic to the cluster.
@@ -1165,3 +1237,19 @@ def _validate_runtime_failure_policy(cls, v):
                 "RESTART_REPLICA policy is not yet implemented. File a GitHub issue if you need this feature."
             )
         return v
+
+
+def _resolve_accelerator_config(
+    value: Union[Dict, AcceleratorConfig, None],
+) -> Optional[AcceleratorConfig]:
+
+    if value is None or isinstance(value, AcceleratorConfig):
+        return value
+    if isinstance(value, dict):
+        kind = value.get("kind")
+        if kind == ACCELERATOR_KIND_TPU:
+            return TPUAcceleratorConfig(**value)
+        raise ValueError(f"Unknown accelerator kind {kind!r}. Supported types: 'tpu'.")
+    raise TypeError(
+        f"accelerator_config must be a dict or AcceleratorConfig, got {type(value)}."
+    )

python/ray/serve/deployment.py

@@ -13,9 +13,11 @@
 from ray.serve._private.usage import ServeUsageTag
 from ray.serve._private.utils import DEFAULT, Default
 from ray.serve.config import (
+    AcceleratorConfig,
     AutoscalingConfig,
     DeploymentActorConfig,
     GangSchedulingConfig,
+    _resolve_accelerator_config,
 )
 from ray.serve.schema import DeploymentSchema, LoggingConfig, RayActorOptionsSchema
 from ray.util.annotations import PublicAPI
@@ -257,6 +259,9 @@ def options(
         deployment_actors: Default[
             Optional[List[Union[Dict, DeploymentActorConfig]]]
         ] = DEFAULT.VALUE,
+        accelerator_config: Default[
+            Union[Dict, AcceleratorConfig, None]
+        ] = DEFAULT.VALUE,
     ) -> "Deployment":
         """Return a copy of this deployment with updated options.
 
@@ -408,6 +413,19 @@ def options(
         if gang_scheduling_config is not DEFAULT.VALUE:
             new_deployment_config.gang_scheduling_config = gang_scheduling_config
 
+        if accelerator_config is not DEFAULT.VALUE:
+            if accelerator_config is not None:
+                accelerator_config = _resolve_accelerator_config(accelerator_config)
+            new_deployment_config.accelerator_config = accelerator_config
+
+        ac = new_deployment_config.accelerator_config
+        gc = new_deployment_config.gang_scheduling_config
+        if ac is not None and gc is not None:
+            raise ValueError(
+                "Cannot specify both `accelerator_config` and `gang_scheduling_config`. "
+                "Accelerator configurations automatically manage their own gang scheduling."
+            )
+
         if deployment_actors is not DEFAULT.VALUE:
             new_deployment_config.deployment_actors = deployment_actors
 
@@ -513,6 +531,7 @@ def deployment_to_schema(d: Deployment) -> DeploymentSchema:
         "gang_scheduling_config": d._deployment_config.gang_scheduling_config,
         "deployment_actors": d._deployment_config.deployment_actors,
         "rolling_update_percentage": d._deployment_config.rolling_update_percentage,
+        "accelerator_config": d._deployment_config.accelerator_config,
     }
 
     # Let non-user-configured options be set to defaults. If the schema
@@ -577,6 +596,7 @@ def schema_to_deployment(s: DeploymentSchema) -> Deployment:
         gang_scheduling_config=s.gang_scheduling_config,
         deployment_actors=s.deployment_actors,
         rolling_update_percentage=s.rolling_update_percentage,
+        accelerator_config=s.accelerator_config,
     )
     deployment_config.user_configured_option_names = (
         s._get_user_configured_option_names()