Revert "Revert "Install tensorflow-gcs-config from the python-tensorflow-whl (#788)""

This reverts commit 354a4eb.

Author: mcollins42

Dockerfile

@@ -1,7 +1,7 @@
 ARG BASE_TAG=latest
 ARG TENSORFLOW_VERSION=2.1.0
 
-FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37-2 as tensorflow_whl
+FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37-3 as tensorflow_whl
 FROM gcr.io/deeplearning-platform-release/base-cpu:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh
@@ -59,6 +59,12 @@ RUN pip install /tmp/tensorflow_cpu/tensorflow*.whl && \
     rm -rf /tmp/tensorflow_cpu && \
     /tmp/clean-layer.sh
 
+# Install tensorflow-gcs-config from a pre-built wheel
+COPY --from=tensorflow_whl /tmp/tensorflow_gcs_config/*.whl /tmp/tensorflow_gcs_config/
+RUN pip install /tmp/tensorflow_gcs_config/tensorflow*.whl && \
+    rm -rf /tmp/tensorflow_gcs_config && \
+    /tmp/clean-layer.sh
+
 RUN apt-get install -y libfreetype6-dev && \
     apt-get install -y libglib2.0-0 libxext6 libsm6 libxrender1 libfontconfig1 --fix-missing && \
     pip install gensim && \

tests/test_imports.py

@@ -5,3 +5,4 @@ class TestImport(unittest.TestCase):
     def test_basic(self):
         import bq_helper
         import cleverhans
+        import tensorflow_gcs_config