chore: import upstream snapshot with attribution

python/ray/tests/test_task_events.py

@@ -0,0 +1,531 @@
+import sys
+import threading
+import time
+from collections import defaultdict
+from typing import Dict
+
+import pytest
+
+import ray
+from ray._common.test_utils import (
+    PrometheusTimeseries,
+    SignalActor,
+    wait_for_condition,
+)
+from ray._private.state_api_test_utils import (
+    verify_failed_task,
+)
+from ray._private.test_utils import (
+    raw_metric_timeseries,
+    wait_for_aggregator_agent_if_enabled,
+)
+from ray._private.worker import RayContext
+from ray.exceptions import RuntimeEnvSetupError
+from ray.runtime_env import RuntimeEnv
+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"),
+]
+
+_SYSTEM_CONFIG = {
+    "task_events_report_interval_ms": 100,
+    "metrics_report_interval_ms": 200,
+    "enable_timeline": False,
+    "gcs_mark_task_failed_on_job_done_delay_ms": 1000,
+}
+
+
+def aggregate_task_event_metric(
+    info: RayContext, timeseries: PrometheusTimeseries
+) -> Dict:
+    """
+    Aggregate metrics of task events into:
+        {
+            "REPORTED": ray_gcs_task_manager_task_events_reported,
+            "STORED": ray_gcs_task_manager_task_events_stored,
+            "DROPPED_PROFILE_EVENT":
+                ray_gcs_task_manager_task_events_dropped PROFILE_EVENT,
+            "DROPPED_STATUS_EVENT":
+                ray_gcs_task_manager_task_events_dropped STATUS_EVENT,
+        }
+    """
+    res = raw_metric_timeseries(info, timeseries)
+    task_events_info = defaultdict(int)
+    if "ray_gcs_task_manager_task_events_dropped" in res:
+        for sample in res["ray_gcs_task_manager_task_events_dropped"]:
+            if "Type" in sample.labels and sample.labels["Type"] != "":
+                task_events_info["DROPPED_" + sample.labels["Type"]] += sample.value
+
+    if "ray_gcs_task_manager_task_events_stored" in res:
+        for sample in res["ray_gcs_task_manager_task_events_stored"]:
+            task_events_info["STORED"] += sample.value
+
+    if "ray_gcs_task_manager_task_events_reported" in res:
+        for sample in res["ray_gcs_task_manager_task_events_reported"]:
+            task_events_info["REPORTED"] += sample.value
+
+    return task_events_info
+
+
+def test_status_task_events_metrics(shutdown_only):
+    info = ray.init(num_cpus=1, _system_config=_SYSTEM_CONFIG)
+
+    # Start a task
+    @ray.remote
+    def f():
+        pass
+
+    for _ in range(10):
+        ray.get(f.remote())
+
+    timeseries = PrometheusTimeseries()
+
+    def verify():
+        metric = aggregate_task_event_metric(info, timeseries)
+        assert metric["REPORTED"] >= 10, (
+            "At least 10 tasks events should be reported. "
+            "Could be more than 10 with multiple flush."
+        )
+        assert metric["STORED"] == 11, "10 task + 1 driver's events should be stored."
+
+        return True
+
+    wait_for_condition(
+        verify,
+        timeout=20,
+        retry_interval_ms=100,
+    )
+
+
+def test_failed_task_error(shutdown_only):
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    # Test failed task with TASK_EXECUTION_EXCEPTION
+    error_msg_str = "fail is expected to fail"
+
+    @ray.remote
+    def fail(x=None):
+        if x is not None:
+            time.sleep(x)
+        raise ValueError(error_msg_str)
+
+    with pytest.raises(ray.exceptions.RayTaskError):
+        ray.get(fail.options(name="fail").remote())
+
+    wait_for_condition(
+        verify_failed_task,
+        name="fail",
+        error_type="TASK_EXECUTION_EXCEPTION",
+        error_message=error_msg_str,
+    )
+
+    # Test canceled tasks with TASK_CANCELLED
+    @ray.remote
+    def not_running():
+        raise ValueError("should not be run")
+
+    with pytest.raises(ray.exceptions.TaskCancelledError):
+        t = not_running.options(name="cancel-before-running").remote()
+        ray.cancel(t)
+        ray.get(t)
+
+    # Cancel task doesn't have additional error message
+    wait_for_condition(
+        verify_failed_task,
+        name="cancel-before-running",
+        error_type="TASK_CANCELLED",
+        error_message="",
+    )
+
+    # Test task failed when worker killed :WORKER_DIED
+    @ray.remote(max_retries=0)
+    def die():
+        exit(27)
+
+    with pytest.raises(ray.exceptions.WorkerCrashedError):
+        ray.get(die.options(name="die-worker").remote())
+
+    wait_for_condition(
+        verify_failed_task,
+        name="die-worker",
+        error_type="WORKER_DIED",
+        error_message="Worker exits with an exit code 27",
+    )
+
+    # Test actor task failed with actor dead: ACTOR_DIED
+    @ray.remote
+    class Actor:
+        def f(self):
+            time.sleep(999)
+
+        def ready(self):
+            pass
+
+    a = Actor.remote()
+    ray.get(a.ready.remote())
+
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.kill(a)
+        ray.get(a.f.options(name="actor-killed").remote())
+
+    wait_for_condition(
+        verify_failed_task,
+        name="actor-killed",
+        error_type="ACTOR_DIED",
+        error_message="The actor is dead because it was killed by `ray.kill`",
+    )
+
+
+def test_failed_task_failed_due_to_node_failure(ray_start_cluster):
+    cluster = ray_start_cluster
+    head_node_id = cluster.add_node(num_cpus=1).node_id
+    ray.init(address=cluster.address)
+    node = cluster.add_node(num_cpus=2)
+
+    signal = SignalActor.options(
+        label_selector={"ray.io/node-id": head_node_id}
+    ).remote()
+
+    @ray.remote(num_cpus=2, max_retries=0)
+    def sleep():
+        ray.get(signal.send.remote())
+        while True:
+            time.sleep(1)
+
+    obj_ref = sleep.options(name="node-killed").remote()
+    ray.get(signal.wait.remote())
+
+    cluster.remove_node(node)
+    with pytest.raises(ray.exceptions.NodeDiedError):
+        ray.get(obj_ref)
+
+    wait_for_condition(
+        verify_failed_task,
+        name="node-killed",
+        error_type="NODE_DIED",
+        error_message="Task failed because the node it was running on is dead or unavailable",
+    )
+
+
+def test_failed_task_unschedulable(shutdown_only):
+    ray.init(num_cpus=1, _system_config=_SYSTEM_CONFIG)
+
+    node_id = ray.get_runtime_context().get_node_id()
+    policy = ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+        node_id=node_id,
+        soft=False,
+    )
+
+    @ray.remote
+    def task():
+        pass
+
+    task.options(
+        scheduling_strategy=policy,
+        name="task-unschedulable",
+        num_cpus=2,
+    ).remote()
+
+    wait_for_condition(
+        verify_failed_task,
+        name="task-unschedulable",
+        error_type="TASK_UNSCHEDULABLE_ERROR",
+        error_message=(
+            "The node specified via NodeAffinitySchedulingStrategy"
+            " doesn't exist any more or is infeasible"
+        ),
+    )
+
+
+def test_failed_task_runtime_env_setup(shutdown_only):
+    @ray.remote
+    def f():
+        pass
+
+    bad_env = RuntimeEnv(
+        pip=["nonexistent-package"],
+    )
+    with pytest.raises(
+        RuntimeEnvSetupError,
+    ):
+        ray.get(f.options(runtime_env=bad_env, name="task-runtime-env-failed").remote())
+
+    wait_for_condition(
+        verify_failed_task,
+        name="task-runtime-env-failed",
+        error_type="RUNTIME_ENV_SETUP_FAILED",
+        # Verify that the pip error content is propagated to task event.
+        error_message="nonexistent-package",
+    )
+
+
+def test_parent_task_id_threaded_task(shutdown_only):
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    # Task starts a thread
+    @ray.remote
+    def main_task():
+        def thd_task():
+            @ray.remote
+            def thd_task():
+                pass
+
+            ray.get(thd_task.remote())
+
+        thd = threading.Thread(target=thd_task)
+        thd.start()
+        thd.join()
+
+    ray.get(main_task.remote())
+
+    def verify():
+        tasks = list_tasks()
+        assert len(tasks) == 2
+        expect_parent_task_id = None
+        actual_parent_task_id = None
+        for task in tasks:
+            if task["name"] == "main_task":
+                expect_parent_task_id = task["task_id"]
+            elif task["name"] == "thd_task":
+                actual_parent_task_id = task["parent_task_id"]
+        assert actual_parent_task_id is not None
+        assert expect_parent_task_id == actual_parent_task_id
+
+        return True
+
+    wait_for_condition(verify)
+
+
+def test_parent_task_id_non_concurrent_actor(shutdown_only):
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    def run_task_in_thread():
+        def thd_task():
+            @ray.remote
+            def thd_task():
+                pass
+
+            ray.get(thd_task.remote())
+
+        thd = threading.Thread(target=thd_task)
+        thd.start()
+        thd.join()
+
+    @ray.remote
+    class Actor:
+        def main_task(self):
+            run_task_in_thread()
+
+    a = Actor.remote()
+    ray.get(a.main_task.remote())
+
+    def verify():
+        tasks = list_tasks()
+        expect_parent_task_id = None
+        actual_parent_task_id = None
+        for task in tasks:
+            if "main_task" in task["name"]:
+                expect_parent_task_id = task["task_id"]
+            elif "thd_task" in task["name"]:
+                actual_parent_task_id = task["parent_task_id"]
+        print(tasks)
+        assert actual_parent_task_id is not None
+        assert expect_parent_task_id == actual_parent_task_id
+
+        return True
+
+    wait_for_condition(verify)
+
+
+@pytest.mark.parametrize("actor_concurrency", [3, 10])
+def test_parent_task_id_concurrent_actor(shutdown_only, actor_concurrency):
+    # Test tasks runs in user started thread from actors have a parent_task_id
+    # as the actor's creation task.
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    def run_task_in_thread(name, i):
+        def thd_task():
+            @ray.remote
+            def thd_task():
+                pass
+
+            ray.get(thd_task.options(name=f"{name}_{i}").remote())
+
+        thd = threading.Thread(target=thd_task)
+        thd.start()
+        thd.join()
+
+    @ray.remote
+    class AsyncActor:
+        async def main_task(self, i):
+            run_task_in_thread("async_thd_task", i)
+
+    @ray.remote
+    class ThreadedActor:
+        def main_task(self, i):
+            run_task_in_thread("threaded_thd_task", i)
+
+    def verify(actor_method_name, actor_class_name):
+        tasks = list_tasks()
+        print(tasks)
+        expect_parent_task_id = None
+        actual_parent_task_id = None
+        for task in tasks:
+            if f"{actor_class_name}.__init__" in task["name"]:
+                expect_parent_task_id = task["task_id"]
+
+        assert expect_parent_task_id is not None
+        for task in tasks:
+            if f"{actor_method_name}" in task["name"]:
+                actual_parent_task_id = task["parent_task_id"]
+                assert expect_parent_task_id == actual_parent_task_id, task
+
+        return True
+
+    async_actor = AsyncActor.options(max_concurrency=actor_concurrency).remote()
+    ray.get([async_actor.main_task.remote(i) for i in range(20)])
+    wait_for_condition(
+        verify, actor_class_name="AsyncActor", actor_method_name="async_thd_task"
+    )
+
+    thd_actor = ThreadedActor.options(max_concurrency=actor_concurrency).remote()
+    ray.get([thd_actor.main_task.remote(i) for i in range(20)])
+    wait_for_condition(
+        verify, actor_class_name="ThreadedActor", actor_method_name="threaded_thd_task"
+    )
+
+
+def test_is_debugger_paused(shutdown_only):
+    ray_context = ray.init(num_cpus=1, _system_config=_SYSTEM_CONFIG)
+    address = ray_context.address_info["address"]
+    node_id = ray_context.address_info["node_id"]
+    wait_for_aggregator_agent_if_enabled(address, node_id)
+
+    @ray.remote(max_retries=0)
+    def f():
+        import time
+
+        # Pause 5 seconds inside debugger
+        with ray._private.worker.global_worker.task_paused_by_debugger():
+            time.sleep(5)
+
+    def verify(num_paused):
+        tasks = list_tasks(filters=[("is_debugger_paused", "=", "True")])
+        return len(tasks) == num_paused
+
+    f_task = f.remote()  # noqa
+
+    wait_for_condition(
+        verify,
+        timeout=20,
+        retry_interval_ms=100,
+        num_paused=1,
+    )
+
+    wait_for_condition(
+        verify,
+        timeout=20,
+        retry_interval_ms=100,
+        num_paused=0,
+    )
+
+
+@pytest.mark.parametrize("actor_concurrency", [1, 3])
+def test_is_debugger_paused_actor(shutdown_only, actor_concurrency):
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    @ray.remote
+    class TestActor:
+        def main_task(self, i):
+            if i == 0:
+                import time
+
+                # Pause 5 seconds inside debugger
+                with ray._private.worker.global_worker.task_paused_by_debugger():
+                    time.sleep(5)
+
+    def verify(expected_task_name):
+        tasks = list_tasks(filters=[("is_debugger_paused", "=", "True")])
+        return len(tasks) == 1 and f"{expected_task_name}_0" in tasks[0]["name"]
+
+    test_actor = TestActor.options(max_concurrency=actor_concurrency).remote()
+    refs = [  # noqa
+        test_actor.main_task.options(name=f"TestActor.main_task_{i}").remote(i)
+        for i in range(20)
+    ]
+
+    wait_for_condition(verify, expected_task_name="TestActor.main_task")
+
+
+@pytest.mark.parametrize("actor_concurrency", [1, 3])
+def test_is_debugger_paused_threaded_actor(shutdown_only, actor_concurrency):
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    @ray.remote
+    class ThreadedActor:
+        def main_task(self, i):
+            def thd_task():
+                @ray.remote
+                def thd_task():
+                    if i == 0:
+                        import time
+
+                        # Pause 5 seconds inside debugger
+                        with ray._private.worker.global_worker.task_paused_by_debugger():  # noqa: E501
+                            time.sleep(5)
+
+                ray.get(thd_task.options(name=f"ThreadedActor.main_task_{i}").remote())
+
+            thd = threading.Thread(target=thd_task)
+            thd.start()
+            thd.join()
+
+    def verify(expected_task_name):
+        tasks = list_tasks(filters=[("is_debugger_paused", "=", "True")])
+        return len(tasks) == 1 and f"{expected_task_name}_0" in tasks[0]["name"]
+
+    threaded_actor = ThreadedActor.options(max_concurrency=actor_concurrency).remote()
+    refs = [  # noqa
+        threaded_actor.main_task.options(name=f"ThreadedActor.main_task_{i}").remote(i)
+        for i in range(20)
+    ]
+
+    wait_for_condition(verify, expected_task_name="ThreadedActor.main_task")
+
+
+@pytest.mark.parametrize("actor_concurrency", [1, 3])
+def test_is_debugger_paused_async_actor(shutdown_only, actor_concurrency):
+    ray.init(_system_config=_SYSTEM_CONFIG)
+
+    @ray.remote
+    class AsyncActor:
+        async def main_task(self, i):
+            if i == 0:
+                import time
+
+                # Pause 5 seconds inside debugger
+                print()
+                with ray._private.worker.global_worker.task_paused_by_debugger():
+                    time.sleep(5)
+
+    def verify(expected_task_name):
+        tasks = list_tasks(filters=[("is_debugger_paused", "=", "True")])
+        print(tasks)
+        return len(tasks) == 1 and f"{expected_task_name}_0" in tasks[0]["name"]
+
+    async_actor = AsyncActor.options(max_concurrency=actor_concurrency).remote()
+    refs = [  # noqa
+        async_actor.main_task.options(name=f"AsyncActor.main_task_{i}").remote(i)
+        for i in range(20)
+    ]
+
+    wait_for_condition(verify, expected_task_name="AsyncActor.main_task")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))