Build TensorFlow addons wheels.

TFA must be build from source to support CustomOps. This is necessary
since TensorFlow C++ API is not ABI compatible. See:

- tensorflow/addons#987.

tensorflow-whl/CHANGELOG.md

@@ -14,4 +14,5 @@
 * `2.1.0-py37`: TensorFlow 2.1.0 with Python 3.7
 * `2.1.0-py37-2`: TensorFlow 2.1.0 with Python 3.7 & DLVM base image.
 * `2.1.0-py37-3`: TensorFlow 2.1.0 with Python 3.7, DLVM base image, tensorflow-gcs-config.
-* `2.2.0-py37`: TensorFlow 2.2.0 with Python 3.7.
+* `2.2.0-py37`: TensorFlow 2.2.0 with Python 3.7.
+* `2.2.0-py37-2`: TensorFlow 2.2.0 with Python 3.7 & TFA.

tensorflow-whl/Dockerfile

@@ -62,15 +62,15 @@ RUN apt-get install -y gnupg zip openjdk-8-jdk && \
  dpkg -i bazel_*.deb && \
  rm bazel_*.deb
 
-# Fetch tensorflow & install dependencies.
+# Fetch TensorFlow & install dependencies.
 RUN cd /usr/local/src && \
  git clone https://github.com/tensorflow/tensorflow && \
  cd tensorflow && \
  git checkout tags/v2.2.0 && \
  pip install keras_applications --no-deps && \
  pip install keras_preprocessing --no-deps
 
-# Create a tensorflow wheel for CPU
+# Create a TensorFlow wheel for CPU
 RUN cd /usr/local/src/tensorflow && \
  cat /dev/null | ./configure && \
  bazel build --config=opt \
@@ -80,6 +80,26 @@ RUN cd /usr/local/src/tensorflow && \
  bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_cpu && \
  bazel clean
 
+# Install TensorFlow CPU wheel which is required to build the GCS & addons libraries.
+RUN pip install /tmp/tensorflow_cpu/tensorflow*.whl
+
+# Build TensorFlow addons library against TensorFlow CPU.
+RUN cd /usr/local/src/ && \
+ git clone https://github.com/tensorflow/addons && \
+ cd addons && \
+ git checkout tags/v0.10.0 && \
+ python ./configure.py && \
+ bazel build --enable_runfiles build_pip_pkg && \
+ bazel-bin/build_pip_pkg /tmp/tfa_cpu && \
+ bazel clean
+
+# Build tensorflow_gcs_config library against TensorFlow CPU.
+ADD tensorflow-gcs-config /usr/local/src/tensorflow_gcs_config/
+RUN cd /usr/local/src/tensorflow_gcs_config && \
+ apt-get install -y libcurl4-openssl-dev && \
+ python setup.py bdist_wheel -d /tmp/tensorflow_gcs_config && \
+ bazel clean
+
 # Create a tensorflow wheel for GPU/cuda
 ENV TF_NEED_CUDA=1
 ENV TF_CUDA_VERSION=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION
@@ -103,13 +123,16 @@ RUN cd /usr/local/src/tensorflow && \
  bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_gpu && \
  bazel clean
 
-ADD tensorflow-gcs-config /usr/local/src/tensorflow_gcs_config/
+# Install TensorFlow GPU wheel which to build addons against.
+RUN pip install /tmp/tensorflow_gpu/tensorflow*.whl
 
-# Build tensorflow_gcs_config library against the tensorflow_cpu build
-RUN cd /usr/local/src/tensorflow_gcs_config && \
- apt-get install -y libcurl4-openssl-dev && \
- pip install /tmp/tensorflow_cpu/tensorflow*.whl && \
- python setup.py bdist_wheel -d /tmp/tensorflow_gcs_config && \
+# Build TensorFlow addons library against TensorFlow GPU.
+ENV CUDA_TOOLKIT_PATH=/usr/local/cuda
+ENV CUDNN_INSTALL_PATH=/usr/lib/x86_64-linux-gnu
+RUN cd /usr/local/src/addons && \
+ python ./configure.py && \
+ bazel build --enable_runfiles build_pip_pkg && \
+ bazel-bin/build_pip_pkg /tmp/tfa_gpu && \
  bazel clean
 
 # Print out the built .whl files