diff --git a/.github/container/install-nsight.sh b/.github/container/install-nsight.sh index dc0ef92cb..976abd366 100755 --- a/.github/container/install-nsight.sh +++ b/.github/container/install-nsight.sh @@ -12,7 +12,7 @@ export DEBIAN_FRONTEND=noninteractive export TZ=America/Los_Angeles apt-get update -apt-get install -y nsight-compute nsight-systems-cli-2024.6.1 +apt-get install -y nsight-compute nsight-systems-cli apt-get clean rm -rf /var/lib/apt/lists/* diff --git a/.github/container/nsys_jax/nsys_jax/scripts/patch_nsys.py b/.github/container/nsys_jax/nsys_jax/scripts/patch_nsys.py index 8b732283c..72e982b82 100644 --- a/.github/container/nsys_jax/nsys_jax/scripts/patch_nsys.py +++ b/.github/container/nsys_jax/nsys_jax/scripts/patch_nsys.py @@ -3,7 +3,7 @@ import shutil import subprocess -patch_content = r"""diff --git a/nsys_recipe/lib/nvtx.py b/nsys_recipe/lib/nvtx.py +nsys_2024_5_and_6_patch_content = r"""diff --git a/nsys_recipe/lib/nvtx.py b/nsys_recipe/lib/nvtx.py index 2470043..7abf892 100644 --- a/nsys_recipe/lib/nvtx.py +++ b/nsys_recipe/lib/nvtx.py @@ -29,6 +29,36 @@ "rangeId": "Range ID", """ +nsys_2024_7_patch_content = r'''diff --git a/nsys_recipe/lib/nvtx.py b/nsys_recipe/lib/nvtx.py +index 1e958f8..1b6138d 100644 +--- a/nsys_recipe/lib/nvtx.py ++++ b/nsys_recipe/lib/nvtx.py +@@ -162,7 +162,7 @@ def _compute_gpu_projection_df(nvtx_df, cuda_df, cuda_nvtx_index_map): + starts.append(start) + ends.append(end) + +- return ( ++ df = ( + pd.DataFrame( + { + "text": nvtx_df.loc[list(nvtx_gpu_end_dict.keys()) + indices, "text"], +@@ -172,10 +172,13 @@ def _compute_gpu_projection_df(nvtx_df, cuda_df, cuda_nvtx_index_map): + "tid": nvtx_df.loc[list(nvtx_gpu_end_dict.keys()) + indices, "tid"], + } + ) +- .sort_values(by=["start", "end"], ascending=[True, False]) +- .reset_index(drop=True) ++ .reset_index() + ) + ++ return df.sort_values( ++ by=["start", "end", "index"], ascending=[True, False, True] ++ ).drop(columns=["index"]) ++ + + def _find_cuda_nvtx_ranges(nvtx_df, cuda_df): + """Find the NVTX ranges that enclose each CUDA operation. +''' def main(): """ @@ -41,7 +71,14 @@ def main(): r"^NVIDIA Nsight Systems version (\d+\.\d+\.\d+)\.\d+-\d+v\d+$", nsys_version ) assert m is not None, f"Could not parse: {nsys_version}" - if m.group(1) in {"2024.5.1", "2024.6.1"}: + match m.group(1): + case "2024.5.1" | "2024.6.1": + patch_content = nsys_2024_5_and_6_patch_content + case "2024.7.1": + patch_content = nsys_2024_7_patch_content + case _: + patch_content = None + if patch_content is not None: print(f"Patching Nsight Systems version {m.group(1)}") # e.g. /opt/nvidia/nsight-systems-cli/2024.7.1/target-linux-x64 tdir = os.path.dirname(os.path.realpath(nsys))