# Use rocm image
FROM rocm/pytorch:rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1
CMD rocm-smi

# Ignore interactive questions during `docker build`
ENV DEBIAN_FRONTEND noninteractive

# Versions
# available options 3.10
ARG PYTHON_VERSION=3.10

# Bash shell
RUN chsh -s /bin/bash
SHELL ["/bin/bash", "-c"]

# Install and update tools to minimize security vulnerabilities
RUN apt-get update
RUN apt-get install -y software-properties-common wget apt-utils patchelf git libprotobuf-dev protobuf-compiler cmake \
    bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 mercurial subversion libopenmpi-dev ffmpeg && \
    apt-get clean
RUN apt-get autoremove -y

ARG PYTHON_EXE=/opt/conda/envs/py_$PYTHON_VERSION/bin/python

# (Optional) Intall test dependencies
RUN $PYTHON_EXE -m pip install -U pip
RUN $PYTHON_EXE -m pip install git+https://github.com/huggingface/transformers
RUN $PYTHON_EXE -m pip install datasets accelerate evaluate coloredlogs absl-py rouge_score seqeval scipy sacrebleu nltk scikit-learn parameterized sentencepiece --no-cache-dir
RUN $PYTHON_EXE -m pip install deepspeed --no-cache-dir
RUN conda install -y mpi4py

# PyTorch
RUN $PYTHON_EXE -m pip install onnx ninja

# ORT Module
RUN $PYTHON_EXE -m pip install --pre onnxruntime-training -f https://download.onnxruntime.ai/onnxruntime_nightly_rocm57.html
RUN $PYTHON_EXE -m pip install torch-ort
RUN $PYTHON_EXE -m pip install --upgrade protobuf==3.20.2
RUN $PYTHON_EXE -m torch_ort.configure

WORKDIR .

CMD ["/bin/bash"]