[core] [1/2] Topology aware scheduling public API

python/ray/_raylet.pyx

@@ -723,6 +723,17 @@ cdef int prepare_actor_concurrency_groups(
     return 1
 
 
+cdef CPlacementStrategy prepare_c_strategy(c_string strategy) except *:
+    if strategy == b"PACK":
+        return PLACEMENT_STRATEGY_PACK
+    elif strategy == b"SPREAD":
+        return PLACEMENT_STRATEGY_SPREAD
+    elif strategy == b"STRICT_PACK":
+        return PLACEMENT_STRATEGY_STRICT_PACK
+    else:
+        return PLACEMENT_STRATEGY_STRICT_SPREAD
+
+
 def raise_sys_exit_with_custom_error_message(
         ray_terminate_msg: str,
         exit_code: int = 0) -> None:
@@ -3693,23 +3704,22 @@ cdef class CoreWorker:
                             c_string strategy,
                             c_bool is_detached,
                             soft_target_node_id,
-                            c_vector[unordered_map[c_string, c_string]] bundle_label_selector):
+                            c_vector[unordered_map[c_string, c_string]] bundle_label_selector,
+                            topology_strategy):
         cdef:
             CPlacementGroupID c_placement_group_id
             CPlacementStrategy c_strategy
             CNodeID c_soft_target_node_id = CNodeID.Nil()
+            c_vector[unordered_map[c_string, CPlacementStrategy]] c_topology_strategy
+            unordered_map[c_string, CPlacementStrategy] c_level
 
-        if strategy == b"PACK":
-            c_strategy = PLACEMENT_STRATEGY_PACK
-        elif strategy == b"SPREAD":
-            c_strategy = PLACEMENT_STRATEGY_SPREAD
-        elif strategy == b"STRICT_PACK":
-            c_strategy = PLACEMENT_STRATEGY_STRICT_PACK
-        else:
-            if strategy == b"STRICT_SPREAD":
-                c_strategy = PLACEMENT_STRATEGY_STRICT_SPREAD
-            else:
-                raise TypeError(strategy)
+        c_strategy = prepare_c_strategy(strategy)
+
+        for level in topology_strategy:
+            c_level.clear()
+            for label, level_strategy in level.items():
+                c_level[label] = prepare_c_strategy(level_strategy)
+            c_topology_strategy.push_back(c_level)
 
         if soft_target_node_id is not None:
             c_soft_target_node_id = CNodeID.FromHex(soft_target_node_id)
@@ -3724,7 +3734,8 @@ cdef class CoreWorker:
                                 bundles,
                                 is_detached,
                                 c_soft_target_node_id,
-                                bundle_label_selector),
+                                bundle_label_selector,
+                                c_topology_strategy),
                             &c_placement_group_id))
 
         return PlacementGroupID(c_placement_group_id.Binary())

python/ray/includes/common.pxd

@@ -432,6 +432,7 @@ cdef extern from "ray/core_worker/common.h" nogil:
             c_bool is_detached,
             CNodeID soft_target_node_id,
             const c_vector[unordered_map[c_string, c_string]] &bundle_label_selector,
+            const c_vector[unordered_map[c_string, CPlacementStrategy]] &topology_strategy,
         )
 
     cdef cppclass CObjectLocation "ray::core::ObjectLocation":

python/ray/tests/BUILD.bazel

@@ -481,6 +481,7 @@ py_test_module_list(
         "test_minimal_install.py",
         "test_path_utils.py",
         "test_runtime_env_ray_minimal.py",
+        "test_topology_strategy.py",
     ],
     tags = [
         "exclusive",

python/ray/tests/test_state_api.py

@@ -138,17 +138,17 @@ def generate_actor_data(id, state=ActorTableData.ActorState.ALIVE, class_name="c
 def generate_pg_data(
     id,
     name="abc",
-    label_domain_key="",
-    label_domain_assignments=None,
+    topology_strategy=None,
+    topology_assignments=None,
 ):
     return PlacementGroupTableData(
         placement_group_id=id,
         state=PlacementGroupTableData.PlacementGroupState.CREATED,
         name=name,
         creator_job_dead=True,
         creator_actor_dead=False,
-        label_domain_key=label_domain_key,
-        label_domain_assignments=label_domain_assignments or {},
+        topology_strategy=topology_strategy or [],
+        topology_assignments=topology_assignments or [],
     )
 
 

python/ray/tests/test_state_api_manager.py

@@ -18,15 +18,18 @@
 )
 from ray._raylet import NodeID
 from ray.core.generated.common_pb2 import (
+    PlacementStrategy,
     TaskInfoEntry,
     TaskStatus,
     TaskType,
+    TopologyStrategyLevel,
 )
 from ray.core.generated.gcs_pb2 import (
     ActorTableData,
     GcsNodeInfo,
     TaskEvents,
     TaskStateUpdate,
+    TopologyAssignmentLevel,
 )
 from ray.core.generated.gcs_service_pb2 import (
     GetAllActorInfoReply,
@@ -293,20 +296,26 @@ async def test_api_manager_list_pgs(state_api_manager):
 
 
 @pytest.mark.asyncio
-async def test_api_manager_list_pgs_label_domain(state_api_manager):
+async def test_api_manager_list_pgs_topology(state_api_manager):
     data_source_client = state_api_manager.data_source_client
 
-    pg_with_domain = generate_pg_data(
+    pg_with_topology = generate_pg_data(
         b"1",
-        name="gpu-domain-pg",
-        label_domain_key="ray.io/gpu-domain",
-        label_domain_assignments={"ray.io/gpu-domain": "rack-1"},
+        name="topology-pg",
+        topology_strategy=[
+            TopologyStrategyLevel(
+                entries={"ray.io/gpu-domain": PlacementStrategy.STRICT_PACK},
+            ),
+        ],
+        topology_assignments=[
+            TopologyAssignmentLevel(assignments={"ray.io/gpu-domain": "rack-1"}),
+        ],
     )
-    pg_without_domain = generate_pg_data(b"2")
+    pg_without_topology = generate_pg_data(b"2")
 
     data_source_client.get_all_placement_group_info.return_value = (
         GetAllPlacementGroupReply(
-            placement_group_table_data=[pg_with_domain, pg_without_domain],
+            placement_group_table_data=[pg_with_topology, pg_without_topology],
             total=2,
         )
     )
@@ -316,23 +325,27 @@ async def test_api_manager_list_pgs_label_domain(state_api_manager):
     )
     assert len(result.result) == 2
 
-    pg_domain = next(r for r in result.result if r["name"] == "gpu-domain-pg")
-    verify_schema(PlacementGroupState, pg_domain, detail=True)
-    assert pg_domain["label_domain_key"] == "ray.io/gpu-domain"
-    assert pg_domain["label_domain_assignments"] == {"ray.io/gpu-domain": "rack-1"}
+    pg_topology = next(r for r in result.result if r["name"] == "topology-pg")
+    verify_schema(PlacementGroupState, pg_topology, detail=True)
+    assert pg_topology["topology_strategy"] == [
+        {"entries": {"ray.io/gpu-domain": "STRICT_PACK"}}
+    ]
+    assert pg_topology["topology_assignments"] == [
+        {"assignments": {"ray.io/gpu-domain": "rack-1"}}
+    ]
 
-    pg_no_domain = next(r for r in result.result if r["name"] == "abc")
-    verify_schema(PlacementGroupState, pg_no_domain, detail=True)
-    assert pg_no_domain["label_domain_key"] == ""
-    assert pg_no_domain["label_domain_assignments"] == {}
+    pg_no_topology = next(r for r in result.result if r["name"] == "abc")
+    verify_schema(PlacementGroupState, pg_no_topology, detail=True)
+    assert pg_no_topology["topology_strategy"] == []
+    assert pg_no_topology["topology_assignments"] == []
 
     # Verify the fields are excluded from non-detail responses.
     result = await state_api_manager.list_placement_groups(
         option=create_api_options(detail=False)
     )
     for pg in result.result:
-        assert "label_domain_key" not in pg
-        assert "label_domain_assignments" not in pg
+        assert "topology_strategy" not in pg
+        assert "topology_assignments" not in pg
 
 
 @pytest.mark.asyncio

python/ray/tests/test_topology_strategy.py

@@ -0,0 +1,119 @@
+import os
+import sys
+
+import pytest
+
+import ray
+from ray.util.placement_group import placement_group, placement_group_table
+
+NODE_ID_LABEL = "ray.io/node-id"
+RACK_LABEL = "rack_id"
+ONE = "1"
+TWO = "2"
+rack1_labels = {RACK_LABEL: ONE}
+rack2_labels = {RACK_LABEL: TWO}
+
+
+def test_topology_strategy_feasible_after_rack_kill(ray_start_cluster):
+    """Verify topology-aware rescheduling after total rack failure.
+
+    Creates a PG on rack 1 (the only available rack at the time). After
+    removing one rack 1 node, the PG enters RESCHEDULING but stays pinned to
+    rack 1, so it remains infeasible even though rack 2 has capacity. Once
+    all rack 1 nodes are removed (total failure), the topology assignment
+    is cleared and the PG reschedules onto rack 2.
+    """
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)
+    ray.init(address=cluster.address)
+
+    rack1_nodes = [cluster.add_node(num_cpus=1, labels=rack1_labels) for _ in range(4)]
+
+    def assert_pg_nodes_label_value(cluster_nodes, pg, label, value):
+        node_id_to_labels = {node["NodeID"]: node["Labels"] for node in cluster_nodes}
+        for node_id in placement_group_table(pg)["bundles_to_node_id"].values():
+            assert node_id_to_labels[node_id].get(label) == value
+
+    bundles = [{"CPU": 1}] * 4
+
+    pg = placement_group(
+        bundles=bundles,
+        topology_strategy=[{RACK_LABEL: "STRICT_PACK"}],
+    )
+    ray.get(pg.ready(), timeout=30)
+    assert placement_group_table(pg)["state"] == "CREATED"
+    assert_pg_nodes_label_value(ray.nodes(), pg, RACK_LABEL, ONE)
+
+    # Bring rack 2 online; PG should still be pinned to rack 1.
+    for _ in range(4):
+        cluster.add_node(num_cpus=1, labels=rack2_labels)
+
+    # Drop one rack 1 node -> partial failure -> RESCHEDULING + infeasible.
+    cluster.remove_node(rack1_nodes[0])
+    with pytest.raises(ray.exceptions.GetTimeoutError):
+        ray.get(pg.ready(), timeout=5)
+    assert placement_group_table(pg)["state"] == "RESCHEDULING"
+
+    # Total rack 1 failure -> clears assignment -> reschedules onto rack 2.
+    for node in rack1_nodes[1:]:
+        cluster.remove_node(node)
+
+    ray.get(pg.ready(), timeout=30)
+    assert placement_group_table(pg)["state"] == "CREATED"
+    assert_pg_nodes_label_value(ray.nodes(), pg, RACK_LABEL, TWO)
+
+
+def test_topology_strategy_strict_pack(ray_start_cluster):
+    """Testing STRICT_PACK on the node level and STRICT_PACK on the rack level"""
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)
+    ray.init(address=cluster.address)
+
+    cluster.add_node(num_cpus=4, labels=rack1_labels)
+
+    def assert_pg_nodes_label_value(cluster_nodes, pg, label, value):
+        node_id_to_labels = {node["NodeID"]: node["Labels"] for node in cluster_nodes}
+        for node_id in placement_group_table(pg)["bundles_to_node_id"].values():
+            assert node_id_to_labels[node_id].get(label) == value
+
+    bundles = [{"CPU": 1}] * 4
+
+    pg = placement_group(
+        bundles=bundles,
+        topology_strategy=[{NODE_ID_LABEL: "STRICT_PACK", RACK_LABEL: "STRICT_PACK"}],
+    )
+    ray.get(pg.ready(), timeout=30)
+    assert placement_group_table(pg)["state"] == "CREATED"
+    assert_pg_nodes_label_value(ray.nodes(), pg, RACK_LABEL, ONE)
+
+
+def test_topology_strategy_strict_spread(ray_start_cluster):
+    """Testing STRICT_SPREAD on the node level and STRICT_PACK on the rack level"""
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)
+    ray.init(address=cluster.address)
+
+    for _ in range(4):
+        cluster.add_node(num_cpus=1, labels=rack1_labels)
+
+    def assert_pg_nodes_label_value(cluster_nodes, pg, label, value):
+        node_id_to_labels = {node["NodeID"]: node["Labels"] for node in cluster_nodes}
+        for node_id in placement_group_table(pg)["bundles_to_node_id"].values():
+            assert node_id_to_labels[node_id].get(label) == value
+
+    bundles = [{"CPU": 1}] * 4
+
+    pg = placement_group(
+        bundles=bundles,
+        topology_strategy=[{NODE_ID_LABEL: "STRICT_SPREAD", RACK_LABEL: "STRICT_PACK"}],
+    )
+    ray.get(pg.ready(), timeout=30)
+    assert placement_group_table(pg)["state"] == "CREATED"
+    assert_pg_nodes_label_value(ray.nodes(), pg, RACK_LABEL, ONE)
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))