<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 MNIST end to end on Kubeflow on GKE
This example guides you through
Step1:  Install the required libraries
Run the next cell to import the libraries required to train this model.
Step2:  Wait for the message Configure docker credentials before moving on to the next cell.
Step3:  Configure a Docker registry for Kubeflow Fairing
In order to build Docker images from your notebook, you need a Docker registry to store the images.
Below you set some variables specifying a Container Registry.
Kubeflow Fairing provides a utility function to guess the name of your GCP project.
Step4:  Use Kubeflow Fairing to build the Docker image
This notebook uses Kubeflow Fairing's kaniko builder to build a Docker image that includes all your dependencies.
  * You use kaniko because you want to be able to run pip to install dependencies.
  * Kaniko gives you the flexibility to build images from Dockerfiles.
Step5:  Run the next cell and wait until you see a message like Built image gcr.io/&lt;your-project&gt;/fairing-job/mnist
Step6:  Create a Cloud Storage bucket
Run the next cell to create a Google Cloud Storage (GCS) bucket to store your models and other results.
Since this notebook is running in Python, the cell uses the GCS Python client libraries, but you can use the gsutil command line instead.
Step8:  Distributed training
To train the model, this example uses TFJob to run a distributed training job. Run the next cell to set up the YAML specification for the job
Step9:  Create the training job
To submit the training job, you could write the spec to a YAML file and then do kubectl apply -f {FILE}.
However, because you are running in a Jupyter notebook, you use the TFJob client. 
* You run the TFJob in a namespace created by a Kubeflow profile.
* The namespace is the same as the namespace where you are running the notebook.
* Creating a profile ensures that the namespace is provisioned with service accounts and other resources needed for Kubeflow.
Step10:  Check the job using kubectl
Above you used the Python SDK for TFJob to check the status. You can also use kubectl get the status of your job. 
The job conditions will tell you whether the job is running, succeeded or failed.
Step12:  Get the training logs
There are two ways to get the logs for the training job
Step16:  Deploy TensorBoard
The next step is to create a Kubernetes deployment to run TensorBoard.
TensorBoard will be accessible behind the Kubeflow IAP endpoint.
Step17:  Set a variable defining your endpoint
Set endpoint to https
Step18:  Access the TensorBoard UI
Run the cell below to find the endpoint for the TensorBoard UI.
Step19:  Wait for the training job to finish
You can use the TFJob client to wait for the job to finish
Step23:  Serve the model
Now you can deploy the model using TensorFlow Serving.
You need to create the following
Step27:  Deploy the UI for the MNIST web app
Deploy the UI to visualize the MNIST prediction results.
This example uses a prebuilt and public Docker image for the UI.
Step28:  Access the MNIST web UI
A reverse proxy route is automatically added to the Kubeflow IAP endpoint. The MNIST endpoint is
<ASSISTANT_TASK:>
Python Code:
import logging
import os
import uuid
from importlib import reload
from oauth2client.client import GoogleCredentials
credentials = GoogleCredentials.get_application_default()
Explanation: MNIST end to end on Kubeflow on GKE
This example guides you through:
Taking an example TensorFlow model and modifying it to support distributed training.
Serving the resulting model using TFServing.
Deploying and using a web app that sends prediction requests to the model.
Requirements
You must be running Kubeflow 1.0 on Kubernetes Engine (GKE) with Cloud Identity-Aware Proxy (Cloud IAP). See the guide to deploying Kubeflow on GCP.
Run this notebook within your Kubeflow cluster. See the guide to setting up your Kubeflow notebooks.
Prepare model
There is a delta between existing distributed MNIST examples and what's needed to run well as a TFJob.
Basically, you must:
Add options in order to make the model configurable.
Use tf.estimator.train_and_evaluate to enable model exporting and serving.
Define serving signatures for model serving.
This tutorial provides a Python program that's already prepared for you: model.py.
Verify that you have a Google Cloud Platform (GCP) account
The cell below checks that this notebook was spawned with credentials to access GCP.
End of explanation
import notebook_setup
reload(notebook_setup)
notebook_setup.notebook_setup()
Explanation: Install the required libraries
Run the next cell to import the libraries required to train this model.
End of explanation
import k8s_util
# Force a reload of Kubeflow. Since Kubeflow is a multi namespace module,
# doing the reload in notebook_setup may not be sufficient.
import kubeflow
reload(kubeflow)
from kubernetes import client as k8s_client
from kubernetes import config as k8s_config
from kubeflow.tfjob.api import tf_job_client as tf_job_client_module
from IPython.core.display import display, HTML
import yaml
Explanation: Wait for the message Configure docker credentials before moving on to the next cell.
End of explanation
from kubernetes import client as k8s_client
from kubernetes.client import rest as k8s_rest
from kubeflow import fairing   
from kubeflow.fairing import utils as fairing_utils
from kubeflow.fairing.builders import append
from kubeflow.fairing.deployers import job
from kubeflow.fairing.preprocessors import base as base_preprocessor
# Setting up Google Container Registry (GCR) for storing output containers.
# You can use any Docker container registry instead of GCR.
GCP_PROJECT = fairing.cloud.gcp.guess_project_name()
DOCKER_REGISTRY = 'gcr.io/{}/fairing-job'.format(GCP_PROJECT)
namespace = fairing_utils.get_current_k8s_namespace()
logging.info(f"Running in project {GCP_PROJECT}")
logging.info(f"Running in namespace {namespace}")
logging.info(f"Using Docker registry {DOCKER_REGISTRY}")
Explanation: Configure a Docker registry for Kubeflow Fairing
In order to build Docker images from your notebook, you need a Docker registry to store the images.
Below you set some variables specifying a Container Registry.
Kubeflow Fairing provides a utility function to guess the name of your GCP project.
End of explanation
# TODO(https://github.com/kubeflow/fairing/issues/426): We should get rid of this once the default 
# Kaniko image is updated to a newer image than 0.7.0.
from kubeflow.fairing import constants
constants.constants.KANIKO_IMAGE = "gcr.io/kaniko-project/executor:v0.14.0"
from kubeflow.fairing.builders import cluster
# output_map is a map of extra files to add to the notebook.
# It is a map from source location to the location inside the context.
output_map =  {
    "Dockerfile.model": "Dockerfile",
    "model.py": "model.py"
}
preprocessor = base_preprocessor.BasePreProcessor(
    command=["python"], # The base class will set this.
    input_files=[],
    path_prefix="/app", # irrelevant since we aren't preprocessing any files
    output_map=output_map)
preprocessor.preprocess()
Explanation: Use Kubeflow Fairing to build the Docker image
This notebook uses Kubeflow Fairing's kaniko builder to build a Docker image that includes all your dependencies.
  * You use kaniko because you want to be able to run pip to install dependencies.
  * Kaniko gives you the flexibility to build images from Dockerfiles.
End of explanation
# Use a Tensorflow image as the base image
# We use a custom Dockerfile 
cluster_builder = cluster.cluster.ClusterBuilder(registry=DOCKER_REGISTRY,
                                                 base_image="", # base_image is set in the Dockerfile
                                                 preprocessor=preprocessor,
                                                 image_name="mnist",
                                                 dockerfile_path="Dockerfile",
                                                 pod_spec_mutators=[fairing.cloud.gcp.add_gcp_credentials_if_exists],
                                                 context_source=cluster.gcs_context.GCSContextSource())
cluster_builder.build()
logging.info(f"Built image {cluster_builder.image_tag}")
Explanation: Run the next cell and wait until you see a message like Built image gcr.io/&lt;your-project&gt;/fairing-job/mnist:&lt;1234567&gt;.
End of explanation
from google.cloud import storage
bucket = f"{GCP_PROJECT}-mnist"
client = storage.Client()
b = storage.Bucket(client=client, name=bucket)
if not b.exists():
    logging.info(f"Creating bucket {bucket}")
    b.create()
else:
    logging.info(f"Bucket {bucket} already exists")    
Explanation: Create a Cloud Storage bucket
Run the next cell to create a Google Cloud Storage (GCS) bucket to store your models and other results.
Since this notebook is running in Python, the cell uses the GCS Python client libraries, but you can use the gsutil command line instead.
End of explanation
train_name = f"mnist-train-{uuid.uuid4().hex[:4]}"
num_ps = 1
num_workers = 2
model_dir = f"gs://{bucket}/mnist"
export_path = f"gs://{bucket}/mnist/export" 
train_steps = 200
batch_size = 100
learning_rate = .01
image = cluster_builder.image_tag
train_spec = fapiVersion: kubeflow.org/v1
kind: TFJob
metadata:
  name: {train_name}  
spec:
  tfReplicaSpecs:
    Ps:
      replicas: {num_ps}
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          serviceAccount: default-editor
          containers:
          - name: tensorflow
            command:
            - python
            - /opt/model.py
            - --tf-model-dir={model_dir}
            - --tf-export-dir={export_path}
            - --tf-train-steps={train_steps}
            - --tf-batch-size={batch_size}
            - --tf-learning-rate={learning_rate}
            image: {image}
            workingDir: /opt
          restartPolicy: OnFailure
    Chief:
      replicas: 1
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          serviceAccount: default-editor
          containers:
          - name: tensorflow
            command:
            - python
            - /opt/model.py
            - --tf-model-dir={model_dir}
            - --tf-export-dir={export_path}
            - --tf-train-steps={train_steps}
            - --tf-batch-size={batch_size}
            - --tf-learning-rate={learning_rate}
            image: {image}
            workingDir: /opt
          restartPolicy: OnFailure
    Worker:
      replicas: 1
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          serviceAccount: default-editor
          containers:
          - name: tensorflow
            command:
            - python
            - /opt/model.py
            - --tf-model-dir={model_dir}
            - --tf-export-dir={export_path}
            - --tf-train-steps={train_steps}
            - --tf-batch-size={batch_size}
            - --tf-learning-rate={learning_rate}
            image: {image}
            workingDir: /opt
          restartPolicy: OnFailure
           
Explanation: Distributed training
To train the model, this example uses TFJob to run a distributed training job. Run the next cell to set up the YAML specification for the job:
End of explanation
tf_job_client = tf_job_client_module.TFJobClient()
tf_job_body = yaml.safe_load(train_spec)
tf_job = tf_job_client.create(tf_job_body, namespace=namespace)  
logging.info(f"Created job {namespace}.{train_name}")
Explanation: Create the training job
To submit the training job, you could write the spec to a YAML file and then do kubectl apply -f {FILE}.
However, because you are running in a Jupyter notebook, you use the TFJob client. 
* You run the TFJob in a namespace created by a Kubeflow profile.
* The namespace is the same as the namespace where you are running the notebook.
* Creating a profile ensures that the namespace is provisioned with service accounts and other resources needed for Kubeflow.
End of explanation
!kubectl get tfjobs -o yaml {train_name}
Explanation: Check the job using kubectl
Above you used the Python SDK for TFJob to check the status. You can also use kubectl get the status of your job. 
The job conditions will tell you whether the job is running, succeeded or failed.
End of explanation
from urllib.parse import urlencode
for replica in ["chief", "worker", "ps"]:    
    logs_filter = fresource.type="k8s_container"    
    labels."k8s-pod/tf-job-name" = "{train_name}"
    labels."k8s-pod/tf-replica-type" = "{replica}"    
    resource.labels.container_name="tensorflow" 
    new_params = {'project': GCP_PROJECT,
                  # Logs for last 7 days
                  'interval': 'P7D',
                  'advancedFilter': logs_filter}
    query = urlencode(new_params)
    url = "https://console.cloud.google.com/logs/viewer?" + query
    display(HTML(f"Link to: <a href='{url}'>{replica} logs</a>"))
Explanation: Get the training logs
There are two ways to get the logs for the training job:
Using kubectl to fetch the pod logs. These logs are ephemeral; they will be unavailable when the pod is garbage collected to free up resources.
Using Stackdriver.
Kubernetes logs are automatically available in Stackdriver.
You can use labels to locate the logs for a specific pod.
In the cell below, you use labels for the training job name and process type to locate the logs for a specific pod.
Run the cell below to get a link to Stackdriver for your logs:
End of explanation
tb_name = "mnist-tensorboard"
tb_deploy = fapiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: mnist-tensorboard
  name: {tb_name}
  namespace: {namespace}
spec:
  selector:
    matchLabels:
      app: mnist-tensorboard
  template:
    metadata:
      labels:
        app: mnist-tensorboard
        version: v1
    spec:
      serviceAccount: default-editor
      containers:
      - command:
        - /usr/local/bin/tensorboard
        - --logdir={model_dir}
        - --port=80
        image: tensorflow/tensorflow:1.15.2-py3
        name: tensorboard
        ports:
        - containerPort: 80
tb_service = fapiVersion: v1
kind: Service
metadata:
  labels:
    app: mnist-tensorboard
  name: {tb_name}
  namespace: {namespace}
spec:
  ports:
  - name: http-tb
    port: 80
    targetPort: 80
  selector:
    app: mnist-tensorboard
  type: ClusterIP
tb_virtual_service = fapiVersion: networking.istio.io/v1alpha3
kind: VirtualService
metadata:
  name: {tb_name}
  namespace: {namespace}
spec:
  gateways:
  - kubeflow/kubeflow-gateway
  hosts:
  - '*'
  http:
  - match:
    - uri:
        prefix: /mnist/{namespace}/tensorboard/
    rewrite:
      uri: /
    route:
    - destination:
        host: {tb_name}.{namespace}.svc.cluster.local
        port:
          number: 80
    timeout: 300s
tb_specs = [tb_deploy, tb_service, tb_virtual_service]
k8s_util.apply_k8s_specs(tb_specs, k8s_util.K8S_CREATE_OR_REPLACE)
Explanation: Deploy TensorBoard
The next step is to create a Kubernetes deployment to run TensorBoard.
TensorBoard will be accessible behind the Kubeflow IAP endpoint.
End of explanation
endpoint = None
if endpoint:
  logging.info(f"endpoint set to {endpoint}")
else:
  logging.info("Warning: You must set {endpoint} in order to print out the URLs where you can access your web apps.")
Explanation: Set a variable defining your endpoint
Set endpoint to https://your-domain (with no slash at the end). Your domain typically has the following pattern: &lt;your-kubeflow-deployment-name&gt;.endpoints.&lt;your-gcp-project&gt;.cloud.goog. You can see your domain in the URL that you're using to access this notebook.
End of explanation
if endpoint:    
    vs = yaml.safe_load(tb_virtual_service)
    path= vs["spec"]["http"][0]["match"][0]["uri"]["prefix"]
    tb_endpoint = endpoint + path
    display(HTML(f"TensorBoard UI is at <a href='{tb_endpoint}'>{tb_endpoint}</a>"))
Explanation: Access the TensorBoard UI
Run the cell below to find the endpoint for the TensorBoard UI.
End of explanation
tf_job = tf_job_client.wait_for_condition(train_name, expected_condition=["Succeeded", "Failed"], namespace=namespace)
if tf_job_client.is_job_succeeded(train_name, namespace):
    logging.info(f"TFJob {namespace}.{train_name} succeeded")
else:
    raise ValueError(f"TFJob {namespace}.{train_name} failed")  
Explanation: Wait for the training job to finish
You can use the TFJob client to wait for the job to finish:
End of explanation
deploy_name = "mnist-model"
model_base_path = export_path
# The web UI defaults to mnist-service so if you change the name, you must
# change it in the UI as well.
model_service = "mnist-service"
deploy_spec = fapiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: mnist
  name: {deploy_name}
  namespace: {namespace}
spec:
  selector:
    matchLabels:
      app: mnist-model
  template:
    metadata:
      # TODO(jlewi): Right now we disable the istio side car because otherwise ISTIO rbac will prevent the
      # UI from sending RPCs to the server. We should create an appropriate ISTIO rbac authorization
      # policy to allow traffic from the UI to the model servier.
      # https://istio.io/docs/concepts/security/#target-selectors
      annotations:        
        sidecar.istio.io/inject: "false"
      labels:
        app: mnist-model
        version: v1
    spec:
      serviceAccount: default-editor
      containers:
      - args:
        - --port=9000
        - --rest_api_port=8500
        - --model_name=mnist
        - --model_base_path={model_base_path}
        - --monitoring_config_file=/var/config/monitoring_config.txt
        command:
        - /usr/bin/tensorflow_model_server
        env:
        - name: modelBasePath
          value: {model_base_path}
        image: tensorflow/serving:1.15.0
        imagePullPolicy: IfNotPresent
        livenessProbe:
          initialDelaySeconds: 30
          periodSeconds: 30
          tcpSocket:
            port: 9000
        name: mnist
        ports:
        - containerPort: 9000
        - containerPort: 8500
        resources:
          limits:
            cpu: "4"
            memory: 4Gi
          requests:
            cpu: "1"
            memory: 1Gi
        volumeMounts:
        - mountPath: /var/config/
          name: model-config
      volumes:
      - configMap:
          name: {deploy_name}
        name: model-config
service_spec = fapiVersion: v1
kind: Service
metadata:
  annotations:    
    prometheus.io/path: /monitoring/prometheus/metrics
    prometheus.io/port: "8500"
    prometheus.io/scrape: "true"
  labels:
    app: mnist-model
  name: {model_service}
  namespace: {namespace}
spec:
  ports:
  - name: grpc-tf-serving
    port: 9000
    targetPort: 9000
  - name: http-tf-serving
    port: 8500
    targetPort: 8500
  selector:
    app: mnist-model
  type: ClusterIP
monitoring_config = fkind: ConfigMap
apiVersion: v1
metadata:
  name: {deploy_name}
  namespace: {namespace}
data:
  monitoring_config.txt: |-
    prometheus_config: {{
      enable: true,
      path: "/monitoring/prometheus/metrics"
    }}
model_specs = [deploy_spec, service_spec, monitoring_config]
k8s_util.apply_k8s_specs(model_specs, k8s_util.K8S_CREATE_OR_REPLACE)     
Explanation: Serve the model
Now you can deploy the model using TensorFlow Serving.
You need to create the following:
* A Kubernetes deployment.
* A Kubernetes service.
* (Optional) A configmap containing the Prometheus monitoring configuration.
End of explanation
ui_name = "mnist-ui"
ui_deploy = fapiVersion: apps/v1
kind: Deployment
metadata:
  name: {ui_name}
  namespace: {namespace}
spec:
  replicas: 1
  selector:
    matchLabels:
      app: mnist-web-ui
  template:
    metadata:
      labels:
        app: mnist-web-ui
    spec:
      containers:
      - image: gcr.io/kubeflow-examples/mnist/web-ui:v20190112-v0.2-142-g3b38225
        name: web-ui
        ports:
        - containerPort: 5000        
      serviceAccount: default-editor
ui_service = fapiVersion: v1
kind: Service
metadata:
  annotations:
  name: {ui_name}
  namespace: {namespace}
spec:
  ports:
  - name: http-mnist-ui
    port: 80
    targetPort: 5000
  selector:
    app: mnist-web-ui
  type: ClusterIP
ui_virtual_service = fapiVersion: networking.istio.io/v1alpha3
kind: VirtualService
metadata:
  name: {ui_name}
  namespace: {namespace}
spec:
  gateways:
  - kubeflow/kubeflow-gateway
  hosts:
  - '*'
  http:
  - match:
    - uri:
        prefix: /mnist/{namespace}/ui/
    rewrite:
      uri: /
    route:
    - destination:
        host: {ui_name}.{namespace}.svc.cluster.local
        port:
          number: 80
    timeout: 300s
ui_specs = [ui_deploy, ui_service, ui_virtual_service]
k8s_util.apply_k8s_specs(ui_specs, k8s_util.K8S_CREATE_OR_REPLACE)     
    
Explanation: Deploy the UI for the MNIST web app
Deploy the UI to visualize the MNIST prediction results.
This example uses a prebuilt and public Docker image for the UI.
End of explanation
if endpoint:    
    vs = yaml.safe_load(ui_virtual_service)
    path= vs["spec"]["http"][0]["match"][0]["uri"]["prefix"]
    ui_endpoint = endpoint + path
    display(HTML(f"mnist UI is at <a href='{ui_endpoint}'>{ui_endpoint}</a>"))
Explanation: Access the MNIST web UI
A reverse proxy route is automatically added to the Kubeflow IAP endpoint. The MNIST endpoint is:
https:/${KUBEFlOW_ENDPOINT}/mnist/${NAMESPACE}/ui/
where NAMESPACE is the namespace where you're running the Jupyter notebook.
End of explanation
<END_TASK>