make tpu tensorflow build conditional

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') {
@@ -155,7 +171,7 @@ pipeline {
  stages {
  stage('Build Tensorflow TPU Image') {
  options {
- timeout(time: 120, unit: 'MINUTES')
+ timeout(time: 20, unit: 'MINUTES')
  }
  steps {
  sh '''#!/bin/bash

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

packages/tpu-tensorflow.Dockerfile

@@ -1,6 +1,6 @@
-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
 
@@ -14,7 +14,7 @@ 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

tpu/Dockerfile

@@ -1,9 +1,9 @@
 ARG BASE_IMAGE_TAG
 ARG LIBTPU_IMAGE_TAG
-ARG TENSORFLOW_WHL_IMAGE_TAG
+ARG TENSORFLOW_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

tpu/build

@@ -12,8 +12,6 @@ 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=''
@@ -49,15 +47,11 @@ for l in `cat ${SRCDIR}/config.txt`; do
 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

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