notebooks/official/custom/custom_training_tensorboard_profiler.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "l2mMvIUG9meX"
      },
      "source": [
        "# Profile model training performance using Cloud Profiler\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_tensorboard_profiler.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",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/custom_training_tensorboard_profiler.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Cloud Profiler lets you monitor and optimize your model training performance by helping you understand the resource consumption of training operations. This tutorial demonstrates how to enable Cloud Profiler so you can debug model training performance for your custom training jobs.\n",
        "\n",
        "Learn more about [Cloud Profiler](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-profiler)."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dmfmQL6w84pS"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to enable Cloud Profiler for custom training jobs.\n",
        "\n",
        "This tutorial uses the following Google Cloud AI services:\n",
        "\n",
        "- `Vertex AI Training`\n",
        "- `Vertex AI TensorBoard`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Setup a service account and a Cloud Storage bucket\n",
        "- Create a TensorBoard instance\n",
        "- Create and run a custom training job\n",
        "- View the Cloud Profiler dashboard\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zfXf0r-K81Y-"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [mnist dataset](https://www.tensorflow.org/datasets/catalog/mnist) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview).\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "I3KFLvpq87rs"
      },
      "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\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ze4-nDLfK4pw"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade google-cloud-aiplatform --quiet"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "aUw6ibN-n5Za"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FM12wbWhn7w0"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "LgFWLeJfoGQu"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### 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 API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "8ckyxpX_oSzD"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zY8DKBoVoVy3"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "mSQjVQmMosMl"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Se9FWWhLotvB"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "IfJRIMBpo5Pg"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "acFN0s3So9-Y"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dQ_mNwuapE5T"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cR_MzpknpGgM"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "h-MuVI_ypJfw"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BeaQlCwMpQUT"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3ivZkPUjpaFz"
      },
      "source": [
        "**4. Setup service account and permissions**\n",
        "\n",
        "A service account will be used to create custom training jobs. If you do not want to use your project's Compute Engine service account, set SERVICE_ACCOUNT to another service account ID. You can create a service account by following the [instructions](https://cloud.google.com/iam/docs/creating-managing-service-accounts#creating)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vYE3b942wza4"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WWIxsCJFCg5Z"
      },
      "outputs": [],
      "source": [
        "# Grant Cloud Storage permission.\n",
        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
        "        --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
        "        --role=\"roles/storage.admin\" \\\n",
        "        --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S_8_5jm-Gk6w"
      },
      "outputs": [],
      "source": [
        "# Grant AI Platform permission.\n",
        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
        "        --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
        "        --role=\"roles/aiplatform.user\" \\\n",
        "        --quiet"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OKtKGmr9pfr6"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "In3aQanwYjFB"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "GOaOsIjxp0oB"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wn5QiIl2p16e"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ankcS-vtp7Wv"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WffSImMvp-Po"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OMrAJ8RGqBQu"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AWRzBFExqERG"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MYzX9ZgjsyyY"
      },
      "source": [
        "### Enable Artifact Registry API\n",
        "\n",
        "First, you must enable the Artifact Registry API service for your project.\n",
        "\n",
        "Learn more about [Enabling service](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bG4Rpt-As55f"
      },
      "outputs": [],
      "source": [
        "! gcloud services enable artifactregistry.googleapis.com --quiet"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-ayTbNdi62_t"
      },
      "source": [
        "### Create a TensorBoard instance\n",
        "\n",
        "A Vertex AI TensorBoard instance, which is a regionalized resource storing your Vertex AI TensorBoard experiments, must be created before the experiments can be visualized. You can create multiple instances in a project. You can use command  `gcloud ai tensorboards list` to get a list of your existing TensorBoard instances."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "9c3QrDTZdaxk"
      },
      "source": [
        "#### Set your TensorBoard instance display name\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "azlwb__AX8gs"
      },
      "outputs": [],
      "source": [
        "TENSORBOARD_NAME = \"your-tensorboard-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "vJrWKK0mY7H7"
      },
      "source": [
        "#### Create a TensorBoard instance\n",
        "\n",
        "If you don't have a TensorBoard instance, create one by running the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JqVNsRFrc_78"
      },
      "outputs": [],
      "source": [
        "tensorboard = aiplatform.Tensorboard.create(\n",
        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
        ")\n",
        "\n",
        "TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name\n",
        "\n",
        "print(\"TensorBoard instance name:\", TENSORBOARD_INSTANCE_NAME)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "yoR29gW2S24w"
      },
      "source": [
        "## Train a model\n",
        "\n",
        "To train a model using your custom training code, choose one of the following options:\n",
        "\n",
        "- **Prebuilt container**: Load your custom training code as a Python package to a prebuilt container image from Google Cloud.\n",
        "\n",
        "- **Custom container**: Create your own container image that contains your custom training code.\n",
        "\n",
        "In this tutorial, we will train a custom model using a custom container."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OrpUIkAIs_uQ"
      },
      "source": [
        "### Create a private Docker repository\n",
        "\n",
        "Your first step is to create your own Docker repository in Google Artifact Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0amu4063tDnG"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "DOCKER_REPOSITORY = f\"{PROJECT_ID}-repo-unique\"\n",
        "\n",
        "! gcloud services enable artifactregistry.googleapis.com\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
        "    ! gcloud components update --quiet\n",
        "\n",
        "! gcloud artifacts repositories create {DOCKER_REPOSITORY} \\\n",
        "    --repository-format=docker \\\n",
        "    --location={REGION} \\\n",
        "    --description=\"Repository for TensorBoard Custom Training Job\" \\\n",
        "    --quiet\n",
        "\n",
        "! gcloud artifacts repositories list"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "MPqbvhO6tK_e"
      },
      "source": [
        "### Configure authentication to your private Docker repository\n",
        "\n",
        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to `Artifact Registry` for your region."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yAeCXZiStPCX"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "\n",
        "if not IS_COLAB:\n",
        "    ! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "pSXrzCWi-bTD"
      },
      "source": [
        "### Create a custom container image and push to your private Docker repository\n",
        "\n",
        "First, you create a training script file and a docker file.\n",
        "\n",
        "Create a directory for all of your training code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZYO55-ZK-lm2"
      },
      "outputs": [],
      "source": [
        "PYTHON_PACKAGE_APPLICATION_DIR = \"trainer\"\n",
        "\n",
        "!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "hyAwgsoQmaYI"
      },
      "source": [
        "#### Prepare the training script\n",
        "\n",
        "Your training code must be configured to write TensorBoard logs to a Cloud Storage bucket, the location of which Vertex AI Training automatically makes available through 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",
        "`AIP_TENSORBOARD_LOG_DIR` is in the `BASE_OUTPUT_DIR` that you provide when creating the custom training job.\n",
        "\n",
        "To enable Cloud Profiler for your training job, add the following to your training script:\n",
        "\n",
        "Add the cloud_profiler import at your top level imports:\n",
        "\n",
        "    from google.cloud.aiplatform.training_utils import cloud_profiler\n",
        "\n",
        "\n",
        "Initialize the cloud_profiler plugin by adding:\n",
        "\n",
        "\n",
        "    cloud_profiler.init()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8JCgWW7Au1w8"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/task.py\n",
        "\n",
        "import tensorflow as tf\n",
        "import argparse\n",
        "import os\n",
        "import sys, traceback\n",
        "from google.cloud.aiplatform.training_utils import cloud_profiler\n",
        "\n",
        "\"\"\"Train an mnist model and use cloud_profiler for profiling.\"\"\"\n",
        "\n",
        "def _create_model():\n",
        "    model = tf.keras.models.Sequential(\n",
        "        [\n",
        "            tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
        "            tf.keras.layers.Dense(128, activation=\"relu\"),\n",
        "            tf.keras.layers.Dropout(0.2),\n",
        "            tf.keras.layers.Dense(10),\n",
        "        ]\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "def main(args):\n",
        "    print('Loading and preprocessing data ...')\n",
        "    mnist = tf.keras.datasets.mnist\n",
        "\n",
        "    (x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
        "    x_train, x_test = x_train / 255.0, x_test / 255.0\n",
        "\n",
        "    print('Creating and training model ...')\n",
        "\n",
        "    model = _create_model()\n",
        "    model.compile(\n",
        "      optimizer=\"adam\",\n",
        "      loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
        "      metrics=[\"accuracy\"],\n",
        "    )\n",
        "\n",
        "    # Initialize the profiler.\n",
        "    print('Initialize the profiler ...')\n",
        "        \n",
        "    try:\n",
        "        cloud_profiler.init()\n",
        "    except:\n",
        "        ex_type, ex_value, ex_traceback = sys.exc_info()\n",
        "        print(\"*** Unexpected:\", ex_type.__name__, ex_value)\n",
        "        traceback.print_tb(ex_traceback, limit=10, file=sys.stdout)\n",
        "    \n",
        "    print('The profiler initiated.')\n",
        "\n",
        "    log_dir = \"logs\"\n",
        "    if 'AIP_TENSORBOARD_LOG_DIR' in os.environ:\n",
        "      log_dir = os.environ['AIP_TENSORBOARD_LOG_DIR']\n",
        "\n",
        "    print('Setting up the TensorBoard callback ...')\n",
        "    tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
        "        log_dir=log_dir,\n",
        "        histogram_freq=1)\n",
        "\n",
        "    print('Training model ...')\n",
        "    model.fit(\n",
        "        x_train,\n",
        "        y_train,\n",
        "        epochs=args.epochs,\n",
        "        verbose=0,\n",
        "        callbacks=[tensorboard_callback],\n",
        "    )\n",
        "    print('Training completed.')\n",
        "\n",
        "    print('Saving model ...')\n",
        "\n",
        "    model_dir = \"model\"\n",
        "    if 'AIP_MODEL_DIR' in os.environ:\n",
        "      model_dir = os.environ['AIP_MODEL_DIR']\n",
        "    tf.saved_model.save(model, model_dir)\n",
        "\n",
        "    print('Model saved at ' + model_dir)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument(\n",
        "        \"--epochs\", type=int, default=100, help=\"Number of epochs to run model.\"\n",
        "    )\n",
        "    \n",
        "    args = parser.parse_args()\n",
        "    main(args)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "LiOEvgkA_90f"
      },
      "source": [
        "#### Prepare the Dockerfile\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aKcS7yzcAXWf"
      },
      "outputs": [],
      "source": [
        "%%writefile Dockerfile\n",
        "# Specifies base image and tag\n",
        "FROM us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-9:latest\n",
        "WORKDIR /root\n",
        "\n",
        "# Installs additional packages as you need.\n",
        "RUN pip3 install google-cloud-aiplatform[cloud_profiler]\n",
        "RUN pip3 install protobuf==3.20.2\n",
        "\n",
        "# Copies the trainer code to the docker image.\n",
        "RUN mkdir /root/trainer\n",
        "COPY trainer/task.py /root/trainer/task.py\n",
        "\n",
        "# Sets up the entry point to invoke the trainer.\n",
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ihYFahRAr6sj"
      },
      "source": [
        "#### Build a custom container image and push to your private Docker repository"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pDs1qDLDA27F"
      },
      "outputs": [],
      "source": [
        "IMAGE_NAME = \"tensorboard-custom-container\"\n",
        "IMAGE_URI = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
        "\n",
        "! gcloud builds submit --project {PROJECT_ID} --region={REGION} --tag {IMAGE_URI} --timeout=60m --quiet"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "k4e6OYmimqTR"
      },
      "source": [
        "### Create and run the custom training job\n",
        "\n",
        "Configure a [custom job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the custom container image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jd6P-lQhFBPg"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = \"tensorboard-job-unique\"\n",
        "\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=JOB_NAME, container_uri=IMAGE_URI\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "51hKGTbU32Eg"
      },
      "source": [
        "#### Run the custom training job\n",
        "\n",
        "Next, you run the custom job to start the training job by invoking the method `run`, with the following parameters:\n",
        "\n",
        "- `args`: The command-line arguments to pass to the training script.\n",
        "   - `--epochs` : The number of epochs for training.\n",
        "- `replica_count`: The number of compute instances for training (replica_count = 1 is single node training).\n",
        "- `machine_type`: The machine type for the compute instances.\n",
        "- `tensorboard`: The TensorBoard instance.\n",
        "- `service_account`: The service account.\n",
        "- `sync`: Whether to block until completion of the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OC6_4KeI4Fit"
      },
      "outputs": [],
      "source": [
        "base_output_dir = \"{}/{}\".format(BUCKET_URI, JOB_NAME)\n",
        "MACHINE_TYPE = \"n1-standard-4\"\n",
        "EPOCHS = 2\n",
        "training_args = [\n",
        "    \"--epochs=\" + str(EPOCHS),\n",
        "]\n",
        "\n",
        "job.run(\n",
        "    args=training_args,\n",
        "    replica_count=1,\n",
        "    machine_type=MACHINE_TYPE,\n",
        "    base_output_dir=base_output_dir,\n",
        "    tensorboard=TENSORBOARD_INSTANCE_NAME,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "JkEe2Nb_85UD"
      },
      "source": [
        "## View the Cloud Profiler dashboard\n",
        "\n",
        "When the custom job state switches to `Running`, you can access the Cloud Profiler dashboard through the Custom jobs page or the Experiments page on the Google Cloud console. \n",
        "\n",
        "The Google Cloud guide to [Profile model training performance using Cloud Profiler](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-profiler) provides detailed instructions for accessing the Cloud Profiler dashboard and capturing a profiling session. \n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "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:\n",
        "\n",
        "- Docker repository\n",
        "- Training job\n",
        "- TensorBoard instance\n",
        "- Cloud Storage bucket\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WR-ZhQ9XwpRI"
      },
      "outputs": [],
      "source": [
        "delete_tensorboard = True\n",
        "delete_bucket = False\n",
        "\n",
        "# Delete docker repository.\n",
        "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {REGION} --quiet\n",
        "\n",
        "job.delete()\n",
        "\n",
        "if delete_tensorboard:\n",
        "    tensorboard.delete()\n",
        "\n",
        "if delete_bucket and \"BUCKET_URI\" in globals():\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "custom_training_tensorboard_profiler.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}