<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Train tensorflow or keras model on GCP or Kubeflow from Notebooks
This notebook introduces you to using Kubeflow Fairing to train the model to Kubeflow on Google Kubernetes Engine (GKE), and Google Cloud AI Platform training. This notebook demonstrate how to
Step1:  Define the model logic
Step2:  Train an Keras model in a notebook
Step3:  Spicify a image registry that will hold the image built by fairing
Step4:  Deploy the training job to kubeflow cluster
Step5:  Deploy distributed training job to kubeflow cluster
Step6:  Deploy the training job as CMLE training job
Doesn’t support CMLE distributed training
Step7:  Inspect training process with tensorboard
Step8:  Deploy the trained model to Kubeflow for predictions
<ASSISTANT_TASK:>
Python Code:
import os
import logging
import tensorflow as tf
import fairing
import numpy as np
from datetime import datetime
from fairing.cloud import gcp
# Setting up google container repositories (GCR) for storing output containers
# You can use any docker container registry istead of GCR
# For local notebook, GCP_PROJECT should be set explicitly
GCP_PROJECT = fairing.cloud.gcp.guess_project_name()
GCP_Bucket = os.environ['GCP_BUCKET'] # e.g., 'gs://kubeflow-demo-g/'
# This is for local notebook instead of that in kubeflow cluster
# os.environ['GOOGLE_APPLICATION_CREDENTIALS']=
Explanation: Train tensorflow or keras model on GCP or Kubeflow from Notebooks
This notebook introduces you to using Kubeflow Fairing to train the model to Kubeflow on Google Kubernetes Engine (GKE), and Google Cloud AI Platform training. This notebook demonstrate how to:
Train an Keras model in a local notebook,
Use Kubeflow Fairing to train an Keras model remotely on Kubeflow cluster,
Use Kubeflow Fairing to train an Keras model remotely on AI Platform training,
Use Kubeflow Fairing to deploy a trained model to Kubeflow, and Call the deployed endpoint for predictions.
You need Python 3.6 to use Kubeflow Fairing.
Setups
Pre-conditions
Deployed a kubeflow cluster through https://deploy.kubeflow.cloud/
Have the following environment variable ready: 
PROJECT_ID # project host the kubeflow cluster or for running AI platform training
DEPLOYMENT_NAME # kubeflow deployment name, the same the cluster name after delpoyed
GCP_BUCKET # google cloud storage bucket
Create service account
bash
export SA_NAME = [service account name]
gcloud iam service-accounts create ${SA_NAME}
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
    --member serviceAccount:${SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
    --role 'roles/editor'
gcloud iam service-accounts keys create ~/key.json \
    --iam-account ${SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com
Authorize for Source Repository
bash
gcloud auth configure-docker
Update local kubeconfig (for submiting job to kubeflow cluster)
bash
export CLUSTER_NAME=${DEPLOYMENT_NAME} # this is the deployment name or the kubenete cluster name
export ZONE=us-central1-c
gcloud container clusters get-credentials ${CLUSTER_NAME} --region ${ZONE}
Set the environmental variable: GOOGLE_APPLICATION_CREDENTIALS
bash
export GOOGLE_APPLICATION_CREDENTIALS = ....
python
os.environ['GOOGLE_APPLICATION_CREDENTIALS']=...
Install the lastest version of fairing
python
pip install git+https://github.com/kubeflow/fairing@master
Please not that the above configuration is required for notebook service running outside Kubeflow environment. And the examples demonstrated in the notebook is fully tested on notebook service outside Kubeflow cluster also.
The environemt variables, e.g. service account, projects and etc, should have been pre-configured while setting up the cluster.
End of explanation
def gcs_copy(src_path, dst_path):
    import subprocess
    print(subprocess.run(['gsutil', 'cp', src_path, dst_path], stdout=subprocess.PIPE).stdout[:-1].decode('utf-8'))
    
def gcs_download(src_path, file_name):
    import subprocess
    print(subprocess.run(['gsutil', 'cp', src_path, file_name], stdout=subprocess.PIPE).stdout[:-1].decode('utf-8'))
class TensorflowModel(object):
    
    def __init__(self):
        self.model_file = "mnist_model.h5"
        self.model = None    
    
    def build(self):
        self.model = tf.keras.models.Sequential([
          tf.keras.layers.Flatten(input_shape=(28, 28)),
          tf.keras.layers.Dense(512, activation=tf.nn.relu),
          tf.keras.layers.Dropout(0.2),
          tf.keras.layers.Dense(10, activation=tf.nn.softmax)
        ])
        self.model.compile(optimizer='adam',
                      loss='sparse_categorical_crossentropy',
                      metrics=['accuracy'])
        print(self.model.summary())
    
    def save_model(self):
        self.model.save(self.model_file)
        gcs_copy(self.model_file, GCP_Bucket + self.model_file)
    
    def train(self):
        self.build()
        
        mnist = tf.keras.datasets.mnist
        (x_train, y_train),(x_test, y_test) = mnist.load_data()
        x_train, x_test = x_train / 255.0, x_test / 255.0
        
        callbacks = [
          # Interrupt training if `val_loss` stops improving for over 2 epochs
          tf.keras.callbacks.EarlyStopping(patience=2, monitor='val_loss'),
          # Write TensorBoard logs to `./logs` directory
          tf.keras.callbacks.TensorBoard(log_dir=GCP_Bucket + 'logs/' 
                                         + datetime.now().date().__str__())
        ]
        self.model.fit(x_train, y_train, batch_size=32, epochs=5, callbacks=callbacks,
                  validation_data=(x_test, y_test))
        self.save_model()
        
    def predict(self, X):
        if not self.model:
            self.model = tf.keras.models.load_model(self.model_file)
        # Do any preprocessing
        prediction = self.model.predict(data=X)
Explanation: Define the model logic
End of explanation
TensorflowModel().train()
Explanation: Train an Keras model in a notebook
End of explanation
# In this demo, I use gsutil, therefore i compile a special image to install GoogleCloudSDK as based image
base_image = 'gcr.io/{}/fairing-predict-example:latest'.format(GCP_PROJECT)
!docker build --build-arg PY_VERSION=3.6.4 . -t {base_image}
!docker push {base_image}
GCP_PROJECT = fairing.cloud.gcp.guess_project_name()
BASE_IMAGE = 'gcr.io/{}/fairing-predict-example:latest'.format(GCP_PROJECT)
DOCKER_REGISTRY = 'gcr.io/{}/fairing-job-tf'.format(GCP_PROJECT)
Explanation: Spicify a image registry that will hold the image built by fairing
End of explanation
from fairing import TrainJob
from fairing.backends import GKEBackend
train_job = TrainJob(TensorflowModel, BASE_IMAGE, input_files=["requirements.txt"],
                     docker_registry=DOCKER_REGISTRY, backend=GKEBackend())
train_job.submit()
Explanation: Deploy the training job to kubeflow cluster
End of explanation
fairing.config.set_builder(name='docker', registry=DOCKER_REGISTRY, 
                           base_image=BASE_IMAGE, push=True)
fairing.config.set_deployer(name='tfjob', worker_count=1, ps_count=1)
run_fn = fairing.config.fn(TensorflowModel)
run_fn()
Explanation: Deploy distributed training job to kubeflow cluster
End of explanation
from fairing import TrainJob
from fairing.backends import GCPManagedBackend
train_job = TrainJob(TensorflowModel, BASE_IMAGE, input_files=["requirements.txt"],
                     docker_registry=DOCKER_REGISTRY, backend=GCPManagedBackend())
train_job.submit()
Explanation: Deploy the training job as CMLE training job
Doesn’t support CMLE distributed training
End of explanation
# ! tensorboard --logdir=gs://kubeflow-demo-g/logs --host=localhost --port=8777
Explanation: Inspect training process with tensorboard
End of explanation
from fairing import PredictionEndpoint
from fairing.backends import KubeflowGKEBackend
# The trained_ames_model.joblib is exported during the above local training
endpoint = PredictionEndpoint(TensorflowModel, BASE_IMAGE, input_files=['mnist_model.h5', "requirements.txt"],
                              docker_registry=DOCKER_REGISTRY, backend=KubeflowGKEBackend())
endpoint.create()
endpoint.delete()
Explanation: Deploy the trained model to Kubeflow for predictions
End of explanation
<END_TASK>