Simplify the build.

tools/docker/build_wheel.Dockerfile

@@ -1,33 +1,60 @@
 #syntax=docker/dockerfile:1.1.5-experimental
 ARG TF_VERSION
-FROM seanpmorgan/tensorflow:2.1.0-custom-op-gpu-ubuntu16-minimal as make_wheel
+FROM seanpmorgan/tensorflow:2.1.0-custom-op-gpu-ubuntu16-minimal as base_install
 ENV TF_NEED_CUDA="1"
 
 RUN apt-get update && apt-get install patchelf
 
 ARG PY_VERSION
-RUN python$PY_VERSION -m pip install --upgrade pip setuptools auditwheel==2.0.0
+RUN ln -sf $(which python$PY_VERSION) /usr/bin/python3
+
+RUN python3 -m pip install --upgrade pip setuptools auditwheel==2.0.0
 
 COPY tools/install_deps/ /install_deps
 ARG TF_VERSION
-RUN python$PY_VERSION -m pip install \
+RUN python3 -m pip install \
         tensorflow==$TF_VERSION \
         -r /install_deps/pytest.txt
 
 COPY requirements.txt .
-RUN python$PY_VERSION -m pip install -r requirements.txt
+RUN python3 -m pip install -r requirements.txt
+
+# Fix presented in
+# https://stackoverflow.com/questions/44967202/pip-is-showing-error-lsb-release-a-returned-non-zero-exit-status-1/44967506
+RUN echo "#! /usr/bin/python2.7" >> /usr/bin/lsb_release2
+RUN cat /usr/bin/lsb_release >> /usr/bin/lsb_release2
+RUN mv /usr/bin/lsb_release2 /usr/bin/lsb_release
 
 COPY ./ /addons
 WORKDIR /addons
+
+# -------------------------------------------------------------------
+FROM base_install as tfa_gpu_tests
+CMD ["bash", "tools/testing/build_and_run_tests.sh"]
+
+# -------------------------------------------------------------------
+FROM base_install as make_wheel
 ARG NIGHTLY_FLAG
 ARG NIGHTLY_TIME
 RUN --mount=type=cache,id=cache_bazel,target=/root/.cache/bazel \
-    bash tools/releases/release_linux.sh $PY_VERSION $NIGHTLY_FLAG
+    bash tools/testing/build_and_run_tests.sh && \
+    bazel clean --expunge && \
+    bazel build \
+        -c opt \
+        --noshow_progress \
+        --noshow_loading_progress \
+        --verbose_failures \
+        --test_output=errors \
+        --crosstool_top=//build_deps/toolchains/gcc7_manylinux2010-nvcc-cuda10.1:toolchain \
+        build_pip_pkg && \
+    # Package Whl
+    bazel-bin/build_pip_pkg artifacts $NIGHTLY_FLAG
 
 RUN bash tools/releases/tf_auditwheel_patch.sh
 RUN auditwheel repair --plat manylinux2010_x86_64 artifacts/*.whl
 RUN ls -al wheelhouse/
 
+# -------------------------------------------------------------------
 FROM scratch as output
 
 COPY --from=make_wheel /addons/wheelhouse/ .

tools/run_gpu_tests.sh

@@ -3,5 +3,10 @@
 set -x -e
 
 export DOCKER_BUILDKIT=1
-docker build -f tools/docker/gpu_tests.Dockerfile -t tfa_gpu_tests ./
-docker run --rm -t --runtime=nvidia tfa_gpu_tests
+docker build \
+       -f tools/docker/build_wheel.Dockerfile \
+       --target tfa_gpu_tests \
+       --build-arg TF_VERSION=2.1.0 \
+       --build-arg PY_VERSION=3.5 \
+       -t tfa_gpu_tests ./
+docker run --rm -t -v cache_bazel:/root/.cache/bazel --runtime=nvidia tfa_gpu_tests