chore: import upstream snapshot with attribution

python/ray/dashboard/modules/reporter/tests/test_gpu_profiler_manager.py

@@ -0,0 +1,377 @@
+"""Unit tests for the GPU profiler manager.
+
+All GPU and dynolog dependencies are mocked out.
+This test just verifies that commands are launched correctly and that
+validations are correctly performed.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from ray.dashboard.modules.reporter.gpu_profile_manager import GpuProfilingManager
+
+
+@pytest.fixture
+def mock_node_has_gpus(monkeypatch):
+    monkeypatch.setattr(GpuProfilingManager, "node_has_gpus", lambda cls: True)
+    yield
+
+
+@pytest.fixture
+def mock_dynolog_binaries(monkeypatch):
+    monkeypatch.setattr("shutil.which", lambda cmd: f"/usr/bin/fake_{cmd}")
+    yield
+
+
+@pytest.fixture
+def mock_subprocess_popen(monkeypatch):
+    mock_popen = MagicMock()
+    mock_proc = MagicMock()
+    mock_popen.return_value = mock_proc
+
+    monkeypatch.setattr("subprocess.Popen", mock_popen)
+    yield (mock_popen, mock_proc)
+
+
+LOCALHOST = "127.0.0.1"
+
+
+@pytest.fixture
+def mock_asyncio_create_subprocess_exec(monkeypatch):
+    mock_create_subprocess_exec = AsyncMock()
+    mock_async_proc = mock_create_subprocess_exec.return_value = AsyncMock()
+    mock_async_proc.communicate.return_value = b"mock stdout", b"mock stderr"
+    mock_async_proc.returncode = 0
+    monkeypatch.setattr("asyncio.create_subprocess_exec", mock_create_subprocess_exec)
+    yield (mock_create_subprocess_exec, mock_async_proc)
+
+
+def test_node_has_gpus_uses_query_gpu_flag(tmp_path, monkeypatch):
+    """node_has_gpus() must use --query-gpu to avoid FabricManager stalls."""
+    captured = {}
+
+    def fake_check_output(cmd, **kwargs):
+        captured["cmd"] = list(cmd)
+        return b""
+
+    monkeypatch.setattr("subprocess.check_output", fake_check_output)
+    # Clear the cache so the monkeypatched check_output is actually called.
+    GpuProfilingManager.node_has_gpus.cache_clear()
+    result = GpuProfilingManager.node_has_gpus()
+    GpuProfilingManager.node_has_gpus.cache_clear()
+
+    assert result is True
+    assert captured["cmd"] == [
+        "nvidia-smi",
+        "--query-gpu=name",
+        "--format=csv,noheader",
+    ]
+
+
+def test_enabled_does_not_call_node_has_gpus_when_dynolog_missing(
+    tmp_path, monkeypatch
+):
+    """enabled must short-circuit on missing dynolog bins before node_has_gpus()."""
+    node_has_gpus_called = []
+
+    def spy_node_has_gpus(self_or_cls):
+        node_has_gpus_called.append(True)
+        return True
+
+    monkeypatch.setattr(GpuProfilingManager, "node_has_gpus", spy_node_has_gpus)
+    # No dynolog binaries on PATH → shutil.which returns None for both.
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    assert not gpu_profiler.enabled
+    assert (
+        not node_has_gpus_called
+    ), "node_has_gpus() must not be called when dynolog binaries are absent"
+
+
+def test_enabled(tmp_path, mock_node_has_gpus, mock_dynolog_binaries):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    assert gpu_profiler.enabled
+
+
+def test_disabled_no_gpus(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        GpuProfilingManager, "node_has_gpus", classmethod(lambda cls: False)
+    )
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    assert not gpu_profiler.enabled
+
+
+def test_disabled_no_dynolog_bin(tmp_path, mock_node_has_gpus):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    assert not gpu_profiler.enabled
+
+
+def test_start_monitoring_daemon(
+    tmp_path, mock_node_has_gpus, mock_dynolog_binaries, mock_subprocess_popen
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+
+    mocked_popen, mocked_proc = mock_subprocess_popen
+    mocked_proc.pid = 123
+    mocked_proc.poll.return_value = None
+
+    gpu_profiler.start_monitoring_daemon()
+    assert gpu_profiler.is_monitoring_daemon_running
+
+    assert mocked_popen.call_count == 1
+    assert mocked_popen.call_args[0][0] == [
+        "/usr/bin/fake_dynolog",
+        "--enable_ipc_monitor",
+        "--port",
+        str(gpu_profiler._DYNOLOG_PORT),
+    ]
+
+    # "Terminate" the daemon
+    mocked_proc.poll.return_value = 0
+    assert not gpu_profiler.is_monitoring_daemon_running
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_disabled(tmp_path):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    assert not gpu_profiler.enabled
+
+    success, output = await gpu_profiler.gpu_profile(pid=123, num_iterations=1)
+    assert not success
+    assert output == gpu_profiler._DISABLED_ERROR_MESSAGE.format(
+        ip_address=gpu_profiler._ip_address
+    )
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_without_starting_daemon(
+    tmp_path, mock_node_has_gpus, mock_dynolog_binaries
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    assert not gpu_profiler.is_monitoring_daemon_running
+
+    with pytest.raises(RuntimeError, match="start_monitoring_daemon"):
+        await gpu_profiler.gpu_profile(pid=123, num_iterations=1)
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_with_dead_daemon(
+    tmp_path, mock_node_has_gpus, mock_dynolog_binaries, mock_subprocess_popen
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    gpu_profiler.start_monitoring_daemon()
+
+    mocked_popen, mocked_proc = mock_subprocess_popen
+    mocked_proc.pid = 123
+    # "Terminate" the daemon
+    mocked_proc.poll.return_value = 0
+    assert not gpu_profiler.is_monitoring_daemon_running
+
+    success, output = await gpu_profiler.gpu_profile(pid=456, num_iterations=1)
+    assert not success
+    print(output)
+    assert "GPU monitoring daemon" in output
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_on_dead_process(
+    tmp_path,
+    monkeypatch,
+    mock_node_has_gpus,
+    mock_dynolog_binaries,
+    mock_subprocess_popen,
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    gpu_profiler.start_monitoring_daemon()
+
+    _, mocked_proc = mock_subprocess_popen
+    mocked_proc.pid = 123
+    mocked_proc.poll.return_value = None
+
+    monkeypatch.setattr(GpuProfilingManager, "is_pid_alive", lambda cls, pid: False)
+
+    success, output = await gpu_profiler.gpu_profile(pid=456, num_iterations=1)
+    assert not success
+    assert output == gpu_profiler._DEAD_PROCESS_ERROR_MESSAGE.format(
+        pid=456, ip_address=gpu_profiler._ip_address
+    )
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_no_matched_processes(
+    tmp_path,
+    monkeypatch,
+    mock_node_has_gpus,
+    mock_dynolog_binaries,
+    mock_subprocess_popen,
+    mock_asyncio_create_subprocess_exec,
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    gpu_profiler.start_monitoring_daemon()
+
+    # Mock the daemon process
+    _, mocked_daemon_proc = mock_subprocess_popen
+    mocked_daemon_proc.pid = 123
+    mocked_daemon_proc.poll.return_value = None
+
+    monkeypatch.setattr(GpuProfilingManager, "is_pid_alive", lambda cls, pid: True)
+
+    # Mock the asyncio.create_subprocess_exec
+    (
+        mocked_create_subprocess_exec,
+        mocked_async_proc,
+    ) = mock_asyncio_create_subprocess_exec
+    mocked_async_proc.communicate.return_value = (
+        f"{gpu_profiler._NO_PROCESSES_MATCHED_ERROR_MESSAGE_PREFIX}".encode(),
+        b"dummy stderr",
+    )
+    process_pid = 456
+    num_iterations = 1
+    success, output = await gpu_profiler.gpu_profile(
+        pid=process_pid, num_iterations=num_iterations
+    )
+
+    assert mocked_create_subprocess_exec.call_count == 1
+
+    assert not success
+    assert output == gpu_profiler._NO_PROCESSES_MATCHED_ERROR_MESSAGE.format(
+        pid=process_pid, ip_address=gpu_profiler._ip_address
+    )
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_timeout(
+    tmp_path,
+    monkeypatch,
+    mock_node_has_gpus,
+    mock_dynolog_binaries,
+    mock_subprocess_popen,
+    mock_asyncio_create_subprocess_exec,
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    gpu_profiler.start_monitoring_daemon()
+
+    # Mock the daemon process
+    _, mocked_daemon_proc = mock_subprocess_popen
+    mocked_daemon_proc.pid = 123
+    mocked_daemon_proc.poll.return_value = None
+
+    monkeypatch.setattr(GpuProfilingManager, "is_pid_alive", lambda cls, pid: True)
+
+    process_pid = 456
+    num_iterations = 1
+    task = asyncio.create_task(
+        gpu_profiler.gpu_profile(
+            pid=process_pid, num_iterations=num_iterations, _timeout_s=0.1
+        )
+    )
+
+    await asyncio.sleep(0.2)
+    success, output = await task
+    assert not success
+    assert "timed out" in output
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_process_dies_during_profiling(
+    tmp_path,
+    monkeypatch,
+    mock_node_has_gpus,
+    mock_dynolog_binaries,
+    mock_subprocess_popen,
+    mock_asyncio_create_subprocess_exec,
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    gpu_profiler.start_monitoring_daemon()
+
+    # Mock the daemon process
+    _, mocked_daemon_proc = mock_subprocess_popen
+    mocked_daemon_proc.pid = 123
+    mocked_daemon_proc.poll.return_value = None
+
+    monkeypatch.setattr(GpuProfilingManager, "is_pid_alive", lambda cls, pid: True)
+
+    process_pid = 456
+    num_iterations = 1
+    task = asyncio.create_task(
+        gpu_profiler.gpu_profile(pid=process_pid, num_iterations=num_iterations)
+    )
+
+    monkeypatch.setattr(GpuProfilingManager, "is_pid_alive", lambda cls, pid: False)
+
+    await asyncio.sleep(0.2)
+
+    success, output = await task
+    assert not success
+    assert output == gpu_profiler._DEAD_PROCESS_ERROR_MESSAGE.format(
+        pid=process_pid, ip_address=gpu_profiler._ip_address
+    )
+
+
+@pytest.mark.asyncio
+async def test_gpu_profile_success(
+    tmp_path,
+    monkeypatch,
+    mock_node_has_gpus,
+    mock_dynolog_binaries,
+    mock_subprocess_popen,
+    mock_asyncio_create_subprocess_exec,
+):
+    gpu_profiler = GpuProfilingManager(tmp_path, ip_address=LOCALHOST)
+    gpu_profiler.start_monitoring_daemon()
+
+    # Mock the daemon process
+    _, mocked_daemon_proc = mock_subprocess_popen
+    mocked_daemon_proc.pid = 123
+    mocked_daemon_proc.poll.return_value = None
+
+    monkeypatch.setattr(GpuProfilingManager, "is_pid_alive", lambda cls, pid: True)
+    monkeypatch.setattr(
+        GpuProfilingManager, "_get_trace_filename", lambda cls: "dummy_trace.json"
+    )
+    dumped_trace_filepath = gpu_profiler._profile_dir_path / "dummy_trace.json"
+    dumped_trace_filepath.touch()
+
+    # Mock the asyncio.create_subprocess_exec
+    (
+        mocked_create_subprocess_exec,
+        mocked_async_proc,
+    ) = mock_asyncio_create_subprocess_exec
+    process_pid = 456
+    num_iterations = 1
+    success, output = await gpu_profiler.gpu_profile(
+        pid=process_pid, num_iterations=num_iterations
+    )
+
+    # Verify the command was launched correctly
+    assert mocked_create_subprocess_exec.call_count == 1
+    profile_launch_args = list(mocked_create_subprocess_exec.call_args[0])
+    assert profile_launch_args[:6] == [
+        "/usr/bin/fake_dyno",
+        "--port",
+        str(gpu_profiler._DYNOLOG_PORT),
+        "gputrace",
+        "--pids",
+        str(process_pid),
+    ]
+
+    assert "--log-file" in profile_launch_args
+    profile_log_file_arg = profile_launch_args[
+        profile_launch_args.index("--log-file") + 1
+    ]
+    assert Path(profile_log_file_arg).is_relative_to(tmp_path)
+
+    assert "--iterations" in profile_launch_args
+    assert profile_launch_args[profile_launch_args.index("--iterations") + 1] == str(
+        num_iterations
+    )
+
+    assert success
+    assert output == str(dumped_trace_filepath.relative_to(tmp_path))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dashboard/modules/reporter/tests/test_gpu_providers.py

@@ -0,0 +1,629 @@
+"""Unit tests for GPU providers."""
+
+import unittest
+from unittest.mock import Mock, patch
+
+from ray.dashboard.modules.reporter.gpu_providers import (
+    MB,
+    AmdGpuProvider,
+    GpuMetricProvider,
+    GpuProvider,
+    GpuProviderType,
+    GpuUtilizationInfo,
+    NvidiaGpuProvider,
+    ProcessGPUInfo,
+)
+
+
+class TestProcessGPUInfo(unittest.TestCase):
+    """Test ProcessGPUInfo TypedDict."""
+
+    def test_creation(self):
+        """Test ProcessGPUInfo creation."""
+        process_info = ProcessGPUInfo(
+            pid=1234, gpu_memory_usage=256, gpu_utilization=None
+        )
+
+        self.assertEqual(process_info["pid"], 1234)
+        self.assertEqual(process_info["gpu_memory_usage"], 256)
+        self.assertIsNone(process_info["gpu_utilization"])
+
+
+class TestGpuUtilizationInfo(unittest.TestCase):
+    """Test GpuUtilizationInfo TypedDict."""
+
+    def test_creation_with_processes(self):
+        """Test GpuUtilizationInfo with process information."""
+        process1 = ProcessGPUInfo(pid=1234, gpu_memory_usage=256, gpu_utilization=None)
+        process2 = ProcessGPUInfo(pid=5678, gpu_memory_usage=512, gpu_utilization=None)
+
+        gpu_info = GpuUtilizationInfo(
+            index=0,
+            name="NVIDIA GeForce RTX 3080",
+            uuid="GPU-12345678-1234-1234-1234-123456789abc",
+            utilization_gpu=75,
+            memory_used=8192,
+            memory_total=10240,
+            processes_pids={1234: process1, 5678: process2},
+        )
+
+        self.assertEqual(gpu_info["index"], 0)
+        self.assertEqual(gpu_info["name"], "NVIDIA GeForce RTX 3080")
+        self.assertEqual(gpu_info["uuid"], "GPU-12345678-1234-1234-1234-123456789abc")
+        self.assertEqual(gpu_info["utilization_gpu"], 75)
+        self.assertEqual(gpu_info["memory_used"], 8192)
+        self.assertEqual(gpu_info["memory_total"], 10240)
+        self.assertEqual(len(gpu_info["processes_pids"]), 2)
+        self.assertIn(1234, gpu_info["processes_pids"])
+        self.assertIn(5678, gpu_info["processes_pids"])
+        self.assertEqual(gpu_info["processes_pids"][1234]["pid"], 1234)
+        self.assertEqual(gpu_info["processes_pids"][1234]["gpu_memory_usage"], 256)
+        self.assertEqual(gpu_info["processes_pids"][5678]["pid"], 5678)
+        self.assertEqual(gpu_info["processes_pids"][5678]["gpu_memory_usage"], 512)
+
+    def test_creation_without_processes(self):
+        """Test GpuUtilizationInfo without process information."""
+        gpu_info = GpuUtilizationInfo(
+            index=1,
+            name="AMD Radeon RX 6800 XT",
+            uuid="GPU-87654321-4321-4321-4321-ba9876543210",
+            utilization_gpu=None,
+            memory_used=4096,
+            memory_total=16384,
+            processes_pids=None,
+        )
+
+        self.assertEqual(gpu_info["index"], 1)
+        self.assertEqual(gpu_info["name"], "AMD Radeon RX 6800 XT")
+        self.assertEqual(gpu_info["uuid"], "GPU-87654321-4321-4321-4321-ba9876543210")
+        self.assertIsNone(gpu_info["utilization_gpu"])  # Should be None, not -1
+        self.assertEqual(gpu_info["memory_used"], 4096)
+        self.assertEqual(gpu_info["memory_total"], 16384)
+        self.assertIsNone(gpu_info["processes_pids"])  # Should be None, not []
+
+
+class TestGpuProvider(unittest.TestCase):
+    """Test abstract GpuProvider class."""
+
+    def test_decode_bytes(self):
+        """Test _decode method with bytes input."""
+        result = GpuProvider._decode(b"test string")
+        self.assertEqual(result, "test string")
+
+    def test_decode_string(self):
+        """Test _decode method with string input."""
+        result = GpuProvider._decode("test string")
+        self.assertEqual(result, "test string")
+
+    def test_abstract_methods_not_implemented(self):
+        """Test that abstract methods raise NotImplementedError."""
+
+        class IncompleteProvider(GpuProvider):
+            pass
+
+        with self.assertRaises(TypeError):
+            IncompleteProvider()
+
+
+class TestNvidiaGpuProvider(unittest.TestCase):
+    """Test NvidiaGpuProvider class."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.provider = NvidiaGpuProvider()
+
+    def test_get_provider_name(self):
+        """Test provider name."""
+        self.assertEqual(self.provider.get_provider_name(), GpuProviderType.NVIDIA)
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_is_available_success(self, mock_pynvml):
+        """Test is_available when NVIDIA GPU is available."""
+        mock_pynvml.nvmlInit.return_value = None
+        mock_pynvml.nvmlShutdown.return_value = None
+
+        # Mock sys.modules to make the import work
+        import sys
+
+        original_modules = sys.modules.copy()
+        sys.modules["ray._private.thirdparty.pynvml"] = mock_pynvml
+
+        try:
+            self.assertTrue(self.provider.is_available())
+            mock_pynvml.nvmlInit.assert_called_once()
+            mock_pynvml.nvmlShutdown.assert_called_once()
+        finally:
+            # Restore original modules
+            sys.modules.clear()
+            sys.modules.update(original_modules)
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_is_available_failure(self, mock_pynvml):
+        """Test is_available when NVIDIA GPU is not available."""
+        mock_pynvml.nvmlInit.side_effect = Exception("NVIDIA driver not found")
+
+        # Mock sys.modules to make the import work but nvmlInit fail
+        import sys
+
+        original_modules = sys.modules.copy()
+        sys.modules["ray._private.thirdparty.pynvml"] = mock_pynvml
+
+        try:
+            self.assertFalse(self.provider.is_available())
+        finally:
+            # Restore original modules
+            sys.modules.clear()
+            sys.modules.update(original_modules)
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_initialize_success(self, mock_pynvml):
+        """Test successful initialization."""
+        # Ensure provider starts fresh
+        self.provider._initialized = False
+
+        mock_pynvml.nvmlInit.return_value = None
+
+        # Mock sys.modules to make the import work
+        import sys
+
+        original_modules = sys.modules.copy()
+        sys.modules["ray._private.thirdparty.pynvml"] = mock_pynvml
+
+        try:
+            self.assertTrue(self.provider._initialize())
+            self.assertTrue(self.provider._initialized)
+            mock_pynvml.nvmlInit.assert_called_once()
+        finally:
+            # Restore original modules
+            sys.modules.clear()
+            sys.modules.update(original_modules)
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_initialize_failure(self, mock_pynvml):
+        """Test failed initialization."""
+        # Ensure provider starts fresh
+        self.provider._initialized = False
+
+        # Make nvmlInit fail
+        mock_pynvml.nvmlInit.side_effect = Exception("Initialization failed")
+
+        # Mock sys.modules to make the import work but nvmlInit fail
+        import sys
+
+        original_modules = sys.modules.copy()
+        sys.modules["ray._private.thirdparty.pynvml"] = mock_pynvml
+
+        try:
+            self.assertFalse(self.provider._initialize())
+            self.assertFalse(self.provider._initialized)
+        finally:
+            # Restore original modules
+            sys.modules.clear()
+            sys.modules.update(original_modules)
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_initialize_already_initialized(self, mock_pynvml):
+        """Test initialization when already initialized."""
+        self.provider._initialized = True
+
+        self.assertTrue(self.provider._initialize())
+        mock_pynvml.nvmlInit.assert_not_called()
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_shutdown(self, mock_pynvml):
+        """Test shutdown."""
+        self.provider._initialized = True
+        self.provider._pynvml = mock_pynvml
+
+        self.provider._shutdown()
+
+        self.assertFalse(self.provider._initialized)
+        mock_pynvml.nvmlShutdown.assert_called_once()
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_shutdown_not_initialized(self, mock_pynvml):
+        """Test shutdown when not initialized."""
+        self.provider._shutdown()
+        mock_pynvml.nvmlShutdown.assert_not_called()
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_get_gpu_utilization_success(self, mock_pynvml):
+        """Test successful GPU utilization retrieval."""
+        # Mock GPU device
+        mock_handle = Mock()
+        mock_memory_info = Mock()
+        mock_memory_info.used = 8 * MB * 1024  # 8GB used
+        mock_memory_info.total = 12 * MB * 1024  # 12GB total
+
+        mock_utilization_info = Mock()
+        mock_utilization_info.gpu = 75
+
+        mock_process = Mock()
+        mock_process.pid = 1234
+        mock_process.usedGpuMemory = 256 * MB
+
+        # Configure mocks
+        mock_pynvml.nvmlInit.return_value = None
+        mock_pynvml.nvmlDeviceGetCount.return_value = 1
+        mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_handle
+        mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_memory_info
+        mock_pynvml.nvmlDeviceGetUtilizationRates.return_value = mock_utilization_info
+        mock_pynvml.nvmlDeviceGetComputeRunningProcesses.return_value = [mock_process]
+        mock_pynvml.nvmlDeviceGetGraphicsRunningProcesses.return_value = []
+        mock_pynvml.nvmlDeviceGetName.return_value = b"NVIDIA GeForce RTX 3080"
+        mock_pynvml.nvmlDeviceGetUUID.return_value = (
+            b"GPU-12345678-1234-1234-1234-123456789abc"
+        )
+        mock_pynvml.nvmlShutdown.return_value = None
+
+        # Set up provider state
+        self.provider._pynvml = mock_pynvml
+        self.provider._initialized = True
+
+        result = self.provider.get_gpu_utilization()
+
+        self.assertEqual(len(result), 1)
+        gpu_info = result[0]
+
+        self.assertEqual(gpu_info["index"], 0)
+        self.assertEqual(gpu_info["name"], "NVIDIA GeForce RTX 3080")
+        self.assertEqual(gpu_info["uuid"], "GPU-12345678-1234-1234-1234-123456789abc")
+        self.assertEqual(gpu_info["utilization_gpu"], 75)
+        self.assertEqual(gpu_info["memory_used"], 8 * 1024)  # 8GB in MB
+        self.assertEqual(gpu_info["memory_total"], 12 * 1024)  # 12GB in MB
+        self.assertEqual(len(gpu_info["processes_pids"]), 1)
+        self.assertEqual(gpu_info["processes_pids"][1234]["pid"], 1234)
+        self.assertEqual(gpu_info["processes_pids"][1234]["gpu_memory_usage"], 256)
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_get_gpu_utilization_with_errors(self, mock_pynvml):
+        """Test GPU utilization retrieval with partial errors."""
+        mock_handle = Mock()
+        mock_memory_info = Mock()
+        mock_memory_info.used = 4 * MB * 1024
+        mock_memory_info.total = 8 * MB * 1024
+
+        # Create mock NVML error class
+        class MockNVMLError(Exception):
+            pass
+
+        mock_pynvml.NVMLError = MockNVMLError
+
+        # Configure mocks with some failures
+        mock_pynvml.nvmlInit.return_value = None
+        mock_pynvml.nvmlDeviceGetCount.return_value = 1
+        mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_handle
+        mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_memory_info
+        mock_pynvml.nvmlDeviceGetUtilizationRates.side_effect = MockNVMLError(
+            "Utilization not available"
+        )
+        mock_pynvml.nvmlDeviceGetComputeRunningProcesses.side_effect = MockNVMLError(
+            "Process info not available"
+        )
+        mock_pynvml.nvmlDeviceGetGraphicsRunningProcesses.side_effect = MockNVMLError(
+            "Process info not available"
+        )
+        mock_pynvml.nvmlDeviceGetName.return_value = b"NVIDIA Tesla V100"
+        mock_pynvml.nvmlDeviceGetUUID.return_value = (
+            b"GPU-87654321-4321-4321-4321-ba9876543210"
+        )
+        mock_pynvml.nvmlShutdown.return_value = None
+
+        # Set up provider state
+        self.provider._pynvml = mock_pynvml
+        self.provider._initialized = True
+
+        result = self.provider.get_gpu_utilization()
+
+        self.assertEqual(len(result), 1)
+        gpu_info = result[0]
+
+        self.assertEqual(gpu_info["index"], 0)
+        self.assertEqual(gpu_info["name"], "NVIDIA Tesla V100")
+        self.assertEqual(gpu_info["utilization_gpu"], -1)  # Should be -1 due to error
+        self.assertEqual(
+            gpu_info["processes_pids"], {}
+        )  # Should be empty dict due to error
+
+    @patch("ray._private.thirdparty.pynvml", create=True)
+    def test_get_gpu_utilization_with_mig(self, mock_pynvml):
+        """Test GPU utilization retrieval with MIG devices."""
+        # Mock regular GPU handle
+        mock_gpu_handle = Mock()
+        mock_memory_info = Mock()
+        mock_memory_info.used = 4 * MB * 1024
+        mock_memory_info.total = 8 * MB * 1024
+
+        # Mock MIG device handle and info
+        mock_mig_handle = Mock()
+        mock_mig_memory_info = Mock()
+        mock_mig_memory_info.used = 2 * MB * 1024
+        mock_mig_memory_info.total = 4 * MB * 1024
+
+        mock_mig_utilization_info = Mock()
+        mock_mig_utilization_info.gpu = 80
+
+        # Configure mocks for MIG-enabled GPU
+        mock_pynvml.nvmlInit.return_value = None
+        mock_pynvml.nvmlDeviceGetCount.return_value = 1
+        mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_gpu_handle
+
+        # MIG mode enabled
+        mock_pynvml.nvmlDeviceGetMigMode.return_value = (
+            True,
+            True,
+        )  # (current, pending)
+        mock_pynvml.nvmlDeviceGetMaxMigDeviceCount.return_value = 1  # Only 1 MIG device
+        mock_pynvml.nvmlDeviceGetMigDeviceHandleByIndex.return_value = mock_mig_handle
+
+        # MIG device info
+        mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_mig_memory_info
+        mock_pynvml.nvmlDeviceGetUtilizationRates.return_value = (
+            mock_mig_utilization_info
+        )
+        mock_pynvml.nvmlDeviceGetComputeRunningProcesses.return_value = []
+        mock_pynvml.nvmlDeviceGetGraphicsRunningProcesses.return_value = []
+        mock_pynvml.nvmlDeviceGetName.return_value = b"NVIDIA A100-SXM4-40GB MIG 1g.5gb"
+        mock_pynvml.nvmlDeviceGetUUID.return_value = (
+            b"MIG-12345678-1234-1234-1234-123456789abc"
+        )
+        mock_pynvml.nvmlShutdown.return_value = None
+
+        # Set up provider state
+        self.provider._pynvml = mock_pynvml
+        self.provider._initialized = True
+
+        result = self.provider.get_gpu_utilization()
+
+        # Should return MIG device info instead of regular GPU
+        self.assertEqual(
+            len(result), 1
+        )  # Only one MIG device due to exception handling
+        gpu_info = result[0]
+
+        self.assertEqual(gpu_info["index"], 0)  # First MIG device (0 * 1000 + 0)
+        self.assertEqual(gpu_info["name"], "NVIDIA A100-SXM4-40GB MIG 1g.5gb")
+        self.assertEqual(gpu_info["uuid"], "MIG-12345678-1234-1234-1234-123456789abc")
+        self.assertEqual(gpu_info["utilization_gpu"], 80)
+        self.assertEqual(gpu_info["memory_used"], 2 * 1024)  # 2GB in MB
+        self.assertEqual(gpu_info["memory_total"], 4 * 1024)  # 4GB in MB
+        self.assertEqual(gpu_info["processes_pids"], {})
+
+
+class TestAmdGpuProvider(unittest.TestCase):
+    """Test AmdGpuProvider class."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.provider = AmdGpuProvider()
+
+    def test_get_provider_name(self):
+        """Test provider name."""
+        self.assertEqual(self.provider.get_provider_name(), GpuProviderType.AMD)
+
+    @patch("ray._private.thirdparty.pyamdsmi", create=True)
+    def test_is_available_success(self, mock_pyamdsmi):
+        """Test is_available when AMD GPU is available."""
+        mock_pyamdsmi.smi_initialize.return_value = None
+        mock_pyamdsmi.smi_shutdown.return_value = None
+
+        self.assertTrue(self.provider.is_available())
+        mock_pyamdsmi.smi_initialize.assert_called_once()
+        mock_pyamdsmi.smi_shutdown.assert_called_once()
+
+    @patch("ray._private.thirdparty.pyamdsmi", create=True)
+    def test_is_available_failure(self, mock_pyamdsmi):
+        """Test is_available when AMD GPU is not available."""
+        mock_pyamdsmi.smi_initialize.side_effect = Exception("AMD driver not found")
+
+        self.assertFalse(self.provider.is_available())
+
+    @patch("ray._private.thirdparty.pyamdsmi", create=True)
+    def test_initialize_success(self, mock_pyamdsmi):
+        """Test successful initialization."""
+        mock_pyamdsmi.smi_initialize.return_value = None
+
+        self.assertTrue(self.provider._initialize())
+        self.assertTrue(self.provider._initialized)
+        mock_pyamdsmi.smi_initialize.assert_called_once()
+
+    @patch("ray._private.thirdparty.pyamdsmi", create=True)
+    def test_get_gpu_utilization_success(self, mock_pyamdsmi):
+        """Test successful GPU utilization retrieval."""
+        mock_process = Mock()
+        mock_process.process_id = 5678
+        mock_process.vram_usage = 512 * MB
+
+        # Configure mocks
+        mock_pyamdsmi.smi_initialize.return_value = None
+        mock_pyamdsmi.smi_get_device_count.return_value = 1
+        mock_pyamdsmi.smi_get_device_id.return_value = "device_0"
+        mock_pyamdsmi.smi_get_device_utilization.return_value = 85
+        mock_pyamdsmi.smi_get_device_compute_process.return_value = [mock_process]
+        mock_pyamdsmi.smi_get_compute_process_info_by_device.return_value = [
+            mock_process
+        ]
+        mock_pyamdsmi.smi_get_device_name.return_value = b"AMD Radeon RX 6800 XT"
+        mock_pyamdsmi.smi_get_device_unique_id.return_value = (
+            "GPU-13579bdf-9abc-def0-0000-000000000000"
+        )
+        mock_pyamdsmi.smi_get_device_memory_used.return_value = 6 * MB * 1024
+        mock_pyamdsmi.smi_get_device_memory_total.return_value = 16 * MB * 1024
+        mock_pyamdsmi.smi_shutdown.return_value = None
+
+        # Set up provider state
+        self.provider._pyamdsmi = mock_pyamdsmi
+        self.provider._initialized = True
+
+        result = self.provider.get_gpu_utilization()
+
+        self.assertEqual(len(result), 1)
+        gpu_info = result[0]
+
+        self.assertEqual(gpu_info["index"], 0)
+        self.assertEqual(gpu_info["name"], "AMD Radeon RX 6800 XT")
+        self.assertEqual(gpu_info["uuid"], "GPU-13579bdf-9abc-def0-0000-000000000000")
+        self.assertEqual(gpu_info["utilization_gpu"], 85)
+        self.assertEqual(gpu_info["memory_used"], 6 * 1024)  # 6GB in MB
+        self.assertEqual(gpu_info["memory_total"], 16 * 1024)  # 16GB in MB
+        self.assertEqual(len(gpu_info["processes_pids"]), 1)
+        self.assertEqual(gpu_info["processes_pids"][5678]["pid"], 5678)
+        self.assertEqual(gpu_info["processes_pids"][5678]["gpu_memory_usage"], 512)
+
+
+class TestGpuMetricProvider(unittest.TestCase):
+    """Test GpuMetricProvider class."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.provider = GpuMetricProvider()
+
+    def test_init(self):
+        """Test GpuMetricProvider initialization."""
+        self.assertIsNone(self.provider._provider)
+        self.assertTrue(self.provider._enable_metric_report)
+        self.assertEqual(len(self.provider._providers), 2)
+        self.assertFalse(self.provider._initialized)
+
+    @patch.object(NvidiaGpuProvider, "is_available", return_value=True)
+    @patch.object(AmdGpuProvider, "is_available", return_value=False)
+    def test_detect_gpu_provider_nvidia(
+        self, mock_amd_available, mock_nvidia_available
+    ):
+        """Test GPU provider detection when NVIDIA is available."""
+        provider = self.provider._detect_gpu_provider()
+
+        self.assertIsInstance(provider, NvidiaGpuProvider)
+        mock_nvidia_available.assert_called_once()
+
+    @patch.object(NvidiaGpuProvider, "is_available", return_value=False)
+    @patch.object(AmdGpuProvider, "is_available", return_value=True)
+    def test_detect_gpu_provider_amd(self, mock_amd_available, mock_nvidia_available):
+        """Test GPU provider detection when AMD is available."""
+        provider = self.provider._detect_gpu_provider()
+
+        self.assertIsInstance(provider, AmdGpuProvider)
+        mock_nvidia_available.assert_called_once()
+        mock_amd_available.assert_called_once()
+
+    @patch.object(NvidiaGpuProvider, "is_available", return_value=False)
+    @patch.object(AmdGpuProvider, "is_available", return_value=False)
+    def test_detect_gpu_provider_none(self, mock_amd_available, mock_nvidia_available):
+        """Test GPU provider detection when no GPUs are available."""
+        provider = self.provider._detect_gpu_provider()
+
+        self.assertIsNone(provider)
+
+    @patch("subprocess.check_output")
+    def test_should_disable_gpu_check_true(self, mock_subprocess):
+        """Test should_disable_gpu_check returns True for specific conditions."""
+        mock_subprocess.return_value = ""  # Empty result means AMD GPU module not live
+
+        class MockNVMLError(Exception):
+            pass
+
+        MockNVMLError.__name__ = "NVMLError_DriverNotLoaded"
+
+        error = MockNVMLError("NVIDIA driver not loaded")
+
+        result = self.provider._should_disable_gpu_check(error)
+        self.assertTrue(result)
+
+    @patch("subprocess.check_output")
+    def test_should_disable_gpu_check_false_wrong_error(self, mock_subprocess):
+        """Test should_disable_gpu_check returns False for wrong error type."""
+        mock_subprocess.return_value = ""
+
+        error = Exception("Some other error")
+
+        result = self.provider._should_disable_gpu_check(error)
+        self.assertFalse(result)
+
+    @patch("subprocess.check_output")
+    def test_should_disable_gpu_check_false_amd_present(self, mock_subprocess):
+        """Test should_disable_gpu_check returns False when AMD GPU is present."""
+        mock_subprocess.return_value = "live"  # AMD GPU module is live
+
+        class MockNVMLError(Exception):
+            pass
+
+        MockNVMLError.__name__ = "NVMLError_DriverNotLoaded"
+
+        error = MockNVMLError("NVIDIA driver not loaded")
+
+        result = self.provider._should_disable_gpu_check(error)
+        self.assertFalse(result)
+
+    def test_get_gpu_usage_disabled(self):
+        """Test get_gpu_usage when GPU usage check is disabled."""
+        self.provider._enable_metric_report = False
+
+        result = self.provider.get_gpu_usage()
+        self.assertEqual(result, [])
+
+    @patch.object(GpuMetricProvider, "_detect_gpu_provider")
+    def test_get_gpu_usage_no_provider(self, mock_detect):
+        """Test get_gpu_usage when no GPU provider is available."""
+        mock_detect.return_value = None
+
+        with patch.object(
+            NvidiaGpuProvider, "_initialize", side_effect=Exception("No GPU")
+        ):
+            result = self.provider.get_gpu_usage()
+
+        self.assertEqual(result, [])
+        self.provider._initialized = False  # Reset for clean test
+        mock_detect.assert_called_once()
+
+    @patch.object(GpuMetricProvider, "_detect_gpu_provider")
+    def test_get_gpu_usage_success(self, mock_detect):
+        """Test successful get_gpu_usage."""
+        mock_provider = Mock()
+        mock_provider.get_gpu_utilization.return_value = [
+            GpuUtilizationInfo(
+                index=0,
+                name="Test GPU",
+                uuid="test-uuid",
+                utilization_gpu=50,
+                memory_used=1024,
+                memory_total=2048,
+                processes_pids={
+                    1234: ProcessGPUInfo(
+                        pid=1234, gpu_memory_usage=1024, gpu_utilization=None
+                    )
+                },
+            )
+        ]
+        mock_detect.return_value = mock_provider
+
+        result = self.provider.get_gpu_usage()
+
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0]["index"], 0)
+        self.assertEqual(result[0]["name"], "Test GPU")
+        mock_provider.get_gpu_utilization.assert_called_once()
+
+    def test_get_provider_name_no_provider(self):
+        """Test get_provider_name when no provider is set."""
+        result = self.provider.get_provider_name()
+        self.assertIsNone(result)
+
+    def test_get_provider_name_with_provider(self):
+        """Test get_provider_name when provider is set."""
+        mock_provider = Mock()
+        mock_provider.get_provider_name.return_value = GpuProviderType.NVIDIA
+        self.provider._provider = mock_provider
+
+        result = self.provider.get_provider_name()
+        self.assertEqual(result, "nvidia")
+
+    def test_is_metric_report_enabled(self):
+        """Test is_metric_report_enabled."""
+        self.assertTrue(self.provider.is_metric_report_enabled())
+
+        self.provider._enable_metric_report = False
+        self.assertFalse(self.provider.is_metric_report_enabled())
+
+
+if __name__ == "__main__":
+    unittest.main()

python/ray/dashboard/modules/reporter/tests/test_healthz.py

@@ -0,0 +1,121 @@
+import signal
+import sys
+
+import pytest
+import requests
+
+import ray._private.ray_constants as ray_constants
+from ray._common.network_utils import find_free_port
+from ray._common.test_utils import wait_for_condition
+from ray.tests.conftest import *  # noqa: F401 F403
+
+
+def test_healthz_head(monkeypatch, ray_start_cluster):
+    dashboard_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_port=dashboard_port)
+    uri = f"http://localhost:{dashboard_port}/api/gcs_healthz"
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    h.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+    # It'll either timeout or just return an error
+    try:
+        wait_for_condition(lambda: requests.get(uri, timeout=1) != 200, timeout=4)
+    except RuntimeError as e:
+        assert "Read timed out" in str(e)
+
+
+def test_healthz_agent_1(monkeypatch, ray_start_cluster):
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/local_raylet_healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+
+    h.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+    # GCS's failure will not lead to healthz failure
+    assert requests.get(uri).status_code == 200
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGSTOP only on posix")
+def test_healthz_agent_2(monkeypatch, ray_start_cluster):
+    monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
+    monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
+    monkeypatch.setenv("RAY_health_check_period_ms", "1000")
+    monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
+
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/local_raylet_healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+
+    h.all_processes[ray_constants.PROCESS_TYPE_RAYLET][0].process.send_signal(
+        signal.SIGSTOP
+    )
+
+    # GCS still think raylet is alive.
+    assert requests.get(uri).status_code == 200
+    # But after heartbeat timeout, it'll think the raylet is down.
+    wait_for_condition(lambda: requests.get(uri).status_code != 200)
+
+
+def test_unified_healthz_head(monkeypatch, ray_start_cluster):
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    resp = requests.get(uri)
+    assert "raylet: success" in resp.text
+    assert "gcs: success" in resp.text
+
+    h.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+    wait_for_condition(lambda: requests.get(uri).status_code == 503)
+    resp = requests.get(uri)
+    assert "gcs: " in resp.text
+    assert "gcs: success" not in resp.text
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGSTOP only on posix")
+def test_unified_healthz_worker(monkeypatch, ray_start_cluster):
+    monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
+    monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
+    monkeypatch.setenv("RAY_health_check_period_ms", "1000")
+    monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
+
+    ray_start_cluster.add_node()
+    agent_port = find_free_port()
+    h = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h.node_ip_address}:{agent_port}/api/healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    resp = requests.get(uri)
+    assert "gcs: success (no local gcs)" in resp.text
+
+    # Stop local raylet and verify this makes /healthz fail.
+    h.all_processes[ray_constants.PROCESS_TYPE_RAYLET][0].process.send_signal(
+        signal.SIGSTOP
+    )
+    wait_for_condition(lambda: requests.get(uri).status_code == 503)
+    resp = requests.get(uri)
+    assert "raylet: Local Raylet failed" in resp.text
+
+
+def test_unified_healthz_worker_gcs_down(monkeypatch, ray_start_cluster):
+    h_head = ray_start_cluster.add_node()
+    agent_port = find_free_port()
+    h_worker = ray_start_cluster.add_node(dashboard_agent_listen_port=agent_port)
+    uri = f"http://{h_worker.node_ip_address}:{agent_port}/api/healthz"
+
+    wait_for_condition(lambda: requests.get(uri).status_code == 200)
+    resp = requests.get(uri)
+    assert "gcs: success (no local gcs)" in resp.text
+
+    # Stop the head GCS server.
+    h_head.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][0].process.kill()
+
+    # Worker health check should still succeed.
+    assert requests.get(uri).status_code == 200
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dashboard/modules/reporter/tests/test_jax_profiler_manager.py

@@ -0,0 +1,106 @@
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray.dashboard.modules.reporter.jax_profile_manager import JaxProfilingManager
+from ray.util.tpu import init_jax_profiler
+
+
+@pytest.fixture
+def mock_profiler_client():
+    mock_client = MagicMock()
+    mock_profiler_module = MagicMock()
+    mock_profiler_module.profiler_client = mock_client
+
+    modules_to_patch = {
+        "tensorflow": MagicMock(),
+        "tensorflow.python": MagicMock(),
+        "tensorflow.python.profiler": mock_profiler_module,
+    }
+
+    with patch.dict("sys.modules", modules_to_patch):
+        yield mock_client
+
+
+@pytest.mark.asyncio
+async def test_jax_profile_success(tmp_path, mock_profiler_client):
+    manager = JaxProfilingManager(tmp_path)
+
+    # Mock success
+    mock_profiler_client.trace.return_value = None
+
+    success, output = await manager.jax_profile(pid=123, port=6000, duration_s=2)
+
+    assert success
+    assert output.startswith("profiles")
+    assert "123_" in output
+
+    mock_profiler_client.trace.assert_called_once()
+    call = mock_profiler_client.trace.call_args
+    assert call.args[0] == "grpc://localhost:6000"
+    assert call.kwargs["logdir"].startswith(str(tmp_path / "profiles"))
+    assert call.kwargs["duration_ms"] == 2000
+
+
+@pytest.mark.asyncio
+async def test_jax_profile_failure(tmp_path, mock_profiler_client):
+    manager = JaxProfilingManager(tmp_path)
+
+    # Mock failure
+    mock_profiler_client.trace.side_effect = Exception("Connection failed")
+
+    success, output = await manager.jax_profile(pid=123, port=6000, duration_s=2)
+
+    assert not success
+    assert "Failed to capture trace: Connection failed" in output
+
+
+@pytest.mark.asyncio
+async def test_jax_profile_no_tensorflow(tmp_path):
+    manager = JaxProfilingManager(tmp_path)
+
+    # Force ImportError on tensorflow
+    with patch.dict("sys.modules", {"tensorflow": None}):
+        success, output = await manager.jax_profile(pid=123, port=6000, duration_s=2)
+
+        assert not success
+        assert "TensorFlow is required" in output
+
+
+@patch("ray.util.tpu.os.getpid")
+@patch("ray.util.tpu.os.getenv")
+def test_setup_jax_profiler_success(mock_getenv, mock_getpid):
+    mock_getenv.return_value = "9999"
+    mock_getpid.return_value = 12345
+
+    mock_jax = MagicMock()
+    mock_worker = MagicMock()
+    mock_worker.node.node_id = "mock_node_id_hex"
+
+    with (
+        patch.dict("sys.modules", {"jax": mock_jax}),
+        patch("ray._private.worker.global_worker", mock_worker),
+        patch("ray.experimental.internal_kv._internal_kv_put") as mock_kv_put,
+    ):
+
+        init_jax_profiler()
+
+        mock_jax.profiler.start_server.assert_called_once_with(9999)
+        import ray
+
+        mock_kv_put.assert_called_once_with(
+            "jax_profiler_port:mock_node_id_hex:12345",
+            b"9999",
+            namespace=ray._private.ray_constants.KV_NAMESPACE_DASHBOARD,
+        )
+
+
+def test_setup_jax_profiler_no_jax():
+    with patch.dict("sys.modules", {"jax": None}):
+        # Should skip starting profiler and not raise error
+        init_jax_profiler()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dashboard/modules/reporter/tests/test_profile_manager.py

@@ -0,0 +1,270 @@
+import os
+import sys
+import tempfile
+import time
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+import ray
+from ray.dashboard.modules.reporter.profile_manager import (
+    CpuProfilingManager,
+    MemoryProfilingManager,
+)
+from ray.dashboard.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def setup_memory_profiler():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        memory_profiler = MemoryProfilingManager(tmpdir)
+
+        @ray.remote
+        class Actor:
+            def getpid(self):
+                return os.getpid()
+
+            def long_run(self):
+                print("Long-running task began.")
+                time.sleep(1000)
+                print("Long-running task completed.")
+
+        actor = Actor.remote()
+
+        yield actor, memory_profiler
+
+
+@pytest.mark.asyncio
+@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.platform == "darwin",
+    reason="Fails on OSX, requires memray & lldb installed in osx image",
+)
+class TestMemoryProfiling:
+    async def test_basic_attach_profiler(self, setup_memory_profiler, shutdown_only):
+        # test basic attach profiler to running process
+        actor, memory_profiler = setup_memory_profiler
+        pid = ray.get(actor.getpid.remote())
+        actor.long_run.remote()
+        success, profiler_filename, message = await memory_profiler.attach_profiler(
+            pid, verbose=True
+        )
+
+        assert success, message
+        assert f"Success attaching memray to process {pid}" in message
+        assert profiler_filename in os.listdir(memory_profiler.profile_dir_path)
+
+    async def test_profiler_multiple_attach(self, setup_memory_profiler, shutdown_only):
+        # test multiple attaches
+        actor, memory_profiler = setup_memory_profiler
+        pid = ray.get(actor.getpid.remote())
+        actor.long_run.remote()
+        success, profiler_filename, message = await memory_profiler.attach_profiler(
+            pid, verbose=True
+        )
+
+        assert success, message
+        assert f"Success attaching memray to process {pid}" in message
+        assert profiler_filename in os.listdir(memory_profiler.profile_dir_path)
+
+        success, _, message = await memory_profiler.attach_profiler(pid)
+        assert success, message
+        assert f"Success attaching memray to process {pid}" in message
+
+    async def test_detach_profiler_successful(
+        self, setup_memory_profiler, shutdown_only
+    ):
+        # test basic detach profiler
+        actor, memory_profiler = setup_memory_profiler
+        pid = ray.get(actor.getpid.remote())
+        actor.long_run.remote()
+        success, _, message = await memory_profiler.attach_profiler(pid, verbose=True)
+        assert success, message
+
+        success, message = await memory_profiler.detach_profiler(pid, verbose=True)
+        assert success, message
+        assert f"Success detaching memray from process {pid}" in message
+
+    async def test_detach_profiler_without_attach(
+        self, setup_memory_profiler, shutdown_only
+    ):
+        # test detach profiler from unattached process
+        actor, memory_profiler = setup_memory_profiler
+        pid = ray.get(actor.getpid.remote())
+
+        success, message = await memory_profiler.detach_profiler(pid)
+        assert not success, message
+        assert "Failed to execute" in message
+        assert "no previous `memray attach`" in message
+
+    async def test_profiler_memray_not_installed(
+        self, setup_memory_profiler, shutdown_only
+    ):
+        # test profiler when memray is not installed
+        actor, memory_profiler = setup_memory_profiler
+        pid = ray.get(actor.getpid.remote())
+
+        with patch("shutil.which", return_value=None):
+            success, _, message = await memory_profiler.attach_profiler(pid)
+            assert not success
+            assert "memray is not installed" in message
+
+    async def test_profiler_attach_process_not_found(
+        self, setup_memory_profiler, shutdown_only
+    ):
+        # test basic attach profiler to non-existing process
+        _, memory_profiler = setup_memory_profiler
+        pid = 123456
+        success, _, message = await memory_profiler.attach_profiler(pid)
+        assert not success, message
+        assert "Failed to execute" in message
+        assert "The given process ID does not exist" in message
+
+    async def test_profiler_get_profiler_result(
+        self, setup_memory_profiler, shutdown_only
+    ):
+        # test get profiler result from running process
+        actor, memory_profiler = setup_memory_profiler
+        pid = ray.get(actor.getpid.remote())
+        actor.long_run.remote()
+        success, profiler_filename, message = await memory_profiler.attach_profiler(
+            pid, verbose=True
+        )
+        assert success, message
+        assert f"Success attaching memray to process {pid}" in message
+
+        # get profiler result in flamegraph and table format
+        supported_formats = ["flamegraph", "table"]
+        unsupported_formats = ["json"]
+        for format in supported_formats + unsupported_formats:
+            success, message = await memory_profiler.get_profile_result(
+                pid, profiler_filename=profiler_filename, format=format
+            )
+            if format in supported_formats:
+                assert success, message
+                assert f"{format} report" in message.decode("utf-8")
+            else:
+                assert not success, message
+                assert f"{format} is not supported" in message
+
+    async def test_profiler_result_not_exist(
+        self, setup_memory_profiler, shutdown_only
+    ):
+        # test get profiler result from unexisting process
+        _, memory_profiler = setup_memory_profiler
+        pid = 123456
+        profiler_filename = "non-existing-file"
+
+        success, message = await memory_profiler.get_profile_result(
+            pid, profiler_filename=profiler_filename, format=format
+        )
+        assert not success, message
+        assert f"process {pid} has not been profiled" in message
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.platform == "win32", reason="No py-spy on Windows.")
+class TestCpuProfiling:
+    async def _capture_pyspy_cmd(self, **cpu_profile_kwargs):
+        """Run cpu_profile with subprocess execution mocked out and return the
+        py-spy command that would have been executed.
+
+        We patch ``asyncio.create_subprocess_exec`` (the same primitive the
+        manager uses) and have the fake process exit non-zero so that
+        ``cpu_profile`` short-circuits before attempting to read the (never
+        created) output file. The command is fully constructed before the
+        subprocess is spawned, so the captured args are valid regardless.
+        """
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cpu_profiler = CpuProfilingManager(tmpdir)
+
+            fake_process = AsyncMock()
+            fake_process.communicate.return_value = (b"", b"boom")
+            fake_process.returncode = 1
+
+            with patch(
+                "ray.dashboard.modules.reporter.profile_manager.shutil.which",
+                return_value="/fake/py-spy",
+            ), patch(
+                "ray.dashboard.modules.reporter.profile_manager."
+                "_can_passwordless_sudo",
+                new=AsyncMock(return_value=False),
+            ), patch(
+                "asyncio.create_subprocess_exec",
+                new=AsyncMock(return_value=fake_process),
+            ) as mock_exec:
+                await cpu_profiler.cpu_profile(pid=12345, **cpu_profile_kwargs)
+
+            assert mock_exec.call_count == 1
+            # create_subprocess_exec(*cmd, ...) -> positional args are the cmd.
+            return list(mock_exec.call_args.args)
+
+    async def test_cpu_profile_idle_flag_added(self):
+        # idle=True should append `--idle` to the py-spy command.
+        cmd = await self._capture_pyspy_cmd(idle=True)
+        assert "--idle" in cmd
+
+    async def test_cpu_profile_idle_not_added_by_default(self):
+        # By default (idle=False) the `--idle` flag should be absent.
+        cmd = await self._capture_pyspy_cmd()
+        assert "--idle" not in cmd
+
+    async def test_cpu_profile_subprocesses_flag_added(self):
+        # subprocesses=True should append `--subprocesses` to the py-spy command.
+        cmd = await self._capture_pyspy_cmd(subprocesses=True)
+        assert "--subprocesses" in cmd
+
+    async def test_cpu_profile_subprocesses_not_added_by_default(self):
+        # By default the `--subprocesses` flag should be absent.
+        cmd = await self._capture_pyspy_cmd()
+        assert "--subprocesses" not in cmd
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.platform == "win32", reason="No py-spy on Windows.")
+class TestTraceDump:
+    async def _capture_pyspy_cmd(self, **trace_dump_kwargs):
+        """Run trace_dump with subprocess execution mocked out and return the
+        py-spy command that would have been executed.
+        """
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cpu_profiler = CpuProfilingManager(tmpdir)
+
+            fake_process = AsyncMock()
+            fake_process.communicate.return_value = (b"", b"boom")
+            fake_process.returncode = 1
+
+            with patch(
+                "ray.dashboard.modules.reporter.profile_manager.shutil.which",
+                return_value="/fake/py-spy",
+            ), patch(
+                "ray.dashboard.modules.reporter.profile_manager."
+                "_can_passwordless_sudo",
+                new=AsyncMock(return_value=False),
+            ), patch(
+                "asyncio.create_subprocess_exec",
+                new=AsyncMock(return_value=fake_process),
+            ) as mock_exec:
+                await cpu_profiler.trace_dump(pid=12345, **trace_dump_kwargs)
+
+            assert mock_exec.call_count == 1
+            return list(mock_exec.call_args.args)
+
+    async def test_trace_dump_subprocesses_flag_added(self):
+        # subprocesses=True should append `--subprocesses` to the py-spy dump command.
+        cmd = await self._capture_pyspy_cmd(subprocesses=True)
+        assert "dump" in cmd
+        assert "--subprocesses" in cmd
+
+    async def test_trace_dump_subprocesses_not_added_by_default(self):
+        # By default the `--subprocesses` flag should be absent.
+        cmd = await self._capture_pyspy_cmd()
+        assert "--subprocesses" not in cmd
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))