Docker image with incremental updates (#1294)

* Add layer pruning and sync tf-io with tf

* Upgrade to tf=2.13.0

Resolved bq_helper package which was broken as a side effect from the upgrade.

* Removed tensorflow-gcs-config package and test.

Package is being used minimally. See b/296908146#comment9

* Revert tf=2.12

Remove test for `tensorflow_gcs_config`

Author: neshdev

Dockerfile.tmpl

@@ -37,6 +37,8 @@ RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/lib
 
 # Keep these variables in sync if base image is updated.
 ENV TENSORFLOW_VERSION=2.12.0
+# See https://github.com/tensorflow/io#tensorflow-version-compatibility
+ENV TENSORFLOW_IO_VERSION=0.32.0
 
 # We need to redefine the ARG here to get the ARG value defined above the FROM instruction.
 # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
@@ -181,15 +183,17 @@ RUN JAXVER=$(pip freeze | grep -e "^jax==") && \
         pandas \
         polars \
         flax \
-        "${JAXVER}"
+        "${JAXVER}" && \
+    /tmp/clean-layer.sh
 
-RUN apt-get install -y default-jre
+RUN apt-get install -y default-jre && \
+    /tmp/clean-layer.sh
 
 RUN pip install -f http://h2o-release.s3.amazonaws.com/h2o/latest_stable_Py.html h2o && /tmp/clean-layer.sh
 
 RUN pip install \
-        "tensorflow-gcs-config<=${TENSORFLOW_VERSION}" \
         "tensorflow==${TENSORFLOW_VERSION}" \
+        "tensorflow-io==${TENSORFLOW_IO_VERSION}"\
         tensorflow-addons \
         tensorflow_decision_forests \
         tensorflow_text && \
@@ -536,6 +540,7 @@ RUN pip install flashtext \
         keras-cv \
         keras-nlp && \
     apt-get -y install libspatialindex-dev
+
 RUN pip install pytorch-ignite \
         qgrid \
         bqplot \
@@ -616,7 +621,14 @@ RUN pip install --upgrade dask && \
     pip install --force-reinstall --no-deps jupyter_server>=2.* && \
     /tmp/clean-layer.sh
 
-RUN pip install setuptools==59.8.0 && pip install -e git+https://github.com/SohierDane/BigQuery_Helper#egg=bq_helper
+# Fix to import bq_helper library without downgrading setuptools
+RUN mkdir -p ~/src && git clone https://github.com/SohierDane/BigQuery_Helper ~/src/BigQuery_Helper && \
+    mkdir -p ~/src/BigQuery_Helper/bq_helper && \
+    mv ~/src/BigQuery_Helper/bq_helper.py ~/src/BigQuery_Helper/bq_helper/__init__.py && \
+    mv ~/src/BigQuery_Helper/test_helper.py ~/src/BigQuery_Helper/bq_helper/ && \
+    sed -i 's/)/packages=["bq_helper"])/g' ~/src/BigQuery_Helper/setup.py && \
+    pip install -e ~/src/BigQuery_Helper && \
+    /tmp/clean-layer.sh
 
 # Add BigQuery client proxy settings
 ENV PYTHONUSERBASE "/root/.local"
@@ -683,4 +695,4 @@ RUN echo "$GIT_COMMIT" > /etc/git_commit && echo "$BUILD_DATE" > /etc/build_date
 {{ if eq .Accelerator "gpu" }}
 # Remove the CUDA stubs.
 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH_NO_STUBS"
-{{ end }}
+{{ end }}

tests/test_imports.py

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