add jenkins setup

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