FROM mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.2-cudnn8-ubuntu20.04:20221010.v1

USER root

ENV AZUREML_CONDA_ENVIRONMENT_PATH /azureml-envs/tensorflow-2.7
# Create conda environment
RUN conda create -p $AZUREML_CONDA_ENVIRONMENT_PATH \
    python=3.8 pip=20.2.4

# Prepend path to AzureML conda environment
ENV PATH $AZUREML_CONDA_ENVIRONMENT_PATH/bin:$PATH

# Install necessary packages to support videos in rllib/gym
RUN apt-get update && apt-get install -y --no-install-recommends \
    python-opengl \
    rsync \
    xvfb && \
    apt-get clean -y && \
    rm -rf /var/lib/apt/lists/* && \
    rm -rf /usr/share/man/*

RUN pip --version
RUN python --version

# Install ray-on-aml
RUN pip install 'ray-on-aml==0.1.6'

RUN pip install ray==0.8.7
RUN pip install gym[atari]==0.19.0
RUN pip install gym[accept-rom-license]==0.19.0

# Install pip dependencies for Tensorflow
RUN pip install 'matplotlib~=3.5.0' \
                'psutil~=5.8.0' \
                'tqdm~=4.62.0' \
                'pandas~=1.3.0' \
                'scipy~=1.7.0' \
                'numpy~=1.21.0' \
                'ipykernel~=6.0' \
                'azureml-core==1.47.0' \
                'azureml-defaults==1.47.0' \
                'azureml-mlflow==1.47.0' \
                'azureml-telemetry==1.47.0' \
                'tensorboard~=2.7.0' \
                'tensorflow-gpu~=2.7.0' \
                'tensorflow-datasets~=4.5.0' \
                'onnxruntime-gpu~=1.9.0' \
                'protobuf~=3.20' \
                'horovod[tensorflow-gpu]~=0.23.0' \
                'debugpy~=1.6.3'

RUN pip install --no-cache-dir \
    azureml-defaults \
    azureml-dataset-runtime[fuse,pandas] \
    azureml-contrib-reinforcementlearning \
    gputil \
    cloudpickle==1.3.0 \
    tabulate \
    dm_tree \
    lz4 \
    psutil \
    setproctitle

# This is required for ray 0.8.7
RUN pip install -U aiohttp==3.7.4

RUN pip install 'msrest<0.7.0'
RUN pip install protobuf==3.20.0

# This is needed for mpi to locate libpython
ENV LD_LIBRARY_PATH $AZUREML_CONDA_ENVIRONMENT_PATH/lib:$LD_LIBRARY_PATH