diff --git a/docker/Dockerfile b/docker/Dockerfile index b9838f432..0daff3995 100644 --- a/docker/Dockerfile +++ b/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 #################### diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 1bb7fd834..89da24f95 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -6,11 +6,12 @@ set -ex # --mode "install" (default) or "wheel" # --pplx-ref pplx-kernels commit hash # --deepep-ref DeepEP commit hash +# --nvshmem-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"