Allow nsys 2024.7 installation.

NVIDIA · Nov 27, 2024 · 9e511ee · 9e511ee
1 parent e9394cc
commit 9e511ee
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 3 deletions.
diff --git 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
@@ -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))