chore: import upstream snapshot with attribution

python/ray/tests/test_debug_tools.py

@@ -0,0 +1,209 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+import ray
+import ray._private.ray_constants as ray_constants
+import ray._private.services as services
+from ray._common.test_utils import wait_for_condition
+
+
+@pytest.fixture
+def ray_gdb_start():
+    # Setup environment and start ray
+    _environ = os.environ.copy()
+    for process_name in ["RAYLET", "PLASMA_STORE"]:
+        os.environ["RAY_{}_GDB".format(process_name)] = "1"
+        os.environ["RAY_{}_TMUX".format(process_name)] = "1"
+
+    yield None
+
+    # Restore original environment and stop ray
+    os.environ.clear()
+    os.environ.update(_environ)
+    ray.shutdown()
+
+
+@pytest.mark.skipif(
+    sys.platform != "linux" and sys.platform != "linux2",
+    reason="This test requires Linux.",
+)
+def test_raylet_gdb(ray_gdb_start):
+    # ray_gdb_start yields the expected process name
+    ray.init(num_cpus=1)
+
+    @ray.remote
+    def f():
+        return 42
+
+    assert ray.get(f.remote()) == 42
+
+    # Check process name in `ps aux | grep gdb`
+    pgrep_command = subprocess.Popen(
+        ["pgrep", "-f", "gdb.*raylet/raylet"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    assert pgrep_command.communicate()[0]
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="memray not supported in win32")
+@pytest.mark.skipif(sys.platform == "darwin", reason="memray not supported in Darwin")
+def test_memory_profiler_command_builder(monkeypatch, tmp_path):
+    session_dir = tmp_path
+    # When there's no env var, command should be just a regular python command.
+    command = services._build_python_executable_command_memory_profileable(
+        ray_constants.PROCESS_TYPE_DASHBOARD, session_dir
+    )
+    assert command == [sys.executable, "-u"]
+
+    with monkeypatch.context() as m:
+        m.setenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV, "dashboard")
+        m.setenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV, "-q")
+        command = services._build_python_executable_command_memory_profileable(
+            ray_constants.PROCESS_TYPE_DASHBOARD, session_dir
+        )
+
+        assert command == [
+            sys.executable,
+            "-u",
+            "-m",
+            "memray",
+            "run",
+            "-o",
+            str(
+                Path(tmp_path)
+                / "profile"
+                / f"{Path(tmp_path).name}_memory_dashboard.bin"
+            ),  # noqa
+            "-q",
+        ]
+
+        # Test with explicit -o path
+        m.delenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV)
+        m.delenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV)
+        m.setenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV, "dashboard")
+        m.setenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV, "-o,/custom/path.bin,-q")
+        command = services._build_python_executable_command_memory_profileable(
+            ray_constants.PROCESS_TYPE_DASHBOARD, session_dir
+        )
+        assert command == [
+            sys.executable,
+            "-u",
+            "-m",
+            "memray",
+            "run",
+            "-o",
+            "/custom/path.bin",
+            "-q",
+        ]
+
+        # Test with explicit --output path
+        m.delenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV)
+        m.delenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV)
+        m.setenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV, "dashboard")
+        m.setenv(
+            services.RAY_MEMRAY_PROFILE_OPTIONS_ENV, "--output,/custom/path.bin,-q"
+        )
+        command = services._build_python_executable_command_memory_profileable(
+            ray_constants.PROCESS_TYPE_DASHBOARD, session_dir
+        )
+        assert command == [
+            sys.executable,
+            "-u",
+            "-m",
+            "memray",
+            "run",
+            "--output",
+            "/custom/path.bin",
+            "-q",
+        ]
+
+        m.delenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV)
+        m.delenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV)
+        m.setenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV, "dashboard,dashboard_agent")
+        m.setenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV, "-q,--live,--live-port,1234")
+        command = services._build_python_executable_command_memory_profileable(
+            ray_constants.PROCESS_TYPE_DASHBOARD_AGENT, session_dir
+        )
+        assert command == [
+            sys.executable,
+            "-u",
+            "-m",
+            "memray",
+            "run",
+            "-q",
+            "--live",
+            "--live-port",
+            "1234",
+        ]
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="memray not supported in win32")
+@pytest.mark.skipif(sys.platform == "darwin", reason="memray not supported in Darwin")
+def test_memory_profile_dashboard_and_agent(monkeypatch, shutdown_only):
+    with monkeypatch.context() as m:
+        m.setenv(services.RAY_MEMRAY_PROFILE_COMPONENT_ENV, "dashboard,dashboard_agent")
+        m.setenv(services.RAY_MEMRAY_PROFILE_OPTIONS_ENV, "-q")
+        addr = ray.init()
+
+        def verify():
+            session_dir = Path(addr["session_dir"])
+            profile_dir = session_dir / "profile"
+            assert profile_dir.exists()
+            files = []
+            for f in profile_dir.iterdir():
+                files.append(f.name)
+            assert len(files) == 2
+            assert f"{session_dir.name}_memory_dashboard.bin" in files
+            assert f"{session_dir.name}_memory_dashboard_agent.bin" in files
+            return True
+
+        wait_for_condition(verify)
+
+
+def test_start_ray_client_server_redis_password_env_updates(monkeypatch):
+    captured = {}
+    expected_process_info = object()
+
+    def fake_start_ray_process(command, process_type, **kwargs):
+        captured["command"] = command
+        captured["process_type"] = process_type
+        captured["kwargs"] = kwargs
+        return expected_process_info
+
+    with monkeypatch.context() as m:
+        m.setattr(services, "start_ray_process", fake_start_ray_process)
+        m.delenv(ray_constants.RAY_REDIS_PASSWORD_ENV, raising=False)
+
+        process_info = services.start_ray_client_server(
+            address="127.0.0.1:6379",
+            ray_client_server_ip="127.0.0.1",
+            ray_client_server_port=10001,
+            redis_username="redis-user",
+            redis_password="secret123",
+            fate_share=False,
+            runtime_env_agent_address="127.0.0.1:12345",
+            node_id="node-1",
+        )
+
+        assert process_info is expected_process_info
+        assert captured["process_type"] == ray_constants.PROCESS_TYPE_RAY_CLIENT_SERVER
+        assert "--redis-username=redis-user" in captured["command"]
+        assert not any(
+            arg.startswith("--redis-password=") for arg in captured["command"]
+        )
+        assert captured["kwargs"]["env_updates"] == {
+            ray_constants.RAY_REDIS_PASSWORD_ENV: "secret123"
+        }
+        assert ray_constants.RAY_REDIS_PASSWORD_ENV not in os.environ
+
+
+if __name__ == "__main__":
+    # Make subprocess happy in bazel.
+    os.environ["LC_ALL"] = "en_US.UTF-8"
+    os.environ["LANG"] = "en_US.UTF-8"
+    sys.exit(pytest.main(["-sv", __file__]))