Files
ray-project--ray/python/ray/tests/test_runtime_env_profiler.py
T
2026-07-13 13:17:40 +08:00

267 lines
8.3 KiB
Python

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__]))