[Serve] Add label_selector and bundle_label_selector to Serve API (ray-project#57694)

<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

This PR adds the `label_selector` option to the supported list of Actor
options for a Serve deployment. Additionally, we add
`bundle_label_selector` to specify label selectors for bundles when
`placement_group_bundles` are specified for the deployment. These two
options are already supported for Tasks/Actors and placement groups
respectively.

Example use case:
```
llm_config = LLMConfig(
    model_loading_config={
        "model_id": "meta-llama/Meta-Llama-3-70B-Instruct",
        "model_source": "huggingface",
    },
    engine_kwargs=tpu_engine_config,
    resources_per_bundle={"TPU": 4},
    runtime_env={"env_vars": {"VLLM_USE_V1": "1"}},
    deployment_config={
        "num_replicas": 4,
        "ray_actor_options": {
            # In a GKE cluster with multiple TPU node-pools, schedule
            # only to the desired slice.
            "label_selector": {
                "ray.io/tpu-topology": "4x4" # added by default by Ray
            }
        }
    }
)
```

The expected behaviors of these new fields is as follows:

**Pack scheduling enabled**
----------------------------------------
**PACK/STRICT_PACK PG strategy:**
- Standard PG without bundle_label_selector or fallback:
- Sorts replicas by resource size (descending). Attempts to find the
"best fit" node (minimizing fragmentation) that has available resources.
Creates a Placement Group on that target node.

- PG node label selector provided:
- Same behavior as regular placement group but filters the list of
candidate nodes to only those matching the label selector before finding
the best fit
  
- PG node label selector and fallback:
  Same as above but when scheduling tries the following:
1. Tries to find a node matching the primary placement_group_bundles and
bundle_label_selector.
2. If no node fits, iterates through the
placement_group_fallback_strategy. For each fallback entry, tries to
find a node matching that entry's bundles and labels.
  3. If a node is found, creates a PG on it.
  
**SPREAD/STRICT_SPREAD PG strategy:**
- If any deployment uses these strategies, the global logic falls back
to "Spread Scheduling" (see below)

**Spread scheduling enabled**
----------------------------------------
- Standard PG without bundle_label_selector or fallback:
- Creates a Placement Group via Ray Core without specifying a
target_node_id. Ray Core decides placement based on the strategy.
- PG node label selector provided:
- Serve passes the bundle_label_selector to the
CreatePlacementGroupRequest. Ray Core handles the soft/hard constraint
logic during PG creation.
- PG node label selector and fallback:
- Serve passes the bundle_label_selector to the
CreatePlacementGroupRequest, fallback_strategy is not yet supported in
the placement group options so this field isn't passed / considered.
It's only used in the "best fit" node selection logic which is skipped
for Spread scheduling.

## Related issue number

ray-project#51564

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] I've run pre-commit jobs to lint the changes in this PR.
([pre-commit
setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting))
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [x] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: Ryan O'Leary <113500783+ryanaoleary@users.noreply.github.com>
Co-authored-by: Cindy Zhang <cindyzyx9@gmail.com>
Co-authored-by: Abrar Sheikh <abrar2002as@gmail.com>
Co-authored-by: Mengjin Yan <mengjinyan3@gmail.com>

Author: 4 people

python/ray/_raylet.pyx

@@ -103,6 +103,8 @@ from ray.includes.common cimport (
     CLabelMatchExpression,
     CLabelIn,
     CLabelNotIn,
+    CLabelSelector,
+    CNodeResources,
     CRayFunction,
     CWorkerType,
     CJobConfig,
@@ -141,6 +143,7 @@ from ray.includes.common cimport (
     PersistPort,
     WaitForPersistedPort,
     CWaitForPersistedPortResult,
+    SetNodeResourcesLabels,
 )
 from ray.includes.unique_ids cimport (
     CActorID,
@@ -597,6 +600,23 @@ cdef int prepare_label_selector(
 
     return 0
 
+def node_labels_match_selector(node_labels: Dict[str, str], selector: Dict[str, str]) -> bool:
+    """
+    Checks if the given node labels satisfy the label selector. This helper function exposes
+    the C++ logic for determining if a node satisfies a label selector to the Python layer.
+    """
+    cdef:
+        CNodeResources c_node_resources
+        CLabelSelector c_label_selector
+        unordered_map[c_string, c_string] c_labels_map
+
+    prepare_labels(node_labels, &c_labels_map)
+    SetNodeResourcesLabels(c_node_resources, c_labels_map)
+    prepare_label_selector(selector, &c_label_selector)
+
+    # Return whether the node resources satisfy the label constraint.
+    return c_node_resources.HasRequiredLabels(c_label_selector)
+
 cdef int prepare_fallback_strategy(
         list fallback_strategy,
         c_vector[CFallbackOption] *fallback_strategy_vector) except -1:

python/ray/includes/common.pxd

@@ -275,6 +275,14 @@ cdef extern from "src/ray/protobuf/common.pb.h" nogil:
         CLineageReconstructionTask()
         const c_string &SerializeAsString() const
 
+cdef extern from "ray/common/scheduling/cluster_resource_data.h" namespace "ray" nogil:
+    cdef cppclass CNodeResources "ray::NodeResources":
+        CNodeResources()
+        unordered_map[c_string, c_string] labels
+        c_bool HasRequiredLabels(const CLabelSelector &label_selector) const
+
+    void SetNodeResourcesLabels(CNodeResources& resources, const unordered_map[c_string, c_string]& labels)
+
 cdef extern from "ray/common/scheduling/label_selector.h" namespace "ray":
     cdef cppclass CLabelSelector "ray::LabelSelector":
         CLabelSelector() nogil except +

python/ray/serve/_private/application_state.py

@@ -1707,6 +1707,14 @@ def override_deployment_info(
         override_max_replicas_per_node = options.pop(
             "max_replicas_per_node", replica_config.max_replicas_per_node
         )
+        override_bundle_label_selector = options.pop(
+            "placement_group_bundle_label_selector",
+            replica_config.placement_group_bundle_label_selector,
+        )
+        override_fallback_strategy = options.pop(
+            "placement_group_fallback_strategy",
+            replica_config.placement_group_fallback_strategy,
+        )
 
         # Record telemetry for container runtime env feature at deployment level
         if override_actor_options.get("runtime_env") and (
@@ -1725,6 +1733,8 @@ def override_deployment_info(
             placement_group_bundles=override_placement_group_bundles,
             placement_group_strategy=override_placement_group_strategy,
             max_replicas_per_node=override_max_replicas_per_node,
+            placement_group_bundle_label_selector=override_bundle_label_selector,
+            placement_group_fallback_strategy=override_fallback_strategy,
         )
         override_options["replica_config"] = replica_config
 

python/ray/serve/_private/cluster_node_info_cache.py

@@ -89,6 +89,10 @@ def get_available_resources_per_node(self) -> Dict[str, Union[float, Dict]]:
 
         return self._cached_available_resources_per_node
 
+    def get_node_labels(self, node_id: str) -> Dict[str, str]:
+        """Get the labels for a specific node from the cache."""
+        return self._cached_node_labels.get(node_id, {})
+
 
 class DefaultClusterNodeInfoCache(ClusterNodeInfoCache):
     def __init__(self, gcs_client: GcsClient):

python/ray/serve/_private/common.py

@@ -855,6 +855,8 @@ class CreatePlacementGroupRequest:
     target_node_id: str
     name: str
     runtime_env: Optional[str] = None
+    bundle_label_selector: Optional[List[Dict[str, str]]] = None
+    fallback_strategy: Optional[List[Dict[str, Any]]] = None
 
 
 # This error is used to raise when a by-value DeploymentResponse is converted to an

python/ray/serve/_private/config.py

@@ -480,6 +480,8 @@ def __init__(
         ray_actor_options: Dict,
         placement_group_bundles: Optional[List[Dict[str, float]]] = None,
         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,
         max_replicas_per_node: Optional[int] = None,
         needs_pickle: bool = True,
     ):
@@ -505,9 +507,14 @@ def __init__(
 
         self.placement_group_bundles = placement_group_bundles
         self.placement_group_strategy = placement_group_strategy
+        self.placement_group_bundle_label_selector = (
+            placement_group_bundle_label_selector
+        )
+        self.placement_group_fallback_strategy = placement_group_fallback_strategy
 
         self.max_replicas_per_node = max_replicas_per_node
 
+        self._normalize_bundle_label_selector()
         self._validate()
 
         # Create resource_dict. This contains info about the replica's resource
@@ -516,6 +523,21 @@ def __init__(
         self.resource_dict = resources_from_ray_options(self.ray_actor_options)
         self.needs_pickle = needs_pickle
 
+    def _normalize_bundle_label_selector(self):
+        """If a single selector is provided for multiple bundles, it is broadcasted
+        uniformly to all bundles.
+        """
+        if (
+            self.placement_group_bundles
+            and self.placement_group_bundle_label_selector
+            and len(self.placement_group_bundle_label_selector) == 1
+            and len(self.placement_group_bundles) > 1
+        ):
+            single_selector = self.placement_group_bundle_label_selector[0]
+            self.placement_group_bundle_label_selector = [
+                single_selector.copy() for _ in range(len(self.placement_group_bundles))
+            ]
+
     def _validate(self):
         self._validate_ray_actor_options()
         self._validate_placement_group_options()
@@ -535,15 +557,22 @@ def update(
         ray_actor_options: dict,
         placement_group_bundles: Optional[List[Dict[str, float]]] = None,
         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,
         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.placement_group_bundle_label_selector = (
+            placement_group_bundle_label_selector
+        )
+        self.placement_group_fallback_strategy = placement_group_fallback_strategy
 
         self.max_replicas_per_node = max_replicas_per_node
 
+        self._normalize_bundle_label_selector()
         self._validate()
 
         self.resource_dict = resources_from_ray_options(self.ray_actor_options)
@@ -557,6 +586,8 @@ def create(
         ray_actor_options: Optional[Dict] = None,
         placement_group_bundles: Optional[List[Dict[str, float]]] = None,
         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,
         max_replicas_per_node: Optional[int] = None,
         deployment_def_name: Optional[str] = None,
     ):
@@ -597,17 +628,23 @@ def create(
                 deployment_def_name = deployment_def.__name__
 
         config = cls(
-            deployment_def_name,
-            pickle_dumps(
+            deployment_def_name=deployment_def_name,
+            serialized_deployment_def=pickle_dumps(
                 deployment_def,
                 f"Could not serialize the deployment {repr(deployment_def)}",
             ),
-            pickle_dumps(init_args, "Could not serialize the deployment init args"),
-            pickle_dumps(init_kwargs, "Could not serialize the deployment init kwargs"),
-            ray_actor_options,
-            placement_group_bundles,
-            placement_group_strategy,
-            max_replicas_per_node,
+            serialized_init_args=pickle_dumps(
+                init_args, "Could not serialize the deployment init args"
+            ),
+            serialized_init_kwargs=pickle_dumps(
+                init_kwargs, "Could not serialize the deployment init kwargs"
+            ),
+            ray_actor_options=ray_actor_options,
+            placement_group_bundles=placement_group_bundles,
+            placement_group_strategy=placement_group_strategy,
+            placement_group_bundle_label_selector=placement_group_bundle_label_selector,
+            placement_group_fallback_strategy=placement_group_fallback_strategy,
+            max_replicas_per_node=max_replicas_per_node,
         )
 
         config._deployment_def = deployment_def
@@ -633,6 +670,8 @@ def _validate_ray_actor_options(self):
             "resources",
             # Other options
             "runtime_env",
+            "label_selector",
+            "fallback_strategy",
         }
 
         for option in self.ray_actor_options:
@@ -674,11 +713,37 @@ def _validate_placement_group_options(self) -> None:
                     "`placement_group_bundles` must also be provided."
                 )
 
+        if self.placement_group_fallback_strategy is not None:
+            if self.placement_group_bundles is None:
+                raise ValueError(
+                    "If `placement_group_fallback_strategy` is provided, "
+                    "`placement_group_bundles` must also be provided."
+                )
+            if not isinstance(self.placement_group_fallback_strategy, list):
+                raise TypeError(
+                    "placement_group_fallback_strategy must be a list of dictionaries. "
+                    f"Got: {type(self.placement_group_fallback_strategy)}."
+                )
+            for i, strategy in enumerate(self.placement_group_fallback_strategy):
+                if not isinstance(strategy, dict):
+                    raise TypeError(
+                        f"placement_group_fallback_strategy entry at index {i} must be a dictionary. "
+                        f"Got: {type(strategy)}."
+                    )
+
+        if self.placement_group_bundle_label_selector is not None:
+            if self.placement_group_bundles is None:
+                raise ValueError(
+                    "If `placement_group_bundle_label_selector` is provided, "
+                    "`placement_group_bundles` must also be provided."
+                )
+
         if self.placement_group_bundles is not None:
             validate_placement_group(
                 bundles=self.placement_group_bundles,
                 strategy=self.placement_group_strategy or "PACK",
                 lifetime="detached",
+                bundle_label_selector=self.placement_group_bundle_label_selector,
             )
 
             resource_error_prefix = (
@@ -772,19 +837,37 @@ def init_kwargs(self) -> Optional[Tuple[Any]]:
     @classmethod
     def from_proto(cls, proto: ReplicaConfigProto, needs_pickle: bool = True):
         return ReplicaConfig(
-            proto.deployment_def_name,
-            proto.deployment_def,
-            proto.init_args if proto.init_args != b"" else None,
-            proto.init_kwargs if proto.init_kwargs != b"" else None,
-            json.loads(proto.ray_actor_options),
-            json.loads(proto.placement_group_bundles)
-            if proto.placement_group_bundles
-            else None,
-            proto.placement_group_strategy
-            if proto.placement_group_strategy != ""
-            else None,
-            proto.max_replicas_per_node if proto.max_replicas_per_node else None,
-            needs_pickle,
+            deployment_def_name=proto.deployment_def_name,
+            serialized_deployment_def=proto.deployment_def,
+            serialized_init_args=(proto.init_args if proto.init_args != b"" else None),
+            serialized_init_kwargs=(
+                proto.init_kwargs if proto.init_kwargs != b"" else None
+            ),
+            ray_actor_options=json.loads(proto.ray_actor_options),
+            placement_group_bundles=(
+                json.loads(proto.placement_group_bundles)
+                if proto.placement_group_bundles
+                else None
+            ),
+            placement_group_strategy=(
+                proto.placement_group_strategy
+                if proto.placement_group_strategy != ""
+                else None
+            ),
+            placement_group_bundle_label_selector=(
+                json.loads(proto.placement_group_bundle_label_selector)
+                if proto.placement_group_bundle_label_selector
+                else None
+            ),
+            placement_group_fallback_strategy=(
+                json.loads(proto.placement_group_fallback_strategy)
+                if proto.placement_group_fallback_strategy
+                else None
+            ),
+            max_replicas_per_node=(
+                proto.max_replicas_per_node if proto.max_replicas_per_node else None
+            ),
+            needs_pickle=needs_pickle,
         )
 
     @classmethod
@@ -793,19 +876,39 @@ def from_proto_bytes(cls, proto_bytes: bytes, needs_pickle: bool = True):
         return cls.from_proto(proto, needs_pickle)
 
     def to_proto(self):
+        placement_group_bundles = (
+            json.dumps(self.placement_group_bundles)
+            if self.placement_group_bundles is not None
+            else ""
+        )
+
+        bundle_label_selector = (
+            json.dumps(self.placement_group_bundle_label_selector)
+            if self.placement_group_bundle_label_selector is not None
+            else ""
+        )
+
+        fallback_strategy = (
+            json.dumps(self.placement_group_fallback_strategy)
+            if self.placement_group_fallback_strategy is not None
+            else ""
+        )
+
+        max_replicas_per_node = (
+            self.max_replicas_per_node if self.max_replicas_per_node is not None else 0
+        )
+
         return ReplicaConfigProto(
             deployment_def_name=self.deployment_def_name,
             deployment_def=self.serialized_deployment_def,
             init_args=self.serialized_init_args,
             init_kwargs=self.serialized_init_kwargs,
             ray_actor_options=json.dumps(self.ray_actor_options),
-            placement_group_bundles=json.dumps(self.placement_group_bundles)
-            if self.placement_group_bundles is not None
-            else "",
+            placement_group_bundles=placement_group_bundles,
             placement_group_strategy=self.placement_group_strategy,
-            max_replicas_per_node=self.max_replicas_per_node
-            if self.max_replicas_per_node is not None
-            else 0,
+            placement_group_bundle_label_selector=bundle_label_selector,
+            placement_group_fallback_strategy=fallback_strategy,
+            max_replicas_per_node=max_replicas_per_node,
         )
 
     def to_proto_bytes(self):
@@ -818,6 +921,8 @@ def to_dict(self):
             "ray_actor_options": self.ray_actor_options,
             "placement_group_bundles": self.placement_group_bundles,
             "placement_group_strategy": self.placement_group_strategy,
+            "placement_group_bundle_label_selector": self.placement_group_bundle_label_selector,
+            "placement_group_fallback_strategy": self.placement_group_fallback_strategy,
             "max_replicas_per_node": self.max_replicas_per_node,
         }
 

python/ray/serve/_private/default_impl.py

@@ -66,6 +66,7 @@ def _default_create_placement_group(
         _soft_target_node_id=request.target_node_id,
         name=request.name,
         lifetime="detached",
+        bundle_label_selector=request.bundle_label_selector,
     )