trfa.Dockerfile

FROM python:3.9.18-bullseye

ENV BAZELISK_VERSION="1.18.0"
ENV USE_BAZEL_VERSION="5.1.1"
ENV TF_VERSION="2.8.3"
ENV OPEN_MPI_VERSION="5.0.0"

ENV MPI_DIR=/opt/ompi
ENV PATH="$MPI_DIR/bin:$HOME/.local/bin:$PATH"
ENV LD_LIBRARY_PATH="$MPI_DIR/lib:$LD_LIBRARY_PATH"
WORKDIR /work

# Install required packages
RUN apt update
RUN apt install -y rsync # rsync required for build_pip_pkg 

# Install OpenMPI as one of the supported approaches for communication in horovod
WORKDIR /work/openmpi
RUN curl --location https://download.open-mpi.org/release/open-mpi/v${OPEN_MPI_VERSION%.*}/openmpi-${OPEN_MPI_VERSION}.tar.bz2 -o openmpi.tar.bz2
RUN tar xf openmpi.tar.bz2
WORKDIR openmpi-${OPEN_MPI_VERSION}
RUN ./configure --prefix=$MPI_DIR
RUN make -j4 all
RUN make install
WORKDIR /work
RUN rm -rf openmpi /tmp/*

# Install bazelisk for bazel version managmenet
RUN curl --location https://github.com/bazelbuild/bazelisk/releases/download/v${BAZELISK_VERSION}/bazelisk-linux-amd64 -o /usr/bin/bazel
RUN chmod +x /usr/bin/bazel
RUN bazel # initiate first download of bazel through bazelisk (masked by bazel exec command)

# Install tensorflow before building tensorflow-recommenders-addons
RUN pip install tensorflow==${TF_VERSION}

# Build tensorflow-recommenders-addons from source
RUN git clone https://github.com/tensorflow/recommenders-addons.git
WORKDIR /work/recommenders-addons
RUN python configure.py
RUN bazel build --local_ram_resources=HOST_RAM*.5 --local_cpu_resources=HOST_CPUS-2 --enable_runfiles build_pip_pkg
RUN bazel-bin/build_pip_pkg artifacts
RUN pip install artifacts/*.whl

# Install horovod
RUN apt install -y cmake # TODO move to the top
RUN HOROVOD_WITH_MPI=1 HOROVOD_WITHOUT_GLOO=1 pip install horovod[tensorflow]

WORKDIR /work

RUN pip install tensorflow-datasets==4.9.0 tensorflow-recommenders==0.6.0

ENTRYPOINT ["bash"]