<jupyter_start><jupyter_text>TF on GKEThis notebook shows how to run the [TensorFlow CIFAR10 sample](https://github.com/tensorflow/models/tree/master/tutorials/image/cifar10_estimator) on GKE using [TFJobs](https://github.com/kubeflow/tf-operator)  Requirements To run this notebook you must have the following installed  * gcloud  * kubectl  * helm  * kubernetes python client library  There is a Docker image based on Datalab suitable for running this notebook.You can start that container as follows```docker run --name=gke-datalab -p "127.0.0.1:8081:8080" \    -v "${HOME}:/content/datalab/home" \    -v /var/run/docker.sock:/var/run/docker.sock -d  -e "PROJECT_ID=" \    gcr.io/tf-on-k8s-dogfood/gke-datalab:v20171103-73616f0```  * You need to map in docker if you want tobuild docker images inside the container.  * Alternatively, you can set "use_gcb" to true in order to build the images using Google Container Builder  Additionally the [py package](https://github.com/kubeflow/tf-operator/tree/master/py) must be a top level package importable as py  * If you cloned [kubeflow/tf-operator](https://github.com/kubeflow/tf-operator) and are running this notebook in place the path with be configured automatically  Preliminaries<jupyter_code># Turn on autoreloading
%load_ext autoreload
%autoreload 2<jupyter_output><empty_output><jupyter_text>import a bunch of modules and set some constants<jupyter_code>from __future__ import print_function

import logging
import os
import sys

# Assumes we are running inside the cloned repo.
# Try to setup the path so we can import py as a top level package
ROOT_DIR = os.path.abspath(os.path.join("../.."))
if os.path.exists(os.path.join(ROOT_DIR, "py")):
  if not ROOT_DIR in sys.path:
    sys.path.append(ROOT_DIR)
  
import kubernetes
from kubernetes import client as k8s_client
from kubernetes import config as k8s_config
from kubernetes.client.rest import ApiException
from kubernetes.client.models.v1_label_selector import V1LabelSelector
import datetime
from googleapiclient import discovery
from googleapiclient import errors
from oauth2client.client import GoogleCredentials
from pprint import pprint
try:
  from py import build_and_push_image
  from py import util
except ImportError:
  raise ImportError("Please ensure the py package in https://github.com/kubeflow/tf-operator is a top level package")
import StringIO
import subprocess
import urllib
import urllib2
import time
import yaml

logging.getLogger().setLevel(logging.INFO)

TF_JOB_GROUP = "kubeflow.org"
TF_JOB_VERSION = "v1alpha1"
TF_JOB_PLURAL = "tfjobs"
TF_JOB_KIND = "TFJob"<jupyter_output><empty_output><jupyter_text>Configure the notebook for your useChange the constants defined below.  1. Change **project** to a project you have access to.     * GKE should be enabled for that project  1. Change **data_dir** and **job_dir**     * Use a GCS bucket that you have access to     * Ensure the service account on your GKE cluster can read/write to this GCS bucket* Optional change the cluster name<jupyter_code>project="cloud-ml-dev"
zone="us-east1-d"
cluster_name="gke-tf-example"
registry = "gcr.io/" + project
data_dir = "gs://cloud-ml-dev_jlewi/cifar10/data"
job_dirs = "gs://cloud-ml-dev_jlewi/cifar10/jobs"
gke = discovery.build("container", "v1")
namespace = "default"

# Whether to build containers using Google Container Builder.
# Set to false it will build by shelling out to docker build.
use_gcb = "false"<jupyter_output><empty_output><jupyter_text>GKE Cluster Setup * The instructions below create a **CPU** cluster* To create a GKE cluster with GPUs sign up for the [GKE GPU Alpha](https://goo.gl/forms/ef7eh2x00hV3hahx1)* To use GPUs set accelerator and accelerator_count* For a full list of cluster options see the [Cluster object](https://cloud.google.com/container-engine/reference/rest/v1/projects.zones.clustersCluster)   in the GKE API docsTo use an existing GKE cluster call **configure_kubectl** but not **create_cluster*** The code below issues a GKE request to create the cluster by calling util.create_cluster  * util.create_cluster uses the GKE python client library* After creating the cluster we call util.configure_kubectl  * This configures your machine to talk to the K8s master of the newly created cluster<jupyter_code>reload(util)
machine_type = "n1-standard-8"
use_gpu = True
if use_gpu:
  accelerator = "nvidia-tesla-k80"
  accelerator_count = 1
else:
  accelerator = None
  accelerator_count = 0

cluster_request = {
    "cluster": {
        "name": cluster_name,
        "description": "A GKE cluster for TF.",
        "initialNodeCount": 1,
        "nodeConfig": {
            "machineType": machine_type,
            "oauthScopes": [
              "https://www.googleapis.com/auth/cloud-platform",
            ],
        },
        # TODO(jlewi): Stop pinning GKE version once 1.8 becomes the default. 
        "initialClusterVersion": "1.8.1-gke.1",
    }
}

if bool(accelerator) != (accelerator_count > 0):
    raise ValueError("If accelerator is set accelerator_count must be  > 0")
    
if accelerator:
  # TODO(jlewi): Stop enabling Alpha once GPUs make it out of Alpha
  cluster_request["cluster"]["enableKubernetesAlpha"] = True

  cluster_request["cluster"]["nodeConfig"]["accelerators"] = [
      {
        "acceleratorCount": accelerator_count,
        "acceleratorType": accelerator,
      },
  ]
util.create_cluster(gke, project, zone, cluster_request)

util.configure_kubectl(project, zone, cluster_name)

k8s_config.load_kube_config()

# Create an API client object to talk to the K8s master.
api_client = k8s_client.ApiClient()<jupyter_output><empty_output><jupyter_text>Install the Operator* We need to deploy the [TFJob](https://github.com/kubeflow/tf-operator) custom resource on our K8s cluster* TFJob is deployed using the [helm](https://github.com/kubernetes/helm) package manager so first we need to setup helm on our cluster<jupyter_code>util.setup_cluster(api_client)<jupyter_output><empty_output><jupyter_text>Now that helm is setup we can deploy the TFJob CRD<jupyter_code>CHART="https://storage.googleapis.com/tf-on-k8s-dogfood-releases/latest/tf-job-operator-chart-latest.tgz"
util.run(["helm", "install", CHART, "-n", "tf-job", "--wait", "--replace", "--set", "rbac.install=true,cloud=gke"])<jupyter_output><empty_output><jupyter_text>Build Docker images To run a TensorFlow program on K8s we need to package our code as Docker images.The [Dockerfile](https://github.com/jlewi/k8s/blob/73616f09f335defc92f9b20225c272862e92e32b/examples/tensorflow-models/Dockerfile.template) for this example starts with the published Docker images for TensorFlow ands the code for our TensorFlow program  * In this example we are using the CIFAR10 example in the [TensorFlow's model zoo](https://github.com/tensorflow/models)  * So our Dockerfile just clones that repo  * Using TF's Docker images ensures we start with a reliable TF environment We need to build separate Docker images for CPU and GPU versions of TensorFlow.  * **modes** controls whether we build images for CPU, GPU or both   * Our Dockerfile is a [Jinja2](http://jinja.pocoo.org/) template, so we can easily    build docker images based on different TensorFlow versions  The base images controls which version of TensorFlow we will use  * Change the base images if you want to use a different version.<jupyter_code>reload(build_and_push_image)

if use_gpu:
  modes = ["cpu", "gpu"]
else:
  modes = ["cpu"]

image = os.path.join(registry, "tf-models")
dockerfile = os.path.join(ROOT_DIR, "examples", "tensorflow-models", "Dockerfile.template")
base_images = {
  "cpu": "gcr.io/tensorflow/tensorflow:1.3.0",
  "gpu": "gcr.io/tensorflow/tensorflow:1.3.0-gpu",
}
images = build_and_push_image.build_and_push(dockerfile, image, modes=modes, base_images=base_images)<jupyter_output><empty_output><jupyter_text>Create the CIFAR10 DatasetsWe need to create the cifar10 TFRecord files by running [generate_cifar10_tfrecords.py](https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10_estimator/generate_cifar10_tfrecords.py)  * We submit a K8s job to run this program  * You can skip this step if your data is already available in data_dir<jupyter_code>batch_api = k8s_client.BatchV1Api(api_client)

job_name = "cifar10-data-"+ datetime.datetime.now().strftime("%y%m%d-%H%M%S")

body = {}
body['apiVersion'] = "batch/v1"
body['kind'] = "Job"
body['metadata'] = {}
body['metadata']['name'] = job_name
body['metadata']['namespace'] = namespace

# Note backoffLimit requires K8s >= 1.8
spec = """
backoffLimit: 4
template:
  spec:
    containers:
    - name: cifar10
      image: {image}
      command: ["python",  "/tensorflow_models/tutorials/image/cifar10_estimator/generate_cifar10_tfrecords.py", "--data-dir={data_dir}"]
    restartPolicy: Never
""".format(data_dir=data_dir, image=images["cpu"])

spec_buffer = StringIO.StringIO(spec)
body['spec'] = yaml.load(spec_buffer)

try: 
    # Create a Resource
    api_response = batch_api.create_namespaced_job(namespace, body)
    print("Created job %s" % api_response.metadata.name)
except ApiException as e:
    print(
        "Exception when calling DefaultApi->apis_fqdn_v1_namespaces_namespace_resource_post: %s\n" % 
        e)<jupyter_output><empty_output><jupyter_text>wait for the job to finish<jupyter_code>while True:
  results = batch_api.read_namespaced_job(job_name, namespace)
  if results.status.succeeded >= 1 or results.status.failed >= 3:
    break
  print("Waiting for job %s ...." % results.metadata.name)
  time.sleep(5)

if results.status.succeeded >= 1:
  print("Job completed successfully")
else:
  print("Job failed")<jupyter_output><empty_output><jupyter_text>Create a TFJob To submit a TFJob, we define a TFJob spec and then create it in our cluster<jupyter_code>crd_api = k8s_client.CustomObjectsApi(api_client)

namespace = "default"
job_name = "cifar10-"+ datetime.datetime.now().strftime("%y%m%d-%H%M%S")
job_dir = os.path.join(job_dirs, job_name)
num_steps = 10
body = {}
body['apiVersion'] = TF_JOB_GROUP + "/" + TF_JOB_VERSION
body['kind'] = TF_JOB_KIND
body['metadata'] = {}
body['metadata']['name'] = job_name
body['metadata']['namespace'] = namespace

master_image = images["cpu"]
if use_gpu:
  master_image = images["gpu"]
spec = """
  replicaSpecs:
    - replicas: 1
      tfReplicaType: MASTER
      template:
        spec:
          containers:
            - image: {master_image}
              name: tensorflow
              command:
                - python
                - /tensorflow_models/tutorials/image/cifar10_estimator/cifar10_main.py
                - --data-dir={data_dir}
                - --job-dir={job_dir}
                - --train-steps={num_steps}
                - --num-gpus={num_gpus}
          restartPolicy: OnFailure
  tfImage: {cpu_image}
  tensorBoard:
    logDir: {job_dir}
""".format(master_image=master_image, cpu_image=images["cpu"], data_dir=data_dir, job_dir=job_dir, num_steps=num_steps, num_gpus=accelerator_count)

spec_buffer = StringIO.StringIO(spec)
body['spec'] = yaml.load(spec_buffer)
if use_gpu:
  body['spec']['replicaSpecs'][0]["template"]["spec"]["containers"][0]["resources"] = {
    "limits": {
      "nvidia.com/gpu": accelerator_count,
    }    
  }

try: 
    # Create a Resource
    api_response = crd_api.create_namespaced_custom_object(TF_JOB_GROUP, TF_JOB_VERSION, namespace, TF_JOB_PLURAL, body) 
    logging.info("Created job %s", api_response["metadata"]["name"])
except ApiException as e:
    print(
        "Exception when calling DefaultApi->apis_fqdn_v1_namespaces_namespace_resource_post: %s\n" % 
        e)<jupyter_output><empty_output><jupyter_text>Monitoring your job and waiting for it to finish We can monitor the job a number of ways  * We can poll K8s to get the status of the TFJob  * We can check the TensorFlow logs      * These are available in StackDriver  * We can access TensorBoard if the TFJob was configured to launch TensorBoard  Running the code below will poll K8s for the TFJob status and also print out relevant links for TensorBoard and the StackDriver logsTo access TensorBoard you will need to run **kubectl proxy** to create a proxy connection to your K8s cluster<jupyter_code># Get pod logs
v1 = k8s_client.CoreV1Api(api_client)

k8s_config.load_kube_config()
api_client = k8s_client.ApiClient()
crd_api = k8s_client.CustomObjectsApi(api_client)

master_started = False
runtime_id = None
while True:
  results = crd_api.get_namespaced_custom_object(TF_JOB_GROUP, TF_JOB_VERSION, namespace, TF_JOB_PLURAL, job_name)

  if not runtime_id:
    runtime_id = results["spec"]["RuntimeId"]
    logging.info("Job has runtime id: %s", runtime_id)
    
    tensorboard_url = "http://127.0.0.1:8001/api/v1/proxy/namespaces/{namespace}/services/tensorboard-{runtime_id}:80/".format(
    namespace=namespace, runtime_id=runtime_id)
    logging.info("Tensorboard will be available at job\n %s", tensorboard_url)

  if not master_started:
    # Get the master pod
    # TODO(jlewi): V1LabelSelector doesn't seem to help
    pods = v1.list_namespaced_pod(namespace=namespace, label_selector="runtime_id={0},job_type=MASTER".format(runtime_id))

    # TODO(jlewi): We should probably handle the case where more than 1 pod gets started.
    # TODO(jlewi): Once GKE logs pod labels we can just filter by labels to get all logs for a particular task
    # and not have to identify the actual pod.
    if pods.items:
      pod = pods.items[0]

      logging.info("master pod is %s", pod.metadata.name)
      query={
        'advancedFilter': 'resource.type="container"\nresource.labels.namespace_id="default"\nresource.labels.pod_id="{0}"'.format(pod.metadata.name), 
        'dateRangeStart': pod.metadata.creation_timestamp.isoformat(),
        'expandAll': 'false',
        'interval': 'NO_LIMIT',
        'logName': 'projects/{0}/logs/tensorflow'.format(project),
       'project': project, 
      }
      logging.info("Logs will be available in stackdriver at\n"
                   "https://console.cloud.google.com/logs/viewer?" + urllib.urlencode(query))
      master_started = True

  if results["status"]["phase"] == "Done":
    break
  print("Job status {0}".format(results["status"]["phase"]))
  time.sleep(5)
  
logging.info("Job %s", results["status"]["state"])<jupyter_output><empty_output><jupyter_text>Cleanup* Delete the GKE cluster<jupyter_code>util.delete_cluster(gke, cluster_name, project, zone)<jupyter_output><empty_output><jupyter_text>Appendix<jupyter_code>from kubernetes.client.models.v1_label_selector import V1LabelSelector
import urllib2
# Get pod logs
k8s_config.load_kube_config()
api_client = k8s_client.ApiClient()
v1 = k8s_client.CoreV1Api(api_client)
runtime_id = results["spec"]["RuntimeId"]
# TODO(jlewi): V1LabelSelector doesn't seem to help
pods = v1.list_namespaced_pod(namespace=namespace, label_selector="runtime_id={0},job_type=MASTER".format(runtime_id))

pod = pods.items[0]<jupyter_output><empty_output><jupyter_text>Read the Pod Logs From K8s We can read pod logs directly from K8s and not depend on stackdriver<jupyter_code>ret = v1.read_namespaced_pod_log(namespace=namespace, name=pod.metadata.name)
print(ret)<jupyter_output><empty_output><jupyter_text>Fetch Logs from StackDriver Programmatically  * On GKE pod logs are stored in stackdriver  * These logs will stick around longer than pod logs  * Fetching from stackNote this tends to be a little slow()<jupyter_code>from google.cloud import logging as gcp_logging
pod_filter = 'resource.type="container" AND resource.labels.pod_id="master-hrhh-0-wrh6g"'
client = gcp_logging.Client(project=project)

for entry in client.list_entries(filter_=pod_filter):
  print(entry.payload.strip())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10828673106253177, 'fraction_numerical': 0.013281647178444331, 'mean_word_length': 4.689443239334779, 'pattern_counts': {'":': 15, '<': 61, '<?xml version=': 0, '>': 69, 'https://': 15, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '38667', 'license': 'Apache-2.0', 'n_tokens_mistral': 5234, 'n_tokens_neox': 5064, 'n_words': 1359, 'path': 'examples/gke/TF on GKE.ipynb', 'repo_name': 'ChanYiLin/tf-operator-extend', 'size': '15736'}