wehub-resource-sync
472 lines
15 KiB
Python
472 lines
15 KiB
Python
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__]))
|