nsys-jax: use pytest tmp_path[_factory] to allow saving profiles (#1748)

This aims to save all of the original `.zip`, but neither the extracted
versions of them nor any derived `.zip` files (as those should be
deterministically derived from the original `.zip`s)

Author: olupton

.github/container/nsys_jax/tests/nsys_jax_test_helpers/__init__.py

@@ -7,25 +7,25 @@
 import zipfile
 
 
-def nsys_jax_with_result(command):
+def nsys_jax_with_result(command, *, out_dir):
     """
     Helper to run nsys-jax with a unique output file that will be automatically
     cleaned up on destruction. Explicitly returns the `subprocess.CompletedProcess`
     instance.
     """
-    output = tempfile.NamedTemporaryFile(suffix=".zip")
+    output = tempfile.NamedTemporaryFile(delete=False, dir=out_dir, suffix=".zip")
     result = subprocess.run(
         ["nsys-jax", "--force-overwrite", "--output", output.name] + command,
     )
     return output, result
 
 
-def nsys_jax(command):
+def nsys_jax(command, *, out_dir):
     """
     Helper to run nsys-jax with a unique output file that will be automatically
     cleaned up on destruction. Throws if running `nsys-jax` does not succeed.
     """
-    output, result = nsys_jax_with_result(command)
+    output, result = nsys_jax_with_result(command, out_dir=out_dir)
     result.check_returncode()
     return output
 
@@ -42,11 +42,11 @@ def extract(archive):
     return tmpdir
 
 
-def nsys_jax_archive(command):
+def nsys_jax_archive(command, *, out_dir):
     """
     Helper to run nsys-jax and automatically extract the output, yielding a directory.
     """
-    archive = nsys_jax(command)
+    archive = nsys_jax(command, out_dir=out_dir)
     tmpdir = extract(archive)
     # Make sure the protobuf bindings can be imported, the generated .py will go into
     # a temporary directory that is not currently cleaned up. The bindings cannot be
@@ -58,13 +58,17 @@ def nsys_jax_archive(command):
 
 
 def multi_process_nsys_jax(
-    num_processes: int, command: typing.Callable[[int], list[str]]
+    num_processes: int,
+    command: typing.Callable[[int], list[str]],
+    *,
+    out_dir,
 ):
     """
     Helper to run a multi-process test under nsys-jax and yield several .zip
     """
     child_outputs = [
-        tempfile.NamedTemporaryFile(suffix=".zip") for _ in range(num_processes)
+        tempfile.NamedTemporaryFile(delete=False, dir=out_dir, suffix=".zip")
+        for _ in range(num_processes)
     ]
     children = [
         subprocess.Popen(

.github/container/nsys_jax/tests/test_basics.py

@@ -9,14 +9,14 @@
 from nsys_jax_test_helpers import nsys_jax  # noqa: E402
 
 
-def test_program_without_gpu_activity():
+def test_program_without_gpu_activity(tmp_path):
     """
     Profiling a program that doesn't do anything should succeed.
     """
-    nsys_jax([sys.executable, "-c", "print('Hello world!')"])
+    nsys_jax([sys.executable, "-c", "print('Hello world!')"], out_dir=tmp_path)
 
 
-def test_stacktrace_entry_with_file():
+def test_stacktrace_entry_with_file(tmp_path):
     """
     Test that if a source file appears in the traceback of a JITed JAX function then
     the source file is bundled into the nsys-jax output archive.
@@ -27,19 +27,21 @@ def test_stacktrace_entry_with_file():
         src_code = "import jax\njax.jit(lambda x: x*2)(4)\n"
         with open(src_file, "w") as f:
             f.write(src_code)
-        archive = nsys_jax([sys.executable, src_file])
+        archive = nsys_jax([sys.executable, src_file], out_dir=tmp_path)
     with zipfile.ZipFile(archive) as ifile:
         src_file_in_archive = f"sources{src_file}"
         assert src_file_in_archive in ifile.namelist()
         with ifile.open(src_file_in_archive, "r") as archived_file:
             assert archived_file.read().decode() == src_code
 
 
-def test_stacktrace_entry_without_file():
+def test_stacktrace_entry_without_file(tmp_path):
     """
     Test that tracing code that does not come from a named file works (bug 4931958).
     """
-    archive = nsys_jax(["python", "-c", "import jax; jax.jit(lambda x: x*2)(4)"])
+    archive = nsys_jax(
+        ["python", "-c", "import jax; jax.jit(lambda x: x*2)(4)"], out_dir=tmp_path
+    )
     with zipfile.ZipFile(archive.name) as ifile:
         # The combination of -c and JAX suppressing references to its own source code
         # should mean that no source code files are gathered.

.github/container/nsys_jax/tests/test_example_program.py

@@ -11,12 +11,13 @@
 
 
 @pytest.fixture(scope="module")
-def profiler_data():
+def profiler_data(tmp_path_factory):
     """
     Fixture that yields the loaded result of profiling example_program.py with nsys-jax.
     """
     outdir = nsys_jax_archive(
-        [sys.executable, os.path.join(os.path.dirname(__file__), "example_program.py")]
+        [sys.executable, os.path.join(os.path.dirname(__file__), "example_program.py")],
+        out_dir=tmp_path_factory.mktemp("test_example_program"),
     )
     return load_profiler_data(pathlib.Path(outdir.name))
 

.github/container/nsys_jax/tests/test_exit_code_handling.py

@@ -20,7 +20,7 @@
 
 
 @pytest.mark.parametrize("kill", kill_args.keys())
-def test_program_that_is_killed_by_nsys(kill):
+def test_program_that_is_killed_by_nsys(kill, tmp_path):
     """
     The default --capture-range-end=stop-shutdown behaviour of nsys profile causes nsys
     to kill the profiled process after cudaProfilerStop if
@@ -38,12 +38,13 @@ def test_program_that_is_killed_by_nsys(kill):
             sys.executable,
             cuda_profiler_api,
             "sleep",
-        ]
+        ],
+        out_dir=tmp_path,
     )
     assert os.path.isfile(output_zip.name)
 
 
-def test_program_that_fails_after_cuda_profiler_stop():
+def test_program_that_fails_after_cuda_profiler_stop(tmp_path):
     """
     With --capture-range=stop then nsys-jax should still propagate a failure code from
     the application.
@@ -56,14 +57,17 @@ def test_program_that_fails_after_cuda_profiler_stop():
             sys.executable,
             cuda_profiler_api,
             "exit42",
-        ]
+        ],
+        out_dir=tmp_path,
     )
     assert result.returncode == 42, result
     assert os.path.isfile(output_zip.name)
 
 
 @pytest.mark.parametrize("kill", kill_args.keys())
-def test_program_that_fails_after_cuda_profiler_stop_as_nsys_tries_to_kill_it(kill):
+def test_program_that_fails_after_cuda_profiler_stop_as_nsys_tries_to_kill_it(
+    tmp_path, kill
+):
     """
     The racy case, where either nsys sends SIGTERM or the application exits with 42.
     Also cover the case where --capture-range-end is not passed explicitly.
@@ -76,7 +80,8 @@ def test_program_that_fails_after_cuda_profiler_stop_as_nsys_tries_to_kill_it(ki
             sys.executable,
             cuda_profiler_api,
             "exit42",
-        ]
+        ],
+        out_dir=tmp_path,
     )
     # If nsys sends SIGTERM fast enough, the child process will exit due to that and
     # nsys-jax will return 0 (because that is the expected result). The application

.github/container/nsys_jax/tests/test_hlo_program.py

@@ -18,7 +18,7 @@
 
 
 @pytest.fixture(scope="module")
-def profiler_data():
+def profiler_data(tmp_path_factory):
     """
     Fixture that yields the loaded result of profiling example_program.py with nsys-jax.
     """
@@ -27,7 +27,8 @@ def profiler_data():
             hlo_runner_main,
             f"--num_repeats={num_repeats}",
             os.path.join(os.path.dirname(__file__), "example.hlo"),
-        ]
+        ],
+        out_dir=tmp_path_factory.mktemp("test_hlo_program"),
     )
     return load_profiler_data(pathlib.Path(outdir.name))
 

.github/container/nsys_jax/tests/test_jax_nccl_test.py

@@ -38,7 +38,7 @@ def set_env(monkeypatch):
 
 
 @pytest.mark.parametrize("collection", ["full", "partial"])
-def test_jax_nccl_single_process(monkeypatch, collection):
+def test_jax_nccl_single_process(monkeypatch, tmp_path, collection):
     set_env(monkeypatch)
     nsys_jax(
         capture_args(collection)
@@ -49,7 +49,8 @@ def test_jax_nccl_single_process(monkeypatch, collection):
             "summary",
             "--",
             "jax-nccl-test",
-        ]
+        ],
+        out_dir=tmp_path,
     )
 
 
@@ -65,7 +66,7 @@ def test_jax_nccl_single_process(monkeypatch, collection):
 
 @pytest.mark.parametrize("process_count", process_counts_to_test)
 @pytest.mark.parametrize("collection", ["full", "partial"])
-def test_jax_nccl_multi_process(monkeypatch, process_count, collection):
+def test_jax_nccl_multi_process(monkeypatch, tmp_path, process_count, collection):
     assert device_count % process_count == 0, (device_count, process_count)
     gpus_per_process = device_count // process_count
     set_env(monkeypatch)
@@ -85,6 +86,7 @@ def test_jax_nccl_multi_process(monkeypatch, process_count, collection):
             str(gpus_per_process),
             "--distributed",
         ],
+        out_dir=tmp_path,
     )
     combined_output = tempfile.NamedTemporaryFile(suffix=".zip")
     subprocess.run(

.github/container/nsys_jax/tests/test_multi_process_program.py

@@ -18,7 +18,7 @@
 
 
 @pytest.fixture(scope="module")
-def individual_results():
+def individual_results(tmp_path_factory):
     """
     Fixture that yields the .zip files from individual subprocesses.
     """
@@ -33,6 +33,7 @@ def individual_results():
             "--rank",
             str(rank),
         ],
+        out_dir=tmp_path_factory.mktemp("test_multi_progress_program"),
     )
 
 

.github/container/nsys_jax/tests/test_overlap_program.py

@@ -14,16 +14,19 @@
 
 
 @pytest.fixture(scope="module")
-def profiler_data_full():
+def profiler_data_full(tmp_path_factory):
     """
     Fixture that yields the loaded result of profiling overlap_program.py with nsys-jax.
     """
-    outdir = nsys_jax_archive([sys.executable, overlap_program])
+    outdir = nsys_jax_archive(
+        [sys.executable, overlap_program],
+        out_dir=tmp_path_factory.mktemp("test_overlap_program_full"),
+    )
     return load_profiler_data(pathlib.Path(outdir.name))
 
 
 @pytest.fixture(scope="module")
-def profiler_data_narrow():
+def profiler_data_narrow(tmp_path_factory):
     """
     Fixture that yields the loaded result of profiling overlap_program.py with nsys-jax.
     """
@@ -34,7 +37,8 @@ def profiler_data_narrow():
             "--",
             sys.executable,
             overlap_program,
-        ]
+        ],
+        out_dir=tmp_path_factory.mktemp("test_overlap_program_narrow"),
     )
     return load_profiler_data(pathlib.Path(outdir.name))
 

.github/workflows/_ci.yaml

@@ -343,7 +343,7 @@ jobs:
             pip install pytest-reportlog nsys-jax[test]
             # abuse knowledge that nsys-jax is installed editable, so the tests exist
             test_path=$(python -c 'import importlib.resources; print(importlib.resources.files("nsys_jax").joinpath("..", "tests").resolve())')
-            pytest --report-log=/opt/output/pytest-report.jsonl "${test_path}"
+            pytest --basetemp=/opt/output/pytest-tmp --report-log=/opt/output/pytest-report.jsonl "${test_path}"
         EOF
       STATISTICS_SCRIPT: |
         summary_line=$(tail -n1 test-nsys-jax.log)
@@ -359,6 +359,7 @@ jobs:
         # pytest-driven part
         test-nsys-jax.log
         pytest-report.jsonl
+        pytest-tmp/
     secrets: inherit
 
   # test-nsys-jax generates several fresh .zip archive outputs by running nsys-jax with real GPU hardware; this test