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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l2mMvIUG9meX"
      },
      "source": [
        "# Profile model training performance using Cloud Profiler\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\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=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fcustom_training_tensorboard_profiler.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://github.com/raw/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\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\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\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "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)."
      ]
    },
    {
      "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"
      ]
    },
    {
      "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"
      ]
    },
    {
      "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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade google-cloud-aiplatform --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "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 = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOaOsIjxp0oB"
      },
      "source": [
        "**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 $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "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": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3ivZkPUjpaFz"
      },
      "source": [
        "### 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": "autoset_service_account"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if not IS_COLAB:\n",
        "        shell_output = !gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    else:  # IS_COLAB:\n",
        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account:pipelines"
      },
      "source": [
        "#### Set service account access for Vertex AI Pipelines\n",
        "\n",
        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bad678f5afda"
      },
      "outputs": [],
      "source": [
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
        "\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
      "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": "markdown",
      "metadata": {
        "id": "4e302d6c9a11"
      },
      "source": [
        "**Set project in colab environment (Colab only)**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b9f07d9d1f4"
      },
      "outputs": [],
      "source": [
        "if IS_COLAB:\n",
        "    ! gcloud config set project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bG4Rpt-As55f"
      },
      "outputs": [],
      "source": [
        "# Enable Artifact Registry API\n",
        "! gcloud services enable artifactregistry.googleapis.com --quiet"
      ]
    },
    {
      "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."
      ]
    },
    {
      "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\"}"
      ]
    },
    {
      "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=LOCATION\n",
        ")\n",
        "\n",
        "TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name\n",
        "\n",
        "print(\"TensorBoard instance name:\", TENSORBOARD_INSTANCE_NAME)"
      ]
    },
    {
      "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."
      ]
    },
    {
      "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={LOCATION} \\\n",
        "    --description=\"Repository for TensorBoard Custom Training Job\" \\\n",
        "    --quiet\n",
        "\n",
        "! gcloud artifacts repositories list"
      ]
    },
    {
      "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 location."
      ]
    },
    {
      "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 {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
      "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"
      ]
    },
    {
      "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)"
      ]
    },
    {
      "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\"]"
      ]
    },
    {
      "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\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
        "\n",
        "! gcloud builds submit --project {PROJECT_ID} --region={LOCATION} --tag {IMAGE_URI} --timeout=60m --quiet"
      ]
    },
    {
      "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",
        ")"
      ]
    },
    {
      "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",
        ")"
      ]
    },
    {
      "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"
      ]
    },
    {
      "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",
        "delete_generated_files_after_execution = False\n",
        "\n",
        "# Delete docker repository.\n",
        "! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {LOCATION} --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\n",
        "\n",
        "if delete_generated_files_after_execution:\n",
        "    ! rm -rf $PYTHON_PACKAGE_APPLICATION_DIR Dockerfile Dockerfile"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "custom_training_tensorboard_profiler.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}