Skip to content

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

Merged
KeitaW merged 4 commits into
awslabs:mainfrom
erezzarum:deepep-v2-eks
Aug 29, 2026
Merged

feat(deepep-v2): Add Kubernetes microbenchmarks examples, Refactor Dockerfile#1241
KeitaW merged 4 commits into
awslabs:mainfrom
erezzarum:deepep-v2-eks

Conversation

@erezzarum

@erezzarum erezzarum commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

Purpose

Changes

  • Refactor Dockerfile.
  • Update README.
  • Add DeepEP v2 microbenchmarks examples for running on Kubernetes.

Test Plan

Environment:

  • AWS Service: Amazon EKS
  • Instance type: p6-b300.48xlarge
  • Number of nodes: 2

Test commands:

export IMAGE_URI=<IMAGE URI>
export INSTANCE_TYPE=p6-b300.48xlarge
export GPU_PER_NODE=8
export EFA_PER_NODE=16
export NUM_NODES=2

envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES' \
  < examples/test-intranode.yaml | kubectl apply -f -

envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES' \
  < examples/test-internode.yaml | kubectl apply -f -

Test Results

Directory Structure

3.test_cases/
└── <framework>/                # e.g. pytorch, megatron, jax
    └── <library>/              # e.g. picotron, FSDP, megatron-lm
        └── <model>/            # e.g. SmolLM-1.7B (may be omitted for single-model cases)
            ├── Dockerfile      # Container / environment setup
            ├── README.md       # Overview, prerequisites, usage
            ├── slurm/          # Slurm-specific launch scripts
            ├── kubernetes/     # Kubernetes manifests
            └── hyperpod-eks/   # HyperPod EKS instructions
  • Top-level files (Dockerfile, README.md, training scripts, configs) cover general setup.
  • Subdirectories (slurm/, kubernetes/, hyperpod-eks/) contain service-specific launch instructions.
  • Not all service subdirectories are required — include only the ones relevant to your test case.

Checklist

  • I have read the contributing guidelines.
  • I am working against the latest main branch.
  • I have searched existing open and recently merged PRs to confirm this is not a duplicate.
  • The contribution is self-contained with documentation and scripts.
  • External dependencies are pinned to a specific version or tag (no latest).
  • A README is included or updated with prerequisites, instructions, and known issues.
  • New test cases follow the expected directory structure.

@erezzarum erezzarum changed the title feat(deepep-v2): Add Kubernetes microbenchmarks examples feat(deepep-v2): Add Kubernetes microbenchmarks examples, Refactor Dockerfile Aug 25, 2026
@erezzarum
erezzarum requested a review from KeitaW August 25, 2026 17:12
@erezzarum

Copy link
Copy Markdown
Collaborator Author

This PR merges changes from #1239

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this — the Kubernetes path for the DeepEP V2 benchmark is a genuinely useful
addition, and the MPIJob design underneath it is sound. I took it to a real cluster rather than
reading it: p6-b300.48xlarge on EKS in us-east-1, building the image from this
branch's own Dockerfile and applying the manifests exactly as the README says. Both benchmarks
do pass there — after two changes, one of which is a single line.

To be precise about what was verified at the current head: the image build, the intranode
job (8 ranks, ~720 GB/s over NVLink) and the internode job (16 ranks across 2 nodes,
116–118 GB/s scale-out over EFA-GDA) — all three, with exactly two deltas applied: the
--deepep-repo line removed so the image would build, and /dev/gdrdrv made reachable.

Because the branch moved while I was testing (db4df83c2b537098), a note on what that
changed: the four kubernetes/ files are byte-identical across both commits, so everything
below about the manifests applies as written. The new commit fixed two things I had already hit
— see the closing batch.

Review Batch 1/5 — Build & Image Contract

The image cannot currently be built, and the rewrite removes the build-time checks that used to
make "this image can actually do GIN" non-negotiable.

The local-tarball EFA installer path was removed (nit)

main supported EFA_INSTALLER_TARBALL to build from a tarball staged in the build context;
this rewrite always curls from the internet, so egress-restricted builds that used to work now
fail. Worth calling out in the PR description if the removal is deliberate.

Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/deepep.Dockerfile Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/deepep.Dockerfile Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/deepep.Dockerfile Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/deepep.Dockerfile Outdated

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 2/5 — Device Access & Runtime Environment

This is the one that actually stops the benchmark running on a stock EKS GPU cluster.

NVIDIA_GDRCOPY=enabled was dropped from the image ENV (should fix)

main sets ENV NVIDIA_GDRCOPY=enabled; this rewrite removes it, leaving the two k8s manifests
as its only home. That reverses a convention this repo adopted on purpose (PR #1145 moved this
var out of the launch manifests and into the Dockerfile so it has exactly one home and can't
drift). The slurm path is unaffected — it bind-mounts the device directly — but I'd put the
ENV back rather than have the manifests be the only place it exists.
NCCL_OFI_RDMA_GDR_FLUSH_DISABLE=0 and NCCL_GIN_PLUGIN were dropped in the same rewrite;
NCCL_GIN_PLUGIN's absence is verified harmless (GIN still resolves via NCCL_NET_PLUGIN=ofi,
since one shared object exports the net/rma/gin tables).

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 3/5 — Documentation Consistency

All of these are copy-paste-and-it-fails issues in the new kubernetes/README.md.

Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/kubernetes/README.md Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/kubernetes/README.md Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/kubernetes/README.md Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/kubernetes/README.md Outdated
Comment thread micro-benchmarks/expert-parallelism/deepep-v2-benchmark/README.md Outdated

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 4/5 — Manifest Robustness

GPU_PER_NODE never reaches the benchmark (nit)

It sizes the nvidia.com/gpu request, but test_ep.py --num-processes defaults to 8 and isn't
passed, so GPU_PER_NODE=4 would request 4 GPUs and still spawn 8 ranks. Either thread it
through or note in env_vars.example that it must stay 8.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 5/5 — PR Scope, and Things That Look Great

Scope: the image rewrite probably wants to be its own PR (should fix)

The title says "Add Kubernetes microbenchmarks examples", and the first revision was exactly
that — 4 files, +409. The current head also rewrites deepep.Dockerfile (263 lines changed),
the parent README, and both slurm launchers. The PR comment says it "merges changes from #1239",
but this isn't #1239's Dockerfile — #1239 is still open and its version keeps the
ncclGinPlugin_v14 gate this one drops. So two open PRs now propose different rewrites of the
same file. Whichever lands second will either conflict outright or, if it is rebased onto the
other, carry its own full-file version and undo the first — neither outcome is one you want to
discover at merge time.

Could the image change either move back out into #1239, or #1239 be closed explicitly in favour
of this one? The k8s manifests are independently reviewable and independently useful.

Things That Look Great

  • The MPI topology is right, and I verified it end-to-end at this head. slotsPerWorker: 1
    with -np ${NUM_NODES} -N 1 and the test spawning 8 local ranks matches how test_ep.py
    reads RANK/WORLD_SIZE. Internode ran 16 ranks across 2 p6-b300.48xlarge and passed
    correctness plus bandwidth: 116–118 GB/s scale-out dispatch alongside ~377–385 GB/s
    NVLink, MPIJobSucceeded.
  • The EFA-GDA pin does what it claims — not a CPU-proxy false green. The log shows
    NCCL_GIN_TYPE set by environment to 5., Loaded gin plugin Libfabric_GDAKI (v14), and
    crucially Skipping plugin Libfabric index 3 type 2: NCCL_GIN_TYPE=5 requested — the
    proxy backend is actively rejected. Zero GIN init failures once the device was reachable.
  • The podAffinity design is correct and I confirmed it schedules. Launcher and workers
    co-located as intended, on both the hostname key (intranode) and zone key (internode).
  • The EFA counts in the README and env_vars.example are all accurate — checked each
    against the API rather than assuming: p6-b300.48xlarge = 16, p6-b200.48xlarge = 8,
    p5.48xlarge = 32, and the node reports vpc.amazonaws.com/efa: 16 allocatable.
  • The CUPTI/SYS_ADMIN note is correct and non-obvious. test_ep.py calls bench_kineto
    unconditionally in the perf section for both the intranode and internode paths, so updating
    the v1 wording ("the intranode test does not profile") to cover both was the right call.
  • imagePullPolicy: IfNotPresent on a pinned tag, and an explicit envsubst allow-list —
    both match this repo's conventions.
  • The newest commit fixed two things I'd already hit before I could report them: the image
    now installs openssh-server (mpi-operator injects /usr/sbin/sshd -De into workers that
    declare no command, and the earlier image had no sshd, so every worker died StartError), and
    TORCH_CUDA_ARCH_LIST now genuinely includes 10.3, making the README's sm_103 line true.
    I re-verified both: the head image builds sm_90 + sm_100 + sm_103 clean on CUDA 13.0.2, and
    both benchmarks pass on B300 from the committed manifests. Nice to see the arch list fixed
    properly rather than only the sentence describing it.

Sources

  • NVIDIA k8s-device-plugin gdrdrv auto-injection was a 0.19.0–0.19.2 default-on bug, reverted in
    0.19.3 — issue #1692,
    PR #1837. Cluster under test runs
    nvcr.io/nvidia/k8s-device-plugin:v0.18.1 (verified live, 2026-08-25).
  • mpi-operator injects /usr/sbin/sshd -De for worker containers with no command/args
    mpi_job_controller.go
    (identical at v0.8.2).
  • setup_deepep_gin.sh accepted flags: lines 86–104 of the file in this PR (unmodified here).
  • EFA interface counts: aws ec2 describe-instance-types → p6-b300.48xlarge 16,
    p6-b200.48xlarge 8, p5.48xlarge 32 (verified live, 2026-08-25).
  • nvidia-nccl-cu13==2.31.2 wheel ships nccl_device.h and nccl_device/gin/efa_gda/*
    (verified live from the published wheel, 2026-08-25).
  • https://download.pytorch.org/whl/cu131 → HTTP 403, cu130 → HTTP 200 (verified live, 2026-08-25).
  • setup_deepep_gin.sh:177 still enforces nccl_device.h, so that guarantee survives the
    Dockerfile rewrite (read from the file in this PR).
  • Build failure, gate counts, GIN init trace, /dev/gdrdrv EPERM, and both passing runs:
    reproduced on 2× p6-b300.48xlarge, EKS, us-east-1 (verified live, 2026-08-25).

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you so much for this @erezzarum ! Few minor comments.

Signed-off-by: Erez Zarum <erezz@amazon.com>
Signed-off-by: Erez Zarum <erezz@amazon.com>
Signed-off-by: Erez Zarum <erezz@amazon.com>
Signed-off-by: Erez Zarum <erezz@amazon.com>

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-verified end to end today on ml-clusters-shared-us-east-1, 2x p6-b300. One infra-side note for anyone reproducing this: on a GPU-operator-managed cluster the plugin runs in CDI mode, where the per-pod NVIDIA_GDRCOPY=enabled env is ignored (NVIDIA/k8s-device-plugin#1692), so the README's requirement translates to setting GDRCOPY_ENABLED=true on the device plugin itself. With that in place, test-internode.yaml ran unmodified and unprivileged: /dev/gdrdrv injected via CDI, Loaded gin plugin Libfabric_GDAKI (v14) on 16/16 ranks, dispatch ~115 GB/s (SO), clean finalize. The README note is accurate and sufficient. Approving. Thanks for turning around all 16 threads so quickly.

@KeitaW
KeitaW merged commit 2889ec1 into awslabs:main Aug 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants