The objective is to show GPU based TensorFlow training on ABM cluster on GCE. The summary of tasks which need to be performed:

• Create ABM GPU cluster with NVIDIA Drivers and V100 GPU
• Perform Model Training using TFJob and GPUs
• Perform ResNet TensorFlow model serving using ABM GPU
• Author the README guide

We need to create a Terraform based ABM cluster setup [1] using GPU on Google Compute Engine. We've made upstream contribution to the Google Terraform Provider for GPU support [2] in the instance_template module release [3]. We need to perform few things --

• Provide a config based sample for GPU hardware enabled ABM (anthos-samples repo)
• Install NVIDIA CUDA drivers on ABM cluster (anthos-samples repo)
• Create a setup guide in this repo on how to create the ABM GPU cluster

[1] https://github.com/GoogleCloudPlatform/anthos-samples/tree/master/anthos-bm-gcp-terraform [2] https://github.com/terraform-google-modules/terraform-google-vm/pull/160 [3] https://github.com/terraform-google-modules/terraform-google-vm/pull/158

TensorFlow ML model training on Anthos Bare Metal. Create a basic Hello World training app using TFJob.

TFJob is a Kubernetes custom resource that you can use to run TensorFlow training jobs on Kubernetes. The Kubeflow implementation of TFJob is in tf-operator.

The model is a simple MNIST model with persistent volume. Once we have a running example on ABM, we'll need to create an end to end ML workflow on ABM and serve the model as well.

The above example also provides an exploration into Kubeflow on ABM and why such a workflow makes sense. Even though we can run TensorFlow directly on Kubernetes (as shown in the TF Serving example), the TFJob abstraction makes it easy to define TensorFlow deployments. The standard way of performing multi-node training in TensorFlow is using TF_CONFIG environment variable. This environment variable is a JSON string which provides cluster and task information, and is set for each binary running on the cluster. The nodes in the cluster can be worker and ps nodes. In multi-worker training, there is usually one worker that takes on a little more responsibility like saving checkpoint and writing summary file for TensorBoard in addition to what a regular worker does. Such worker is referred to as the chief worker, and it is customary that the worker with index 0 is appointed as the chief worker. Each node in the cluster also has one of the roles worker, ps, or chief.

The setting of TF_CONFIG environment variable can be a manual process if done outside Kubeflow, and this is the part which TFJob Controller automatically manages.

There has been current work going on with ABM Edge large scale rollout here. One of the common workloads we've been seeing on the Edge is the deployment of ML models (like serving ML models for image classification).

Deploying these pre-trained models on the Edge involves complex software systems. ABM 1.8.0+ provides an Edge profile with minimal vCPU and Mem resource requirements and therefore well suited to deploy appropriate size ML models. Once the rollout work is finished we'd like to replace that work with an ML workload.