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

python/ray/_raylet.pyx

@@ -566,6 +566,19 @@ cdef class Language:
     JAVA = Language.from_native(LANGUAGE_JAVA)
 
 
+cdef CPlacementStrategy prepare_c_strategy(c_string strategy) except *:
+    # Called by CoreWorker.create_placement_group(..., c_string strategy, ...).
+    # The Python placement_group wrapper validates `strategy` to be one of the
+    # strategies below beforehand.
+    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:
@@ -3767,23 +3780,18 @@ 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,
+                            dict topology_strategy):
         cdef:
             CPlacementGroupID c_placement_group_id
             CPlacementStrategy c_strategy
             CNodeID c_soft_target_node_id = CNodeID.Nil()
+            unordered_map[c_string, CPlacementStrategy] c_topology_strategy
 
-        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 label, level_strategy in topology_strategy.items():
+            c_topology_strategy[label] = prepare_c_strategy(level_strategy)
 
         if soft_target_node_id is not None:
             c_soft_target_node_id = CNodeID.FromHex(soft_target_node_id)
@@ -3798,7 +3806,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/dashboard/client/src/components/PlacementGroupTable.component.test.tsx

@@ -107,8 +107,8 @@ describe("PlacementGroupTable", () => {
     const labelSelectorHeaders = screen.getAllByText("Label Selector");
     expect(labelSelectorHeaders.length).toBeGreaterThan(0);
 
-    const labelDomainHeaders = screen.getAllByText("Label Domain");
-    expect(labelDomainHeaders.length).toBeGreaterThan(0);
+    const topologyHeaders = screen.getAllByText("Topology");
+    expect(topologyHeaders.length).toBeGreaterThan(0);
 
     const schedulingDetailHeaders = screen.getAllByText("Scheduling Detail");
     expect(schedulingDetailHeaders.length).toBeGreaterThan(0);
@@ -287,7 +287,7 @@ describe("PlacementGroupTable", () => {
     // Check that null stats are handled gracefully
     expect(screen.getByText("pg-123456789")).toBeInTheDocument();
     const dashes = screen.getAllByText("-");
-    expect(dashes.length).toBeGreaterThan(0); // Null scheduling detail and empty label domain
+    expect(dashes.length).toBeGreaterThan(0); // Null scheduling detail and empty topology
   });
 
   it("renders placement groups with empty name", () => {
@@ -308,7 +308,7 @@ describe("PlacementGroupTable", () => {
     // Check that empty names are handled gracefully
     expect(screen.getByText("pg-123456789")).toBeInTheDocument();
     const dashes = screen.getAllByText("-");
-    expect(dashes.length).toBeGreaterThan(0); // Empty name and empty label domain
+    expect(dashes.length).toBeGreaterThan(0); // Empty name and empty topology
   });
 
   it("renders state counter for placement groups", () => {
@@ -386,8 +386,8 @@ describe("PlacementGroupTable", () => {
     const labelSelectorHeaders = screen.getAllByText("Label Selector");
     expect(labelSelectorHeaders.length).toBeGreaterThan(0);
 
-    const labelDomainHeaders = screen.getAllByText("Label Domain");
-    expect(labelDomainHeaders.length).toBeGreaterThan(0);
+    const topologyHeaders = screen.getAllByText("Topology");
+    expect(topologyHeaders.length).toBeGreaterThan(0);
 
     const schedulingDetailHeaders = screen.getAllByText("Scheduling Detail");
     expect(schedulingDetailHeaders.length).toBeGreaterThan(0);

python/ray/dashboard/client/src/components/PlacementGroupTable.tsx

@@ -89,7 +89,7 @@ const PlacementGroupTable = ({
     { label: "State" },
     { label: "Reserved Resources" },
     { label: "Label Selector" },
-    { label: "Label Domain" },
+    { label: "Topology" },
     { label: "Scheduling Detail" },
   ];
 
@@ -194,8 +194,8 @@ const PlacementGroupTable = ({
                 state,
                 stats,
                 bundles,
-                label_domain_key,
-                label_domain_assignments,
+                topology_strategy,
+                topology_assignments,
               }) => (
                 <TableRow key={placement_group_id}>
                   <TableCell align="center">
@@ -215,16 +215,15 @@ const PlacementGroupTable = ({
                     <LabelSelector bundles={bundles} />
                   </TableCell>
                   <TableCell align="center">
-                    {label_domain_key ? (
+                    {topology_strategy ? (
                       <Box sx={{ textAlign: "left" }}>
                         <div>
-                          <b>key:</b> {label_domain_key}
+                          <b>strategy:</b> {JSON.stringify(topology_strategy)}
                         </div>
                         <div>
                           <b>assignment:</b>{" "}
-                          {label_domain_assignments &&
-                          Object.keys(label_domain_assignments).length > 0
-                            ? JSON.stringify(label_domain_assignments)
+                          {topology_assignments
+                            ? JSON.stringify(topology_assignments)
                             : "-"}
                         </div>
                       </Box>

python/ray/dashboard/client/src/type/placementGroup.ts

@@ -26,8 +26,14 @@ export type PlacementGroup = {
     [key: string]: number | string;
   } | null;
   bundles: Bundle[];
-  label_domain_key?: string | null;
-  label_domain_assignments?: {
+  // Topology strategy: map from label key (e.g. "ray.io/gpu-domain") to
+  // corresponding placement strategy.
+  topology_strategy?: {
+    [key: string]: string;
+  } | null;
+  // Topology assignments: map from topology label key to the value the
+  // scheduler selected for this PG (e.g. {"ray.io/gpu-domain": "rack-1"}).
+  topology_assignments?: {
     [key: string]: string;
   } | null;
 };

python/ray/includes/common.pxd

@@ -436,6 +436,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 unordered_map[c_string, CPlacementStrategy] &topology_strategy,
         )
 
     cdef cppclass CObjectLocation "ray::core::ObjectLocation":

python/ray/includes/common.pxi

@@ -163,6 +163,8 @@ RAY_NODE_ACCELERATOR_TYPE_KEY = kLabelKeyNodeAcceleratorType.decode()
 RAY_NODE_MARKET_TYPE_KEY = kLabelKeyNodeMarketType.decode()
 RAY_NODE_REGION_KEY = kLabelKeyNodeRegion.decode()
 RAY_NODE_ZONE_KEY = kLabelKeyNodeZone.decode()
+# Keep this in sync with NODE_ID_LABEL_KEY in ray.util.placement_group. That
+# module cannot import this exported value because it forms a circular dependency
 RAY_NODE_ID_KEY = kLabelKeyNodeID.decode()
 RAY_NODE_GROUP_KEY = kLabelKeyNodeGroup.decode()
 

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_bundle_label_selector.py

@@ -297,152 +297,6 @@ def test_spread_strategy_bundle_label_selector(ray_start_cluster):
     ray.util.remove_placement_group(fail_pg)
 
 
-def test_gpu_domain_scheduling_reschedule_on_node_failure(ray_start_cluster):
-    """
-    Spins up 6 nodes in a single GPU domain (rack-1). Schedules 4 bundles,
-    then kills 2 nodes containing bundles. Verifies that the PG reschedules
-    the bundles onto remaining nodes that share the same GPU domain.
-    """
-    cluster = ray_start_cluster
-    cluster.add_node(num_cpus=0)
-    ray.init(address=cluster.address)
-
-    rack_labels = {
-        "ray.io/gpu-domain": "rack-1",
-        "ray.io/accelerator-type": "GB300",
-    }
-
-    rack_nodes = []
-    for _ in range(6):
-        rack_nodes.append(cluster.add_node(num_cpus=1, labels=rack_labels))
-    for _ in range(2):
-        cluster.add_node(num_cpus=1)
-
-    bundles = [{"CPU": 1}] * 4
-    label_selector = [{"ray.io/accelerator-type": "GB300"}] * 4
-
-    pg = placement_group(
-        bundles=bundles,
-        bundle_label_selector=label_selector,
-    )
-    ray.get(pg.ready(), timeout=30)
-    assert placement_group_table(pg)["state"] == "CREATED"
-
-    cluster.remove_node(rack_nodes[0])
-    cluster.remove_node(rack_nodes[1])
-
-    ray.get(pg.ready(), timeout=30)
-    assert placement_group_table(pg)["state"] == "CREATED"
-
-    # Verify that the rescheduled bundles are on nodes with the rack-1 gpu-domain label.
-    node_id_to_labels = {node["NodeID"]: node.get("Labels", {}) for node in ray.nodes()}
-    for node_id in placement_group_table(pg)["bundles_to_node_id"].values():
-        labels = node_id_to_labels[node_id]
-        assert labels.get("ray.io/gpu-domain") == "rack-1"
-
-
-def test_gpu_domain_scheduling_infeasible_after_node_kill(ray_start_cluster):
-    """
-    Spins up 2 nodes in a single GPU domain (rack-1). Schedules 2 bundles,
-    then kills one node containing a bundle. The surviving node already contains a bundle so
-    the orphaned bundle can't be rescheduled on it, and is left infeasible.
-    """
-    cluster = ray_start_cluster
-    cluster.add_node(num_cpus=0)
-    ray.init(address=cluster.address)
-
-    rack_labels = {
-        "ray.io/gpu-domain": "rack-1",
-        "ray.io/accelerator-type": "GB300",
-    }
-
-    cluster.add_node(num_cpus=1, labels=rack_labels)
-    node_b = cluster.add_node(num_cpus=1, labels=rack_labels)
-    cluster.add_node(num_cpus=1)
-
-    bundles = [{"CPU": 1}] * 2
-    label_selector = [{"ray.io/accelerator-type": "GB300"}] * 2
-
-    pg = placement_group(
-        bundles=bundles,
-        bundle_label_selector=label_selector,
-    )
-    ray.get(pg.ready(), timeout=10)
-    assert placement_group_table(pg)["state"] == "CREATED"
-
-    cluster.remove_node(node_b)
-
-    with pytest.raises(ray.exceptions.GetTimeoutError):
-        ray.get(pg.ready(), timeout=5)
-
-    state = placement_group_table(pg)["state"]
-    assert state == "RESCHEDULING", f"Expected RESCHEDULING, got {state}"
-
-
-def test_gpu_domain_scheduling_rescheduling_on_gpu_domain_failure(ray_start_cluster):
-    """Verify the PG reschedules onto a new GPU domain after total domain failure.
-
-    Creates a placement group on rack 1. Then removes one rack 1 node so the PG cannot
-    be fully placed on that rack anymore, confirming the PG enters RESCHEDULING. While
-    partial rack 1 remains, the PG stays infeasible because label locality pins it to
-    the original domain. Once all rack 1 nodes are removed (total failure), the gpu
-    domain 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_labels = {
-        "ray.io/gpu-domain": "rack-1",
-        "ray.io/accelerator-type": "GB300",
-    }
-
-    rack2_labels = {
-        "ray.io/gpu-domain": "rack-2",
-        "ray.io/accelerator-type": "GB300",
-    }
-
-    rack1_nodes = []
-    for _ in range(4):
-        rack1_nodes.append(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():
-            labels = node_id_to_labels[node_id]
-            assert labels.get(label) == value
-
-    bundles = [{"CPU": 1}] * 4
-    label_selector = [{"ray.io/accelerator-type": "GB300"}] * 4
-
-    pg = placement_group(
-        bundles=bundles,
-        bundle_label_selector=label_selector,
-    )
-    ray.get(pg.ready(), timeout=30)
-    assert placement_group_table(pg)["state"] == "CREATED"
-
-    assert_pg_nodes_label_value(ray.nodes(), pg, "ray.io/gpu-domain", "rack-1")
-
-    for _ in range(4):
-        cluster.add_node(num_cpus=1, labels=rack2_labels)
-
-    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"
-
-    for node in rack1_nodes[1:]:
-        cluster.remove_node(node)
-
-    ray.get(pg.ready(), timeout=30)
-    assert placement_group_table(pg)["state"] == "CREATED"
-
-    # Verify that the PG is now on rack 2 after all rack 1 nodes are removed
-    assert_pg_nodes_label_value(ray.nodes(), pg, "ray.io/gpu-domain", "rack-2")
-
-
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))