From 49238807cdfc799b5fbb877a8daf50af56be5c9d Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Tue, 10 Sep 2024 13:22:01 -0700
Subject: [PATCH] feat: allow to set python version for tensorflow

---
 tensorflow.dockerfile | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/tensorflow.dockerfile b/tensorflow.dockerfile
index 9a0c1cc..b793def 100644
--- a/tensorflow.dockerfile
+++ b/tensorflow.dockerfile
@@ -1,5 +1,6 @@
 FROM wtongze/rocm-base:5.7.3
 ARG TENSORFLOW_BRANCH="r2.11-rocm-enhanced"
+ARG PYTHON_VERSION="3.11"
 ARG NUMPY_VERSION="1.22.4"
 ARG DEBIAN_FRONTEND="noninteractive"
 
@@ -16,20 +17,32 @@ RUN bash -c 'echo -e "gfx803\ngfx900" | sudo tee /opt/rocm/bin/target.lst'
 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
 
 RUN git clone --depth 1 https://github.com/ROCm/tensorflow-upstream.git -b ${TENSORFLOW_BRANCH} tensorflow
+
 WORKDIR /home/rocm/tensorflow
 
-# > https://github.com/ROCm/ROCm/issues/1880
+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} python${PYTHON_VERSION}-dev
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
 ENV PATH=$PATH:/home/rocm/.local/bin
 
-RUN sudo apt-get update && \
-sudo apt-get install -y python3-dev && \
-sudo rm -rf /var/lib/apt/lists/*
+# TODO: simplify dependencies
+# > https://github.com/ROCm/ROCm/issues/1880
+RUN python${PYTHON_VERSION} -m pip install -U --user pip numpy==${NUMPY_VERSION} wheel packaging requests opt_einsum patchelf
+RUN python${PYTHON_VERSION} -m pip install -U --user keras_preprocessing --no-deps
 
-RUN pip install -U --user pip numpy==${NUMPY_VERSION} wheel packaging requests opt_einsum patchelf
-RUN pip install -U --user keras_preprocessing --no-deps
+RUN sed -i "s/python3/python${PYTHON_VERSION}/g" build_rocm_python3
+RUN sudo ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
 
-RUN sudo ln -s /usr/bin/python3 /usr/bin/python
 RUN ./build_rocm_python3 /opt/rocm
-RUN cp -a /tmp/tensorflow_pkg /home/rocm
 
+WORKDIR /tmp/tensorflow_pkg
+RUN cp (ls -l tensorflow*.whl | cut -d' ' -f9 | tr -d '\n') /home/rocm/tensorflow-rocm.whl
+
+WORKDIR /home/rocm
 CMD [ "/bin/bash", "-l" ]