add jenkins setup

Author: djherbis

Jenkinsfile

@@ -150,7 +150,34 @@ pipeline {
               }
             }
           }
-        } 
+        }
+        stage('TPU VM') {
+          stages {
+            stage('Build Tensorflow TPU Image') {
+              options {
+                timeout(time: 120, 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 +188,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

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

@@ -6,6 +6,8 @@ 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}
 
+ENV ISTPUVM=1
+
 COPY --from=libtpu /libtpu.so /lib
 
 COPY --from=tensorflow_whl /tmp/tensorflow_pkg/tensorflow*.whl /tmp/tensorflow_pkg/

tpu/build

@@ -0,0 +1,63 @@
+#!/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'
+TENSORFLOW_DOCKERFILE='tensorflow.Dockerfile'
+TENSORFLOW_IMAGE_TAG='kaggle/python-tpuvm-tensorflow-build'
+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 TENSORFLOW_DOCKERFILE
+readonly TENSORFLOW_IMAGE_TAG
+readonly DOCKERFILE
+readonly IMAGE_TAG
+readonly BUILD_ARGS
+
+TENSORFLOW_DOCKERFILE_PATH="$SRCDIR/$TENSORFLOW_DOCKERFILE"
+DOCKERFILE_PATH="$SRCDIR/$DOCKERFILE"
+
+set -x
+docker build --rm --pull $CACHE_FLAG -t "$TENSORFLOW_IMAGE_TAG" -f "$TENSORFLOW_DOCKERFILE_PATH" $BUILD_ARGS .
+docker build --rm --pull $CACHE_FLAG -t "$IMAGE_TAG" -f "$DOCKERFILE_PATH" $BUILD_ARGS .

tpu/config.txt

@@ -0,0 +1,3 @@
+BASE_IMAGE_TAG=v108
+LIBTPU_IMAGE_TAG=libtpu_1.1.0_RC00
+TENSORFLOW_WHL_IMAGE_TAG=local