[tensorflow] make python version a build arg, better build cleanup

tensorflow/Dockerfile

@@ -29,12 +29,20 @@ RUN mkdir -p /opt/bazel && \
     /opt/bazel/installer.sh && \
     rm -f /opt/bazel/installer.sh
 
+ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:/usr/local/cuda/lib64/stubs:/usr/include/x64_64-linux-gnu:$LD_LIBRARY_PATH
+ENV TF_NEED_CUDA 1
+ENV TF_NEED_TENSORRT 1
+ENV TF_CUDA_VERSION=${CUDA}
+ENV TF_CUDNN_VERSION=${CUDNN_MAJOR_VERSION}
+
 ARG TF_VERSION=1.15.3
 RUN git clone --recursive --branch ${TF_VERSION} https://git.imp.fu-berlin.de/bioroboticslab/robofish/vendor/tensorflow.git /opt/tensorflow/src
+COPY tf_configure.bazelrc /opt/tensorflow/src/.tf_configure.bazelrc
 
+ARG PYTHON_VERSION=3.8
 RUN mkdir -p /opt/tensorflow/.venv && \
-    python3.8 -m venv /opt/tensorflow/.venv/python3.8 && \
-    . /opt/tensorflow/.venv/python3.8/bin/activate && \
+    python${PYTHON_VERSION} -m venv /opt/tensorflow/.venv && \
+    . /opt/tensorflow/.venv/bin/activate && \
     pip --no-cache-dir install --upgrade pip && \
     pip --no-cache-dir install wheel && \
     pip --no-cache-dir install \
@@ -48,20 +56,15 @@ RUN mkdir -p /opt/tensorflow/.venv && \
         sklearn \
         pandas \
         future \
-        portpicker
-
-ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:/usr/local/cuda/lib64/stubs:/usr/include/x64_64-linux-gnu:$LD_LIBRARY_PATH
-ENV TF_NEED_CUDA 1
-ENV TF_NEED_TENSORRT 1
-ENV TF_CUDA_VERSION=${CUDA}
-ENV TF_CUDNN_VERSION=${CUDNN_MAJOR_VERSION}
-
-COPY tf_configure.bazelrc /opt/tensorflow/src/.tf_configure.bazelrc
-RUN cd /opt/tensorflow/src && \
-    . /opt/tensorflow/.venv/python3.8/bin/activate && \
+        portpicker \
+        && \
+    cd /opt/tensorflow/src && \
+    . /opt/tensorflow/.venv/bin/activate && \
     bazel build \
-        --action_env PYTHONPATH="/opt/tensorflow/.venv/python3.8/lib/python3.8/site-packages" \
+        --action_env PYTHONPATH="/opt/tensorflow/.venv/lib/python${PYTHON_VERSION}/site-packages" \
         --config=cuda --config=opt --config=mkl \
         //tensorflow/tools/pip_package:build_pip_package && \
     ./bazel-bin/tensorflow/tools/pip_package/build_pip_package /opt/tensorflow/build && \
-    rm -rf /root/.cache/bazel
+    deactivate && \
+    rm -rf /root/.cache/bazel && \
+    rm -rf /opt/tensorflow/.venv