install-nsight-systems.sh: run conditionally (#1496)

Author: olupton

.github/container/Dockerfile.base

@@ -4,6 +4,7 @@ ARG GIT_USER_NAME="JAX Toolbox"
 ARG [email protected]
 ARG CLANG_VERSION=18
 ARG JAX_TOOLBOX_REF=main
+ARG NSIGHT_SYSTEMS_VERSION_OVERRIDE
 
 ###############################################################################
 ## Build base image
@@ -15,6 +16,7 @@ ARG GIT_USER_EMAIL
 ARG GIT_USER_NAME
 ARG CLANG_VERSION
 ARG JAX_TOOLBOX_REF
+ARG NSIGHT_SYSTEMS_VERSION_OVERRIDE
 ENV CUDA_BASE_IMAGE=${BASE_IMAGE}
 
 ###############################################################################
@@ -167,7 +169,7 @@ ADD jax-nccl-test parallel-launch /usr/local/bin/
 ##############################################################################
 
 ADD install-nsight-systems.sh /usr/local/bin
-RUN install-nsight-systems.sh
+RUN if [[ -n "${NSIGHT_SYSTEMS_VERSION_OVERRIDE}" ]]; then install-nsight-systems.sh "${NSIGHT_SYSTEMS_VERSION_OVERRIDE}"; fi
 
 ###############################################################################
 ## Install the nsys-jax JAX/XLA-aware profiling scripts, patch Nsight Systems

.github/container/install-nsight-systems.sh

@@ -1,6 +1,12 @@
 #!/bin/bash
 set -exuo pipefail
 
+version=$1
+if [[ -z "${version}" ]]; then
+  echo "Usage: $0 major.minor.patch"
+  exit 1
+fi
+
 # Remove the symlink that makes `nsys` refer to the CUDA-bundled version:
 rm /usr/local/cuda/bin/nsys
 
@@ -15,7 +21,7 @@ export DEBIAN_FRONTEND=noninteractive
 export TZ=America/Los_Angeles
 
 apt-get update
-apt-get install -y nsight-systems-cli-2025.3.1
+apt-get install -y nsight-systems-cli-${version}
 apt-get clean
 
 rm -rf /var/lib/apt/lists/*

.github/workflows/_build_base.yaml

@@ -135,6 +135,7 @@ jobs:
             GIT_USER_EMAIL=${{ inputs.GIT_USER_EMAIL }}
             BUILD_DATE=${{ inputs.BUILD_DATE }}
             JAX_TOOLBOX_REF=${{ github.head_ref || github.sha }}
+            NSIGHT_SYSTEMS_VERSION_OVERRIDE=2025.3.1
             ${{ inputs.BASE_IMAGE != 'latest' && format('BASE_IMAGE={0}', inputs.BASE_IMAGE) || '' }}
         
       - name: Generate sitrep