[Core] Add fallback strategy scheduling logic (#56369)

This PR also updates the cluster resource scheduler logic to account for
the list of `LabelSelector`s specified by the `fallback_strategy`,
falling back to each fallback strategy `LabelSelector` in-order until
one is satisfied when selecting the best node. We're able to support
fallback selectors by considering them in the cluster resource scheduler
in-order using the existing label selector logic in `IsFeasible` and
`IsAvailable`, returning the first valid node returned by
`GetBestSchedulableNode`.

#51564

---------

Signed-off-by: Ryan O'Leary <[email protected]>
Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Mengjin Yan <[email protected]>

Author: ryanaoleary

python/ray/_raylet.pyx

@@ -87,9 +87,11 @@ from ray.includes.common cimport (
     CRayStatus,
     CActorTableData,
     CErrorTableData,
+    CFallbackOption,
     CGcsClientOptions,
     CGcsNodeInfo,
     CJobTableData,
+    CLabelSelector,
     CLogBatch,
     CTaskArg,
     CTaskArgByReference,
@@ -835,12 +837,13 @@ cdef int prepare_labels(
 
 cdef int prepare_label_selector(
         dict label_selector_dict,
-        unordered_map[c_string, c_string] *label_selector) except -1:
+        CLabelSelector *c_label_selector) except -1:
+
+    c_label_selector[0] = CLabelSelector()
 
     if label_selector_dict is None:
         return 0
 
-    label_selector[0].reserve(len(label_selector_dict))
     for key, value in label_selector_dict.items():
         if not isinstance(key, str):
             raise ValueError(f"Label selector key type must be string, but got {type(key)}")
@@ -853,10 +856,39 @@ cdef int prepare_label_selector(
             inner = value[value.index("(")+1:-1].strip()
             if not inner:
                 raise ValueError(f"No values provided for Label Selector '{value[:value.index('(')]}' operator on key '{key}'.")
-        label_selector[0][key.encode("utf-8")] = value.encode("utf-8")
+        # Add key-value constraint to the LabelSelector object.
+        c_label_selector[0].AddConstraint(key.encode("utf-8"), value.encode("utf-8"))
 
     return 0
 
+cdef int prepare_fallback_strategy(
+        list fallback_strategy,
+        c_vector[CFallbackOption] *fallback_strategy_vector) except -1:
+
+    cdef dict label_selector_dict
+    cdef CLabelSelector c_label_selector
+
+    if fallback_strategy is None:
+        return 0
+
+    for strategy_dict in fallback_strategy:
+        if not isinstance(strategy_dict, dict):
+            raise ValueError(
+                "Fallback strategy must be a list of dicts, "
+                f"but got list containing {type(strategy_dict)}")
+
+        label_selector_dict = strategy_dict.get("label_selector")
+
+        if label_selector_dict is not None and not isinstance(label_selector_dict, dict):
+            raise ValueError("Invalid fallback strategy element: invalid 'label_selector'.")
+
+        prepare_label_selector(label_selector_dict, &c_label_selector)
+
+        fallback_strategy_vector.push_back(
+             CFallbackOption(c_label_selector)
+        )
+
+    return 0
 
 cdef int prepare_resources(
         dict resource_dict,
@@ -3644,11 +3676,13 @@ cdef class CoreWorker:
                     int64_t generator_backpressure_num_objects,
                     c_bool enable_task_events,
                     labels,
-                    label_selector):
+                    label_selector,
+                    fallback_strategy):
         cdef:
             unordered_map[c_string, double] c_resources
             unordered_map[c_string, c_string] c_labels
-            unordered_map[c_string, c_string] c_label_selector
+            CLabelSelector c_label_selector
+            c_vector[CFallbackOption] c_fallback_strategy
             CRayFunction ray_function
             CTaskOptions task_options
             c_vector[unique_ptr[CTaskArg]] args_vector
@@ -3675,6 +3709,7 @@ cdef class CoreWorker:
             prepare_resources(resources, &c_resources)
             prepare_labels(labels, &c_labels)
             prepare_label_selector(label_selector, &c_label_selector)
+            prepare_fallback_strategy(fallback_strategy, &c_fallback_strategy)
             ray_function = CRayFunction(
                 language.lang, function_descriptor.descriptor)
             prepare_args_and_increment_put_refs(
@@ -3691,7 +3726,8 @@ cdef class CoreWorker:
                 c_label_selector,
                 # `tensor_transport` is currently only supported in Ray Actor tasks.
                 # For Ray tasks, we always use `OBJECT_STORE`.
-                TENSOR_TRANSPORT_OBJECT_STORE)
+                TENSOR_TRANSPORT_OBJECT_STORE,
+                c_fallback_strategy)
 
             current_c_task_id = current_task.native()
 
@@ -3742,6 +3778,7 @@ cdef class CoreWorker:
                      label_selector,
                      c_bool allow_out_of_order_execution,
                      c_bool enable_tensor_transport,
+                     fallback_strategy,
                      ):
         cdef:
             CRayFunction ray_function
@@ -3755,7 +3792,8 @@ cdef class CoreWorker:
             c_vector[CObjectID] incremented_put_arg_ids
             optional[c_bool] is_detached_optional = nullopt
             unordered_map[c_string, c_string] c_labels
-            unordered_map[c_string, c_string] c_label_selector
+            CLabelSelector c_label_selector
+            c_vector[CFallbackOption] c_fallback_strategy
             c_string call_site
 
         self.python_scheduling_strategy_to_c(
@@ -3770,6 +3808,7 @@ cdef class CoreWorker:
             prepare_resources(placement_resources, &c_placement_resources)
             prepare_labels(labels, &c_labels)
             prepare_label_selector(label_selector, &c_label_selector)
+            prepare_fallback_strategy(fallback_strategy, &c_fallback_strategy)
             ray_function = CRayFunction(
                 language.lang, function_descriptor.descriptor)
             prepare_args_and_increment_put_refs(
@@ -3799,7 +3838,8 @@ cdef class CoreWorker:
                         enable_tensor_transport,
                         enable_task_events,
                         c_labels,
-                        c_label_selector),
+                        c_label_selector,
+                        c_fallback_strategy),
                     extension_data,
                     call_site,
                     &c_actor_id,
@@ -3914,9 +3954,10 @@ cdef class CoreWorker:
             c_string serialized_retry_exception_allowlist
             c_string serialized_runtime_env = b"{}"
             unordered_map[c_string, c_string] c_labels
-            unordered_map[c_string, c_string] c_label_selector
+            CLabelSelector c_label_selector
             c_string call_site
             CTensorTransport c_tensor_transport_val
+            c_vector[CFallbackOption] c_fallback_strategy
 
         serialized_retry_exception_allowlist = serialize_retry_exception_allowlist(
             retry_exception_allowlist,
@@ -3953,7 +3994,8 @@ cdef class CoreWorker:
                         enable_task_events,
                         c_labels,
                         c_label_selector,
-                        c_tensor_transport_val),
+                        c_tensor_transport_val,
+                        c_fallback_strategy),
                     max_retries,
                     retry_exceptions,
                     serialized_retry_exception_allowlist,

python/ray/actor.py

@@ -1807,6 +1807,7 @@ def _remote(self, args=None, kwargs=None, **actor_options) -> ActorProxy[T]:
             enable_task_events=enable_task_events,
             labels=actor_options.get("_labels"),
             label_selector=actor_options.get("label_selector"),
+            fallback_strategy=actor_options.get("fallback_strategy"),
             allow_out_of_order_execution=allow_out_of_order_execution,
             enable_tensor_transport=meta.enable_tensor_transport,
         )

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

@@ -654,7 +654,7 @@ def all_nodes_launched():
         nodes = {node["NodeID"]: node["Labels"] for node in ray.nodes()}
         task_selectors = {
             "task_0": {"accelerator-type": "A100"},
-            "task_1": {"region": "in(us-east1,me-central1)"},
+            "task_1": {"region": "in(me-central1,us-east1)"},
             "task_2": {"accelerator-type": "!in(A100,TPU)"},
             "task_3": {"market-type": "!spot"},
         }
@@ -799,7 +799,7 @@ def all_actors_scheduled():
         nodes = {node["NodeID"]: node["Labels"] for node in ray.nodes()}
         actor_selectors = {
             "actor_0": {"accelerator-type": "A100"},
-            "actor_1": {"region": "in(us-east1,me-central1)"},
+            "actor_1": {"region": "in(me-central1,us-east1)"},
             "actor_2": {"accelerator-type": "!in(A100,TPU)"},
             "actor_3": {"market-type": "!spot"},
         }
@@ -900,9 +900,22 @@ def test_pg_scheduled_on_node_with_bundle_label_selector(autoscaler_v2):
         ray.get(pg.ready())
 
         # Validate the number and types of the auto-scaled nodes are as expected.
-        status = get_cluster_status(gcs_address)
-        assert len(status.active_nodes) == expected_nodes
+        # Add a wait here to avoid flaky test behavior.
+        def check_nodes_active():
+            status = get_cluster_status(gcs_address)
+            return len(status.active_nodes) == expected_nodes
 
+        try:
+            wait_for_condition(check_nodes_active, timeout=30, retry_interval_ms=500)
+        except Exception as e:
+            latest_status = get_cluster_status(gcs_address)
+            raise AssertionError(
+                f"Timed out waiting for {expected_nodes} active nodes. "
+                f"Got: {len(latest_status.active_nodes)}. "
+                f"Full status: {latest_status}"
+            ) from e
+
+        status = get_cluster_status(gcs_address)
         actual_node_types = {node.ray_node_type_name for node in status.active_nodes}
         expected_node_types = {"a100_node", "tpu_node"}
         assert actual_node_types == expected_node_types

python/ray/includes/common.pxd

@@ -239,6 +239,17 @@ cdef extern from "src/ray/protobuf/common.pb.h" nogil:
         CLineageReconstructionTask()
         const c_string &SerializeAsString() const
 
+cdef extern from "ray/common/scheduling/label_selector.h" namespace "ray":
+    cdef cppclass CLabelSelector "ray::LabelSelector":
+        CLabelSelector() nogil except +
+        void AddConstraint(const c_string& key, const c_string& value) nogil except +
+
+cdef extern from "ray/common/scheduling/fallback_strategy.h" namespace "ray":
+    cdef cppclass CFallbackOption "ray::FallbackOption":
+        CLabelSelector label_selector
+
+        CFallbackOption() nogil except +
+        CFallbackOption(CLabelSelector) nogil except +
 
 # This is a workaround for C++ enum class since Cython has no corresponding
 # representation.
@@ -346,8 +357,9 @@ cdef extern from "ray/core_worker/common.h" nogil:
                      c_string serialized_runtime_env,
                      c_bool enable_task_events,
                      const unordered_map[c_string, c_string] &labels,
-                     const unordered_map[c_string, c_string] &label_selector,
-                     CTensorTransport tensor_transport)
+                     CLabelSelector label_selector,
+                     CTensorTransport tensor_transport,
+                     c_vector[CFallbackOption] fallback_strategy)
 
     cdef cppclass CActorCreationOptions "ray::core::ActorCreationOptions":
         CActorCreationOptions()
@@ -368,7 +380,8 @@ cdef extern from "ray/core_worker/common.h" nogil:
             c_bool enable_tensor_transport,
             c_bool enable_task_events,
             const unordered_map[c_string, c_string] &labels,
-            const unordered_map[c_string, c_string] &label_selector)
+            CLabelSelector label_selector,
+            c_vector[CFallbackOption] fallback_strategy)
 
     cdef cppclass CPlacementGroupCreationOptions \
             "ray::core::PlacementGroupCreationOptions":

python/ray/remote_function.py

@@ -58,6 +58,7 @@ class RemoteFunction:
         _memory: The heap memory request in bytes for this task/actor,
             rounded down to the nearest integer.
         _label_selector: The label requirements on a node for scheduling of the task or actor.
+        _fallback_strategy: Soft constraints of a list of decorator options to fall back on when scheduling on a node.
         _resources: The default custom resource requirements for invocations of
             this remote function.
         _num_returns: The default number of return values for invocations
@@ -471,6 +472,7 @@ def _remote(
         enable_task_events = task_options.get("enable_task_events")
         labels = task_options.get("_labels")
         label_selector = task_options.get("label_selector")
+        fallback_strategy = task_options.get("fallback_strategy")
 
         def invocation(args, kwargs):
             if self._is_cross_language:
@@ -503,6 +505,7 @@ def invocation(args, kwargs):
                 enable_task_events,
                 labels,
                 label_selector,
+                fallback_strategy,
             )
             # Reset worker's debug context from the last "remote" command
             # (which applies only to this .remote call).

python/ray/tests/test_label_scheduling.py

@@ -96,5 +96,92 @@ def test_label_selector_multiple(cluster_with_labeled_nodes):
     assert ray.get(actor.get_node_id.remote(), timeout=3) == node_3
 
 
+def test_fallback_strategy(cluster_with_labeled_nodes):
+    # Create a RayCluster with labelled nodes.
+    gpu_node, _, _ = cluster_with_labeled_nodes
+
+    # Define an unsatisfiable label selector.
+    infeasible_label_selector = {"ray.io/accelerator-type": "does-not-exist"}
+
+    # Create a fallback strategy with multiple accelerator options.
+    accelerator_fallbacks = [
+        {"label_selector": {"ray.io/accelerator-type": "A100"}},
+        {"label_selector": {"ray.io/accelerator-type": "TPU"}},
+    ]
+
+    # Attempt to schedule the actor. The scheduler should fail to find a node with the
+    # primary `label_selector` and fall back to the first available option, 'A100'.
+    label_selector_actor = MyActor.options(
+        label_selector=infeasible_label_selector,
+        fallback_strategy=accelerator_fallbacks,
+    ).remote()
+
+    # Assert that the actor was scheduled on the expected node.
+    assert ray.get(label_selector_actor.get_node_id.remote(), timeout=5) == gpu_node
+
+
+def test_empty_selector_fallback_strategy(cluster_with_labeled_nodes):
+    node_1, node_2, node_3 = cluster_with_labeled_nodes
+
+    # Define an unsatisfiable label selector.
+    infeasible_label_selector = {"ray.io/accelerator-type": "does-not-exist"}
+
+    # Create a fallback strategy with multiple label selector fallbacks. The
+    # first fallback option is unsatisfiable, so it falls back to the empty label
+    # selector option. This fallback should match any node.
+    accelerator_fallbacks = [
+        {"label_selector": {"ray.io/accelerator-type": "also-does-not-exist"}},
+        {"label_selector": {}},
+    ]
+
+    label_selector_actor = MyActor.options(
+        label_selector=infeasible_label_selector,
+        fallback_strategy=accelerator_fallbacks,
+    ).remote()
+
+    # Assert that the actor was scheduled on the expected node.
+    assert ray.get(label_selector_actor.get_node_id.remote(), timeout=5) in {
+        node_1,
+        node_2,
+        node_3,
+    }
+
+
+def test_infeasible_fallback_strategy(cluster_with_labeled_nodes):
+    # Define an unsatisfiable label selector and fallback strategy.
+    label_selector = {"ray.io/accelerator-type": "does-not-exist"}
+    fallback_strategy = [
+        {"label_selector": {"ray.io/accelerator-type": "does-not-exist-either"}},
+        {"label_selector": {"ray.io/accelerator-type": "also-nonexistant"}},
+    ]
+
+    # Attempt to schedule the actor, but it should timeout since none of
+    # the nodes match any label selector.
+    label_selector_actor = MyActor.options(
+        label_selector=label_selector, fallback_strategy=fallback_strategy
+    ).remote()
+    with pytest.raises(TimeoutError):
+        ray.get(label_selector_actor.get_node_id.remote(), timeout=3)
+
+
+def test_fallback_with_feasible_primary_selector(cluster_with_labeled_nodes):
+    gpu_node, _, _ = cluster_with_labeled_nodes
+
+    feasible_label_selector = {"ray.io/accelerator-type": "A100"}
+    feasible_fallback_strategy = [
+        {"label_selector": {"ray.io/accelerator-type": "B200"}},
+    ]
+
+    # Attempt to schedule the actor. The scheduler should use the
+    # primary selector and ignore the fallback.
+    label_selector_actor = MyActor.options(
+        label_selector=feasible_label_selector,
+        fallback_strategy=feasible_fallback_strategy,
+    ).remote()
+
+    # Assert that the actor was scheduled on the expected node and not the fallback.
+    assert ray.get(label_selector_actor.get_node_id.remote(), timeout=5) == gpu_node
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

src/ray/common/BUILD.bazel

@@ -230,6 +230,7 @@ ray_cc_library(
         ":ray_config",
         ":ray_object",
         ":runtime_env",
+        "//src/ray/common/scheduling:fallback_strategy",
         "//src/ray/common/scheduling:label_selector",
         "//src/ray/common/scheduling:resource_set",
         "//src/ray/common/scheduling:scheduling_class_util",