feat(deepep-v2): Add Kubernetes microbenchmarks examples, Refactor Dockerfile

micro-benchmarks/expert-parallelism/deepep-v2-benchmark/README.md

@@ -19,8 +19,10 @@ installation works:
 - **[`deepep.Dockerfile`](./deepep.Dockerfile)** — a reference image showing the full stack the
   script needs (the EFA installer, which bundles the EFA-GDA-capable aws-ofi-nccl plugin, a
   GIN-capable NCCL, torch), built from a bare CUDA base; it can also be used directly for testing.
-- **[`slurm/`](./slurm/)** — launchers that run DeepEP's own `tests/elastic/test_ep.py` to
-  validate the installation:
+- **[`slurm/`](./slurm/)** - launchers that run DeepEP's own `tests/elastic/test_ep.py` to
+  validate the installation.
+- **[`kubernetes/`](./kubernetes/)** - launchers that run DeepEP's own `tests/elastic/test_ep.py` to
+  validate the installation.
 
 | Benchmark | Nodes | Measures |
 |-----------|-------|----------|
@@ -30,15 +32,15 @@ installation works:
 ## ⚠️ Version constraints
 
 > - **NCCL >= 2.31** — the GIN device API (`nccl_device.h`) DeepEP V2 links against.
->   The Dockerfile builds NCCL from source and fails the build if the header is missing.
+>   The Dockerfile install NVIDIA's precompiled NCCL from PyPi as a python package.
 > - **EFA installer >= 1.50** — provides everything EFA-GDA needs in the container (the
 >   userspace runtime and the aws-ofi-nccl plugin with the EFA-GDA GIN backend, type 5,
 >   this benchmark measures). Installers up to 1.49.0 will not work; the image build fails
 >   loudly on a too-old installer. DeepEP's kernels can also run over the CPU-proxy GIN
 >   backend, which has no such floor.
 > - A host EFA kernel driver **>= 3.3.0**, the first release with the completion-counter API
 >   (`efadv_create_comp_cntr`). Run the EFA installer **on the node itself** (not in the
->   container) to install/upgrade the kernel driver; the container ships only the userspace
+>   container) to install/upgrade the kernel driver; the provided Dockerfile ships only the userspace
 >   stack. Check with `modinfo efa | grep ^version`; stock AMIs may ship older.
 > - The **gdrcopy kernel module (`gdrdrv`) loaded on compute nodes** — the launchers
 >   bind-mount `/dev/gdrdrv`, and the plugin's GIN initialization opens a gdr handle.
@@ -58,18 +60,16 @@ it, if left unset, NCCL selects the CPU-proxy backend when running on EFA.
 | Variable | Value | Purpose |
 |----------|-------|---------|
 | `NCCL_GIN_TYPE` | `5` | Pins the EFA-GDA GIN backend; unset falls back to the CPU-proxy backend. |
-| `EP_NCCL_ROOT_DIR` | `$NCCL_HOME` | Points DeepEP's build/runtime checks at the GIN-capable NCCL. |
-| `LD_PRELOAD` | `$NCCL_HOME/lib/libnccl.so.2` | Guarantees the GIN-capable NCCL wins over any other on the loader path. |
 
 ## Prerequisites
 
 - EFA-enabled GPU nodes with the EFA kernel driver installed (section below).
-- **GPU architecture:** the image builds for **Hopper (`sm_90`) and Blackwell (`sm_100`)** by
+- **GPU architecture:** the image builds for **Hopper (`sm_90`) and Blackwell (`sm_100`, `sm_103`)** by
   default, so one image runs on `p5`/`p5en` and `p6`. Earlier DeepEP revisions required
   **CUDA >= 13.1** on p6/Blackwell (ptxas 13.0.88 rejects the 32-bit `st.bulk` size operand);
   DeepEP main carries the fix ([amazon-contributing/DeepEP#3](https://github.com/amazon-contributing/DeepEP/pull/3)),
-  so CUDA 13.0 works as well. Override
-  `--build-arg TORCH_CUDA_ARCH_LIST` / `--build-arg NVCC_GENCODE` for a single-arch image.
+  so CUDA 13.0 works.
+  Override `--build-arg TORCH_CUDA_ARCH_LIST` / `--build-arg NVCC_GENCODE` for a single-arch image.
 - Docker with BuildKit, and enroot + pyxis on the cluster for the Slurm flow.
 
 ## Using the script in your own container
@@ -81,11 +81,10 @@ it, if left unset, NCCL selects the CPU-proxy backend when running on EFA.
 
 The script refuses to run if the NCCL at `--nccl-root` lacks the GIN device API, so it cannot
 silently produce a DeepEP that falls back to a slower path. Runtime still requires the
-EFA-GDA-capable aws-ofi-nccl GIN plugin (bundled with EFA installer >= 1.50): run the EFA
-installer **inside your container** too (with `--skip-kmod`, as the reference Dockerfile does)
-to provide the userspace stack and the plugin, then point `NCCL_NET_PLUGIN` and
-`NCCL_GIN_PLUGIN` at `/opt/amazon/ofi-nccl/lib/libnccl-net-ofi.so`. The host installer run
-(previous section) provides only the kernel driver.
+EFA-GDA-capable aws-ofi-nccl GIN plugin (bundled with EFA installer >= 1.50).
+
+> NOTE: The aws-ofi-nccl plugin has to be in your container, run the EFA Installer inside your container with the `-y --skip-kmod --skip-limit-conf --no-verify` arguments. \
+> For installing the EFA Driver on the host, the EFA Installer bundles the driver, run the EFA Installer on the host with the `--minimal -y` arguments.
 
 ## Building the reference image
 
@@ -96,11 +95,11 @@ DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
   -t deepep-v2:gin .
 ```
 
-Override `NCCL_REF` or `EFA_INSTALLER_VERSION` for different versions:
+Override `NCCL_VERSION` or `EFA_INSTALLER_VERSION` for different versions:
 
 ```bash
 DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
-  --build-arg NCCL_REF=v2.31.2-1 \
+  --build-arg NCCL_VERSION=2.31.2 \
   --build-arg EFA_INSTALLER_VERSION=1.50.0 \
   -t deepep-v2:gin .
 ```

micro-benchmarks/expert-parallelism/deepep-v2-benchmark/deepep.Dockerfile

@@ -1,164 +1,128 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: MIT-0
-#
-# DeepEP V2 image for AWS EFA clusters.
-# See README.md for build and run instructions.
-
-# CUDA 13.1 is the default. Earlier DeepEP revisions needed >= 13.1 on
-# p6/Blackwell (ptxas 13.0.88 rejects the 32-bit st.bulk size operand); DeepEP
-# main carries the 64-bit fix (amazon-contributing/DeepEP#3), so 13.0 works too.
-ARG CUDA_VERSION=13.1.2
+ARG CUDA_VERSION=13.0.2
+
 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
 
-ARG TORCH_VERSION=2.11.0
-# Torch wheel index. Kept explicit rather than derived from nvcc: pytorch.org has
-# no cu131 index, and cu130 wheels run on a 13.1 toolkit (same CUDA major).
-ARG TORCH_CUDA_INDEX=cu130
 ARG GDRCOPY_VERSION=v2.5.2
-
-# EFA installer: provides the complete EFA userspace stack (OpenMPI, the EFA
-# runtime, and the aws-ofi-nccl plugin). EFA-GDA requires installer >= 1.50;
-# the build fails loudly on a too-old installer (plugin gate below).
-# EFA_INSTALLER_TARBALL overrides the download with a local tarball staged next
-# to this Dockerfile.
 ARG EFA_INSTALLER_VERSION=1.50.0
-ARG EFA_INSTALLER_TARBALL=
-
-# NCCL: defaults to the public NVIDIA release. For internal use, override with
-# the upstream fork + staging branch (see build instructions in README.md).
-ARG NCCL_REPO=https://github.com/NVIDIA/nccl.git
-ARG NCCL_REF=v2.31.2-1
-
-# DeepEP V2. The repo is pinned to amazon-contributing/DeepEP inside
-# setup_deepep_gin.sh; only the ref is overridable.
-ARG DEEPEP_REF=main
+ARG DEEPEP_REF="main"
+ARG NCCL_VERSION=2.31.2
+ARG PYTHON_VERSION=3.10
+ARG TORCH_VERSION=2.11.0
+ARG TORCH_CUDA_INDEX=cu130
 
 # CUDA architecture(s), semicolon-separated:
-# 9.0 = Hopper (H100/H200, sm_90); 10.0/10.3 = Blackwell (B200/B300, sm_100/sm_103).
-# Defaults to Hopper + Blackwell so one image runs on p5/p5en and p6.
-ARG TORCH_CUDA_ARCH_LIST="9.0;10.0"
-ARG NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_100,code=sm_100"
+# 9.0 = Hopper (H100, sm_90), 10.0 and 10.3 = Blackwell (B200/B300, sm_100,sm_103). Defaults to
+# both so one image runs on Hopper and Blackwell; override with e.g.
+# --build-arg TORCH_CUDA_ARCH_LIST=9.0 \
+# --build-arg NVCC_GENCODE=-gencode=arch=compute_90,code=sm_90 to build a smaller Hopper-only image.
+ARG TORCH_CUDA_ARCH_LIST="9.0;10.0;10.3"
+ARG NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_100,code=sm_100 -gencode=arch=compute_103,code=sm_103"
+
 ARG CUDA_HOME="/usr/local/cuda"
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-## Remove conflicting distro packages
-RUN apt-get update -y \
-    && apt-get remove -y --allow-change-held-packages \
+RUN apt-get update -y && \
+    apt-get remove -y --allow-change-held-packages \
     ibverbs-utils \
     libibverbs-dev \
     libibverbs1 \
     libmlx5-1 \
     libnccl2 \
-    libnccl-dev || true
-RUN rm -rf /opt/hpcx /usr/local/mpi /etc/ld.so.conf.d/hpcx.conf && ldconfig
-
-## Toolchain
-RUN apt-get update -y && apt-get install -y \
-    apt-utils autoconf automake build-essential check cmake curl debhelper \
-    devscripts git gcc gdb libsubunit-dev libtool ninja-build meson pandoc \
-    pkg-config vim wget \
-    python3.10-dev python3.10-venv python3-distutils \
-    cython3 \
-    libnl-3-dev libnl-route-3-dev libudev-dev libsystemd-dev \
-    libhwloc-dev \
-    && rm -rf /var/lib/apt/lists/*
-RUN apt-get purge -y cuda-compat-* || true
+    libnccl-dev
+
+RUN rm -rf /opt/hpcx \
+    && rm -rf /usr/local/mpi \
+    && rm -f /etc/ld.so.conf.d/hpcx.conf \
+    && ldconfig
+
+RUN apt-get install -y --no-install-recommends \
+    apt-utils \
+    autoconf \
+    automake \
+    build-essential \
+    check \
+    cmake \
+    curl \
+    debhelper \
+    devscripts \
+    git \
+    gcc \
+    gdb \
+    libsubunit-dev \
+    libtool \
+    openssh-client \
+    openssh-server \
+    pkg-config \
+    python3-distutils \
+    vim \
+    python${PYTHON_VERSION}-dev \
+    python${PYTHON_VERSION}-venv
+
+RUN apt-get purge -y cuda-compat-*
+
+RUN mkdir -p /var/run/sshd
+RUN sed -i 's/[ #]\(.*StrictHostKeyChecking \).*/ \1no/g' /etc/ssh/ssh_config && \
+    echo "    UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \
+    sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config
 
 RUN curl https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py \
     && python3 /tmp/get-pip.py \
-    && pip3 install awscli nvidia-ml-py ninja Cython
-
-## GDRCopy
-# Required by aws-ofi-nccl: the plugin's GIN initialization opens a gdr handle
-# regardless of backend, so building or running without GDRCopy leaves the GIN
-# backends unable to initialize (verified: DeepEP then fails with "NCCL GIN is
-# unavailable"). Built from source because no distro package ships it. The
-# matching gdrdrv kernel module must be loaded on the host (see README).
+    && pip3 install awscli nvidia-ml-py ninja
+
+#################################################
+## Install NVIDIA GDRCopy
 ARG GDRCOPY_PREFIX="/opt/gdrcopy"
 RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \
     && cd /tmp/gdrcopy \
     && make prefix="${GDRCOPY_PREFIX}" install \
     && rm -rf /tmp/gdrcopy
+
 ENV LD_LIBRARY_PATH="${GDRCOPY_PREFIX}/lib:$LD_LIBRARY_PATH"
 ENV LIBRARY_PATH="${GDRCOPY_PREFIX}/lib:$LIBRARY_PATH"
 ENV PATH="${GDRCOPY_PREFIX}/bin:$PATH"
 
-## EFA installer (OpenMPI + EFA runtime; local tarball wins when set)
-RUN --mount=type=bind,target=/ctx \
-    apt-get update -y \
-    && cd /tmp \
-    && if [ -n "${EFA_INSTALLER_TARBALL}" ]; then \
-         echo "=== EFA installer: LOCAL tarball /ctx/${EFA_INSTALLER_TARBALL} ==="; \
-         if [ ! -f "/ctx/${EFA_INSTALLER_TARBALL}" ]; then \
-           echo "ERROR: ${EFA_INSTALLER_TARBALL} not found in build context" >&2; exit 1; \
-         fi; \
-         tar -xf "/ctx/${EFA_INSTALLER_TARBALL}"; \
-       else \
-         echo "=== EFA installer: downloading v${EFA_INSTALLER_VERSION} ==="; \
-         curl -fsSL -O "https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz" \
-         && tar -xf "aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz"; \
-       fi \
-    && cd aws-efa-installer \
-    && ./efa_installer.sh --disable-ngc -y --skip-kmod --skip-limit-conf --no-verify \
-    && ldconfig \
-    && { nm -D /opt/amazon/ofi-nccl/lib/libnccl-net-ofi.so 2>/dev/null | grep -qw ncclGinPlugin_v14 \
-         && echo "OK: bundled plugin is EFA-GDA-capable (ncclGinPlugin_v14 present)" \
-         || { echo "ERROR: installer's plugin has no ncclGinPlugin_v14 export: EFA-GDA requires installer >= 1.50" >&2; exit 1; }; } \
-    && rm -rf /tmp/aws-efa-installer* /var/lib/apt/lists/*
+#################################################
+## Install EFA installer
+RUN cd $HOME && \
+    curl --retry 3 --retry-delay 2 -fsSL -o aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \
+        https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz && \
+    tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz && \
+    cd aws-efa-installer && \
+    apt-get update && \
+    ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify && \
+    cd .. && rm -rf aws-efa-installer* && \
+    ldconfig
+
 ENV LD_LIBRARY_PATH=/opt/amazon/openmpi/lib:$LD_LIBRARY_PATH
-ENV PATH=/opt/amazon/openmpi/bin:/opt/amazon/efa/bin:$PATH
-ENV EFA_PREFIX=/opt/amazon/efa
-
-## NCCL
-ENV NCCL_HOME=/opt/nccl/build
-RUN git clone ${NCCL_REPO} /opt/nccl \
-    && cd /opt/nccl \
-    && git checkout ${NCCL_REF} \
-    && make -j"$(nproc)" src.build CUDA_HOME="${CUDA_HOME}" NVCC_GENCODE="${NVCC_GENCODE}" \
-    && test -f ${NCCL_HOME}/include/nccl_device.h \
-       || (echo "ERROR: nccl_device.h missing -- this NCCL is not GIN-capable" >&2 && exit 1)
-ENV LD_LIBRARY_PATH="${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
-
-## aws-ofi-nccl: the EFA installer (>= 1.50) bundles the plugin with EFA-GDA
-## support; the installer layer above gates on the plugin's ncclGinPlugin_v14
-## export (v11/v13 op-tables are exported by every build; the v14 table is
-## EFA-GDA-only), so a too-old installer fails that layer immediately.
-ENV OFI_HOME=/opt/amazon/ofi-nccl
-ENV LD_LIBRARY_PATH="${OFI_HOME}/lib:${LD_LIBRARY_PATH}"
-ENV NCCL_NET_PLUGIN="${OFI_HOME}/lib/libnccl-net-ofi.so"
-ENV NCCL_GIN_PLUGIN="${OFI_HOME}/lib/libnccl-net-ofi.so"
-
-# Keep NCCL's bootstrap off the docker/loopback/veth interfaces.
-ENV NCCL_SOCKET_IFNAME=^docker,lo,veth
-
-## PyTorch.
-# The torch wheel drags in pip's stock NCCL (nvidia-nccl-cu*), which predates GIN.
-# We remove it rather than keep both: DeepEP refuses to import when it detects two
-# NCCL runtimes, and with both on the loader path which one wins is load-order
-# luck. The GIN capability itself is enforced above at NCCL build time
-# (nccl_device.h check errors the build), so this cannot silently downgrade: the
-# image ends up with exactly one NCCL, the GIN-capable one it was built and
-# tested with.
-RUN pip3 install torch==${TORCH_VERSION} numpy --index-url https://download.pytorch.org/whl/${TORCH_CUDA_INDEX} \
-    && pip3 uninstall -y nvidia-nccl-cu13 nvidia-nccl-cu12 nvidia-nccl 2>/dev/null || true
-
-## No NVSHMEM: DeepEP V2's transport is NCCL GIN, so the NVSHMEM backend is not
-## built (the DeepEP fork's build only links NVSHMEM when NVSHMEM_DIR points at an
-## install). To additionally build the legacy NVSHMEM backend, install the
-## libnvshmem3-*-cuda-13 packages and set NVSHMEM_DIR before running the setup script.
-
-## DeepEP V2.
-ARG DEEPEP_PREFIX="/opt/amazon/deepep"
+ENV PATH=/opt/amazon/openmpi/bin/:/opt/amazon/efa/bin:$PATH
+
+RUN rm -rf /var/lib/apt/lists/*
+
+## Set Open MPI variables to exclude network interface and conduit.
+ENV OMPI_MCA_pml=^ucx            \
+    OMPI_MCA_btl=tcp,self           \
+    OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent\
+    OPAL_PREFIX=/opt/amazon/openmpi \
+    NCCL_SOCKET_IFNAME=^docker,lo,veth
+
+## Turn off PMIx Error https://github.com/open-mpi/ompi/issues/7516
+ENV PMIX_MCA_gds=hash
+
+################################ PyTorch ########################################
+RUN pip3 install torch==${TORCH_VERSION} numpy --index-url https://download.pytorch.org/whl/${TORCH_CUDA_INDEX}
+
+################################ DeepEP v2 ########################################
 RUN --mount=type=bind,source=setup_deepep_gin.sh,target=/tmp/setup_deepep_gin.sh \
-    set -e; \
-    TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
+    pip3 uninstall -y nvidia-nccl-cu13 nvidia-nccl-cu12 nvidia-nccl 2>/dev/null || true && \
+    CUDA_MAJOR=$(nvcc --version | grep -oP 'release \K[0-9]+') && \
+    pip3 install --no-deps nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION} && \
     /tmp/setup_deepep_gin.sh \
-        --deepep-ref "${DEEPEP_REF}" \
-        --deepep-prefix "${DEEPEP_PREFIX}" \
-        --nccl-root "${NCCL_HOME}"
+      --deepep-ref $DEEPEP_REF \
+      --nccl-root "/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvidia/nccl"
 
-## Runtime env
 ENV NVIDIA_GDRCOPY=enabled
-ENV NCCL_OFI_RDMA_GDR_FLUSH_DISABLE=0
-WORKDIR /root
+ENV NCCL_NET_PLUGIN=ofi
+ENV NCCL_TUNER_PLUGIN=ofi