wehub-resource-sync
275 lines
7.4 KiB
Python
275 lines
7.4 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import (
|
|
wait_for_condition,
|
|
)
|
|
from ray.util.state import list_tasks
|
|
|
|
# Run every test in this module twice: default and with core-worker to aggregator feature flag enabled
|
|
pytestmark = [
|
|
pytest.mark.parametrize(
|
|
"event_routing_config", ["default", "aggregator"], indirect=True
|
|
),
|
|
pytest.mark.usefixtures("event_routing_config"),
|
|
]
|
|
|
|
|
|
def test_globally_disable_task_events(
|
|
shutdown_only,
|
|
):
|
|
"""Test that globally disabled task events overrides per-actor/task options."""
|
|
system_config = {
|
|
"task_events_report_interval_ms": 0,
|
|
"metrics_report_interval_ms": 200,
|
|
"enable_timeline": False,
|
|
}
|
|
ray.init(
|
|
num_cpus=1,
|
|
_system_config=system_config,
|
|
)
|
|
|
|
@ray.remote(enable_task_events=True)
|
|
def noop():
|
|
pass
|
|
|
|
@ray.remote(num_cpus=0, enable_task_events=True)
|
|
class Actor:
|
|
def f(self):
|
|
# This should always be traced.
|
|
ray.get(noop.options(name="inner-task-traced").remote())
|
|
|
|
@ray.method(enable_task_events=True)
|
|
def g(self):
|
|
pass
|
|
|
|
a = Actor.remote()
|
|
|
|
for _ in range(10):
|
|
ray.get(a.f.remote())
|
|
ray.get(a.g.remote())
|
|
|
|
assert len(list_tasks()) == 0
|
|
|
|
|
|
def test_enable_task_events_actor(
|
|
shutdown_only,
|
|
):
|
|
"""Test enabling/disabling actor task events via actor options.
|
|
|
|
- enable_task_events=True should enable all tasks from the actor.
|
|
- enable_task_events=False should disable all tasks from the actor.
|
|
- The above can be overridden per method.
|
|
"""
|
|
system_config = {
|
|
"task_events_report_interval_ms": 100,
|
|
"metrics_report_interval_ms": 200,
|
|
"enable_timeline": False,
|
|
}
|
|
|
|
ray.init(
|
|
num_cpus=1,
|
|
_system_config=system_config,
|
|
)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class A:
|
|
def default_method(self):
|
|
pass
|
|
|
|
@ray.method(enable_task_events=True)
|
|
def enabled_method(self):
|
|
pass
|
|
|
|
@ray.method(enable_task_events=False)
|
|
def disabled_method(self):
|
|
pass
|
|
|
|
a_enabled = A.options(enable_task_events=True).remote()
|
|
a_disabled = A.options(enable_task_events=False).remote()
|
|
|
|
ray.get(
|
|
[
|
|
a_enabled.default_method.remote(),
|
|
a_enabled.enabled_method.remote(),
|
|
a_enabled.disabled_method.remote(),
|
|
]
|
|
)
|
|
ray.get(
|
|
[
|
|
a_disabled.default_method.remote(),
|
|
a_disabled.enabled_method.remote(),
|
|
a_disabled.disabled_method.remote(),
|
|
]
|
|
)
|
|
|
|
expected_task_events = {
|
|
# All methods except the override-disabled method should be exported.
|
|
(a_enabled._actor_id.hex(), "A.__init__"),
|
|
(a_enabled._actor_id.hex(), "A.default_method"),
|
|
(a_enabled._actor_id.hex(), "A.enabled_method"),
|
|
# Only the override-enabled method should be exported.
|
|
(a_disabled._actor_id.hex(), "A.enabled_method"),
|
|
}
|
|
|
|
def _verify():
|
|
task_events = {(t.actor_id, t.name) for t in list_tasks()}
|
|
assert task_events == expected_task_events
|
|
return True
|
|
|
|
wait_for_condition(_verify)
|
|
|
|
|
|
def test_task_events_actor_called_from_task(shutdown_only):
|
|
"""Call an actor method from within a task, verify events are/aren't exported."""
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def f(self):
|
|
pass
|
|
|
|
@ray.remote(enable_task_events=False)
|
|
def task(actor):
|
|
actor.f.remote()
|
|
|
|
a_enabled = Actor.options(enable_task_events=True).remote()
|
|
a_disabled = Actor.options(enable_task_events=False).remote()
|
|
ray.get(task.remote(a_enabled))
|
|
ray.get(task.remote(a_disabled))
|
|
|
|
# Expect to only see events from a_enabled, not a_disabled.
|
|
expected_task_events = {
|
|
(a_enabled._actor_id.hex(), "Actor.__init__"),
|
|
(a_enabled._actor_id.hex(), "Actor.f"),
|
|
}
|
|
|
|
def _verify():
|
|
task_events = {(t.actor_id, t.name) for t in list_tasks()}
|
|
assert task_events == expected_task_events
|
|
return True
|
|
|
|
wait_for_condition(_verify)
|
|
|
|
|
|
def test_enable_task_events_invalid_options(shutdown_only):
|
|
"""Test the invalid values for the option."""
|
|
|
|
@ray.remote
|
|
def f():
|
|
pass
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
pass
|
|
|
|
with pytest.raises(TypeError):
|
|
ray.get(f.options(enable_task_events="invalid").remote())
|
|
|
|
with pytest.raises(TypeError):
|
|
ray.get(f.options(enable_task_events=None).remote())
|
|
|
|
with pytest.raises(TypeError):
|
|
ray.get(Actor.options(enable_task_events="invalid").remote())
|
|
|
|
with pytest.raises(TypeError):
|
|
ray.get(Actor.options(enable_task_events=None).remote())
|
|
|
|
|
|
def test_enable_task_events_tasks(shutdown_only):
|
|
"""Test enabling/disabled task events at the task level."""
|
|
system_config = {
|
|
"task_events_report_interval_ms": 100,
|
|
"metrics_report_interval_ms": 200,
|
|
"enable_timeline": False,
|
|
}
|
|
|
|
ray.init(
|
|
num_cpus=1,
|
|
_system_config=system_config,
|
|
)
|
|
|
|
@ray.remote(enable_task_events=True)
|
|
def inner_enabled():
|
|
pass
|
|
|
|
@ray.remote(enable_task_events=False)
|
|
def inner_disabled():
|
|
pass
|
|
|
|
@ray.remote(enable_task_events=True)
|
|
def outer_enabled():
|
|
return ray.get(inner_disabled.remote())
|
|
|
|
@ray.remote(enable_task_events=False)
|
|
def outer_disabled():
|
|
return ray.get(inner_enabled.remote())
|
|
|
|
ray.get([outer_enabled.remote(), outer_disabled.remote()])
|
|
|
|
expected_task_events = {"inner_enabled", "outer_enabled"}
|
|
|
|
def _verify():
|
|
task_events = {t.name for t in list_tasks()}
|
|
assert task_events == expected_task_events
|
|
return True
|
|
|
|
wait_for_condition(_verify)
|
|
|
|
|
|
def test_enable_task_events_nested_actor(shutdown_only):
|
|
"""Test that enabling/disabling is independent for nested actors."""
|
|
ray.init(
|
|
num_cpus=1,
|
|
_system_config={
|
|
"task_events_report_interval_ms": 100,
|
|
"metrics_report_interval_ms": 200,
|
|
"enable_timeline": False,
|
|
},
|
|
)
|
|
|
|
@ray.remote
|
|
class Inner:
|
|
def noop(self):
|
|
pass
|
|
|
|
@ray.remote
|
|
class Outer:
|
|
def __init__(self, inner: ray.actor.ActorHandle):
|
|
self._inner = inner
|
|
|
|
def call_inner(self):
|
|
ray.get(self._inner.noop.remote())
|
|
|
|
# Create one copy of Outer that has task events enabled, but calls a copy of
|
|
# inner that has task events disabled.
|
|
inner_disabled = Inner.options(enable_task_events=False).remote()
|
|
outer_enabled = Outer.options(enable_task_events=True).remote(inner_disabled)
|
|
|
|
# Create another copy of Outer that has task events disabled, but calls a copy of
|
|
# inner that has task events enabled.
|
|
inner_enabled = Inner.options(enable_task_events=True).remote()
|
|
outer_disabled = Outer.options(enable_task_events=False).remote(inner_enabled)
|
|
|
|
ray.get(outer_enabled.call_inner.remote())
|
|
ray.get(outer_disabled.call_inner.remote())
|
|
|
|
expected_task_events = {
|
|
(outer_enabled._actor_id.hex(), "Outer.__init__"),
|
|
(outer_enabled._actor_id.hex(), "Outer.call_inner"),
|
|
(inner_enabled._actor_id.hex(), "Inner.__init__"),
|
|
(inner_enabled._actor_id.hex(), "Inner.noop"),
|
|
}
|
|
|
|
def _verify():
|
|
task_events = {(t.actor_id, t.name) for t in list_tasks()}
|
|
assert task_events == expected_task_events
|
|
return True
|
|
|
|
wait_for_condition(_verify)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|