[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>
This commit is contained in:
Seiji Eicher
2026-01-05 15:56:08 -08:00
committed by GitHub
parent 9513029898
commit 3c98c2d21b
2 changed files with 24 additions and 2 deletions

View File

@@ -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"