[CI/Build] Allow user to configure NVSHMEM version via ENV or command line (#30732)

Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

docker/Dockerfile

@@ -273,6 +273,7 @@ RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
 COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
 ARG PPLX_COMMIT_HASH
 ARG DEEPEP_COMMIT_HASH
+ARG NVSHMEM_VER
 RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p /tmp/ep_kernels_workspace/dist && \
     export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
@@ -280,7 +281,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --workspace /tmp/ep_kernels_workspace \
         --mode wheel \
         ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
-        ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \
+        ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \
+        ${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \
     find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
 #################### EXTENSIONS BUILD IMAGE ####################
 

tools/ep_kernels/install_python_libraries.sh

@@ -6,11 +6,12 @@ set -ex
 #   --mode <mode>        "install" (default) or "wheel"
 #   --pplx-ref <commit>  pplx-kernels commit hash
 #   --deepep-ref <commit> DeepEP commit hash
+#   --nvshmem-ver <ver>  NVSHMEM version 
 
 CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
 PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"}
 DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
-NVSHMEM_VER=3.3.24  # Suppports both CUDA 12 and 13
+NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"}  # Default supports both CUDA 12 and 13
 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
 MODE=${MODE:-install}
 CUDA_VERSION_MAJOR=$(${CUDA_HOME}/bin/nvcc --version | egrep -o "release [0-9]+" | cut -d ' ' -f 2)
@@ -50,6 +51,18 @@ while [[ $# -gt 0 ]]; do
             DEEPEP_COMMIT_HASH="$2"
             shift 2
             ;;
+        --nvshmem-ver)
+            if [[ -z "$2" || "$2" =~ ^- ]]; then
+                echo "Error: --nvshmem-ver requires an argument." >&2
+                exit 1
+            fi
+            if [[ "$2" =~ / ]]; then
+                echo "Error: NVSHMEM version should not contain slashes." >&2
+                exit 1
+            fi
+            NVSHMEM_VER="$2"
+            shift 2
+            ;;
         *)
             echo "Error: Unknown argument '$1'" >&2
             exit 1
@@ -57,6 +70,13 @@ while [[ $# -gt 0 ]]; do
     esac
 done
 
+# Validate NVSHMEM_VER to prevent path traversal attacks
+# Only allow alphanumeric characters, dots, and hyphens (typical version string chars)
+if [[ ! "$NVSHMEM_VER" =~ ^[a-zA-Z0-9.-]+$ ]]; then
+    echo "Error: NVSHMEM_VER contains invalid characters. Only alphanumeric, dots, and hyphens are allowed." >&2
+    exit 1
+fi
+
 mkdir -p "$WORKSPACE"
 
 WHEEL_DIR="$WORKSPACE/dist"