chore: import upstream snapshot with attribution

python/ray/tests/test_dashboard_profiler.py

@@ -0,0 +1,315 @@
+import os
+import subprocess
+import sys
+
+import pytest
+import requests
+
+import ray
+from ray._private.test_utils import (
+    format_web_url,
+    wait_until_server_available,
+    wait_until_succeeded_without_exception,
+)
+
+
+@pytest.fixture(autouse=True)
+def enable_profiling():
+    os.environ["RAY_DASHBOARD_ENABLE_PROFILING"] = "1"
+    yield
+    os.environ.pop("RAY_DASHBOARD_ENABLE_PROFILING", None)
+
+
+@pytest.mark.skipif(
+    os.environ.get("RAY_MINIMAL") == "1",
+    reason="This test is not supposed to work for minimal installation.",
+)
+@pytest.mark.skipif(sys.platform == "win32", reason="No py-spy on Windows.")
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="No py-spy on python 3.12.")
+@pytest.mark.skipif(
+    sys.platform == "darwin",
+    reason="Fails on OSX: https://github.com/ray-project/ray/issues/30114",
+)
+@pytest.mark.parametrize("native", ["0", "1"])
+@pytest.mark.parametrize("node_info", ["node_id", "ip"])
+def test_profiler_endpoints(ray_start_with_dashboard, native, node_info):
+    if native == "1" and sys.platform == "linux":
+        pytest.skip(
+            "py-spy --native 'failed to get os threadid' "
+            "(https://github.com/ray-project/ray/issues/30566); "
+            "see also https://github.com/benfred/py-spy/issues/490; "
+            "disabled as part of Ubuntu 22.04 upgrade"
+        )
+    # Sanity check py-spy are installed.
+    subprocess.check_call(["py-spy", "--version"])
+
+    assert wait_until_server_available(ray_start_with_dashboard["webui_url"]) is True
+    address_info = ray_start_with_dashboard
+    webui_url = address_info["webui_url"]
+    webui_url = format_web_url(webui_url)
+
+    @ray.remote
+    class Actor:
+        def getpid(self):
+            return os.getpid()
+
+        def do_stuff_infinite(self):
+            while True:
+                pass
+
+    a = Actor.remote()
+    pid = ray.get(a.getpid.remote())
+    a.do_stuff_infinite.remote()
+
+    node_id = ray_start_with_dashboard.address_info["node_id"]
+    node_ip = ray_start_with_dashboard.address_info["node_ip_address"]
+
+    def get_node_info():
+        if node_info == "node_id":
+            return f"node_id={node_id}"
+        else:
+            return f"ip={node_ip}"
+
+    def get_actor_stack():
+        url = (
+            f"{webui_url}/worker/traceback?pid={pid}&{get_node_info()}&native={native}"
+        )
+        print("GET URL", url)
+        response = requests.get(url)
+        print("STATUS CODE", response.status_code)
+        print("HEADERS", response.headers)
+        content = response.content.decode("utf-8")
+        print("CONTENT", content)
+        response.raise_for_status()
+        assert "text/plain" in response.headers["Content-Type"], response.headers
+        # Sanity check we got the stack trace text.
+        assert "do_stuff_infinite" in content, content
+        if native == "1":
+            assert "ray::core::CoreWorker" in content, content
+        else:
+            assert "ray::core::CoreWorker" not in content, content
+
+    assert wait_until_succeeded_without_exception(
+        get_actor_stack,
+        (requests.RequestException, AssertionError),
+        timeout_ms=20000,
+        retry_interval_ms=1000,
+    )
+
+    def get_actor_flamegraph():
+        response = requests.get(
+            f"{webui_url}/worker/cpu_profile?pid={pid}&{get_node_info()}&native={native}"
+        )
+        response.raise_for_status()
+        assert response.headers["Content-Type"] == "image/svg+xml", response.headers
+        content = response.content.decode("utf-8")
+        print(content)
+        # Sanity check we got the flame graph SVG.
+        assert "<!DOCTYPE svg" in content, content
+        assert "do_stuff_infinite" in content, content
+        if native == "1":
+            assert "ray::core" in content, content
+        else:
+            assert "ray::core" not in content, content
+
+    assert wait_until_succeeded_without_exception(
+        get_actor_flamegraph,
+        (requests.RequestException, AssertionError),
+        timeout_ms=20000,
+        retry_interval_ms=1000,
+    )
+
+
+@pytest.mark.skipif(
+    os.environ.get("RAY_MINIMAL") == "1",
+    reason="This test is not supposed to work for minimal installation.",
+)
+@pytest.mark.skipif(sys.platform == "win32", reason="No memray on Windows.")
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="No memray on python 3.12")
+@pytest.mark.skipif(
+    sys.platform == "darwin",
+    reason="Fails on OSX, requires memray & lldb installed in osx image",
+)
+@pytest.mark.parametrize("leaks", ["0", "1"])
+@pytest.mark.parametrize("node_info", ["node_id", "ip"])
+def test_memory_profiler_endpoint(ray_start_with_dashboard, leaks, node_info):
+    # Sanity check memray are installed.
+    subprocess.check_call(["memray", "--version"])
+
+    assert wait_until_server_available(ray_start_with_dashboard["webui_url"]) is True
+    address_info = ray_start_with_dashboard
+    webui_url = address_info["webui_url"]
+    webui_url = format_web_url(webui_url)
+
+    @ray.remote
+    class Actor:
+        def getpid(self):
+            return os.getpid()
+
+        def do_stuff_infinite(self):
+            while True:
+                pass
+
+    a = Actor.remote()
+    pid = ray.get(a.getpid.remote())
+    a.do_stuff_infinite.remote()
+
+    node_id = ray_start_with_dashboard.address_info["node_id"]
+    node_ip = ray_start_with_dashboard.address_info["node_ip_address"]
+
+    def get_node_info():
+        if node_info == "node_id":
+            return f"node_id={node_id}"
+        else:
+            return f"ip={node_ip}"
+
+    def get_actor_memory_flamegraph():
+        response = requests.get(
+            f"{webui_url}/memory_profile?pid={pid}&{get_node_info()}&leaks={leaks}&duration=5"
+        )
+        response.raise_for_status()
+
+        assert response.headers["Content-Type"] == "text/html", response.headers
+        content = response.content.decode("utf-8")
+        print(content)
+        assert "memray" in content, content
+
+        if leaks == "1":
+            assert "flamegraph report (memory leaks)" in content, content
+        else:
+            assert "flamegraph report" in content, content
+
+    assert wait_until_succeeded_without_exception(
+        get_actor_memory_flamegraph,
+        (requests.RequestException, AssertionError),
+        timeout_ms=20000,
+        retry_interval_ms=1000,
+        raise_last_ex=True,
+    )
+
+    def get_actor_memory_multiple_flamegraphs():
+        response = requests.get(
+            f"{webui_url}/memory_profile?pid={pid}&{get_node_info()}&leaks={leaks}&duration=5"
+        )
+        response.raise_for_status()
+
+        assert response.headers["Content-Type"] == "text/html", response.headers
+        content = response.content.decode("utf-8")
+        print(content)
+        assert "memray" in content, content
+
+        if leaks == "1":
+            assert "flamegraph report (memory leaks)" in content, content
+        else:
+            assert "flamegraph report" in content, content
+
+    assert wait_until_succeeded_without_exception(
+        get_actor_memory_multiple_flamegraphs,
+        (requests.RequestException, AssertionError),
+        timeout_ms=20000,
+        retry_interval_ms=1000,
+        raise_last_ex=True,
+    )
+
+
+@pytest.mark.skipif(
+    os.environ.get("RAY_MINIMAL") == "1",
+    reason="This test is not supposed to work for minimal installation.",
+)
+@pytest.mark.skipif(sys.platform == "win32", reason="No py-spy on Windows.")
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="No py-spy on python 3.12.")
+@pytest.mark.skipif(
+    sys.platform == "darwin",
+    reason="Fails on OSX, requires memray & lldb installed in osx image",
+)
+@pytest.mark.parametrize("node_info", ["node_id", "ip"])
+def test_profiler_failure_message(ray_start_with_dashboard, node_info):
+    # Sanity check py-spy and memray is installed.
+    subprocess.check_call(["py-spy", "--version"])
+    subprocess.check_call(["memray", "--version"])
+
+    assert wait_until_server_available(ray_start_with_dashboard["webui_url"]) is True
+    address_info = ray_start_with_dashboard
+    webui_url = address_info["webui_url"]
+    webui_url = format_web_url(webui_url)
+
+    @ray.remote
+    class Actor:
+        def getpid(self):
+            return os.getpid()
+
+        def do_stuff_infinite(self):
+            while True:
+                pass
+
+    a = Actor.remote()
+    pid = ray.get(a.getpid.remote())
+    a.do_stuff_infinite.remote()
+
+    node_id = ray_start_with_dashboard.address_info["node_id"]
+    node_ip = ray_start_with_dashboard.address_info["node_ip_address"]
+
+    def get_node_info():
+        if node_info == "node_id":
+            return f"node_id={node_id}"
+        else:
+            return f"ip={node_ip}"
+
+    def get_actor_stack():
+        response = requests.get(
+            f"{webui_url}/worker/traceback?pid={pid}&{get_node_info()}"
+        )
+        response.raise_for_status()
+        content = response.content.decode("utf-8")
+        print("CONTENT", content)
+        assert "do_stuff_infinite" in content, content
+
+    # First ensure dashboard is up, before we test for failure cases.
+    assert wait_until_succeeded_without_exception(
+        get_actor_stack,
+        (requests.RequestException, AssertionError),
+        timeout_ms=20000,
+        retry_interval_ms=1000,
+    )
+
+    # Check we return the right status code and error message on failure.
+    response = requests.get(
+        f"{webui_url}/worker/traceback?pid=1234567&{get_node_info()}"
+    )
+    content = response.content.decode("utf-8")
+    print(content)
+    assert "text/plain" in response.headers["Content-Type"], response.headers
+    assert "Failed to execute" in content, content
+
+    # Check we return the right status code and error message on failure.
+    response = requests.get(
+        f"{webui_url}/worker/cpu_profile?pid=1234567&{get_node_info()}"
+    )
+    content = response.content.decode("utf-8")
+    print(content)
+    assert "text/plain" in response.headers["Content-Type"], response.headers
+    assert "Failed to execute" in content, content
+
+    # Check we return the right status code and error message on failure.
+    response = requests.get(f"{webui_url}/memory_profile?pid=1234567&{get_node_info()}")
+    content = response.content.decode("utf-8")
+    print(content)
+    assert "text/plain" in response.headers["Content-Type"], response.headers
+    assert "Failed to execute" in content, content
+
+    # Check wrong ID/ip failure
+    if node_info == "node_id":
+        wrong_param = "node_id=DUMMY_ID"
+        expect_msg = "Failed to execute: no agent address found for node DUMMY_ID"
+    else:
+        wrong_param = "ip=1.2.3.4"
+        expect_msg = "Failed to execute: no agent address found for node IP 1.2.3.4"
+
+    response = requests.get(f"{webui_url}/memory_profile?pid=1234567&{wrong_param}")
+    content = response.content.decode("utf-8")
+    print(content)
+    assert expect_msg in content, content
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))