Support GPU image for dnntrainer component

.cloudbuild.yaml

@@ -121,6 +121,10 @@ steps:
   entrypoint: '/bin/bash' 
   args: ['-c', 'cd /workspace/components/kubeflow/launcher && ./build_image.sh -p $PROJECT_ID -t $COMMIT_SHA']
   id: 'buildLauncher'
+- name: 'gcr.io/cloud-builders/docker'
+  entrypoint: '/bin/bash' 
+  args: ['-c', 'cd /workspace/components/kubeflow/dnntrainer && ./build_image.sh -p $PROJECT_ID -t $COMMIT_SHA -m gpu']
+  id: 'buildGpuTrainer'
 
 # Build the Dataproc-based pipeline component images
 - name: 'gcr.io/cloud-builders/docker'
@@ -199,6 +203,7 @@ images:
 # Images for the Kubeflow-based pipeline components
 - 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-deployer:$COMMIT_SHA'
 - 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf-trainer:$COMMIT_SHA'
+- 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf-trainer-gpu:$COMMIT_SHA'
 - 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf:$COMMIT_SHA'
 
 # Images for the Dataproc-based pipeline components

.release.cloudbuild.yaml

@@ -185,6 +185,17 @@ steps:
   args: ['tag', 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf-trainer:$COMMIT_SHA', 'gcr.io/ml-pipeline/ml-pipeline-kubeflow-tf-trainer:$COMMIT_SHA']
   id:   'tagTrainerCommitSHA'
   waitFor: ['pullTrainer']
+- name: 'gcr.io/cloud-builders/docker'
+  args: ['pull', 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf-trainer-gpu:$COMMIT_SHA']
+  id:   'pullGpuTrainer'
+- name: 'gcr.io/cloud-builders/docker'
+  args: ['tag', 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf-trainer-gpu:$COMMIT_SHA', 'gcr.io/ml-pipeline/ml-pipeline-kubeflow-tf-trainer-gpu:$TAG_NAME']
+  id:   'tagGpuTrainerVersionNumber'
+  waitFor: ['pullGpuTrainer']
+- name: 'gcr.io/cloud-builders/docker'
+  args: ['tag', 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf-trainer-gpu:$COMMIT_SHA', 'gcr.io/ml-pipeline/ml-pipeline-kubeflow-tf-trainer-gpu:$COMMIT_SHA']
+  id:   'tagGpuTrainerCommitSHA'
+  waitFor: ['pullGpuTrainer']
 - name: 'gcr.io/cloud-builders/docker'
   args: ['pull', 'gcr.io/$PROJECT_ID/ml-pipeline-kubeflow-tf:$COMMIT_SHA']
   id:   'pullLauncher'

components/kubeflow/dnntrainer/Dockerfile

@@ -12,18 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-FROM ubuntu:16.04
+ARG TF_VERSION
+FROM tensorflow/tensorflow:$TF_VERSION
 
 RUN apt-get update -y
 
 RUN apt-get install --no-install-recommends -y -q ca-certificates python-dev python-setuptools \
                                                   wget unzip git
 
-RUN easy_install pip
-
 RUN apt-get install --no-install-recommends -y -q build-essential && \
     pip install pyyaml==3.12 six==1.11.0 \
-        tensorflow==1.6.0 \
         tensorflow-transform==0.6.0 \
         tensorflow-model-analysis==0.6.0 && \
     apt-get --purge autoremove -y build-essential

components/kubeflow/dnntrainer/build_image.sh

@@ -13,11 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-while getopts ":hp:t:i:" opt; do
+while getopts ":hp:t:i:m:" opt; do
   case "${opt}" in
     h) echo "-p: project name"
         echo "-t: tag name"
         echo "-i: image name. If provided, project name and tag name are not necessary"
+        echo "-m: mode. Optional. The value can be 'cpu' or 'gpu'. Defaults to 'cpu'."
         exit
       ;;
     p) PROJECT_ID=${OPTARG}
@@ -26,12 +27,19 @@ while getopts ":hp:t:i:" opt; do
       ;;
     i) IMAGE_NAME=${OPTARG}
       ;;
-    \? ) echo "Usage: cmd [-p] project [-t] tag [-i] image"
+    m) MODE=${OPTARG}
+      if [ "${MODE}" != "cpu" ] && [ "${MODE}" != "gpu" ]; then
+        echo "mode value can only be 'cpu' or 'gpu'."
+        exit
+      fi
+      ;;
+    \? ) echo "Usage: cmd [-p] project [-t] tag [-i] image [-m] mode"
       exit
       ;;
   esac
 done
 
+set -x
 LOCAL_IMAGE_NAME=ml-pipeline-kubeflow-tf-trainer
 
 if [ -z "${PROJECT_ID}" ]; then
@@ -42,12 +50,18 @@ if [ -z "${TAG_NAME}" ]; then
   TAG_NAME=$(date +v%Y%m%d)-$(git describe --tags --always --dirty)-$(git diff | shasum -a256 | cut -c -6)
 fi
 
+TF_VERSION=1.6.0
+if [ "${MODE}" = "gpu" ]; then
+  TF_VERSION="${TF_VERSION}-gpu"
+  LOCAL_IMAGE_NAME="${LOCAL_IMAGE_NAME}-gpu"
+fi
+
 mkdir -p ./build
 rsync -arvp ./src/ ./build/
 cp ../../license.sh ./build
 cp ../../third_party_licenses.csv ./build
 
-docker build -t ${LOCAL_IMAGE_NAME} .
+docker build --build-arg TF_VERSION=${TF_VERSION} -t ${LOCAL_IMAGE_NAME} .
 if [ -z "${IMAGE_NAME}" ]; then
   docker tag ${LOCAL_IMAGE_NAME} gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME}:${TAG_NAME}
   docker push gcr.io/${PROJECT_ID}/${LOCAL_IMAGE_NAME}:${TAG_NAME}

components/third_party_licenses.csv

@@ -169,3 +169,11 @@ tensorflow-tensorboard,https://raw.githubusercontent.com/tensorflow/tensorboard/
 tensorflow-data-validation,https://raw.githubusercontent.com/tensorflow/data-validation/master/LICENSE,Apache 2.0
 tensorflow-metadata,https://raw.githubusercontent.com/tensorflow/metadata/master/LICENSE,Apache 2.0
 defusedxml,https://raw.githubusercontent.com/tiran/defusedxml/master/LICENSE,PSF
+backports.functools-lru-cache,https://raw.githubusercontent.com/jaraco/backports.functools_lru_cache/master/LICENSE,MIT
+cycler,https://raw.githubusercontent.com/matplotlib/cycler/master/LICENSE,MIT
+h5py,https://raw.githubusercontent.com/h5py/h5py/07f3c362887563c1040b10d2cbd8cca49ceac1b9/licenses/license.txt
+matplotlib,https://raw.githubusercontent.com/matplotlib/matplotlib/master/LICENSE/LICENSE
+Pillow,https://raw.githubusercontent.com/python-pillow/Pillow/master/LICENSE
+sklearn,https://raw.githubusercontent.com/scikit-learn/scikit-learn/master/COPYING,BSD
+tensorflow-gpu,https://raw.githubusercontent.com/tensorflow/tensorflow/master/LICENSE,Apache 2.0
+webencodings,https://raw.githubusercontent.com/gsnedders/python-webencodings/master/LICENSE,BSD