[Core] Update V2 Autoscaler to support scheduling using Node labels and LabelSelector API

python/ray/autoscaler/v2/scheduler.py

@@ -23,6 +23,7 @@
 from ray.core.generated.autoscaler_pb2 import (
     ClusterResourceConstraint,
     GangResourceRequest,
+    LabelOperator,
     ResourceRequest,
     ResourceRequestByCount,
 )
@@ -275,8 +276,10 @@ def new(
                 # Available resources for scheduling requests of different
                 # sources.
                 available_resources=dict(instance.ray_node.available_resources),
-                # Use ray node's dynamic labels.
-                labels=dict(instance.ray_node.dynamic_labels),
+                labels={
+                    **(instance.ray_node.labels or {}),
+                    **(instance.ray_node.dynamic_labels or {}),
+                },
                 status=SchedulingNodeStatus.SCHEDULABLE,
                 im_instance_id=instance.im_instance.instance_id,
                 im_instance_status=instance.im_instance.status,
@@ -437,7 +440,7 @@ def _compute_score(
         A "higher" score means that this node is more suitable for scheduling the
         current scheduled resource requests.
 
-        The score is a tuple of 4 values:
+        The score is a tuple of 5 values:
             1. Whether this node is a GPU node and the current resource request has
                 GPU requirements:
                     0: if this node is a GPU node and the current resource request
@@ -447,6 +450,11 @@ def _compute_score(
             2. The number of resource types being scheduled.
             3. The minimum utilization rate across all resource types.
             4. The average utilization rate across all resource types.
+            5. Whether this node has labels matching the current resource request's
+                label_selector requirements:
+                    0: if this node does not satisfy the label_selector requirements.
+                    1: if this node satisfies the label_selector requirements (or no
+                        requirements provided).
 
         NOTE:
             This function is adapted from  _resource_based_utilization_scorer from
@@ -499,6 +507,9 @@ def _compute_score(
             if is_gpu_node and not any_gpu_requests:
                 gpu_ok = False
 
+        # Check if node satisfies label requirements.
+        matches_labels = self._satisfies_label_constraints(sched_requests)
+
         # Prioritize avoiding gpu nodes for non-gpu workloads first,
         # then prioritize matching multiple resource types,
         # then prioritize using all resources,
@@ -510,8 +521,40 @@ def _compute_score(
             float(sum(util_by_resources)) / len(util_by_resources)
             if util_by_resources
             else 0,
+            matches_labels,
         )
 
+    def _satisfies_label_constraints(
+        self, sched_requests: List[ResourceRequest]
+    ) -> int:
+        """Returns 1 if this node satisfies at least one label selector, 0 otherwise."""
+        for req in sched_requests:
+            for selector in req.label_selectors:
+                # A label selector passes only if all constraints are satisfied
+                all_constraints_pass = True
+                for constraint in selector.label_constraints:
+                    key = constraint.label_key
+                    values = set(constraint.label_values)
+                    op = constraint.operator
+                    node_val = self.labels.get(key)
+
+                    if op == LabelOperator.LABEL_OPERATOR_IN:
+                        if node_val not in values:
+                            all_constraints_pass = False
+                            break
+                    elif op == LabelOperator.LABEL_OPERATOR_NOT_IN:
+                        if node_val in values:
+                            all_constraints_pass = False
+                            break
+                    else:
+                        all_constraints_pass = False
+                        break
+
+                if all_constraints_pass:
+                    return 1  # One label selector matched
+
+        return 0  # No label selectors are satisfied by SchedulingNode
+
     def _try_schedule_one(
         self, request: ResourceRequest, resource_request_source: ResourceRequestSource
     ) -> bool:
@@ -528,6 +571,36 @@ def _try_schedule_one(
             True if the resource request is scheduled on this node.
         """
 
+        # Enforce label selector constraints
+        if request.label_selectors:
+            selector_satisfied = False
+            for selector in request.label_selectors:
+                all_constraints_pass = True
+                for constraint in selector.label_constraints:
+                    key = constraint.label_key
+                    values = set(constraint.label_values)
+                    op = constraint.operator
+                    node_val = self.labels.get(key)
+
+                    if op == LabelOperator.LABEL_OPERATOR_IN:
+                        if node_val not in values:
+                            all_constraints_pass = False
+                            break
+                    elif op == LabelOperator.LABEL_OPERATOR_NOT_IN:
+                        if node_val in values:
+                            all_constraints_pass = False
+                            break
+                    else:
+                        all_constraints_pass = False
+                        break
+
+                if all_constraints_pass:
+                    selector_satisfied = True
+                    break  # At least one selector matched
+
+            if not selector_satisfied:
+                return False  # Node doesn't satisfy any label selector
+
         # Check if there's placement constraints that are not satisfied.
         for constraint in request.placement_constraints:
             if constraint.HasField("anti_affinity"):

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

@@ -25,6 +25,7 @@
 from ray.core.generated.autoscaler_pb2 import (
     ClusterResourceConstraint,
     GangResourceRequest,
+    LabelOperator,
     NodeState,
     NodeStatus,
     ResourceRequest,
@@ -1930,64 +1931,67 @@ def try_schedule(node_resources: Dict, requests: List[Dict]) -> Tuple:
         infeasible, score = node.try_schedule(requests, source)
         return ResourceRequestUtil.to_resource_maps(infeasible), score
 
-    assert try_schedule({"CPU": 1}, [{"CPU": 1}]) == ([], (True, 1, 1.0, 1.0))
+    assert try_schedule({"CPU": 1}, [{"CPU": 1}]) == ([], (True, 1, 1.0, 1.0, 0))
 
-    assert try_schedule({"GPU": 4}, [{"GPU": 2}]) == ([], (True, 1, 0.5, 0.5))
+    assert try_schedule({"GPU": 4}, [{"GPU": 2}]) == ([], (True, 1, 0.5, 0.5, 0))
     assert try_schedule({"GPU": 4}, [{"GPU": 1}, {"GPU": 1}]) == (
         [],
-        (True, 1, 0.5, 0.5),
+        (True, 1, 0.5, 0.5, 0),
+    )
+    assert try_schedule({"GPU": 2}, [{"GPU": 2}]) == ([], (True, 1, 2, 2, 0))
+    assert try_schedule({"GPU": 2}, [{"GPU": 1}, {"GPU": 1}]) == (
+        [],
+        (True, 1, 2, 2, 0),
     )
-    assert try_schedule({"GPU": 2}, [{"GPU": 2}]) == ([], (True, 1, 2, 2))
-    assert try_schedule({"GPU": 2}, [{"GPU": 1}, {"GPU": 1}]) == ([], (True, 1, 2, 2))
     assert try_schedule({"GPU": 1}, [{"GPU": 1, "CPU": 1}, {"GPU": 1}]) == (
         [{"GPU": 1, "CPU": 1}],
-        (True, 1, 1, 1),
+        (True, 1, 1, 1, 0),
     )
     assert try_schedule({"GPU": 1, "CPU": 1}, [{"GPU": 1, "CPU": 1}, {"GPU": 1}]) == (
         [{"GPU": 1}],
-        (True, 2, 1, 1),
+        (True, 2, 1, 1, 0),
     )
-    assert try_schedule({"GPU": 2, "TPU": 1}, [{"GPU": 2}]) == ([], (True, 1, 0, 1))
-    assert try_schedule({"CPU": 64}, [{"CPU": 64}]) == ([], (True, 1, 64, 64))
-    assert try_schedule({"CPU": 64}, [{"CPU": 32}]) == ([], (True, 1, 8, 8))
+    assert try_schedule({"GPU": 2, "TPU": 1}, [{"GPU": 2}]) == ([], (True, 1, 0, 1, 0))
+    assert try_schedule({"CPU": 64}, [{"CPU": 64}]) == ([], (True, 1, 64, 64, 0))
+    assert try_schedule({"CPU": 64}, [{"CPU": 32}]) == ([], (True, 1, 8, 8, 0))
     assert try_schedule({"CPU": 64}, [{"CPU": 16}, {"CPU": 16}]) == (
         [],
-        (True, 1, 8, 8),
+        (True, 1, 8, 8, 0),
     )
 
     # GPU Scores
     assert try_schedule({"GPU": 1, "CPU": 1}, [{"CPU": 1}]) == (
         [],
-        (False, 1, 0.0, 0.5),
+        (False, 1, 0.0, 0.5, 0),
     )
     assert try_schedule({"GPU": 1, "CPU": 1}, [{"CPU": 1, "GPU": 1}]) == (
         [],
-        (True, 2, 1.0, 1.0),
+        (True, 2, 1.0, 1.0, 0),
     )
     assert try_schedule({"GPU": 1, "CPU": 1}, [{"GPU": 1}]) == (
         [],
-        (True, 1, 0.0, 0.5),
+        (True, 1, 0.0, 0.5, 0),
     )
 
     # Zero resources
     assert try_schedule({"CPU": 0, "custom": 1}, [{"custom": 1}]) == (
         [],
-        (True, 1, 1, 1),
+        (True, 1, 1, 1, 0),
     )
     assert try_schedule({"CPU": 0, "custom": 1}, [{"CPU": 1}]) == (
         [{"CPU": 1}],
-        (True, 0, 0.0, 0.0),
+        (True, 0, 0.0, 0.0, 0),
     )
 
     # Implicit resources
     implicit_resource = ray._raylet.IMPLICIT_RESOURCE_PREFIX + "a"
     assert try_schedule({"CPU": 1}, [{implicit_resource: 1}]) == (
         [],
-        (True, 0, 0.0, 0.0),
+        (True, 0, 0.0, 0.0, 0),
     )
     assert try_schedule({"CPU": 1}, [{implicit_resource: 1}] * 2) == (
         [{implicit_resource: 1}],
-        (True, 0, 0.0, 0.0),
+        (True, 0, 0.0, 0.0, 0),
     )
 
 
@@ -2341,6 +2345,163 @@ def get_nodes_for(gang_resource_requests) -> Tuple[Dict, List[List[Dict]]]:
     ) == ({"p2.8xlarge": 1}, [])
 
 
+def test_schedule_node_with_matching_labels():
+    """
+    Test that a node with matching labels is considered schedulable and used to satisfy a request
+    with a label_selector.
+    """
+    scheduler = ResourceDemandScheduler(event_logger)
+    node_type_configs = {
+        "labelled_node": NodeTypeConfig(
+            name="labelled_node",
+            resources={"CPU": 1},
+            min_worker_nodes=0,
+            max_worker_nodes=10,
+            labels={"accelerator": "A100"},
+        ),
+    }
+
+    # The existing instance has matching dynamic label.
+    instance = make_autoscaler_instance(
+        im_instance=Instance(
+            instance_type="labelled_node",
+            status=Instance.RAY_RUNNING,
+            instance_id="1",
+            node_id=b"r-1",
+        ),
+        ray_node=NodeState(
+            node_id=b"r-1",
+            ray_node_type_name="labelled_node",
+            available_resources={"CPU": 1},
+            total_resources={"CPU": 1},
+            labels={"accelerator": "A100"},
+            status=NodeStatus.RUNNING,
+        ),
+        cloud_instance_id="c-1",
+    )
+
+    # No new nodes should be launched if the existing node satisfies the request.
+    resource_request = ResourceRequestUtil.make(
+        {"CPU": 1},
+        label_selectors=[[("accelerator", LabelOperator.LABEL_OPERATOR_IN, ["A100"])]],
+    )
+
+    request = sched_request(
+        node_type_configs=node_type_configs,
+        resource_requests=[resource_request],
+        instances=[instance],
+    )
+    reply = scheduler.schedule(request)
+    to_launch, _ = _launch_and_terminate(reply)
+    assert to_launch == {}
+
+
+def test_scale_up_node_to_satisfy_labels():
+    """
+    Test that a resource request with a label selector scales up a new node with
+    labels to satisfy the constraint.
+    """
+    scheduler = ResourceDemandScheduler(event_logger)
+
+    node_type_configs = {
+        "tpu_node": NodeTypeConfig(
+            name="tpu_node",
+            resources={"CPU": 1},
+            labels={"accelerator": "TPU"},
+            min_worker_nodes=0,
+            max_worker_nodes=10,
+        ),
+        "gpu_node": NodeTypeConfig(
+            name="gpu_node",
+            resources={"CPU": 1},
+            labels={"accelerator": "A100"},
+            min_worker_nodes=0,
+            max_worker_nodes=10,
+        ),
+    }
+
+    # Request: want a node with label "accelerator: A100"
+    resource_request = ResourceRequestUtil.make(
+        {"CPU": 1},
+        label_selectors=[[("accelerator", LabelOperator.LABEL_OPERATOR_IN, ["A100"])]],
+    )
+
+    request = sched_request(
+        node_type_configs=node_type_configs,
+        resource_requests=[resource_request],
+    )
+
+    reply = scheduler.schedule(request)
+    to_launch, _ = _launch_and_terminate(reply)
+
+    assert to_launch == {"gpu_node": 1}
+
+
+def test_pg_with_bundle_infeasible_label_selectors():
+    """
+    Test that placement group scheduling honors bundle_label_selectors.
+    """
+    scheduler = ResourceDemandScheduler(event_logger)
+    AFFINITY = ResourceRequestUtil.PlacementConstraintType.AFFINITY
+
+    node_type_configs = {
+        "gpu_node": NodeTypeConfig(
+            name="gpu_node",
+            resources={"CPU": 4, "GPU": 1},
+            min_worker_nodes=0,
+            max_worker_nodes=5,
+            labels={"accelerator": "A100"},
+        ),
+        "tpu_node": NodeTypeConfig(
+            name="tpu_node",
+            resources={"CPU": 4},
+            min_worker_nodes=0,
+            max_worker_nodes=5,
+            labels={"accelerator": "TPU"},
+        ),
+    }
+
+    # Create ResourceRequests for a placement group where each bundle has different label selectors
+    gpu_request = ResourceRequestUtil.make(
+        {"CPU": 2, "GPU": 1},
+        constraints=[(AFFINITY, "pg-1", "")],
+        label_selectors=[[("accelerator", LabelOperator.LABEL_OPERATOR_IN, ["A100"])]],
+    )
+    tpu_request = ResourceRequestUtil.make(
+        {"CPU": 2},
+        constraints=[(AFFINITY, "pg-1", "")],
+        label_selectors=[[("accelerator", LabelOperator.LABEL_OPERATOR_IN, ["TPU"])]],
+    )
+
+    request = sched_request(
+        node_type_configs=node_type_configs,
+        gang_resource_requests=[[gpu_request, tpu_request]],
+    )
+
+    reply = scheduler.schedule(request)
+    to_launch, _ = _launch_and_terminate(reply)
+
+    assert sorted(to_launch) == sorted({"gpu_node": 1, "tpu_node": 1})
+
+    # Both bundles require A100, but no node has enough resources -> infeasible
+    infeasbile_gpu_request = ResourceRequestUtil.make(
+        {"CPU": 3, "GPU": 1},
+        constraints=[(AFFINITY, "pg-2", "")],
+        label_selectors=[[("accelerator", LabelOperator.LABEL_OPERATOR_IN, ["A100"])]],
+    )
+
+    request = sched_request(
+        node_type_configs=node_type_configs,
+        gang_resource_requests=[[infeasbile_gpu_request, infeasbile_gpu_request]],
+    )
+
+    reply = scheduler.schedule(request)
+    to_launch, _ = _launch_and_terminate(reply)
+
+    assert to_launch == {}
+    assert len(reply.infeasible_gang_resource_requests) == 1
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))

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

@@ -357,7 +357,9 @@ def verify():
             state,
             [
                 ExpectedNodeState(
-                    head_node_id, NodeStatus.RUNNING, labels={f"_PG_{pg_id}": ""}
+                    head_node_id,
+                    NodeStatus.RUNNING,
+                    labels={f"_PG_{pg_id}": ""},
                 ),
             ],
         )