diff --git a/docker/Dockerfile b/docker/Dockerfile index 84a1802dbe03..aa3aad21d6c0 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -244,9 +244,15 @@ 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 # Install EP kernels(pplx-kernels and DeepEP) +ARG PPLX_COMMIT_HASH +ARG DEEPEP_COMMIT_HASH RUN --mount=type=cache,target=/root/.cache/uv \ export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \ - /tmp/install_python_libraries.sh /tmp/ep_kernels_workspace wheel && \ + /tmp/install_python_libraries.sh \ + --workspace /tmp/ep_kernels_workspace \ + --mode wheel \ + ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \ + ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \ find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete # Check the size of the wheel if RUN_WHEEL_CHECK is true diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 1cea1bef8dbc..88be5cd778ff 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -1,22 +1,68 @@ #!/usr/bin/env bash set -ex -# usage: ./build.sh [workspace_dir] [mode] -# mode: "install" (default) → install directly into current Python env -# "wheel" → build wheels into WORKSPACE/dist +# usage: ./install_python_libraries.sh [options] +# --workspace workspace directory (default: ./ep_kernels_workspace) +# --mode "install" (default) or "wheel" +# --pplx-ref pplx-kernels commit hash +# --deepep-ref DeepEP commit hash + +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.9 +WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace} +MODE=${MODE:-install} + +# Parse arguments +while [[ $# -gt 0 ]]; do + case $1 in + --workspace) + if [[ -z "$2" || "$2" =~ ^- ]]; then + echo "Error: --workspace requires an argument." >&2 + exit 1 + fi + WORKSPACE="$2" + shift 2 + ;; + --mode) + if [[ -z "$2" || "$2" =~ ^- ]]; then + echo "Error: --mode requires an argument." >&2 + exit 1 + fi + MODE="$2" + shift 2 + ;; + --pplx-ref) + if [[ -z "$2" || "$2" =~ ^- ]]; then + echo "Error: --pplx-ref requires an argument." >&2 + exit 1 + fi + PPLX_COMMIT_HASH="$2" + shift 2 + ;; + --deepep-ref) + if [[ -z "$2" || "$2" =~ ^- ]]; then + echo "Error: --deepep-ref requires an argument." >&2 + exit 1 + fi + DEEPEP_COMMIT_HASH="$2" + shift 2 + ;; + *) + echo "Error: Unknown argument '$1'" >&2 + exit 1 + ;; + esac +done -WORKSPACE=${1:-$(pwd)/ep_kernels_workspace} -MODE=${2:-install} mkdir -p "$WORKSPACE" WHEEL_DIR="$WORKSPACE/dist" mkdir -p "$WHEEL_DIR" -NVSHMEM_VER=3.3.9 pushd "$WORKSPACE" -CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} - # install dependencies if not installed if [ -z "$VIRTUAL_ENV" ]; then uv pip install --system cmake torch ninja @@ -133,7 +179,7 @@ do_build \ "https://github.com/ppl-ai/pplx-kernels" \ "pplx-kernels" \ "setup.py" \ - "12cecfd" \ + "$PPLX_COMMIT_HASH" \ "" # build DeepEP @@ -141,7 +187,7 @@ do_build \ "https://github.com/deepseek-ai/DeepEP" \ "DeepEP" \ "setup.py" \ - "73b6ea4" \ + "$DEEPEP_COMMIT_HASH" \ "export NVSHMEM_DIR=$WORKSPACE/nvshmem; " if [ "$MODE" = "wheel" ]; then