Build Kaggle TPU image (#1116)

TODO:
- Currently fails on Cloud Build but works locally on Cloud Top: https://pantheon.corp.google.com/cloud-build/builds;region=global/904c778d-45c4-4646-9183-f10b18a4910a?project=kkb-infra. May be related to not enough RAM. Will require more investigation.
- TensorFlow related libraries (tensorflow-addons, tensorflow-gcs-config) will need to be recompiled against the new TPU binary.
- Install TPU support for JAX & Pytorch.

http://b/152075195

Author: rosbo

tpu/Dockerfile

@@ -0,0 +1,13 @@
+ARG BASE_IMAGE_TAG
+ARG LIBTPU_IMAGE_TAG
+ARG TENSORFLOW_WHL_IMAGE_TAG
+
+FROM gcr.io/cloud-tpu-v2-images/libtpu:${LIBTPU_IMAGE_TAG} as libtpu
+FROM gcr.io/kaggle-images/python-tpu-tensorflow-whl:${TENSORFLOW_WHL_IMAGE_TAG} AS tensorflow_whl
+FROM gcr.io/kaggle-images/python:${BASE_IMAGE_TAG}
+
+COPY --from=libtpu /libtpu.so /lib
+
+COPY --from=tensorflow_whl /tmp/tensorflow_pkg/tensorflow*.whl /tmp/tensorflow_pkg/
+RUN pip install /tmp/tensorflow_pkg/tensorflow*.whl && \
+    rm -rf /tmp/tensorflow_pkg

tpu/README.md

@@ -0,0 +1,9 @@
+# Build the Kaggle TPU image
+
+NOTE: Building a new Kaggle TPU image can only be done by members of the Kaggle team.
+
+1. Set the `_BASE_IMAGE_TAG` substitution in [cloudbuild.yaml](cloudbuild.yaml) to the desired version.
+1. Submit the build to Google Cloud Build by running:
+    ```
+    gcloud builds submit --async
+    ```

tpu/cloudbuild.yaml

@@ -0,0 +1,41 @@
+steps:
+- id: 'build-tensorflow-whl'
+  name: 'gcr.io/cloud-builders/docker'
+  args:
+  - build
+  - --rm
+  - --tag=gcr.io/kaggle-images/python-tpu-tensorflow-whl:$BUILD_ID
+  - --file=tensorflow.Dockerfile
+  - --build-arg=BASE_IMAGE_TAG=$_BASE_IMAGE_TAG
+  - .
+
+# TODO(b/152075195): Build JAX & Pytorch TPU enabled packages.
+
+- id: 'build-tpu-image'
+  waitFor: ['build-tensorflow-whl']
+  name: 'gcr.io/cloud-builders/docker'
+  args:
+  - build
+  - --rm
+  - --tag=gcr.io/kaggle-images/python-tpu:$BUILD_ID
+  - --file=Dockerfile
+  - --build-arg=BASE_IMAGE_TAG=$_BASE_IMAGE_TAG
+  - --build-arg=LIBTPU_IMAGE_TAG=$_LIBTPU_IMAGE_TAG
+  - --build-arg=TENSORFLOW_WHL_IMAGE_TAG=$BUILD_ID
+  - .
+
+options:
+  machineType: E2_HIGHCPU_32
+  diskSizeGb: 1000
+
+timeout: 86400s
+
+substitutions:
+  _BASE_IMAGE_TAG: v107
+  _LIBTPU_IMAGE_TAG: libtpu_1.1.0_RC00
+
+images:
+- gcr.io/kaggle-images/python-tpu:$BUILD_ID
+- gcr.io/kaggle-images/python-tpu-tensorflow-whl:$BUILD_ID
+
+tags: ['python', 'tpu']

tpu/tensorflow.Dockerfile

@@ -0,0 +1,39 @@
+ARG BASE_IMAGE_TAG
+
+FROM gcr.io/kaggle-images/python:${BASE_IMAGE_TAG} AS builder
+
+# Use Bazelisk to ensure the proper bazel version is used.
+RUN cd /usr/local/src && \
+    wget --no-verbose "https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-linux-amd64" && \
+    mv bazelisk-linux-amd64 /usr/local/bin/bazel && \
+    chmod u+x /usr/local/bin/bazel
+
+# Fetch TensorFlow & install dependencies.
+RUN cd /usr/local/src && \
+    git clone https://github.com/tensorflow/tensorflow && \
+    cd tensorflow && \
+    git checkout tags/v${TENSORFLOW_VERSION} && \
+    # TODO(rosbo): Is it really needed?
+    pip install keras_applications --no-deps && \
+    pip install keras_preprocessing --no-deps
+
+# Create a TensorFlow wheel for CPU
+RUN cd /usr/local/src/tensorflow && \
+    cat /dev/null | ./configure && \
+    bazel build \
+            --config=opt \
+            --distinct_host_configuration=true \
+            --define=framework_shared_object=true \
+            --define=with_tpu_support=true \
+            --copt=-DLIBTPU_ON_GCE \
+        //tensorflow/tools/pip_package:build_pip_package \
+            --local_ram_resources=HOST_RAM*.5
+
+RUN cd /usr/local/src/tensorflow && \
+    bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
+
+# TODO(b/152075195): Will likely need to install custom build for TFA & tensorflow-gcs-config
+
+# Use multi-stage builds to minimize image output size.
+FROM alpine:latest
+COPY --from=builder /tmp/tensorflow_pkg/tensorflow*.whl /tmp/tensorflow_pkg/