chore: import upstream snapshot with attribution

python/ray/serve/tests/test_enable_task_events.py

@@ -0,0 +1,87 @@
+import sys
+
+import httpx
+import pytest
+from starlette.requests import Request
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.constants import RAY_SERVE_ENABLE_TASK_EVENTS
+from ray.util.state import list_tasks
+
+
+@ray.remote
+def some_task():
+    return "hi from task"
+
+
+@serve.deployment
+class Deployment:
+    async def __call__(self, request: Request) -> str:
+        if (await request.json())["call_task"] is True:
+            msg = await some_task.remote()
+        else:
+            msg = "hi from deployment"
+
+        return msg
+
+
+@pytest.mark.skipif(RAY_SERVE_ENABLE_TASK_EVENTS, reason="Task events enabled.")
+def test_task_events_disabled_by_default(serve_instance):
+    # Deploy simple Serve app and call it over HTTP.
+    # Check that none of the Serve-managed actors or tasks generate events.
+    serve.run(Deployment.bind())
+
+    assert (
+        httpx.request("GET", "http://localhost:8000/", json={"call_task": False}).text
+        == "hi from deployment"
+    )
+    for _ in range(100):
+        assert len(list_tasks()) == 0
+
+    # Now call a Ray task from within the deployment.
+    # A task event should be generated.
+    assert (
+        httpx.request("GET", "http://localhost:8000/", json={"call_task": True}).text
+        == "hi from task"
+    )
+    wait_for_condition(lambda: len(list_tasks()) == 1)
+    [task_state] = list_tasks()
+    assert task_state.name == "some_task"
+    assert task_state.type == "NORMAL_TASK"
+
+
+@pytest.mark.skipif(not RAY_SERVE_ENABLE_TASK_EVENTS, reason="Task events disabled.")
+def test_enable_task_events(serve_instance):
+    # Deploy simple Serve app and call it over HTTP.
+    # Check that the Serve-managed actors generate events.
+    serve.run(Deployment.bind())
+
+    assert (
+        httpx.request("GET", "http://localhost:8000/", json={"call_task": False}).text
+        == "hi from deployment"
+    )
+
+    def check_for_expected_actor_tasks():
+        found_proxy = False
+        found_replica = False
+        found_controller = False
+        for task_state in list_tasks():
+            if "ProxyActor" in task_state.name:
+                found_proxy = True
+            if "ServeReplica" in task_state.name:
+                found_replica = True
+            if "ServeController" in task_state.name:
+                found_controller = True
+
+        assert found_proxy, "No proxy tasks found."
+        assert found_replica, "No replica tasks found."
+        assert found_controller, "No controller tasks found."
+        return True
+
+    wait_for_condition(check_for_expected_actor_tasks, timeout=5)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))