format

.cloud-build/execute_changed_notebooks_helper.py

@@ -24,6 +24,7 @@ import re
 import subprocess
 from typing import List, Optional
 
+import execute_notebook_helper
 import execute_notebook_remote
 import nbformat
 from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
@@ -287,14 +288,15 @@ def process_and_execute_notebooks(
         timeout (str):
             Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
     """
-    notebook_execution_results: List[NotebookExecutionResult] = []
 
     # Calculate deadline
     deadline = datetime.datetime.now() + datetime.timedelta(
         seconds=max(timeout - WORKER_TIMEOUT_BUFFER_IN_SECONDS, 0)
     )
 
-    if len(notebooks) > 0:
+    if len(notebooks) > 1:
+        notebook_execution_results: List[NotebookExecutionResult] = []
+
         print(f"Found {len(notebooks)} modified notebooks: {notebooks}")
 
         if should_parallelize and len(notebooks) > 1:
@@ -333,43 +335,56 @@ def process_and_execute_notebooks(
                 )
                 for notebook in notebooks
             ]
+
+        print("\n=== RESULTS ===\n")
+
+        results_sorted = sorted(
+            notebook_execution_results,
+            key=lambda result: result.is_pass,
+            reverse=True,
+        )
+
+        # Print results
+        print(
+            tabulate(
+                [
+                    [
+                        result.name,
+                        "PASSED" if result.is_pass else "FAILED",
+                        format_timedelta(result.duration),
+                        result.log_url,
+                        result.output_uri,
+                    ]
+                    for result in results_sorted
+                ],
+                headers=["build_tag", "status", "duration", "log_url", "output_url"],
+            )
+        )
+
+        print("\n=== END RESULTS===\n")
+
+        total_notebook_duration = functools.reduce(
+            operator.add,
+            [datetime.timedelta(seconds=0)]
+            + [result.duration for result in results_sorted],
+        )
+
+        print(
+            f"Cumulative notebook duration: {format_timedelta(total_notebook_duration)}"
+        )
+
+        # Raise error if any notebooks failed
+        if not all([result.is_pass for result in results_sorted]):
+            raise RuntimeError("Notebook failures detected. See logs for details")
+
+    elif len(notebooks) == 1:
+        notebook = notebooks[0]
+        execute_notebook_helper.execute_notebook(
+            notebook_source=notebook,
+            output_file_or_uri="/".join(
+                [artifacts_bucket, pathlib.Path(notebook).name]
+            ),
+            should_log_output=True,
+        )
     else:
         print("No notebooks modified in this pull request.")
-
-    print("\n=== RESULTS ===\n")
-
-    results_sorted = sorted(
-        notebook_execution_results,
-        key=lambda result: result.is_pass,
-        reverse=True,
-    )
-
-    # Print results
-    print(
-        tabulate(
-            [
-                [
-                    result.name,
-                    "PASSED" if result.is_pass else "FAILED",
-                    format_timedelta(result.duration),
-                    result.log_url,
-                ]
-                for result in results_sorted
-            ],
-            headers=["build_tag", "status", "duration", "log_url"],
-        )
-    )
-
-    print("\n=== END RESULTS===\n")
-
-    total_notebook_duration = functools.reduce(
-        operator.add,
-        [datetime.timedelta(seconds=0)]
-        + [result.duration for result in results_sorted],
-    )
-
-    print(f"Cumulative notebook duration: {format_timedelta(total_notebook_duration)}")
-
-    # Raise error if any notebooks failed
-    if not all([result.is_pass for result in results_sorted]):
-        raise RuntimeError("Notebook failures detected. See logs for details")

.cloud-build/test_folders.txt

@@ -1,3 +1,2 @@
 notebooks/official
 notebooks/notebook_template.ipynb
-notebooks/community/ml_ops

notebooks/official/CODEOWNERS

@@ -12,6 +12,7 @@
 /migration @aferlitsch
 /explainabl_ai @aferlitsch
 /pipelines @aferlitsch
+/tensorboard @zbl94
 
 /model_monitoring/model_monitoring.ipynb @mco-gh
 /ml_metadata/sdk-metric-parameter-tracking-for-custom-jobs.ipynb @jialuzh

notebooks/official/tensorboard/vertex_tensorboard_custom_training_with_prebuilt_container.ipynb

@@ -0,0 +1,871 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JICNqM6hx1XC"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2022 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eBEO2w9My9py"
+      },
+      "source": [
+        "# Vertex AI TensorBoard with Custom Training using Prebuilt Containers\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/ai-platform/notebooks/deploy-notebook?name=Model%20Monitoring&download_url=https%3A%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmaster%2Fnotebooks%2Fcommunity%2Ftensorboard%2Fvertex_tensorboard_custom_training_with_prebuilt_container.ipynb\">\n",
+        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Google Cloud Notebooks\">Open in Cloud Notebook\n",
+        "    </a>\n",
+        "  </td> \n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/tensorboard/vertex_tensorboard_custom_training_with_prebuilt_container.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/tensorboard/vertex_tensorboard_custom_training_with_prebuilt_container.ipynb\">\n",
+        "        <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GRbmi2_tlzAv"
+      },
+      "source": [
+        "## Overview"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VL7XCFV7yCBU"
+      },
+      "source": [
+        "### What is Vertex AI TensorBoard\n",
+        "\n",
+        "[Open source TensorBoard](https://www.tensorflow.org/tensorboard/get_started)\n",
+        "(TB) is a Google open source project for machine learning experiment\n",
+        "visualization. Vertex AI TensorBoard is an enterprise-ready managed\n",
+        "version of TensorBoard.\n",
+        "\n",
+        "Vertex AI TensorBoard provides various detailed visualizations, that\n",
+        "includes:\n",
+        "\n",
+        "*   Tracking and visualizing metrics such as loss and accuracy over time\n",
+        "*   Visualizing model computational graphs (ops and layers)\n",
+        "*   Viewing histograms of weights, biases, or other tensors as they change over time\n",
+        "*   Projecting embeddings to a lower dimensional space\n",
+        "*   Displaying image, text, and audio samples\n",
+        "\n",
+        "In addition to the powerful visualizations from\n",
+        "TensorBoard, Vertex AI TensorBoard provides:\n",
+        "\n",
+        "*  A persistent, shareable link to your experiment's dashboard\n",
+        "\n",
+        "*  A searchable list of all experiments in a project\n",
+        "\n",
+        "*  Tight integrations with Vertex AI services for model training\n",
+        "\n",
+        "*  Enterprise-grade security, privacy, and compliance\n",
+        "\n",
+        "With Vertex AI TensorBoard, you can track, visualize, and compare\n",
+        "ML experiments and share them with your team."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fbF2lF8rlp3I"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "Dataset used in this tutorial will be the [flower dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) provided by TensorFlow. No other datasets required.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UE8vLw7SlpwE"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn how to create a custom training job using prebuilt containers, and monitor your training process on Vertex AI TensorBoard in near real time.\n",
+        "\n",
+        "The steps performed include:\n",
+        "\n",
+        "* Setup service account and Google Cloud Storage buckets.\n",
+        "* Write your customized training code.\n",
+        "* Package and upload your training code to Google Cloud Storage.\n",
+        "* Create & launch your custom training job with Tensorboard enabled for near real time monitorning."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "yajq2millpnu"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "lAcDYZfslpeF"
+      },
+      "source": [
+        "### Set up your local development environment\n",
+        "\n",
+        "**If you are using Colab or Vertex AI Workbench**, your environment already meets all the requirements to run this notebook. You can skip this step.\n",
+        "\n",
+        "Otherwise, make sure your environment meets this notebook's requirements. You need the following:\n",
+        "\n",
+        "- The Cloud Storage SDK\n",
+        "- Git\n",
+        "- Python 3\n",
+        "- virtualenv\n",
+        "- Jupyter notebook running in a virtual environment with Python 3\n",
+        "\n",
+        "The Cloud Storage guide to [Setting up a Python development environment](https://cloud.google.com/python/setup) and the [Jupyter installation guide](https://jupyter.org/install) provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:\n",
+        "\n",
+        "1. [Install and initialize the SDK](https://cloud.google.com/sdk/docs/).\n",
+        "\n",
+        "2. [Install Python 3](https://cloud.google.com/python/setup#installing_python).\n",
+        "\n",
+        "3. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3.  Activate the virtual environment.\n",
+        "\n",
+        "4. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.\n",
+        "\n",
+        "5. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
+        "\n",
+        "6. Open this notebook in the Jupyter Notebook Dashboard."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "s3moH5AexXpk"
+      },
+      "source": [
+        "### Install additional packages\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pJwxNvcA8V_c"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "! pip3 install google-cloud-aiplatform {USER_FLAG} -q\n",
+        "\n",
+        "# Automatically restart kernel after installs\n",
+        "\n",
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cczxMMkYK9a4"
+      },
+      "source": [
+        "## Before You Begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "64BZ1jL5GEi0"
+      },
+      "source": [
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
+        "\n",
+        "3. [Enable the Vertex AI APIs and Cloud Storage.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,storage-component.googleapis.com)\n",
+        "\n",
+        "4. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "\n",
+        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "Note: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$`."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UIl_kn1pGH_T"
+      },
+      "source": [
+        "### Set your project ID"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "75C2px4XtS5l"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6yPUOueCF3pI"
+      },
+      "source": [
+        "### Set your region"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gzEibeHGF1jb"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jetPysDMtV1-"
+      },
+      "source": [
+        "### Login to your Google Cloud account"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IK8eOR8tt6TI"
+      },
+      "outputs": [],
+      "source": [
+        "# The Google Cloud Notebook product has specific requirements\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
+        "\n",
+        "# If on Google Cloud Notebooks, then don't execute this code\n",
+        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        from google.colab import auth as google_auth\n",
+        "\n",
+        "        google_auth.authenticate_user()\n",
+        "\n",
+        "    # If you are running this notebook locally, replace the string below with the\n",
+        "    # path to your service account key and run this cell to authenticate your GCP\n",
+        "    # account.\n",
+        "    elif not os.getenv(\"IS_TESTING\"):\n",
+        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KIPcg_Xhwvsn"
+      },
+      "source": [
+        "### Import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "p4w8c1pHw2Yt"
+      },
+      "outputs": [],
+      "source": [
+        "import google.cloud.aiplatform as aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hwXxa4Qgnh4Y"
+      },
+      "source": [
+        "## Setup Service Account and Permissions"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7qXFUiHLoFRw"
+      },
+      "source": [
+        "Create the service account (if you have not already done it) and grant permissions for AI Platform and Cloud Storage."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "mpKjfsXumuNV"
+      },
+      "outputs": [],
+      "source": [
+        "USER_SA_NAME = \"your-serivce-account-name\"  # @param {type:\"string\"}\n",
+        "SA_EMAIL = \"{}@{}.iam.gserviceaccount.com\".format(USER_SA_NAME, PROJECT_ID)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "RdWLe8S_tlee"
+      },
+      "outputs": [],
+      "source": [
+        "# Create service account.\n",
+        "! gcloud --project={PROJECT_ID} iam service-accounts create {USER_SA_NAME}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "UlDhuciOt5vo"
+      },
+      "outputs": [],
+      "source": [
+        "# Grant Cloud Storage permission.\n",
+        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
+        "   --member=serviceAccount:$SA_EMAIL \\\n",
+        "   --role=roles/storage.admin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "lTKVB71soRyr"
+      },
+      "outputs": [],
+      "source": [
+        "# Grant AI Platform permission.\n",
+        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
+        "   --member=serviceAccount:$SA_EMAIL \\\n",
+        "   --role=roles/aiplatform.user"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BpOvmA4VugWu"
+      },
+      "source": [
+        "## Create Cloud Storage Buckets.\n",
+        "Two Cloud Storage buckets will be used to a) store your training code distribution (details below), and b) the outputs (including Tensorboard logs) your training code generates. The buckets must be regional that is, not multi-region or dual-region, and the following resources must be in same region:\n",
+        "\n",
+        "* the Cloud Storage bucket\n",
+        "* the Vertex AI training job\n",
+        "* the Vertex AI TensorBoard instance\n",
+        "\n",
+        "These created buckets will be deleted in the cleaning up section in the end. "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Dt3fkuz1uw6e"
+      },
+      "outputs": [],
+      "source": [
+        "GCS_BUCKET_OUTPUT = \"{}-tensorboard-out-{}\".format(PROJECT_ID, REGION)\n",
+        "GCS_BUCKET_TRAINING = \"{}-tensorboard-train-{}\".format(PROJECT_ID, REGION)\n",
+        "\n",
+        "! gsutil mb -p $PROJECT_ID -l $REGION gs://$GCS_BUCKET_OUTPUT\n",
+        "! gsutil mb -p $PROJECT_ID -l $REGION gs://$GCS_BUCKET_TRAINING"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dR2mOCllvlqN"
+      },
+      "source": [
+        "## Write Your Training Code\n",
+        "Your training code must be configured to write TensorBoard logs to the Cloud Storage bucket, the location of which the Vertex AI Training Service will automatically make available via a predefined environment variable `AIP_TENSORBOARD_LOG_DIR`.\n",
+        "\n",
+        "This can usually be done by providing `os.environ['AIP_TENSORBOARD_LOG_DIR']` as the log directory to the open source TensorBoard log writing APIs.\n",
+        "\n",
+        "For example, in TensorFlow 2.x, you can use following code to create a `tensorboard_callback`:\n",
+        "```\n",
+        "tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
+        "  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],\n",
+        "  histogram_freq=1)\n",
+        "```\n",
+        "\n",
+        "`AIP_TENSORBOARD_LOG_DIR` will be in the `BASE_OUTPUT_DIR` that you provided below when creating the custom training job.\n",
+        "\n",
+        "We will use the following sample code as an example:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "balswZevBQt_"
+      },
+      "outputs": [],
+      "source": [
+        "# Download the sample code\n",
+        "! gsutil cp gs://cloud-samples-data/ai-platform/hello-custom/hello-custom-sample-v1.tar.gz - | tar -xzv\n",
+        "%cd hello-custom-sample/"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JGEJkNvHBmc8"
+      },
+      "outputs": [],
+      "source": [
+        "# The training code we want to edit is:\n",
+        "! cat trainer/task.py"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vjdLIqYyDZFS"
+      },
+      "source": [
+        "In `trainer/task.py`, create the `tensorboard_callback` and add the callback to `model.fit(...)`\n",
+        "\n",
+        "Sample code: \n",
+        "```\n",
+        "# previous things\n",
+        "model.compile(...)\n",
+        "\n",
+        "tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
+        "  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],\n",
+        "  histogram_freq=1)\n",
+        "  \n",
+        "model.fit(dataset, epochs=10, callbacks=[tensorboard_callback])\n",
+        "```\n",
+        "\n",
+        "Update your `trainer/task.py` with `tensorboard_callback`. You can use the following sample code."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HvI4Y5edurEp"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile trainer/task.py\n",
+        "\n",
+        "import logging\n",
+        "import os\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_datasets as tfds\n",
+        "\n",
+        "IMG_WIDTH = 128\n",
+        "\n",
+        "\n",
+        "def normalize_img(image):\n",
+        "    \"\"\"Normalizes image.\n",
+        "\n",
+        "    * Resizes image to IMG_WIDTH x IMG_WIDTH pixels\n",
+        "    * Casts values from `uint8` to `float32`\n",
+        "    * Scales values from [0, 255] to [0, 1]\n",
+        "\n",
+        "    Returns:\n",
+        "      A tensor with shape (IMG_WIDTH, IMG_WIDTH, 3). (3 color channels)\n",
+        "    \"\"\"\n",
+        "    image = tf.image.resize_with_pad(image, IMG_WIDTH, IMG_WIDTH)\n",
+        "    return image / 255.\n",
+        "\n",
+        "\n",
+        "def normalize_img_and_label(image, label):\n",
+        "    \"\"\"Normalizes image and label.\n",
+        "\n",
+        "    * Performs normalize_img on image\n",
+        "    * Passes through label unchanged\n",
+        "\n",
+        "    Returns:\n",
+        "      Tuple (image, label) where\n",
+        "      * image is a tensor with shape (IMG_WIDTH, IMG_WIDTH, 3). (3 color\n",
+        "        channels)\n",
+        "      * label is an unchanged integer [0, 4] representing flower type\n",
+        "    \"\"\"\n",
+        "    return normalize_img(image), label\n",
+        "\n",
+        "\n",
+        "if 'AIP_MODEL_DIR' not in os.environ:\n",
+        "    raise KeyError(\n",
+        "        'The `AIP_MODEL_DIR` environment variable has not been' +\n",
+        "        'set. See https://cloud.google.com/ai-platform-unified/docs/tutorials/image-recognition-custom/training'\n",
+        "    )\n",
+        "output_directory = os.environ['AIP_MODEL_DIR']\n",
+        "\n",
+        "logging.info('Loading and preprocessing data ...')\n",
+        "dataset = tfds.load('tf_flowers:3.*.*',\n",
+        "                    split='train',\n",
+        "                    try_gcs=True,\n",
+        "                    shuffle_files=True,\n",
+        "                    as_supervised=True)\n",
+        "dataset = dataset.map(normalize_img_and_label,\n",
+        "                      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
+        "dataset = dataset.cache()\n",
+        "dataset = dataset.shuffle(1000)\n",
+        "dataset = dataset.batch(128)\n",
+        "dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\n",
+        "\n",
+        "logging.info('Creating and training model ...')\n",
+        "model = tf.keras.Sequential([\n",
+        "    tf.keras.layers.Conv2D(16,\n",
+        "                           3,\n",
+        "                           padding='same',\n",
+        "                           activation='relu',\n",
+        "                           input_shape=(IMG_WIDTH, IMG_WIDTH, 3)),\n",
+        "    tf.keras.layers.MaxPooling2D(),\n",
+        "    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),\n",
+        "    tf.keras.layers.MaxPooling2D(),\n",
+        "    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),\n",
+        "    tf.keras.layers.MaxPooling2D(),\n",
+        "    tf.keras.layers.Flatten(),\n",
+        "    tf.keras.layers.Dense(512, activation=\"relu\"),\n",
+        "    tf.keras.layers.Dense(5)  # 5 classes\n",
+        "])\n",
+        "model.compile(\n",
+        "    optimizer='adam',\n",
+        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "    metrics=['accuracy'])\n",
+        "\n",
+        "### Create a tensorboard call back and write to the gcs path provided by AIP_TENSORBOARD_LOG_DIR\n",
+        "tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
+        "  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],\n",
+        "  histogram_freq=1)\n",
+        "\n",
+        "### Train the model with tensorboard_callback\n",
+        "model.fit(dataset, epochs=14, callbacks=[tensorboard_callback])\n",
+        "\n",
+        "logging.info(f'Exporting SavedModel to: {output_directory}')\n",
+        "# Add softmax layer for intepretability\n",
+        "probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])\n",
+        "probability_model.save(output_directory)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "L067Jw_QFcZ3"
+      },
+      "source": [
+        "## Upload your training code to Cloud Storage\n",
+        "\n",
+        "You must package the training code as a source distribution and upload it to Cloud Storage in order for Vertex AI to run the code in a custom training pipeline.\n",
+        "\n",
+        "1. Run the following commands to create a source distribution in the gzipped tarball format. The command uses the `setup.py` file included in the sample code."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bNEr950xGbGT"
+      },
+      "outputs": [],
+      "source": [
+        "! python3 setup.py sdist --formats=gztar"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a7Gt9WPFG9V7"
+      },
+      "source": [
+        "2. Run the following command to upload the source distribution you just created, `dist/hello-custom-training-3.0.tar.gz`, to the Cloud Storage bucket that you created before."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wTv1gReWJqq0"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil cp dist/hello-custom-training-3.0.tar.gz gs://$GCS_BUCKET_TRAINING"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IaQjIPvuKLwW"
+      },
+      "source": [
+        "## Create a Custom Training Job"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IbN8D3BUWgK3"
+      },
+      "source": [
+        "Setup the endpoint we will talk to."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4C-yG7K3We4k"
+      },
+      "outputs": [],
+      "source": [
+        "ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "svUGBOow_Obj"
+      },
+      "source": [
+        "If no existing Tensorboard instance for this project and region, create one."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "OAe1xJeS_X3F"
+      },
+      "outputs": [],
+      "source": [
+        "TENSORBOARD_DISPLAY_NAME = \"your-tensorboard-display-name\"  # @param {type:\"string\"}\n",
+        "\n",
+        "tensorboard = aiplatform.Tensorboard.create(\n",
+        "    display_name=TENSORBOARD_DISPLAY_NAME, project=PROJECT_ID, location=REGION\n",
+        ")\n",
+        "tensorboard_resource_name = tensorboard.gca_resource.name\n",
+        "print(\"TensorBoard resource name:\", tensorboard_resource_name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nnOnCPPEDNUw"
+      },
+      "source": [
+        "If you already have a Tensorboard for this `PROJECT_ID` and `REGION`, you can get your `Tensorboard_ID` either from Google Cloud Console, Vertex AI > Experiments > Tensorboard Instance, or from the command below:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "k2wqpkEEyDCD"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.Tensorboard.list(project=PROJECT_ID, location=REGION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kbnRdlEIV4rV"
+      },
+      "source": [
+        "Prepare your `Tensorboard_ID` and `TENSORBOARD_INSTANCE_NAME`.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "KFQTI2z2CX1T"
+      },
+      "outputs": [],
+      "source": [
+        "TENSORBOARD_ID = \"255807977271853056\"  # @param {type:\"string\"}\n",
+        "TENSORBOARD_INSTANCE_NAME = \"projects/{}/locations/{}/tensorboards/{}\".format(\n",
+        "    PROJECT_ID, REGION, TENSORBOARD_ID\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mudxBDal_a_k"
+      },
+      "source": [
+        "Run following example request to create your own custom training job and stream the training results to Tensorboard."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "noqBo6p_1y7E"
+      },
+      "outputs": [],
+      "source": [
+        "from datetime import datetime\n",
+        "\n",
+        "INVOCATION_TIMESTAMP = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "JOB_NAME = \"tensorboard-example-job-{}\".format(INVOCATION_TIMESTAMP)\n",
+        "BASE_OUTPUT_DIR = \"gs://{}/{}\".format(GCS_BUCKET_OUTPUT, JOB_NAME)\n",
+        "\n",
+        "# The AI Platform services require regional API endpoints.\n",
+        "client_options = {\"api_endpoint\": f\"{REGION}-aiplatform.googleapis.com\"}\n",
+        "# Initialize client that will be used to create and send requests.\n",
+        "# This client only needs to be created once, and can be reused for multiple requests.\n",
+        "client = aiplatform.gapic.JobServiceClient(client_options=client_options)\n",
+        "custom_job = {\n",
+        "    \"display_name\": JOB_NAME,\n",
+        "    \"job_spec\": {\n",
+        "        \"worker_pool_specs\": [\n",
+        "            {\n",
+        "                \"machine_spec\": {\n",
+        "                    \"machine_type\": \"n1-standard-8\",\n",
+        "                },\n",
+        "                \"replica_count\": 1,\n",
+        "                \"python_package_spec\": {\n",
+        "                    \"executor_image_uri\": \"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-3:latest\",\n",
+        "                    \"package_uris\": [\n",
+        "                        f\"gs://{GCS_BUCKET_TRAINING}/hello-custom-training-3.0.tar.gz\"\n",
+        "                    ],\n",
+        "                    \"python_module\": \"trainer.task\",\n",
+        "                },\n",
+        "            }\n",
+        "        ],\n",
+        "        \"service_account\": SA_EMAIL,\n",
+        "        \"tensorboard\": TENSORBOARD_INSTANCE_NAME,\n",
+        "        \"base_output_directory\": {\"output_uri_prefix\": BASE_OUTPUT_DIR},\n",
+        "    },\n",
+        "}\n",
+        "parent = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
+        "response = client.create_custom_job(parent=parent, custom_job=custom_job)\n",
+        "print(\"response:\", response)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VfMsn_RnEtnj"
+      },
+      "source": [
+        "In Google Cloud Console, you can monitor your training job at Vertex AI > Training > Custom Jobs. In each custom training job, near real time updated TensorBoard is available at `OPEN TENSORBOARD` button."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TFEriiywMZga"
+      },
+      "source": [
+        "## Cleaning up\n",
+        "\n",
+        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "O1Je2JZTMoMg"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil -m rm -r gs://$GCS_BUCKET_OUTPUT\n",
+        "! gsutil -m rm -r gs://$GCS_BUCKET_TRAINING"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "vertex_tensorboard_custom_training_with_prebuilt_container.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}