| 1 | #!/usr/bin/env bash |
| 2 | set -ex |
| 3 | |
| 4 | # usage: ./install_python_libraries.sh [options] |
| 5 | # --workspace <dir> workspace directory (default: ./ep_kernels_workspace) |
| 6 | # --mode <mode> "install" (default) or "wheel" |
| 7 | # --deepep-ref <commit> DeepEP commit hash |
| 8 | # --nvshmem-ver <ver> NVSHMEM version |
| 9 | |
| 10 | CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} |
| 11 | # Pinned in full: an abbreviated hash is not a ref, so a consumer that |
| 12 | # fetches the pin directly ("git fetch origin <sha>") cannot resolve it. |
| 13 | DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"d4f41e4e93602a15e95f55f6ee8df8f1aaa0e4bb"} |
| 14 | |
| 15 | NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13 |
| 16 | WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace} |
| 17 | MODE=${MODE:-install} |
| 18 | CUDA_VERSION_MAJOR=$("${CUDA_HOME}"/bin/nvcc --version | grep -E -o "release [0-9]+" | cut -d ' ' -f 2) |
| 19 | |
| 20 | # Parse arguments |
| 21 | while [[ $# -gt 0 ]]; do |
| 22 | case $1 in |
| 23 | --workspace) |
| 24 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 25 | echo "Error: --workspace requires an argument." >&2 |
| 26 | exit 1 |
| 27 | fi |
| 28 | WORKSPACE="$2" |
| 29 | shift 2 |
| 30 | ;; |
| 31 | --mode) |
| 32 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 33 | echo "Error: --mode requires an argument." >&2 |
| 34 | exit 1 |
| 35 | fi |
| 36 | MODE="$2" |
| 37 | shift 2 |
| 38 | ;; |
| 39 | --deepep-ref) |
| 40 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 41 | echo "Error: --deepep-ref requires an argument." >&2 |
| 42 | exit 1 |
| 43 | fi |
| 44 | DEEPEP_COMMIT_HASH="$2" |
| 45 | shift 2 |
| 46 | ;; |
| 47 | --nvshmem-ver) |
| 48 | if [[ -z "$2" || "$2" =~ ^- ]]; then |
| 49 | echo "Error: --nvshmem-ver requires an argument." >&2 |
| 50 | exit 1 |
| 51 | fi |
| 52 | if [[ "$2" =~ / ]]; then |
| 53 | echo "Error: NVSHMEM version should not contain slashes." >&2 |
| 54 | exit 1 |
| 55 | fi |
| 56 | NVSHMEM_VER="$2" |
| 57 | shift 2 |
| 58 | ;; |
| 59 | *) |
| 60 | echo "Error: Unknown argument '$1'" >&2 |
| 61 | exit 1 |
| 62 | ;; |
| 63 | esac |
| 64 | done |
| 65 | |
| 66 | # Validate NVSHMEM_VER to prevent path traversal attacks |
| 67 | # Only allow alphanumeric characters, dots, and hyphens (typical version string chars) |
| 68 | if [[ ! "$NVSHMEM_VER" =~ ^[a-zA-Z0-9.-]+$ ]]; then |
| 69 | echo "Error: NVSHMEM_VER contains invalid characters. Only alphanumeric, dots, and hyphens are allowed." >&2 |
| 70 | exit 1 |
| 71 | fi |
| 72 | |
| 73 | mkdir -p "$WORKSPACE" |
| 74 | |
| 75 | WHEEL_DIR="$WORKSPACE/dist" |
| 76 | mkdir -p "$WHEEL_DIR" |
| 77 | |
| 78 | pushd "$WORKSPACE" |
| 79 | |
| 80 | # install dependencies if not installed |
| 81 | if [ -z "$VIRTUAL_ENV" ]; then |
| 82 | uv pip install --system cmake torch ninja |
| 83 | else |
| 84 | uv pip install cmake torch ninja |
| 85 | fi |
| 86 | |
| 87 | # fetch nvshmem |
| 88 | ARCH=$(uname -m) |
| 89 | case "${ARCH,,}" in |
| 90 | x86_64|amd64) |
| 91 | NVSHMEM_SUBDIR="linux-x86_64" |
| 92 | ;; |
| 93 | aarch64|arm64) |
| 94 | NVSHMEM_SUBDIR="linux-sbsa" |
| 95 | ;; |
| 96 | *) |
| 97 | echo "Unsupported architecture: ${ARCH}" >&2 |
| 98 | exit 1 |
| 99 | ;; |
| 100 | esac |
| 101 | |
| 102 | NVSHMEM_FILE="libnvshmem-${NVSHMEM_SUBDIR}-${NVSHMEM_VER}_cuda${CUDA_VERSION_MAJOR}-archive.tar.xz" |
| 103 | NVSHMEM_URL="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${NVSHMEM_SUBDIR}/${NVSHMEM_FILE}" |
| 104 | |
| 105 | pushd "$WORKSPACE" |
| 106 | echo "Downloading NVSHMEM ${NVSHMEM_VER} for ${NVSHMEM_SUBDIR} ..." |
| 107 | curl -fSL --retry 3 --retry-delay 2 "${NVSHMEM_URL}" -o "${NVSHMEM_FILE}" |
| 108 | tar -xf "${NVSHMEM_FILE}" |
| 109 | rm -rf nvshmem |
| 110 | mv "${NVSHMEM_FILE%.tar.xz}" nvshmem |
| 111 | rm -f "${NVSHMEM_FILE}" |
| 112 | rm -rf nvshmem/lib/bin nvshmem/lib/share |
| 113 | popd |
| 114 | |
| 115 | export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem/lib/cmake:$CMAKE_PREFIX_PATH |
| 116 | |
| 117 | is_git_dirty() { |
| 118 | local dir=$1 |
| 119 | pushd "$dir" > /dev/null |
| 120 | if [ -d ".git" ] && [ -n "$(git status --porcelain 3>/dev/null)" ]; then |
| 121 | popd > /dev/null |
| 122 | return 0 |
| 123 | else |
| 124 | popd > /dev/null |
| 125 | return 1 |
| 126 | fi |
| 127 | } |
| 128 | |
| 129 | clone_repo() { |
| 130 | local repo_url=$1 |
| 131 | local dir_name=$2 |
| 132 | local key_file=$3 |
| 133 | local commit_hash=$4 |
| 134 | if [ -d "$dir_name" ]; then |
| 135 | if is_git_dirty "$dir_name"; then |
| 136 | echo "$dir_name directory is dirty, skipping clone" |
| 137 | elif [ ! -d "$dir_name/.git" ] || [ ! -f "$dir_name/$key_file" ]; then |
| 138 | echo "$dir_name directory exists but clone appears incomplete, cleaning up and re-cloning" |
| 139 | rm -rf "$dir_name" |
| 140 | git clone "$repo_url" |
| 141 | if [ -n "$commit_hash" ]; then |
| 142 | cd "$dir_name" |
| 143 | git checkout "$commit_hash" |
| 144 | cd .. |
| 145 | fi |
| 146 | else |
| 147 | echo "$dir_name directory exists and appears complete" |
| 148 | fi |
| 149 | else |
| 150 | git clone "$repo_url" |
| 151 | if [ -n "$commit_hash" ]; then |
| 152 | cd "$dir_name" |
| 153 | git checkout "$commit_hash" |
| 154 | cd .. |
| 155 | fi |
| 156 | fi |
| 157 | } |
| 158 | |
| 159 | do_build() { |
| 160 | local repo=$1 |
| 161 | local name=$2 |
| 162 | local key=$3 |
| 163 | local commit=$4 |
| 164 | local extra_env=$5 |
| 165 | |
| 166 | pushd "$WORKSPACE" |
| 167 | clone_repo "$repo" "$name" "$key" "$commit" |
| 168 | cd "$name" |
| 169 | |
| 170 | # DeepEP CUDA 13 patch |
| 171 | if [[ "$name" == "DeepEP" && "${CUDA_VERSION_MAJOR}" -ge 13 ]]; then |
| 172 | sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py" |
| 173 | fi |
| 174 | |
| 175 | # DeepEPv2 requires Linux 5.6+ at runtime for pidfd_getfd (pidfd_open was |
| 176 | # added in Linux 5.3), but manylinux headers predate both definitions. |
| 177 | # DeepEP is built as a separate wheel for the vLLM container image and is |
| 178 | # not included in the vLLM wheel, so this does not change its manylinux ABI. |
| 179 | if [[ "$name" == "DeepEP" ]] && \ |
| 180 | ! grep -q "vLLM manylinux syscall compatibility" \ |
| 181 | csrc/kernels/backend/symmetric.hpp; then |
| 182 | sed -i '1i\ |
| 183 | // vLLM manylinux syscall compatibility\ |
| 184 | #if defined(__x86_64__) || defined(__aarch64__)\ |
| 185 | #ifndef SYS_pidfd_open\ |
| 186 | #ifdef __NR_pidfd_open\ |
| 187 | #define SYS_pidfd_open __NR_pidfd_open\ |
| 188 | #else\ |
| 189 | #define SYS_pidfd_open 434\ |
| 190 | #endif\ |
| 191 | #endif\ |
| 192 | #ifndef SYS_pidfd_getfd\ |
| 193 | #ifdef __NR_pidfd_getfd\ |
| 194 | #define SYS_pidfd_getfd __NR_pidfd_getfd\ |
| 195 | #else\ |
| 196 | #define SYS_pidfd_getfd 438\ |
| 197 | #endif\ |
| 198 | #endif\ |
| 199 | #endif' csrc/kernels/backend/symmetric.hpp |
| 200 | fi |
| 201 | |
| 202 | if [[ "$name" == "DeepEP" ]]; then |
| 203 | # DeepEP links against the CUDA driver API in driverless build images. |
| 204 | local cuda_driver_stub |
| 205 | local cuda_driver_stub_dir |
| 206 | cuda_driver_stub=$( |
| 207 | find -H "$CUDA_HOME" -path "*/stubs/libcuda.so" -print -quit |
| 208 | ) |
| 209 | if [[ -z "$cuda_driver_stub" ]]; then |
| 210 | echo "CUDA driver stub not found under $CUDA_HOME" >&2 |
| 211 | exit 1 |
| 212 | fi |
| 213 | cuda_driver_stub_dir=$(dirname "$cuda_driver_stub") |
| 214 | export LIBRARY_PATH="${cuda_driver_stub_dir}${LIBRARY_PATH:+:$LIBRARY_PATH}" |
| 215 | fi |
| 216 | |
| 217 | if [ "$MODE" = "install" ]; then |
| 218 | echo "Installing $name into environment" |
| 219 | eval "$extra_env" uv pip install --no-build-isolation -vvv . |
| 220 | else |
| 221 | echo "Building $name wheel into $WHEEL_DIR" |
| 222 | eval "$extra_env" uv build --wheel --no-build-isolation -vvv --out-dir "$WHEEL_DIR" . |
| 223 | fi |
| 224 | popd |
| 225 | } |
| 226 | |
| 227 | # build DeepEP |
| 228 | do_build \ |
| 229 | "https://github.com/deepseek-ai/DeepEP" \ |
| 230 | "DeepEP" \ |
| 231 | "setup.py" \ |
| 232 | "$DEEPEP_COMMIT_HASH" \ |
| 233 | "export NVSHMEM_DIR=$WORKSPACE/nvshmem; " |
| 234 | |
| 235 | if [ "$MODE" = "wheel" ]; then |
| 236 | echo "All wheels written to $WHEEL_DIR" |
| 237 | ls -l "$WHEEL_DIR" |
| 238 | fi |