jax_paxml.Dockerfile.orig

FROM nvcr.io/nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04

ARG EFA_INSTALLER_VERSION=1.29.1
ARG NCCL_VERSION=v2.18.6-1
ARG AWS_OFI_NCCL_VERSION=v1.7.4-aws
ARG JAX_VERSION=0.4.18
ARG PRAXIS_VERSION=1.2.0
ARG PAXML_VERSION=1.2.0

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHON_VERSION=3.10
ENV LD_LIBRARY_PATH=/opt/amazon/openmpi/lib:/opt/nccl/build/lib:/opt/aws-ofi-nccl/install/lib:/usr/local/cuda-12/lib64:$LD_LIBRARY_PATH
ENV PATH=/opt/amazon/openmpi/bin/:/opt/amazon/efa/bin:/usr/local/cuda-12/bin:$PATH
ENV CUDA_HOME=/usr/local/cuda-12


#########################
# Packages and Pre-reqs #
RUN apt-get update -y && \
    apt-get purge -y --allow-change-held-packages libmlx5-1 ibverbs-utils libibverbs-dev libibverbs1 libnccl-dev libnccl2
RUN apt-get install -y --allow-unauthenticated \
    autoconf \
    automake \
    bash \
    build-essential \
    ca-certificates \
    curl \
    debianutils \
    dnsutils \
    g++ \
    git \
    libtool \
    libhwloc-dev \
    netcat \
    openssh-client \
    openssh-server \
    openssl \
    python3-distutils \
    python"${PYTHON_VERSION}"-dev \
    python-is-python3 \
    util-linux

RUN update-ca-certificates

###########################
# Python/Pip dependencies #
RUN curl https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py \
    && python"${PYTHON_VERSION}" /tmp/get-pip.py
RUN pip"${PYTHON_VERSION}" install numpy wheel build

######################################
# Install EFA Libfabric and Open MPI #
RUN cd /tmp \
    && 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 -d --skip-kmod --skip-limit-conf --no-verify

############################
# Compile and Install NCCL #
RUN git clone -b "${NCCL_VERSION}" https://github.com/NVIDIA/nccl.git /opt/nccl \
    && cd /opt/nccl \
    && make -j src.build CUDA_HOME=${CUDA_HOME} \
    && cp -R /opt/nccl/build/* /usr/

###############################
# Compile AWS OFI NCCL Plugin #
RUN git clone -b "${AWS_OFI_NCCL_VERSION}" https://github.com/aws/aws-ofi-nccl.git /opt/aws-ofi-nccl \
    && cd /opt/aws-ofi-nccl \
    && ./autogen.sh \
    && ./configure --prefix=/opt/aws-ofi-nccl/install \
       --with-libfabric=/opt/amazon/efa/ \
       --with-cuda=${CUDA_HOME} \
       --with-mpi=/opt/amazon/openmpi/ \
       --with-nccl=/opt/nccl/build \
       --enable-platform-aws \
    && make -j && make install

###############
# Install JAX #
RUN pip install --upgrade "jax[cuda12_pip]==${JAX_VERSION}" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
RUN pip install "orbax-checkpoint>=0.4.0,<0.5.0"

##################
# Install Praxis #
RUN pip install praxis==${PRAXIS_VERSION}

#################
# Install Paxml #
RUN pip install paxml==${PAXML_VERSION}

#####################################
# Allow unauthenticated SSH for MPI #
RUN mkdir -p /var/run/sshd \
    && sed -i 's/[ #]\(.*StrictHostKeyChecking \).*/ \1no/g' /etc/ssh/ssh_config \
    && echo "    UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config \
    && sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config

COPY run_paxml.sh /run_paxml.sh
RUN pip install scipy==1.12.0
RUN pip install protobuf==3.20.3
RUN cp /usr/local/lib/python3.10/dist-packages/google/protobuf/internal/builder.py .
RUN pip3 install protobuf==3.19.6
RUN cp /builder.py /usr/local/lib/python3.10/dist-packages/google/protobuf/internal/builder.py