content: FROM nvcr.io/nvidia/pytorch:22.12-py3
  
  # Ubuntu: 20.04
  # Python: 3.8
  # CUDA: 11.8.0
  # cuDNN: 8.7.0.84
  # NCCL: v2.16.2-1 + FP8 Support
  # PyTorch: 1.14.0a0+410ce96
  
  LABEL maintainer="MS-AMP"
  
  ENV DEBIAN_FRONTEND=noninteractive
  RUN apt-get update && \
      apt-get install -y --no-install-recommends \
      autoconf \
      automake \
      bc \
      build-essential \
      curl \
      dmidecode \
      git \
      iproute2 \
      jq \
      moreutils \
      net-tools \
      openssh-client \
      openssh-server \
      sudo \
      util-linux \
      vim \
      wget \
      python3-mpi4py \
      && \
      apt-get autoremove && \
      apt-get clean && \
      rm -rf /var/lib/apt/lists/* /tmp/*
  
  ARG NUM_MAKE_JOBS=
  ENV PATH="${PATH}" \
      LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}"
  
  WORKDIR /opt/msamp
  
  ADD third_party third_party
  RUN cd third_party/msccl && \
      make -j ${NUM_MAKE_JOBS} src.build NVCC_GENCODE="\
      -gencode=arch=compute_70,code=sm_70 \
      -gencode=arch=compute_80,code=sm_80 \
      -gencode=arch=compute_90,code=sm_90" && \
      make install
  # cache TE build to save time in CI
  RUN python3 -m pip install --upgrade pip && \
      python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@v1.1
  
  ADD . .
  RUN python3 -m pip install . && \
      make postinstall
  
  ENV LD_PRELOAD="/usr/local/lib/libmsamp_dist.so:/usr/local/lib/libnccl.so:${LD_PRELOAD}"
  
  # Set up entrypoint
  COPY dockerfile/entrypoint.sh /entrypoint.sh
  RUN chmod +x /entrypoint.sh
  ENTRYPOINT ["/entrypoint.sh"]
  