[Autoscaler] Add bundle_label_selector to request_resources sdk (#54843)

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
Signed-off-by: Ryan O'Leary <113500783+ryanaoleary@users.noreply.github.com>
Co-authored-by: Jiajun Yao <jeromeyjj@gmail.com>
Co-authored-by: Mengjin Yan <mengjinyan3@gmail.com>

Author: 3 people

python/ray/autoscaler/_private/commands.py

@@ -184,9 +184,16 @@ def debug_status(
 
 
 def request_resources(
-    num_cpus: Optional[int] = None, bundles: Optional[List[dict]] = None
+    num_cpus: Optional[int] = None,
+    bundles: Optional[List[dict]] = None,
+    bundle_label_selectors: Optional[List[dict]] = None,
 ) -> None:
-    """Remotely request some CPU or GPU resources from the autoscaler.
+    """Remotely request some CPU or GPU resources from the autoscaler. Optionally
+    specify label selectors for nodes with the requested resources.
+
+    If `bundle_label_selectors` is provided, `bundles` must also be provided.
+    Both must be lists of the same length, and `bundle_label_selectors` expects a list
+    of string dictionaries.
 
     This function is to be called e.g. on a node before submitting a bunch of
     ray.remote calls to ensure that resources rapidly become available.
@@ -198,14 +205,24 @@ def request_resources(
         bundles (List[ResourceDict]): Scale the cluster to ensure this set of
             resource shapes can fit. This request is persistent until another
             call to request_resources() is made.
+        bundle_label_selectors (List[Dict[str,str]]): Optional label selectors
+            that new nodes must satisfy. (e.g. [{"accelerator-type": "A100"}])
+            The elements in the bundle_label_selectors should be one-to-one mapping
+            to the elements in bundles.
     """
     if not ray.is_initialized():
         raise RuntimeError("Ray is not initialized yet")
     to_request = []
-    if num_cpus:
-        to_request += [{"CPU": 1}] * num_cpus
+    for _ in range(num_cpus or 0):
+        to_request.append({"resources": {"CPU": 1}, "label_selector": {}})
+    assert not bundle_label_selectors or (
+        bundles and len(bundles) == len(bundle_label_selectors)
+    ), "If bundle_label_selectors is provided, bundles must also be provided and have the same length."
     if bundles:
-        to_request += bundles
+        for i, bundle in enumerate(bundles):
+            selector = bundle_label_selectors[i] if bundle_label_selectors else {}
+            to_request.append({"resources": bundle, "label_selector": selector})
+
     _internal_kv_put(
         AUTOSCALER_RESOURCE_REQUEST_CHANNEL, json.dumps(to_request), overwrite=True
     )

python/ray/autoscaler/sdk/sdk.py

@@ -6,6 +6,7 @@
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, Iterator, List, Optional, Union
 
+from ray._private.label_utils import validate_label_selector
 from ray.autoscaler._private import commands
 from ray.autoscaler._private.cli_logger import cli_logger
 from ray.autoscaler._private.event_system import (
@@ -21,7 +22,7 @@ def create_or_update_cluster(
     *,
     no_restart: bool = False,
     restart_only: bool = False,
-    no_config_cache: bool = False
+    no_config_cache: bool = False,
 ) -> Dict[str, Any]:
     """Create or updates an autoscaling Ray cluster from a config json.
 
@@ -87,7 +88,7 @@ def run_on_cluster(
     stop: bool = False,
     no_config_cache: bool = False,
     port_forward: Optional[commands.Port_forward] = None,
-    with_output: bool = False
+    with_output: bool = False,
 ) -> Optional[str]:
     """Runs a command on the specified cluster.
 
@@ -133,7 +134,7 @@ def rsync(
     ip_address: Optional[str] = None,
     use_internal_ip: bool = False,
     no_config_cache: bool = False,
-    should_bootstrap: bool = True
+    should_bootstrap: bool = True,
 ):
     """Rsyncs files to or from the cluster.
 
@@ -206,7 +207,9 @@ def get_worker_node_ips(cluster_config: Union[dict, str]) -> List[str]:
 
 @DeveloperAPI
 def request_resources(
-    num_cpus: Optional[int] = None, bundles: Optional[List[dict]] = None
+    num_cpus: Optional[int] = None,
+    bundles: Optional[List[dict]] = None,
+    bundle_label_selectors: Optional[List[dict]] = None,
 ) -> None:
     """Command the autoscaler to scale to accommodate the specified requests.
 
@@ -230,6 +233,11 @@ def request_resources(
         bundles (List[ResourceDict]): Scale the cluster to ensure this set of
             resource shapes can fit. This request is persistent until another
             call to request_resources() is made to override.
+        bundle_label_selectors: A list of label selectors, applied per-bundle to the same
+            index in the `bundles` list. For bundles without a label requirement, the
+            corresponding item in the list is an empty dictionary. For each bundle.
+            Label selectors consist of zero or more key-value pairs where the key is
+            a label and the value is a operator (in, !in, etc.) and label value.
 
     Examples:
         >>> from ray.autoscaler.sdk import request_resources
@@ -241,6 +249,13 @@ def request_resources(
         >>> # Same as requesting num_cpus=3.
         >>> request_resources( # doctest: +SKIP
         ...     bundles=[{"CPU": 1}, {"CPU": 1}, {"CPU": 1}])
+        >>> # Requests 2 num_cpus=1 bundles, the first with
+        >>> # label_selector={"accelerator-type": "in(A100)"} and second with
+        >>> # label_selector={"market-type": "spot"}.
+        >>> request_resources( # doctest: +SKIP
+        ...     bundles=[{"CPU": 1}, {"CPU": 1}]),
+        ...     bundle_label_selectors=[{"accelerator-type": "in(A100)"},
+        ...                            {"market-type": "spot"}])
     """
     if num_cpus is not None and not isinstance(num_cpus, int):
         raise TypeError("num_cpus should be of type int.")
@@ -257,8 +272,34 @@ def request_resources(
                     raise TypeError("each bundle should be a Dict.")
         else:
             raise TypeError("bundles should be of type List")
-
-    return commands.request_resources(num_cpus, bundles)
+    if bundle_label_selectors is not None:
+        if bundles is None:
+            raise ValueError(
+                "`bundles` must be provided when `bundle_label_selectors` is specified."
+            )
+        if len(bundle_label_selectors) != len(bundles):
+            raise ValueError(
+                "`bundle_label_selector` must be a list with length equal to the number of bundles."
+            )
+        for label_selector in bundle_label_selectors:
+            if (
+                not isinstance(label_selector, dict)
+                or not all(isinstance(k, str) for k in label_selector.keys())
+                or not all(isinstance(v, str) for v in label_selector.values())
+            ):
+                raise ValueError(
+                    "Bundle label selector must be a list of string dictionary"
+                    " label selectors. For example: "
+                    '`[{ray.io/market_type": "spot"}, {"ray.io/accelerator-type": "A100"}]`.'
+                )
+            error_message = validate_label_selector(label_selector)
+            if error_message:
+                raise ValueError(
+                    f"Invalid label selector provided in bundle_label_selectors list."
+                    f" Detailed error: '{error_message}'"
+                )
+
+    return commands.request_resources(num_cpus, bundles, bundle_label_selectors)
 
 
 @DeveloperAPI

python/ray/autoscaler/v2/sdk.py

@@ -1,6 +1,6 @@
 import time
-from collections import defaultdict
-from typing import List
+from collections import Counter
+from typing import List, NamedTuple
 
 from ray._raylet import GcsClient
 from ray.autoscaler.v2.schema import ClusterStatus, Stats
@@ -14,8 +14,15 @@
 DEFAULT_RPC_TIMEOUT_S = 10
 
 
+class ResourceRequest(NamedTuple):
+    resources: dict
+    label_selector: dict
+
+
 def request_cluster_resources(
-    gcs_address: str, to_request: List[dict], timeout: int = DEFAULT_RPC_TIMEOUT_S
+    gcs_address: str,
+    to_request: List[dict],
+    timeout: int = DEFAULT_RPC_TIMEOUT_S,
 ):
     """Request resources from the autoscaler.
 
@@ -28,28 +35,46 @@ def request_cluster_resources(
 
     Args:
         gcs_address: The GCS address to query.
-        to_request: A list of resource bundles to request the cluster to have.
-            Each bundle is a dict of resource name to resource quantity, e.g:
-            [{"CPU": 1}, {"GPU": 1}].
+        to_request: A list of resource requests to request the cluster to have.
+            Each resource request is a tuple of resources and a label_selector
+            to apply per-bundle. e.g.: [{"resources": {"CPU": 1, "GPU": 1}, "label_selector": {"accelerator-type": "A100"}}]
         timeout: Timeout in seconds for the request to be timeout
 
     """
     assert len(gcs_address) > 0, "GCS address is not specified."
 
-    # Aggregate bundle by shape.
-    resource_requests_by_count = defaultdict(int)
-    for request in to_request:
-        bundle = frozenset(request.items())
-        resource_requests_by_count[bundle] += 1
+    # Convert bundle dicts to ResourceRequest tuples.
+    normalized: List[ResourceRequest] = []
+    for r in to_request:
+        assert isinstance(
+            r, dict
+        ), f"Internal Error: Expected a dict, but got {type(r)}"
+        resources = r.get("resources", {})
+        selector = r.get("label_selector", {})
+        normalized.append(ResourceRequest(resources, selector))
+
+    to_request = normalized
+
+    # Aggregate bundle by shape
+    def keyfunc(r):
+        return (
+            frozenset(r.resources.items()),
+            frozenset(r.label_selector.items()),
+        )
+
+    grouped_requests = Counter(keyfunc(r) for r in to_request)
+
+    bundles: List[dict] = []
+    label_selectors: List[dict] = []
+    counts: List[int] = []
 
-    bundles = []
-    counts = []
-    for bundle, count in resource_requests_by_count.items():
+    for (bundle, selector), count in grouped_requests.items():
         bundles.append(dict(bundle))
+        label_selectors.append(dict(selector))
         counts.append(count)
 
     GcsClient(gcs_address).request_cluster_resource_constraint(
-        bundles, counts, timeout_s=timeout
+        bundles, label_selectors, counts, timeout_s=timeout
     )
 
 

python/ray/autoscaler/v2/tests/test_autoscaler.py

@@ -121,7 +121,9 @@ def test_basic_scaling(make_autoscaler):
 
     # Resource requests
     print("=================== Test scaling up constraint 1/2====================")
-    request_cluster_resources(gcs_address, [{"CPU": 1}, {"GPU": 1}])
+    request_cluster_resources(
+        gcs_address, [{"resources": {"CPU": 1}}, {"resources": {"GPU": 1}}]
+    )
 
     def verify():
         autoscaler.update_autoscaling_state()

python/ray/autoscaler/v2/tests/test_sdk.py

@@ -17,7 +17,10 @@
     NodeInfo,
     ResourceRequestByCount,
 )
-from ray.autoscaler.v2.sdk import get_cluster_status, request_cluster_resources
+from ray.autoscaler.v2.sdk import (
+    get_cluster_status,
+    request_cluster_resources,
+)
 from ray.autoscaler.v2.tests.util import (
     get_available_resources,
     get_cluster_resource_state,
@@ -26,6 +29,7 @@
 )
 from ray.core.generated import autoscaler_pb2, autoscaler_pb2_grpc
 from ray.core.generated.autoscaler_pb2 import ClusterResourceState, NodeStatus
+from ray.core.generated.common_pb2 import LabelSelectorOperator
 from ray.util.state.api import list_nodes
 
 
@@ -247,7 +251,7 @@ def test_request_cluster_resources_basic(shutdown_only):
     gcs_address = ctx.address_info["gcs_address"]
 
     # Request one
-    request_cluster_resources(gcs_address, [{"CPU": 1}])
+    request_cluster_resources(gcs_address, [{"resources": {"CPU": 1}}])
 
     def verify():
         state = get_cluster_resource_state(stub)
@@ -257,7 +261,9 @@ def verify():
     wait_for_condition(verify)
 
     # Request another overrides the previous request
-    request_cluster_resources(gcs_address, [{"CPU": 2, "GPU": 1}, {"CPU": 1}])
+    request_cluster_resources(
+        gcs_address, [{"resources": {"CPU": 2, "GPU": 1}}, {"resources": {"CPU": 1}}]
+    )
 
     def verify():
         state = get_cluster_resource_state(stub)
@@ -267,7 +273,7 @@ def verify():
         return True
 
     # Request multiple is aggregated by shape.
-    request_cluster_resources(gcs_address, [{"CPU": 1}] * 100)
+    request_cluster_resources(gcs_address, [{"resources": {"CPU": 1}}] * 100)
 
     def verify():
         state = get_cluster_resource_state(stub)
@@ -277,6 +283,65 @@ def verify():
     wait_for_condition(verify)
 
 
+def test_request_cluster_resources_with_label_selectors(shutdown_only):
+    ctx = ray.init(num_cpus=1)
+    stub = _autoscaler_state_service_stub()
+    gcs_address = ctx.address_info["gcs_address"]
+
+    # Define two bundles, each with its own label_selector, to request.
+    bundles = [
+        {"CPU": 1},
+        {"GPU": 1, "CPU": 2},
+    ]
+    bundle_label_selectors = [
+        {"region": "us-west1"},
+        {"accelerator-type": "!in(A100)"},
+    ]
+    to_request = [
+        {"resources": b, "label_selector": s}
+        for b, s in zip(bundles, bundle_label_selectors)
+    ]
+
+    # Send the request for these resource bundles
+    request_cluster_resources(gcs_address, to_request)
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        # Validate shape and resource request count
+        assert_cluster_resource_constraints(state, bundles, [1, 1])
+
+        # Check that requests carry expected label selectors
+        requests = state.cluster_resource_constraints[0].resource_requests
+
+        # First resource request
+        label_selectors_0 = requests[0].request.label_selectors
+        selector_0 = label_selectors_0[0]
+        constraints_0 = {
+            c.label_key: list(c.label_values) for c in selector_0.label_constraints
+        }
+        assert constraints_0 == {"region": ["us-west1"]}
+        assert (
+            selector_0.label_constraints[0].operator
+            == LabelSelectorOperator.LABEL_OPERATOR_IN
+        )
+
+        # Second resource request
+        label_selectors_1 = requests[1].request.label_selectors
+        selector_1 = label_selectors_1[0]
+        constraints_1 = {
+            c.label_key: list(c.label_values) for c in selector_1.label_constraints
+        }
+        assert constraints_1 == {"accelerator-type": ["A100"]}
+        assert (
+            selector_1.label_constraints[0].operator
+            == LabelSelectorOperator.LABEL_OPERATOR_NOT_IN
+        )
+
+        return True
+
+    wait_for_condition(verify)
+
+
 def test_node_info_basic(shutdown_only, monkeypatch):
     with monkeypatch.context() as m:
         m.setenv("RAY_CLOUD_INSTANCE_ID", "instance-id")
@@ -646,7 +711,7 @@ def verify_task_demands():
 
     # Request resources through SDK
     request_cluster_resources(
-        gcs_address=cluster.address, to_request=[{"GPU": 1, "CPU": 2}]
+        gcs_address=cluster.address, to_request=[{"resources": {"GPU": 1, "CPU": 2}}]
     )
 
     def verify_cluster_constraint_demand():

python/ray/includes/common.pxd

@@ -561,7 +561,8 @@ cdef extern from "ray/gcs_rpc_client/accessor.h" nogil:
         CRayStatus RequestClusterResourceConstraint(
             int64_t timeout_ms,
             const c_vector[unordered_map[c_string, double]] &bundles,
-            const c_vector[int64_t] &count_array
+            const c_vector[unordered_map[c_string, c_string]] &label_selectors,
+            const c_vector[int64_t] &count_array,
         )
 
         CRayStatus GetClusterResourceState(