Multistage docker container for building tensorflow from source Ubuntu 20.04

Dockerfile

FROM ubuntu:20.04 as cuda
ENV TZ=Europe/Moscow
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
RUN apt-get update && apt-get install --no-install-recommends -y build-essential git \
  wget make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl \
  llvm libncurses5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev ca-certificates gnupg
RUN git clone https://github.com/pyenv/pyenv.git /root/pyenv
RUN /root/pyenv/bin/pyenv install 3.8.0
RUN ln -s /root/.pyenv/versions/3.8.0/bin/python /usr/bin/python
RUN ln -s /root/.pyenv/versions/3.8.0/bin/pip /usr/bin/pip
RUN apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
RUN sh -c 'echo "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64 /" > /etc/apt/sources.list.d/cuda.list'
RUN apt-get update
RUN apt-get install -y cuda-toolkit-11-1
RUN wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64/nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
RUN apt install ./nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb && apt-get update
RUN apt-get install --no-install-recommends \
    libcudnn8=8.0.5.39-1+cuda11.1  \
    libcudnn8-dev=8.0.5.39-1+cuda11.1
COPY ./tensorrt /usr/local/tensorrt

FROM cuda as tensorbuild
RUN git clone https://github.com/tensorflow/tensorflow.git
WORKDIR /tensorflow
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/extras/CUPTI/lib64"
RUN pip install -U --user pip "numpy<1.19.0" wheel && pip install -U --user keras_preprocessing --no-deps
RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add - \
  && echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list \
  && apt update && apt install -y bazel
# RUN perl -pi.bak -e 's%, CompareUFunc%, (PyUFuncGenericFunction) CompareUFunc%g' tensorflow/python/lib/core/bfloat16.cc \
#   && sed -i 's/cuda_config.cudart_version/"11.0"/g' third_party/gpus/cuda_configure.bzl
RUN PYTHON_BIN_PATH=$(which python3) \
  PYTHON_LIB_PATH="$(python -c 'import site; print(site.getsitepackages()[0])')" \
  TF_NEED_OPENCL_SYCL=0 \
  TF_NEED_OPENCL=0 \
  TF_NEED_ROCM=0 \
  TF_NEED_CUDA=1 \
  TF_NEED_TENSORRT=1 \
  TF_CUDA_VERSION=11 \
  TF_CUDNN_VERSION=8 \
  TF_TENSORRT_VERSION=7 \
  CUDA_TOOLKIT_PATH=/usr/local/cuda \
  CUDNN_INSTALL_PATH=/usr/lib/x86_64-linux-gnu \
  TENSORRT_INSTALL_PATH=/usr/local/tensorrt \
  TF_CUDA_COMPUTE_CAPABILITIES=8.6 \
  TF_CUDA_CLANG=0 \
  TF_NEED_IGNITE=0 \
  TF_ENABLE_XLA=0 \
  TF_NEED_MPI=0 \
  GCC_HOST_COMPILER_PATH=$(which gcc) \
  CC_OPT_FLAGS="-march=native" \
  TF_SET_ANDROID_WORKSPACE=0 \
  ./configure
RUN bazel build --config=cuda //tensorflow/tools/pip_package:build_pip_package
RUN ./bazel-bin/tensorflow/tools/pip_package/build_pip_package --nightly_flag /tmp/tensorflow_pkg

FROM cuda
COPY --from=tensorbuild /tmp/tensorflow_pkg /tmp/tensorflow_pkg
RUN pip install /tmp/tensorflow_pkg/tf_nightly-2.5.0-cp38-cp38-linux_x86_64.whl