apiVersion: kubeflow.org/v1alpha2
kind: MPIJob
metadata:
  name: pytorch-ddp-efa
spec:
  slotsPerWorker: 1
  cleanPodPolicy: Running
  mpiReplicaSpecs:
    Launcher:
      replicas: 1
      template:
         spec:
          imagePullPolicy: Always
          restartPolicy: OnFailure
          containers:
          - image: public.ecr.aws/w6p6i9i7/aws-efa-nccl-rdma:megatron-cu11.1
            name: megatron-training-launcher
            env:
             - name: GPUS_PER_NODE
               value: "8"
             - name: MASTER_ADDR
               value: pytorch-ddp-efa-worker-0
             - name: MASTER_PORT
               value: "6002"
             - name: NNODES
               value: $OMPI_MCA_orte_num_nodes
             - name: NODE_RANK
               value: $OMPI_COMM_WORLD_NODE_RANK
             - name: WORLD_SIZE
               value: $OMPI_COMM_WORLD_SIZE
             - name: LD_LIBRARY_PATH
               value: /opt/amazon/openmpi/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/aws-ofi-nccl/install/lib:/usr/local/nvidia/lib:$LD_LIBRARY_PATH
             - name: PATH
               value: $PATH:/opt/amazon/efa/bin:/usr/local/bin:/usr/bin:/opt/conda/bin
             - name: XLA_FLAGS
               value: "--xla_gpu_cuda_data_dir=/usr/local/cuda"
             - name: TF_XLA_FLAGS
               value: "--tf_xla_cpu_global_jit"
             - name: DISTRIBUTED_ARGS
               value: "--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
    
            command:
            - /opt/amazon/openmpi/bin/mpirun
            - --allow-run-as-root
            - --tag-output
            - -np
            - "2"
            - -bind-to
            - none
            - -map-by
            - slot
            - -x 
            - PATH
            - -x 
            - LD_LIBRARY_PATH
            - -x
            - MASTER_ADDR
            - -x
            - MASTER_PORT
            - --mca
            - pml
            - ^cm
            - /tmp/torch_ddp_efa.sh

    Worker:
      replicas: 2
      template:
        spec:
          imagePullPolicy: Always
          containers:
          - image: public.ecr.aws/w6p6i9i7/aws-efa-nccl-rdma:megatron-cu11.1
            name: pytorch-efa-worker
            resources:
              limits:
                nvidia.com/gpu: 8
                hugepages-2Mi: 5120Mi
                memory: 384000Mi
                vpc.amazonaws.com/efa: 4
              requests:
                memory: 384000Mi