chore: import upstream snapshot with attribution

python/ray/tests/test_runtime_env_profiler.py

@@ -0,0 +1,266 @@
+import glob
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray._private.runtime_env.nsight import parse_nsight_config
+from ray.exceptions import RuntimeEnvSetupError
+
+
+def wait_for_report(profilers_dir, num_reports):
+    assert len(os.listdir(profilers_dir)) == num_reports
+    return True
+
+
+NSIGHT_FAKE_DIR = str(Path(os.path.realpath(__file__)).parent / "nsight_fake")
+
+
+@pytest.fixture(scope="class")
+def nsight_fake_dependency():
+    subprocess.check_call(
+        ["pip", "install", NSIGHT_FAKE_DIR],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    yield
+    subprocess.check_call(
+        ["pip", "uninstall", "nsys", "--y"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+
+class TestNsightProfiler:
+    @pytest.mark.skipif(
+        os.environ.get("CI") and sys.platform != "linux",
+        reason="Requires PR wheels built in CI, so only run on linux CI machines.",
+    )
+    @pytest.mark.usefixtures("nsight_fake_dependency")
+    def test_nsight_basic(
+        self,
+        shutdown_only,
+    ):
+        """Test Nsight profile generate report on logs dir"""
+        ray.init()
+        session_dir = ray.worker._global_node.get_session_dir_path()
+        profilers_dir = Path(session_dir) / "logs" / "nsight"
+
+        # test nsight default config
+        @ray.remote(runtime_env={"nsight": "default"})
+        def test_generate_report():
+            return 0
+
+        ray.get(test_generate_report.remote())
+
+        wait_for_condition(wait_for_report, profilers_dir=profilers_dir, num_reports=1)
+        nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+        assert len(nsys_reports) == 1
+        for report in nsys_reports:
+            assert "worker_process_" in report
+            os.remove(report)
+
+    @pytest.mark.skipif(
+        os.environ.get("CI") and sys.platform != "linux",
+        reason="Requires PR wheels built in CI, so only run on linux CI machines.",
+    )
+    @pytest.mark.usefixtures("nsight_fake_dependency")
+    def test_nsight_custom_option(
+        self,
+        shutdown_only,
+    ):
+        """Test Nsight profile generate report on logs dir"""
+        ray.init()
+        session_dir = ray.worker._global_node.get_session_dir_path()
+        profilers_dir = Path(session_dir) / "logs" / "nsight"
+
+        # test nsight custom filename
+        CUSTOM_REPORT = "custom_report"
+
+        @ray.remote(
+            runtime_env={
+                "nsight": {
+                    "t": "cuda,cublas,cudnn",
+                    "stop-on-exit": "true",
+                    "o": CUSTOM_REPORT,
+                }
+            }
+        )
+        def test_generate_custom_report():
+            return 0
+
+        ray.get(test_generate_custom_report.remote())
+
+        wait_for_condition(wait_for_report, profilers_dir=profilers_dir, num_reports=1)
+        nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+        assert len(nsys_reports) == 1
+        for report in nsys_reports:
+            assert f"{CUSTOM_REPORT}.nsys-rep" in report
+            os.remove(report)
+
+    @pytest.mark.skipif(
+        os.environ.get("CI") and sys.platform != "linux",
+        reason="Requires PR wheels built in CI, so only run on linux CI machines.",
+    )
+    @pytest.mark.usefixtures("nsight_fake_dependency")
+    def test_nsight_multiple_tasks(
+        self,
+        shutdown_only,
+    ):
+        """Test Nsight profile on multiple ray tasks/actors"""
+        ray.init()
+        session_dir = ray.worker._global_node.get_session_dir_path()
+        profilers_dir = Path(session_dir) / "logs" / "nsight"
+
+        # test ray task
+        @ray.remote(
+            runtime_env={
+                "nsight": {"o": "ray_task_%p"},
+            }
+        )
+        def test_generate_report():
+            return 0
+
+        # test ray actor
+        @ray.remote(runtime_env={"nsight": {"o": "ray_actor"}})
+        class NsightActor:
+            def run(self):
+                return 0
+
+        ray_actor = NsightActor.remote()
+        ray.get(
+            [
+                test_generate_report.remote(),
+                test_generate_report.remote(),
+                ray_actor.run.remote(),
+                ray_actor.run.remote(),
+            ]
+        )
+        ray.kill(ray_actor)
+
+        wait_for_condition(wait_for_report, profilers_dir=profilers_dir, num_reports=3)
+        nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+        assert len(nsys_reports) == 3  # ray actor used only one worker process
+        for report in nsys_reports:
+            assert "ray_task" in report or "ray_actor" in report
+            os.remove(report)
+
+    @pytest.mark.skipif(
+        os.environ.get("CI") and sys.platform != "linux",
+        reason="Requires PR wheels built in CI, so only run on linux CI machines.",
+    )
+    @pytest.mark.usefixtures("nsight_fake_dependency")
+    def test_nsight_invalid_option(
+        self,
+        shutdown_only,
+    ):
+        """Test Nsight profile for invalid option"""
+        ray.init()
+        session_dir = ray.worker._global_node.get_session_dir_path()
+        profilers_dir = Path(session_dir) / "logs" / "nsight"
+
+        # test nsight invalid config
+        @ray.remote(
+            runtime_env={
+                "nsight": {
+                    "not-option": "random",
+                }
+            }
+        )
+        def test_invalid_nsight():
+            return 0
+
+        with pytest.raises(
+            RuntimeEnvSetupError,
+            match="nsight profile failed to run with the following error message",
+        ):
+            ray.get(test_invalid_nsight.remote())
+
+        nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+        assert len(nsys_reports) == 0
+
+        """Test Nsight profile for unavailable string option"""
+
+        # test nsight not supported string config
+        @ray.remote(runtime_env={"nsight": "not_default"})
+        def test_wrong_config_nsight():
+            return 0
+
+        with pytest.raises(
+            RuntimeEnvSetupError, match="Unsupported nsight config: not_default."
+        ):
+            ray.get(test_wrong_config_nsight.remote())
+
+        nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+        assert len(nsys_reports) == 0
+
+    @pytest.mark.skipif(
+        sys.platform == "linux",
+        reason="For non linux system, Nsight CLI is not supported",
+    )
+    def test_nsight_unsupported(
+        self,
+        shutdown_only,
+    ):
+        """Test Nsight profile for unsupported OS"""
+        ray.init()
+        session_dir = ray.worker._global_node.get_session_dir_path()
+        profilers_dir = Path(session_dir) / "logs" / "nsight"
+
+        # test nsight default config
+        @ray.remote(runtime_env={"nsight": "default"})
+        def test_nsight_not_supported():
+            return 0
+
+        with pytest.raises(
+            RuntimeEnvSetupError, match="Nsight CLI is only available in Linux."
+        ):
+            ray.get(test_nsight_not_supported.remote())
+
+        nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+        assert len(nsys_reports) == 0
+
+    def test_parse_nsight_config(self):
+        """Test parse nsight config into nsight command prefix"""
+        nsight_config = {
+            "o": "single_dash",
+            "two-dash": "double_dash",
+        }
+        nsight_cmd = parse_nsight_config(nsight_config)
+        assert nsight_cmd == [
+            "nsys",
+            "profile",
+            "-o",
+            "single_dash",
+            "--two-dash=double_dash",
+        ]
+
+
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform != "linux",
+    reason="Test only for compatible OS (linux) with minmial installation",
+)
+def test_nsight_not_installed():
+    """Test Nsight profile not installed"""
+    ray.init()
+    session_dir = ray.worker._global_node.get_session_dir_path()
+    profilers_dir = Path(session_dir) / "logs" / "nsight"
+
+    # test nsight default config
+    @ray.remote(runtime_env={"nsight": "default"})
+    def test_nsight_not_supported():
+        return 0
+
+    with pytest.raises(RuntimeEnvSetupError, match="nsight is not installed"):
+        ray.get(test_nsight_not_supported.remote())
+
+    nsys_reports = glob.glob(os.path.join(f"{profilers_dir}/*.nsys-rep"))
+    assert len(nsys_reports) == 0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))