content: # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  # SPDX-License-Identifier: MIT-0
  
  ARG AWS_REGION=us-west-2
  
  FROM 763104351884.dkr.ecr.${AWS_REGION}.amazonaws.com/pytorch-training:2.0.1-gpu-py310-cu118-ubuntu20.04-ec2
  
  ARG DEBIAN_FRONTEND=noninteractive
  
  RUN apt-get update && apt-get install -y libpmix-dev libpmix2
  
  # reinstall EFA, to restore the openmpi
  ENV EFA_INSTALLER_VERSION=latest
  RUN apt-get update && \
      cd /tmp && \
      rm -fr /opt/amazon/openmpi && \
      curl -O 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 && \
      ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf && \
      ldconfig && \
      rm -rf /tmp/aws-efa-installer /var/lib/apt/lists/* && \
      /opt/amazon/efa/bin/fi_info --version
  
  # Repeat this from base Dockefile
  # Install AWS OFI NCCL plug-in
  ENV AWS_OFI_NCCL_VERSION=1.7.1
  RUN apt-get update && apt-get install -y autoconf
  RUN mkdir /tmp/efa-ofi-nccl \
   && cd /tmp/efa-ofi-nccl \
   && git clone https://github.com/aws/aws-ofi-nccl.git -b v${AWS_OFI_NCCL_VERSION}-aws \
   && cd aws-ofi-nccl \
   && ./autogen.sh \
   && ./configure --with-libfabric=/opt/amazon/efa \
    --with-mpi=/opt/amazon/openmpi \
    --with-cuda=/opt/conda --prefix=/usr/local \
   && make \
   && make install \
   && rm -rf /tmp/efa-ofi-nccl \
   && rm -rf /var/lib/apt/lists/* \
   && apt-get clean
  
  ENV LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH
  ENV PATH=/opt/amazon/efa/bin:$PATH
  
  RUN cd /opt && git clone https://github.com/NVIDIA/nccl-tests \
      && cd nccl-tests \
      && make MPI=1 MPI_HOME=/opt/amazon/openmpi
  
  
  #################################################
  ## Install NVIDIA GDRCopy
  RUN apt-get update && apt-get install -y check libsubunit0 libsubunit-dev pkg-config \
      && git clone https://github.com/NVIDIA/gdrcopy.git /opt/gdrcopy \
      && cd /opt/gdrcopy \
      && CUDA=/opt/conda make lib_install \
      # Optional: tests tool. Need to point to the stub libcuda.so
      # See: https://gitlab.com/nvidia/container-images/cuda/-/blob/master/dist/11.7.1/ubuntu2004/devel/Dockerfile#L68
      && LIBRARY_PATH=/opt/conda/lib/stubs CUDA=/opt/conda make exes_install
  