Build and release sgl-deep-ep wheels (#33498)

2026-08-06 02:07:36 -07:00
parent 8b29c90218
commit ba9074035a
4 changed files with 615 additions and 0 deletions
@@ -0,0 +1,75 @@
+ARG BASE_IMAGE=pytorch/manylinux2_28-builder
+ARG CUDA_VERSION=13.0
+
+FROM ${BASE_IMAGE}:cuda${CUDA_VERSION}
+
+ARG ARCHITECTURE=x86_64
+ARG CUDA_TAG=cu130
+ARG CUDA_VERSION=13.0
+ARG GDRCOPY_VERSION=2.5.1
+ARG PYTHON_TAG=cp312-cp312
+ARG TORCH_VERSION=2.11.0
+
+ENV CUDA_HOME=/usr/local/cuda
+ENV GDRCOPY_HOME=/usr/local
+ENV LD_LIBRARY_PATH=/usr/local/lib:/usr/local/lib64:/usr/local/cuda/lib64:${LD_LIBRARY_PATH}
+ENV PATH=/opt/python/${PYTHON_TAG}/bin:${PATH}
+ENV PYTHON_BIN=/opt/python/${PYTHON_TAG}/bin/python
+
+# These mirror the build and RDMA dependencies used by ci_install_deepep.sh.
+# The image builds only GDRCopy's user-space library: the host owns gdrdrv and
+# passes /dev/gdrdrv through at runtime.
+RUN yum install -y --nogpgcheck \
+        cmake \
+        curl \
+        gcc \
+        gcc-c++ \
+        git \
+        infiniband-diags \
+        libfabric \
+        libfabric-devel \
+        libibverbs \
+        libibverbs-devel \
+        libibverbs-utils \
+        librdmacm \
+        librdmacm-devel \
+        make \
+        openssh-server \
+        patchelf \
+        perftest \
+        pkgconfig \
+        rdma-core \
+        wget \
+    && yum clean all \
+    && rm -rf /var/cache/yum
+
+RUN set -eux; \
+    if [ "${ARCHITECTURE}" = aarch64 ]; then cuda_target=sbsa; else cuda_target="${ARCHITECTURE}"; fi; \
+    cuda_stub="/usr/local/cuda-${CUDA_VERSION}/targets/${cuda_target}-linux/lib/stubs/libcuda.so"; \
+    test -f "${cuda_stub}"; \
+    mkdir -p /usr/lib64 "/usr/lib/${ARCHITECTURE}-linux-gnu"; \
+    ln -sf "${cuda_stub}" /usr/lib64/libcuda.so; \
+    ln -sf "${cuda_stub}" "/usr/lib/${ARCHITECTURE}-linux-gnu/libcuda.so"
+
+RUN git clone --depth 1 --branch "v${GDRCOPY_VERSION}" \
+        https://github.com/NVIDIA/gdrcopy.git /opt/gdrcopy \
+    && make -C /opt/gdrcopy CUDA="${CUDA_HOME}" prefix=/usr/local lib_install \
+    && ldconfig \
+    && test -f /usr/local/include/gdrapi.h \
+    && ldconfig -p | grep -q libgdrapi
+
+RUN --mount=type=cache,id=sgl-deep-ep-pip-${CUDA_TAG}-${PYTHON_TAG}-${ARCHITECTURE},target=/root/.cache/pip \
+    set -eux; \
+    "${PYTHON_BIN}" -m pip uninstall -y deep-ep sgl-deep-ep || true; \
+    "${PYTHON_BIN}" -m pip install --upgrade pip; \
+    "${PYTHON_BIN}" -m pip install --force-reinstall \
+        "torch==${TORCH_VERSION}" \
+        --index-url "https://download.pytorch.org/whl/${CUDA_TAG}"; \
+    "${PYTHON_BIN}" -m pip install \
+        "auditwheel>=6.0" \
+        build \
+        ninja \
+        packaging \
+        setuptools \
+        wheel; \
+    "${PYTHON_BIN}" -c 'import torch; assert torch.__version__.startswith("2.11.0"); print(torch.__version__, torch.version.cuda)'