biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

4
docker/Dockerfile
View File

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

22
tools/ep_kernels/install_python_libraries.sh
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"