chore: import upstream snapshot with attribution

python/ray/tests/test_runtime_context.py

@@ -0,0 +1,471 @@
+import os
+import signal
+import sys
+import time
+import warnings
+
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from ray.util.state import list_tasks
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Fails on windows")
+def test_was_current_actor_reconstructed(shutdown_only):
+    ray.init()
+
+    @ray.remote(max_restarts=10)
+    class A(object):
+        def __init__(self):
+            self._was_reconstructed = (
+                ray.get_runtime_context().was_current_actor_reconstructed
+            )
+
+        def get_was_reconstructed(self):
+            return self._was_reconstructed
+
+        def update_was_reconstructed(self):
+            return ray.get_runtime_context().was_current_actor_reconstructed
+
+        def get_pid(self):
+            return os.getpid()
+
+    a = A.remote()
+    # `was_reconstructed` should be False when it's called in actor.
+    assert ray.get(a.get_was_reconstructed.remote()) is False
+    # `was_reconstructed` should be False when it's called in a remote method
+    # and the actor never fails.
+    assert ray.get(a.update_was_reconstructed.remote()) is False
+
+    pid = ray.get(a.get_pid.remote())
+    os.kill(pid, signal.SIGKILL)
+    time.sleep(2)
+    # These 2 methods should be return True because
+    # this actor failed and restored.
+    assert ray.get(a.get_was_reconstructed.remote()) is True
+    assert ray.get(a.update_was_reconstructed.remote()) is True
+
+    @ray.remote(max_restarts=10)
+    class A(object):
+        def current_job_id(self):
+            return ray.get_runtime_context().get_job_id()
+
+        def current_actor_id(self):
+            return ray.get_runtime_context().get_actor_id()
+
+    @ray.remote
+    def f():
+        assert ray.get_runtime_context().get_actor_id() is None
+        assert ray.get_runtime_context().get_task_id() is not None
+        assert ray.get_runtime_context().get_node_id() is not None
+        assert ray.get_runtime_context().get_job_id() is not None
+        context = ray.get_runtime_context().get()
+        assert "actor_id" not in context
+        assert context["task_id"].hex() == ray.get_runtime_context().get_task_id()
+        assert context["node_id"].hex() == ray.get_runtime_context().get_node_id()
+        assert context["job_id"].hex() == ray.get_runtime_context().get_job_id()
+
+    a = A.remote()
+    assert ray.get(a.current_job_id.remote()) is not None
+    assert ray.get(a.current_actor_id.remote()) is not None
+    ray.get(f.remote())
+
+
+def test_get_context_dict(ray_start_regular):
+    context_dict = ray.get_runtime_context().get()
+    assert context_dict["node_id"] is not None
+    assert context_dict["job_id"] is not None
+    assert "actor_id" not in context_dict
+    assert "task_id" not in context_dict
+
+    @ray.remote
+    class Actor:
+        def check(self, node_id, job_id):
+            context_dict = ray.get_runtime_context().get()
+            assert context_dict["node_id"] == node_id
+            assert context_dict["job_id"] == job_id
+            assert context_dict["actor_id"] is not None
+            assert context_dict["task_id"] is not None
+            assert context_dict["actor_id"] != "not an ActorID"
+
+    a = Actor.remote()
+    ray.get(a.check.remote(context_dict["node_id"], context_dict["job_id"]))
+
+    @ray.remote
+    def task(node_id, job_id):
+        context_dict = ray.get_runtime_context().get()
+        assert context_dict["node_id"] == node_id
+        assert context_dict["job_id"] == job_id
+        assert context_dict["task_id"] is not None
+        assert "actor_id" not in context_dict
+
+    ray.get(task.remote(context_dict["node_id"], context_dict["job_id"]))
+
+
+def test_current_actor(ray_start_regular):
+    @ray.remote
+    class Echo:
+        def __init__(self):
+            pass
+
+        def echo(self, s):
+            self_actor = ray.get_runtime_context().current_actor
+            return self_actor.echo2.remote(s)
+
+        def echo2(self, s):
+            return s
+
+    e = Echo.remote()
+    obj = e.echo.remote("hello")
+    assert ray.get(ray.get(obj)) == "hello"
+
+
+def test_get_assigned_resources(ray_start_10_cpus):
+    @ray.remote
+    class Echo:
+        def check(self):
+            return ray.get_runtime_context().get_assigned_resources()
+
+    e = Echo.remote()
+    result = e.check.remote()
+    print(ray.get(result))
+    assert ray.get(result).get("CPU") is None
+    ray.kill(e)
+
+    e = Echo.options(num_cpus=4).remote()
+    result = e.check.remote()
+    assert ray.get(result)["CPU"] == 4.0
+    ray.kill(e)
+
+    @ray.remote
+    def check():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    result = check.remote()
+    assert ray.get(result)["CPU"] == 1.0
+
+    result = check.options(num_cpus=2).remote()
+    assert ray.get(result)["CPU"] == 2.0
+
+
+# Use default filterwarnings behavior for this test
+@pytest.mark.filterwarnings("default")
+def test_ids(ray_start_regular):
+    rtc = ray.get_runtime_context()
+    # node id
+    assert isinstance(rtc.get_node_id(), str)
+    with warnings.catch_warnings(record=True) as w:
+        assert rtc.get_node_id() == rtc.node_id.hex()
+        assert any("Use get_node_id() instead" in str(warning.message) for warning in w)
+
+    # session name
+    assert isinstance(rtc.get_session_name(), str)
+    session_name = rtc.get_session_name()
+    assert len(session_name) > 0
+
+    @ray.remote
+    def verify_session_name(expected_session_name):
+        rtc = ray.get_runtime_context()
+        assert isinstance(rtc.get_session_name(), str)
+        assert rtc.get_session_name() == expected_session_name
+        return True
+
+    ray.get(verify_session_name.remote(session_name))
+
+    # job id
+    assert isinstance(rtc.get_job_id(), str)
+    with warnings.catch_warnings(record=True) as w:
+        assert rtc.get_job_id() == rtc.job_id.hex()
+        assert any("Use get_job_id() instead" in str(warning.message) for warning in w)
+
+    assert rtc.get_actor_name() is None
+
+    # placement group id
+    # Driver doesn't belong to any placement group.
+    assert rtc.get_placement_group_id() is None
+    pg = ray.util.placement_group(
+        name="bar",
+        strategy="PACK",
+        bundles=[
+            {"CPU": 1, "GPU": 0},
+        ],
+    )
+    ray.get(pg.ready())
+
+    @ray.remote
+    def foo_pg():
+        rtc = ray.get_runtime_context()
+        assert isinstance(rtc.get_placement_group_id(), str)
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            assert rtc.get_placement_group_id() == rtc.current_placement_group_id.hex()
+            assert any(
+                "Use get_placement_group_id() instead" in str(warning.message)
+                for warning in w
+            )
+
+    ray.get(
+        foo_pg.options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+        ).remote()
+    )
+    ray.util.remove_placement_group(pg)
+
+    # task id
+    assert rtc.get_task_id() is None
+
+    @ray.remote
+    def foo_task():
+        rtc = ray.get_runtime_context()
+        assert isinstance(rtc.get_task_id(), str)
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            assert rtc.get_task_id() == rtc.task_id.hex()
+            assert any(
+                "Use get_task_id() instead" in str(warning.message) for warning in w
+            )
+
+    ray.get(foo_task.remote())
+
+    # actor id
+    assert rtc.get_actor_id() is None
+
+    @ray.remote
+    class FooActor:
+        def foo(self):
+            rtc = ray.get_runtime_context()
+            assert isinstance(rtc.get_actor_id(), str)
+            with warnings.catch_warnings(record=True) as w:
+                warnings.simplefilter("always")
+                assert rtc.get_actor_id() == rtc.actor_id.hex()
+                assert any(
+                    "Use get_actor_id() instead" in str(warning.message)
+                    for warning in w
+                )
+
+    actor = FooActor.remote()
+    ray.get(actor.foo.remote())
+
+    # actor name
+    @ray.remote
+    class NamedActor:
+        def name(self):
+            return ray.get_runtime_context().get_actor_name()
+
+    ACTOR_NAME = "actor_name"
+    named_actor = NamedActor.options(name=ACTOR_NAME).remote()
+    assert ray.get(named_actor.name.remote()) == ACTOR_NAME
+
+    # unnamed actor name
+    unnamed_actor = NamedActor.options().remote()
+    assert ray.get(unnamed_actor.name.remote()) == ""
+
+    # task actor name
+    @ray.remote
+    def task_actor_name():
+        ray.get_runtime_context().get_actor_name()
+
+    assert ray.get(task_actor_name.remote()) is None
+
+    # driver actor name
+    assert rtc.get_actor_name() is None
+
+
+def test_auto_init(shutdown_only):
+    assert not ray.is_initialized()
+    ray.get_runtime_context()
+    assert ray.is_initialized()
+
+
+def test_get_task_name(shutdown_only):
+    ray.init()
+
+    # for a normal task
+    @ray.remote
+    def get_task_name_for_normal_task():
+        return ray.get_runtime_context().get_task_name()
+
+    expected_task_name = "normal_task_name"
+    task_name = ray.get(
+        get_task_name_for_normal_task.options(name=expected_task_name).remote()
+    )
+    assert (
+        task_name == expected_task_name
+    ), f"Check normal task name failed. expected={expected_task_name}, \
+actual={task_name}"
+
+    # for an actor task
+    @ray.remote
+    class Actor:
+        def get_task_name_for_actor_task(self):
+            return ray.get_runtime_context().get_task_name()
+
+    expected_task_name = "Actor.get_task_name_for_actor_task"
+    actor = Actor.remote()
+    task_name = ray.get(actor.get_task_name_for_actor_task.remote())
+    assert (
+        task_name == expected_task_name
+    ), f"Check actor task name failed. expected={expected_task_name}, \
+actual={task_name}"
+
+    # for a threaded actor task
+    @ray.remote
+    class ThreadedActor:
+        def get_task_name_for_threaded_actor_task(self):
+            return ray.get_runtime_context().get_task_name()
+
+    expected_task_name = "ThreadedActor.get_task_name_for_threaded_actor_task"
+    threaded_actor = ThreadedActor.options(max_concurrency=2).remote()
+    task_name = ray.get(threaded_actor.get_task_name_for_threaded_actor_task.remote())
+    assert (
+        task_name == expected_task_name
+    ), f"Check actor task name failed. expected={expected_task_name}, \
+actual={task_name}"
+
+    # for a async actor task
+    @ray.remote
+    class AsyncActor:
+        async def get_task_name_for_async_actor_task(self):
+            return ray.get_runtime_context().get_task_name()
+
+    expected_task_name = "AsyncActor.get_task_name_for_async_actor_task"
+    async_actor = AsyncActor.remote()
+    task_name = ray.get(async_actor.get_task_name_for_async_actor_task.remote())
+    assert (
+        task_name == expected_task_name
+    ), f"Check actor task name failed. expected={expected_task_name}, \
+actual={task_name}"
+
+
+def test_get_task_function_name(shutdown_only):
+    ray.init()
+
+    # for a normal task
+    @ray.remote
+    def get_task_function_name_for_normal_task():
+        return ray.get_runtime_context().get_task_function_name()
+
+    expected_task_function_name = __name__ + ".get_task_function_name_for_normal_task"
+    task_function_name = ray.get(get_task_function_name_for_normal_task.remote())
+    assert (
+        task_function_name == expected_task_function_name
+    ), f"Check normal task function failed. expected={expected_task_function_name}, \
+actual={task_function_name}"
+
+    # for an actor task
+    @ray.remote
+    class Actor:
+        def get_task_function_name_for_actor_task(self):
+            return ray.get_runtime_context().get_task_function_name()
+
+    expected_task_function_name = (
+        __name__ + ".Actor.get_task_function_name_for_actor_task"
+    )
+    actor = Actor.remote()
+    task_function_name = ray.get(actor.get_task_function_name_for_actor_task.remote())
+    assert (
+        task_function_name == expected_task_function_name
+    ), f"Check actor task function failed. expected={expected_task_function_name}, \
+actual={task_function_name}"
+
+    # for a threaded actor task
+    @ray.remote
+    class ThreadedActor:
+        def get_task_function_name_for_threaded_actor_task(self):
+            return ray.get_runtime_context().get_task_function_name()
+
+    expected_task_function_name = (
+        __name__ + ".ThreadedActor.get_task_function_name_for_threaded_actor_task"
+    )
+    threaded_actor = ThreadedActor.options(max_concurrency=2).remote()
+    task_function_name = ray.get(
+        threaded_actor.get_task_function_name_for_threaded_actor_task.remote()
+    )
+    assert (
+        task_function_name == expected_task_function_name
+    ), f"Check actor task function failed. expected={expected_task_function_name}, \
+actual={task_function_name}"
+
+    # for a async actor task
+    @ray.remote
+    class AsyncActor:
+        async def get_task_function_name_for_async_actor_task(self):
+            return ray.get_runtime_context().get_task_function_name()
+
+    expected_task_function_name = (
+        __name__
+        + ".test_get_task_function_name.<locals>.AsyncActor.\
+get_task_function_name_for_async_actor_task"
+    )
+    async_actor = AsyncActor.remote()
+    task_function_name = ray.get(
+        async_actor.get_task_function_name_for_async_actor_task.remote()
+    )
+    assert (
+        task_function_name == expected_task_function_name
+    ), f"Check actor task function failed. expected={expected_task_function_name}, \
+actual={task_function_name}"
+
+
+def test_async_actor_task_id(shutdown_only):
+    ray.init()
+
+    @ray.remote
+    class A:
+        async def f(self):
+            task_id = ray.get_runtime_context().get_task_id()
+            return task_id
+
+    a = A.remote()
+    task_id = ray.get(a.f.remote())
+
+    def verify():
+        tasks = list_tasks(filters=[("name", "=", "A.f")])
+        assert len(tasks) == 1
+        assert tasks[0].task_id == task_id
+        return True
+
+    wait_for_condition(verify)
+
+
+def test_get_node_labels(ray_start_cluster_head):
+    cluster = ray_start_cluster_head
+    cluster.add_node(
+        resources={"worker1": 1},
+        num_cpus=1,
+        labels={
+            "ray.io/accelerator-type": "A100",
+            "ray.io/availability-region": "us-west4",
+            "ray.io/market-type": "spot",
+        },
+    )
+
+    @ray.remote
+    class Actor:
+        def get_node_id(self):
+            return ray.get_runtime_context().get_node_id()
+
+        def get_node_labels(self):
+            return ray.get_runtime_context().get_node_labels()
+
+    expected_node_labels = {
+        "ray.io/accelerator-type": "A100",
+        "ray.io/availability-region": "us-west4",
+        "ray.io/market-type": "spot",
+    }
+
+    # Check node labels from Actor runtime context
+    a = Actor.options(label_selector={"ray.io/accelerator-type": "A100"}).remote()
+    node_labels = ray.get(a.get_node_labels.remote())
+    expected_node_labels["ray.io/node-id"] = ray.get(a.get_node_id.remote())
+    assert expected_node_labels == node_labels
+
+    # Check node labels from driver runtime context (none are set except default)
+    driver_labels = ray.get_runtime_context().get_node_labels()
+    assert {"ray.io/node-id": ray.get_runtime_context().get_node_id()} == driver_labels
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))