Merge pull request #1132 from Kaggle/tpu-1vm-image

Get tensorflow, jax, and pytorch working on TPU1VM

Author: djherbis

Jenkinsfile

@@ -56,6 +56,22 @@ pipeline {
             '''
           }
         }
+        stage('tensorflow TPU') {
+          options {
+            timeout(time: 180, unit: 'MINUTES')
+          }
+          steps {
+            sh '''#!/bin/bash
+              set -exo pipefail
+              source tpu/config.txt
+              cd packages/
+              ./build_package --base-image gcr.io/kaggle-images/python:${BASE_IMAGE_TAG} \
+                --package tpu-tensorflow \
+                --version $TENSORFLOW_VERSION \
+                --push
+            '''
+          }
+        }
       }
     }
     stage('Build/Test/Diff') {
@@ -150,7 +166,34 @@ pipeline {
               }
             }
           }
-        } 
+        }
+        stage('TPU VM') {
+          stages {
+            stage('Build Tensorflow TPU Image') {
+              options {
+                timeout(time: 20, unit: 'MINUTES')
+              }
+              steps {
+                sh '''#!/bin/bash
+                  set -exo pipefail
+
+                  ./tpu/build | ts
+                  ./push --tpu ${PRETEST_TAG}
+                '''
+              }
+            }
+            stage('Diff TPU VM Image') {
+              steps {
+                sh '''#!/bin/bash
+                set -exo pipefail
+
+                docker pull gcr.io/kaggle-private-byod/python-tpuvm:${PRETEST_TAG}
+                ./diff --tpu --target gcr.io/kaggle-private-byod/python-tpuvm:${PRETEST_TAG}
+              '''
+              }
+            }
+          }
+        }
       }
     }
 
@@ -161,6 +204,7 @@ pipeline {
 
           gcloud container images add-tag gcr.io/kaggle-images/python:${PRETEST_TAG} gcr.io/kaggle-images/python:${STAGING_TAG}
           gcloud container images add-tag gcr.io/kaggle-private-byod/python:${PRETEST_TAG} gcr.io/kaggle-private-byod/python:${STAGING_TAG}
+          gcloud container images add-tag gcr.io/kaggle-private-byod/python-tpuvm:${PRETEST_TAG} gcr.io/kaggle-private-byod/python-tpuvm:${STAGING_TAG}
         '''
       }
     }

diff

@@ -32,6 +32,10 @@ while :; do
             BASE_IMAGE_TAG='gcr.io/kaggle-private-byod/python:latest'
             TARGET_IMAGE_TAG='kaggle/python-gpu-build'
             ;;
+        -x|--tpu)
+            BASE_IMAGE_TAG='gcr.io/kaggle-private-byod/python-tpuvm:latest'
+            TARGET_IMAGE_TAG='kaggle/python-tpuvm-build'
+            ;;
         -b|--base)
             if [[ -z "$2" ]]; then
                 usage

packages/build_package

@@ -117,6 +117,8 @@ fi
 
 # Keep only `tf2-gpu.2-6:m80` in `gcr.io/deeplearning-platform-release/tf2-gpu.2-6:m80` 
 TAG=${BASE_IMAGE/gcr.io\/deeplearning-platform-release\//}
+# Keep only `python:v108` in `gcr.io/kaggle-images/python:v108`
+TAG=${TAG/gcr.io\/kaggle-images\//}
 # Replace the `:` in `tf2-gpu.2-6:m80` by `-`
 TAG=${TAG/:/-}
 # Append the package version

tpu/tensorflow.Dockerfile renamed to packages/tpu-tensorflow.Dockerfile

@@ -1,6 +1,8 @@
-ARG BASE_IMAGE_TAG
+ARG BASE_IMAGE
 
-FROM gcr.io/kaggle-images/python:${BASE_IMAGE_TAG} AS builder
+FROM ${BASE_IMAGE} AS builder
+
+ARG PACKAGE_VERSION
 
 # Use Bazelisk to ensure the proper bazel version is used.
 RUN cd /usr/local/src && \
@@ -12,12 +14,12 @@ RUN cd /usr/local/src && \
 RUN cd /usr/local/src && \
     git clone https://github.com/tensorflow/tensorflow && \
     cd tensorflow && \
-    git checkout tags/v${TENSORFLOW_VERSION} && \
+    git checkout tags/v${PACKAGE_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
+# Create a TensorFlow wheel for TPU
 RUN cd /usr/local/src/tensorflow && \
     cat /dev/null | ./configure && \
     bazel build \
@@ -32,7 +34,22 @@ RUN cd /usr/local/src/tensorflow && \
 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
+# 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.12.1 && \
+#    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
 
 # Use multi-stage builds to minimize image output size.
 FROM alpine:latest

push

@@ -8,6 +8,7 @@ Push a newly-built image with the given LABEL to gcr.io and DockerHub.
 
 Options:
     -g, --gpu                   Push the image with GPU support.
+    -t, --tpu                   Push the image with GPU support.
     -s, --source-image IMAGE    Tag for the source image. 
 EOF
 }
@@ -26,6 +27,10 @@ while :; do
             SOURCE_IMAGE_TAG='kaggle/python-gpu-build:latest'
             TARGET_IMAGE='gcr.io/kaggle-private-byod/python'
             ;;
+        -t|--tpu)
+            SOURCE_IMAGE_TAG='kaggle/python-tpuvm-build:latest'
+            TARGET_IMAGE='gcr.io/kaggle-private-byod/python-tpuvm'
+            ;;
         -s|--source-image)
             if [[ -z $2 ]]; then
                 usage

tests/common.py

@@ -4,3 +4,4 @@
 import unittest
 
 gpu_test = unittest.skipIf(len(os.environ.get('CUDA_VERSION', '')) == 0, 'Not running GPU tests')
+tpu_test = unittest.skipIf(len(os.environ.get('ISTPUVM', '')) == 0, 'Not running TPU tests')

tpu/Dockerfile

@@ -1,13 +1,27 @@
 ARG BASE_IMAGE_TAG
 ARG LIBTPU_IMAGE_TAG
-ARG TENSORFLOW_WHL_IMAGE_TAG
+ARG TENSORFLOW_VERSION
+ARG TORCH_VERSION
 
 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-tpu-tensorflow-whl:python-${BASE_IMAGE_TAG}-${TENSORFLOW_VERSION} AS tensorflow_whl
 FROM gcr.io/kaggle-images/python:${BASE_IMAGE_TAG}
 
+ENV ISTPUVM=1
+
 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
+    rm -rf /tmp/tensorflow_pkg && \
+    /tmp/clean-layer.sh
+
+# https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm#changing_pytorch_version
+RUN pip uninstall -y torch && \
+    pip install torch==${TORCH_VERSION} && \
+    pip install torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION}-cp37-cp37m-linux_x86_64.whl && \
+    /tmp/clean-layer.sh
+
+# https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm
+RUN pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html && \
+    /tmp/clean-layer.sh

tpu/build

@@ -0,0 +1,57 @@
+#!/bin/bash
+set -e
+
+usage() {
+cat << EOF
+Usage: $0 [OPTIONS]
+Build a new Python TPU 1VM Docker image.
+
+Options:
+    -c, --use-cache           Use layer cache when building a new image.
+EOF
+}
+
+CACHE_FLAG='--no-cache'
+DOCKERFILE='Dockerfile'
+IMAGE_TAG='kaggle/python-tpuvm-build'
+BUILD_ARGS=''
+
+while :; do
+    case "$1" in 
+        -h|--help)
+            usage
+            exit
+            ;;
+        -c|--use-cache)
+            CACHE_FLAG=''
+            ;;
+        -?*)
+            usage
+            printf 'ERROR: Unknown option: %s\n' "$1" >&2
+            exit
+            ;;
+        *)            
+            break
+    esac
+
+    shift
+done
+
+BUILD_ARGS+=" --build-arg GIT_COMMIT=$(git rev-parse HEAD)"
+BUILD_ARGS+=" --build-arg BUILD_DATE=$(date '+%Y%m%d-%H%M%S')"
+
+# Read build args from config.txt file.
+SRCDIR=$(dirname "${BASH_SOURCE[0]}")
+for l in `cat ${SRCDIR}/config.txt`; do
+    BUILD_ARGS+=" --build-arg $l"
+done
+
+readonly CACHE_FLAG
+readonly DOCKERFILE
+readonly IMAGE_TAG
+readonly BUILD_ARGS
+
+DOCKERFILE_PATH="$SRCDIR/$DOCKERFILE"
+
+set -x
+docker build --rm --pull $CACHE_FLAG -t "$IMAGE_TAG" -f "$DOCKERFILE_PATH" $BUILD_ARGS .

tpu/config.txt

@@ -0,0 +1,5 @@
+# TODO(b/213335159): Use ci-pretest for BASE_IMAGE_TAG once stable.
+BASE_IMAGE_TAG=v108
+LIBTPU_IMAGE_TAG=libtpu_1.1.0_RC00
+TENSORFLOW_VERSION=2.8.0
+TORCH_VERSION=1.11.0