ARG ROCM_VERSION="5.7.3"

FROM wtongze/rocm-base:${ROCM_VERSION}

ARG ROCM_VERSION="5.7.3"
ARG TENSORFLOW_BRANCH="r2.12-rocm-enhanced"
ARG PYTHON_VERSION="3.11"
ARG NUMPY_VERSION="1.24.3"
ARG HIP_ARCH="gfx803;gfx900"
ARG DEBIAN_FRONTEND="noninteractive"

# Add GPU target
RUN echo "${HIP_ARCH}" | sed 's/;/\n/g' | sudo tee /opt/rocm/bin/target.lst

# Install bazelisk
RUN sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.21.0/bazelisk-linux-amd64 && sudo chmod 755 /usr/local/bin/bazel

WORKDIR /home/rocm
RUN git clone --depth 1 https://github.com/ROCm/tensorflow-upstream.git -b ${TENSORFLOW_BRANCH} tensorflow

# For ROCm 4.5+
RUN <<EOF
sudo apt-get update
sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF

RUN <<EOF
sudo apt-get update
sudo apt install -y software-properties-common
sudo add-apt-repository -y ppa:deadsnakes/ppa
sudo apt-get update
sudo apt-get install -y python${PYTHON_VERSION}-full python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-distutils
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF

WORKDIR /home/rocm/tensorflow

# TODO: simplify dependencies
# > Ref: https://github.com/ROCm/ROCm/issues/1880#issuecomment-1547838227
RUN python${PYTHON_VERSION} -m ensurepip --upgrade
RUN python${PYTHON_VERSION} -m pip install -U --user numpy==${NUMPY_VERSION} wheel packaging requests opt_einsum patchelf
RUN python${PYTHON_VERSION} -m pip install -U --user keras_preprocessing --no-deps

RUN sed -i -e "s/python3/python${PYTHON_VERSION}/g" -e "s/pip3/python${PYTHON_VERSION} -m pip/g" build_rocm_python3
RUN sudo ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python

ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
ENV ROCM_TOOLKIT_PATH=${ROCM_PATH}
RUN ./build_rocm_python3 ${ROCM_PATH}

WORKDIR /tmp/tensorflow_pkg
RUN cp $(ls -l tensorflow*.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/tensorflow-rocm.whl

WORKDIR /home/rocm
CMD [ "/bin/bash", "-l" ]