[Serve] Add label_selector and bundle_label_selector to Serve API

python/ray/serve/_private/config.py

@@ -459,6 +459,7 @@ def __init__(
         ray_actor_options: Dict,
         placement_group_bundles: Optional[List[Dict[str, float]]] = None,
         placement_group_strategy: Optional[str] = None,
+        bundle_label_selector: Optional[List[Dict[str, str]]] = None,
         max_replicas_per_node: Optional[int] = None,
         needs_pickle: bool = True,
     ):
@@ -484,6 +485,7 @@ def __init__(
 
         self.placement_group_bundles = placement_group_bundles
         self.placement_group_strategy = placement_group_strategy
+        self.bundle_label_selector = bundle_label_selector
 
         self.max_replicas_per_node = max_replicas_per_node
 
@@ -514,12 +516,14 @@ def update(
         ray_actor_options: dict,
         placement_group_bundles: Optional[List[Dict[str, float]]] = None,
         placement_group_strategy: Optional[str] = None,
+        bundle_label_selector: Optional[List[Dict[str, str]]] = None,
         max_replicas_per_node: Optional[int] = None,
     ):
         self.ray_actor_options = ray_actor_options
 
         self.placement_group_bundles = placement_group_bundles
         self.placement_group_strategy = placement_group_strategy
+        self.bundle_label_selector = bundle_label_selector
 
         self.max_replicas_per_node = max_replicas_per_node
 
@@ -536,6 +540,7 @@ def create(
         ray_actor_options: Optional[Dict] = None,
         placement_group_bundles: Optional[List[Dict[str, float]]] = None,
         placement_group_strategy: Optional[str] = None,
+        bundle_label_selector: Optional[List[Dict[str, str]]] = None,
         max_replicas_per_node: Optional[int] = None,
         deployment_def_name: Optional[str] = None,
     ):
@@ -586,6 +591,7 @@ def create(
             ray_actor_options,
             placement_group_bundles,
             placement_group_strategy,
+            bundle_label_selector,
             max_replicas_per_node,
         )
 
@@ -612,6 +618,7 @@ def _validate_ray_actor_options(self):
             "resources",
             # Other options
             "runtime_env",
+            "label_selector",
         }
 
         for option in self.ray_actor_options:
@@ -658,6 +665,7 @@ def _validate_placement_group_options(self) -> None:
                 bundles=self.placement_group_bundles,
                 strategy=self.placement_group_strategy or "PACK",
                 lifetime="detached",
+                bundle_label_selector=self.bundle_label_selector,
             )
 
             resource_error_prefix = (
@@ -763,6 +771,9 @@ def from_proto(cls, proto: ReplicaConfigProto, needs_pickle: bool = True):
             if proto.placement_group_strategy != ""
             else None,
             proto.max_replicas_per_node if proto.max_replicas_per_node else None,
+            json.loads(proto.bundle_label_selector)
+            if proto.bundle_label_selector
+            else None,
             needs_pickle,
         )
 
@@ -785,6 +796,9 @@ def to_proto(self):
             max_replicas_per_node=self.max_replicas_per_node
             if self.max_replicas_per_node is not None
             else 0,
+            bundle_label_selector=json.dumps(self.bundle_label_selector)
+            if self.bundle_label_selector is not None
+            else "",
         )
 
     def to_proto_bytes(self):

python/ray/serve/api.py

@@ -339,6 +339,7 @@ def deployment(
     ray_actor_options: Default[Dict] = DEFAULT.VALUE,
     placement_group_bundles: Default[List[Dict[str, float]]] = DEFAULT.VALUE,
     placement_group_strategy: Default[str] = DEFAULT.VALUE,
+    bundle_label_selector: Default[List[Dict[str, str]]] = DEFAULT.VALUE,
     max_replicas_per_node: Default[int] = DEFAULT.VALUE,
     user_config: Default[Optional[Any]] = DEFAULT.VALUE,
     max_ongoing_requests: Default[int] = DEFAULT.VALUE,
@@ -378,7 +379,7 @@ class MyDeployment:
         route_prefix: Route prefix for HTTP requests. Defaults to '/'. Deprecated.
         ray_actor_options: Options to pass to the Ray Actor decorator, such as
             resource requirements. Valid options are: `accelerator_type`, `memory`,
-            `num_cpus`, `num_gpus`, `resources`, and `runtime_env`.
+            `num_cpus`, `num_gpus`, `resources`, `runtime_env`, and `label_selector`.
         placement_group_bundles: Defines a set of placement group bundles to be
             scheduled *for each replica* of this deployment. The replica actor will
             be scheduled in the first bundle provided, so the resources specified in
@@ -389,6 +390,8 @@ class MyDeployment:
             This cannot be set together with max_replicas_per_node.
         placement_group_strategy: Strategy to use for the replica placement group
             specified via `placement_group_bundles`. Defaults to `PACK`.
+        bundle_label_selector: A list of label selectors to apply to the
+            placement group on a per-bundle level.
         max_replicas_per_node: The max number of replicas of this deployment that can
             run on a single node. Valid values are None (default, no limit)
             or an integer in the range of [1, 100].
@@ -507,6 +510,11 @@ def decorator(_func_or_class):
                 if placement_group_strategy is not DEFAULT.VALUE
                 else None
             ),
+            bundle_label_selector=(
+                bundle_label_selector
+                if bundle_label_selector is not DEFAULT.VALUE
+                else None
+            ),
             max_replicas_per_node=(
                 max_replicas_per_node
                 if max_replicas_per_node is not DEFAULT.VALUE

python/ray/serve/tests/conftest.py

@@ -344,3 +344,38 @@ def metrics_start_shutdown(request):
     serve.shutdown()
     ray.shutdown()
     reset_ray_address()
+
+
+# Helper function to return the node ID of a remote worker.
+@ray.remote(num_cpus=0)
+def _get_node_id():
+    return ray.get_runtime_context().get_node_id()
+
+
+# Test fixture to start a Serve instance in a RayCluster with two labelled nodes
+@pytest.fixture
+def serve_instance_with_labeled_nodes():
+    cluster = Cluster()
+    node1_config = {
+        "num_cpus": 1,
+        "resources": {"worker1": 1},
+        "labels": {"region": "us-west", "gpu-type": "A100"},
+    }
+    cluster.add_node(**node1_config)
+    ray.init(address=cluster.address)
+    node_1_id = ray.get(_get_node_id.options(resources={"worker1": 1}).remote())
+
+    node2_config = {
+        "num_cpus": 1,
+        "resources": {"worker2": 1},
+        "labels": {"region": "us-east", "gpu-type": "H100"},
+    }
+    cluster.add_node(**node2_config)
+    node_2_id = ray.get(_get_node_id.options(resources={"worker2": 1}).remote())
+
+    serve.start()
+    yield _get_global_client(), node_1_id, node_2_id
+
+    serve.shutdown()
+    ray.shutdown()
+    cluster.shutdown()

python/ray/serve/tests/test_api.py

@@ -1219,6 +1219,35 @@ def __init__(self, counter):
     wait_for_condition(lambda: ray.get(counter.get_count.remote()) == 9)
 
 
+async def test_serve_label_selector_api(serve_instance_with_labeled_nodes):
+    serve_instance, node_1_id, node_2_id = serve_instance_with_labeled_nodes
+
+    # Test `label_selector` in `ray_actor_options` targets the desired node.
+    @serve.deployment(ray_actor_options={"label_selector": {"region": "us-west"}})
+    class Deployment1:
+        def get_node_id(self):
+            return ray.get_runtime_context().get_node_id()
+
+    handle1 = serve.run(Deployment1.bind(), name="app1")
+    replica_node_id_1 = await handle1.get_node_id.remote()
+    assert replica_node_id_1 == node_1_id
+    serve.delete("app1")
+
+    # Test `bundle_label_selector` targets the desired node.
+    @serve.deployment(
+        placement_group_bundles=[{"CPU": 1}],
+        bundle_label_selector=[{"gpu-type": "H100"}],
+    )
+    class Deployment2:
+        def get_node_id(self):
+            return ray.get_runtime_context().get_node_id()
+
+    handle2 = serve.run(Deployment2.bind(), name="app2")
+    replica_node_id_2 = await handle2.get_node_id.remote()
+    assert replica_node_id_2 == node_2_id
+    serve.delete("app2")
+
+
 if __name__ == "__main__":
     import sys
 

src/ray/protobuf/serve.proto

@@ -255,6 +255,7 @@ message ReplicaConfig {
   string placement_group_bundles = 6;
   string placement_group_strategy = 7;
   int32 max_replicas_per_node = 8;
+  string bundle_label_selector = 9;
 }
 
 enum TargetCapacityDirection {