chore: import upstream snapshot with attribution

python/ray/tests/test_node_label_scheduling_strategy.py

@@ -0,0 +1,341 @@
+import sys
+
+import pytest
+
+import ray
+from ray.util.scheduling_strategies import (
+    DoesNotExist,
+    Exists,
+    In,
+    NodeLabelSchedulingStrategy,
+    NotIn,
+)
+
+
+@ray.remote
+class MyActor:
+    def __init__(self):
+        self.value = 0
+
+    def value(self):
+        return self.value
+
+    def get_node_id(self):
+        return ray.get_runtime_context().get_node_id()
+
+
+@ray.remote
+def get_node_id():
+    return ray.get_runtime_context().get_node_id()
+
+
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ['ray start --head --labels={"gpu_type":"A100","region":"us"}'],
+    indirect=True,
+)
+def test_node_label_scheduling_basic(call_ray_start):
+    ray.init(address=call_ray_start)
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            {"gpu_type": In("A100", "T100"), "region": Exists()}
+        )
+    ).remote()
+    assert ray.get(actor.value.remote(), timeout=3) == 0
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": NotIn("A100")})
+    ).remote()
+    with pytest.raises(TimeoutError):
+        assert ray.get(actor.value.remote(), timeout=3) == 0
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={"gpu_type": DoesNotExist()},
+            soft={"gpu_type": In("A100")},
+        )
+    ).remote()
+    with pytest.raises(TimeoutError):
+        assert ray.get(actor.value.remote(), timeout=3) == 0
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={"gpu_type": In("T100")},
+        )
+    ).remote()
+    with pytest.raises(TimeoutError):
+        assert ray.get(actor.value.remote(), timeout=3) == 0
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={},
+            soft={"gpu_type": In("T100ssss")},
+        )
+    ).remote()
+    assert ray.get(actor.value.remote(), timeout=3) == 0
+
+
+def test_node_label_scheduling_in_cluster(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        resources={"worker1": 1},
+        num_cpus=3,
+        labels={"gpu_type": "A100", "azone": "azone-1"},
+    )
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+    node_1 = ray.get(get_node_id.options(resources={"worker1": 1}).remote())
+    cluster.add_node(
+        resources={"worker2": 1},
+        num_cpus=3,
+        labels={"gpu_type": "T100", "azone": "azone-1"},
+    )
+    node_2 = ray.get(get_node_id.options(resources={"worker2": 1}).remote())
+    cluster.add_node(
+        resources={"worker3": 1},
+        num_cpus=3,
+        labels={"gpu_type": "T100", "azone": "azone-2"},
+    )
+    node_3 = ray.get(get_node_id.options(resources={"worker3": 1}).remote())
+    cluster.add_node(resources={"worker4": 1}, num_cpus=3)
+    node_4 = ray.get(get_node_id.options(resources={"worker4": 1}).remote())
+    cluster.wait_for_nodes()
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": In("A100")})
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) == node_1
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy({"ray.io/node-id": In(node_4)})
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) == node_4
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": In("T100")})
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_2, node_3)
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            {"azone": In("azone-1", "azone-2")}
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_1, node_2, node_3)
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            {"gpu_type": In("T100"), "azone": In("azone-1")}
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_2)
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            {
+                "gpu_type": NotIn("A100"),
+            }
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_2, node_3, node_4)
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            {
+                "gpu_type": DoesNotExist(),
+            }
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_4)
+
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            {
+                "gpu_type": Exists(),
+            }
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_1, node_2, node_3)
+
+
+def test_node_label_scheduling_with_soft(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        resources={"worker1": 1},
+        num_cpus=3,
+        labels={"gpu_type": "A100", "azone": "azone-1"},
+    )
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+    node_1 = ray.get(get_node_id.options(resources={"worker1": 1}).remote())
+    cluster.add_node(
+        resources={"worker2": 1},
+        num_cpus=3,
+        labels={"gpu_type": "T100", "azone": "azone-1"},
+    )
+    node_2 = ray.get(get_node_id.options(resources={"worker2": 1}).remote())
+    cluster.add_node(
+        resources={"worker3": 1},
+        num_cpus=3,
+        labels={"gpu_type": "T100", "azone": "azone-2"},
+    )
+    node_3 = ray.get(get_node_id.options(resources={"worker3": 1}).remote())
+    cluster.add_node(resources={"worker4": 1}, num_cpus=3)
+    node_4 = ray.get(get_node_id.options(resources={"worker4": 1}).remote())
+    cluster.wait_for_nodes()
+
+    # hard match and soft match
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={"azone": In("azone-1")}, soft={"gpu_type": In("T100")}
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) == node_2
+
+    # hard match and soft don't match
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={"azone": In("azone-1")}, soft={"gpu_type": In("H100")}
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_1, node_2)
+
+    # no hard and  soft match
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={}, soft={"gpu_type": Exists()}
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (node_1, node_2, node_3)
+
+    # no hard and soft don't match
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={}, soft={"gpu_type": In("H100")}
+        )
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) in (
+        node_1,
+        node_2,
+        node_3,
+        node_4,
+    )
+
+    # hard don't match and soft match
+    actor = MyActor.options(
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={"azone": In("azone-3")}, soft={"gpu_type": In("T100")}
+        )
+    ).remote()
+    with pytest.raises(TimeoutError):
+        ray.get(actor.get_node_id.remote(), timeout=3)
+
+
+def test_node_not_available(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(resources={"worker1": 1}, num_cpus=1, labels={"gpu_type": "A100"})
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+    node_1 = ray.get(get_node_id.options(resources={"worker1": 1}).remote())
+    cluster.add_node(resources={"worker2": 1}, num_cpus=1)
+    node_2 = ray.get(get_node_id.options(resources={"worker2": 1}).remote())
+    cluster.wait_for_nodes()
+
+    # Infeasible
+    actor = MyActor.options(
+        num_cpus=2,
+        scheduling_strategy=NodeLabelSchedulingStrategy(hard={"gpu_type": In("A100")}),
+    ).remote()
+    with pytest.raises(TimeoutError):
+        ray.get(actor.get_node_id.remote(), timeout=3)
+
+    actor = MyActor.options(
+        num_cpus=1,
+        scheduling_strategy=NodeLabelSchedulingStrategy(hard={"gpu_type": In("A100")}),
+    ).remote()
+    assert ray.get(actor.get_node_id.remote(), timeout=3) == node_1
+
+    # Soft match node is not available, sheduling to other available node.
+    actor_2 = MyActor.options(
+        num_cpus=1,
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={}, soft={"gpu_type": In("A100")}
+        ),
+    ).remote()
+    assert ray.get(actor_2.get_node_id.remote(), timeout=3) == node_2
+
+    # No available nodes.
+    actor_3 = MyActor.options(
+        num_cpus=1,
+        scheduling_strategy=NodeLabelSchedulingStrategy(
+            hard={}, soft={"gpu_type": In("A100")}
+        ),
+    ).remote()
+    with pytest.raises(TimeoutError):
+        ray.get(actor_3.get_node_id.remote(), timeout=3)
+
+    # node_1 change to available
+    ray.kill(actor)
+    assert ray.get(actor_3.get_node_id.remote(), timeout=3) == node_1
+
+
+def test_node_label_scheduling_invalid_paramter(call_ray_start):
+    ray.init(address=call_ray_start)
+    with pytest.raises(
+        ValueError, match="Type of value in position 0 for the In operator must be str"
+    ):
+        MyActor.options(
+            scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": In(123)})
+        )
+
+    with pytest.raises(
+        ValueError,
+        match="Type of value in position 0 for the NotIn operator must be str",
+    ):
+        MyActor.options(
+            scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": NotIn(123)})
+        )
+
+    with pytest.raises(
+        ValueError,
+        match="The variadic parameter of the In operator must be a non-empty tuple",
+    ):
+        MyActor.options(
+            scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": In()})
+        )
+
+    with pytest.raises(
+        ValueError,
+        match="The variadic parameter of the NotIn operator must be a non-empty tuple",
+    ):
+        MyActor.options(
+            scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": NotIn()})
+        )
+
+    with pytest.raises(ValueError, match="The soft parameter must be a map"):
+        MyActor.options(
+            scheduling_strategy=NodeLabelSchedulingStrategy(hard=None, soft=["1"])
+        )
+
+    with pytest.raises(
+        ValueError, match="The map key of the hard parameter must be of type str"
+    ):
+        MyActor.options(scheduling_strategy=NodeLabelSchedulingStrategy({111: "1111"}))
+
+    with pytest.raises(
+        ValueError, match="must be one of the `In`, `NotIn`, `Exists` or `DoesNotExist`"
+    ):
+        MyActor.options(
+            scheduling_strategy=NodeLabelSchedulingStrategy({"gpu_type": "1111"})
+        )
+
+    with pytest.raises(
+        ValueError,
+        match="The `hard` and `soft` parameter "
+        "of NodeLabelSchedulingStrategy cannot both be empty.",
+    ):
+        MyActor.options(scheduling_strategy=NodeLabelSchedulingStrategy(hard={}))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))