Running multi-node workloads¶

apiVersion: kubeflow.org/v2beta1
kind: MPIJob
metadata:
  generateName: nccl-test-2n-
spec:
  slotsPerWorker: 8
  runPolicy:
    cleanPodPolicy: Running
  mpiReplicaSpecs:
    Launcher:
      replicas: 1
      template:
        spec:
          containers:
            - name: launcher
              image: vccr.io/nccl-tests/nccl-tests:cuda13.1.1-nccl2.29.3-1-v2.17.9
              env:
                - name: OMPI_ALLOW_RUN_AS_ROOT
                  value: "1"
                - name: OMPI_ALLOW_RUN_AS_ROOT_CONFIRM
                  value: "1"
              command: ["/bin/bash", "-c"]
              args:
                - |
                  echo "=== NCCL 16-GPU Test (2 nodes) ==="

                  # Wait for MPI hostfile
                  echo "Waiting for MPI hostfile..."
                  while [ ! -f /etc/mpi/hostfile ] || [ ! -s /etc/mpi/hostfile ]; do
                    sleep 2
                  done
                  echo "Hostfile:"
                  cat /etc/mpi/hostfile

                  # Wait for workers to be reachable via SSH
                  echo "Waiting for workers..."
                  for worker in $(awk '{print $1}' /etc/mpi/hostfile); do
                    retries=0
                    until ssh -o ConnectTimeout=2 "$worker" hostname >/dev/null 2>&1; do
                      retries=$((retries + 1))
                      if [ "$retries" -ge 60 ]; then
                        echo "TIMEOUT: $worker not reachable after 5 minutes"
                        exit 1
                      fi
                      echo "  Waiting for $worker... (attempt $retries)"
                      sleep 5
                    done
                    echo "  $worker ready"
                  done
                  echo "All workers ready"

                  echo ""
                  echo "=========================================="
                  echo "Running: all_reduce_perf"
                  echo "=========================================="
                  mpirun \
                    -np 16 \
                    -bind-to none \
                    -x NCCL_IB_PKEY=1 \
                    /opt/nccl-tests/build/all_reduce_perf -b 512M -e 8G -f 2 -g 1

                  echo ""
                  echo "=== NCCL test completed ==="
              resources:
                requests:
                  cpu: 2
                  memory: 256Mi
    Worker:
      replicas: 2
      template:
        metadata:
          labels:
            app: nccl-test
        spec:
          containers:
            - name: worker
              image: vccr.io/nccl-tests/nccl-tests:cuda13.1.1-nccl2.29.3-1-v2.17.9
              securityContext:
                capabilities:
                  add:
                    - IPC_LOCK
              resources:
                requests:
                  cpu: 32
                  memory: 128Gi
                  nvidia.com/gpu: 8
                  rdma/rdma_shared_device_a: 1
                limits:
                  nvidia.com/gpu: 8
                  rdma/rdma_shared_device_a: 1
              volumeMounts:
                - mountPath: /dev/shm
                  name: dshm
          affinity:
            podAntiAffinity:
              requiredDuringSchedulingIgnoredDuringExecution:
                - labelSelector:
                    matchLabels:
                      app: nccl-test
                  topologyKey: kubernetes.io/hostname
          volumes:
            - name: dshm
              emptyDir:
                medium: Memory
                sizeLimit: 64Gi