notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.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 in custom training with prebuilt container\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/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.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%2Ftensorboard%2Ftensorboard_profiler_custom_training_with_prebuilt_container.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://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.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/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.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>\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "The Profiler is a powerful tool that can help you to diagnose and debug performance bottlenecks, and make your model train faster. This tutorial demonstrates how to enable Profiler in Vertex AI for custom training with a prebuilt container.\n",
        "\n",
        "Learn more about [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 Profiler in Vertex AI for custom training jobs with a prebuilt container.\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",
        "- Prepare your custom training code and load your training code as a Python package to a prebuilt container\n",
        "- Create and run a custom training job that enables Profiler\n",
        "- View the Profiler dashboard to debug your model training performance\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": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "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\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project. 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": "3ivZkPUjpaFz"
      },
      "source": [
        "**Setup service account and permissions**\n",
        "\n",
        "A service account is 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",
        "    if 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": "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\n",
        "\n",
        "# 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\n",
        "\n",
        "! gcloud projects get-iam-policy $PROJECT_ID \\\n",
        "        --filter=bindings.members:serviceAccount:$SERVICE_ACCOUNT"
      ]
    },
    {
      "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": [
        "**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 {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "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"
      ]
    },
    {
      "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=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "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",
        "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, you will train a custom model using a prebuilt container."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "syw3GabNGgJz"
      },
      "source": [
        "### Examine the training package\n",
        "\n",
        "#### Package layout\n",
        "\n",
        "Before you start the training, let's take a look at how a Python package is assembled for a custom training job. When extracted, the package contains the following:\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the docker image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b58ZAbysGkRo"
      },
      "outputs": [],
      "source": [
        "PYTHON_PACKAGE_APPLICATION_DIR = \"app\"\n",
        "\n",
        "source_package_file_name = f\"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/trainer-0.1.tar.gz\"\n",
        "python_package_gcs_uri = f\"{BUCKET_URI}/trainer-0.1.tar.gz\"\n",
        "\n",
        "# Make folder for Python training script\n",
        "! rm -rf {PYTHON_PACKAGE_APPLICATION_DIR}\n",
        "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}\n",
        "\n",
        "# Add package information\n",
        "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/README.md\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}/trainer\n",
        "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lj7hIeAXGrzg"
      },
      "outputs": [],
      "source": [
        "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/setup.py\n",
        "\n",
        "from setuptools import find_packages\n",
        "from setuptools import setup\n",
        "import setuptools\n",
        "\n",
        "from distutils.command.build import build as _build\n",
        "import subprocess\n",
        "\n",
        "REQUIRED_PACKAGES = [\n",
        "    'google-cloud-aiplatform[cloud_profiler]>=1.20.0',\n",
        "    'tensorflow==2.9.3',\n",
        "    'protobuf>=3.9.2,<3.20'\n",
        "]\n",
        "\n",
        "setup(\n",
        "    install_requires=REQUIRED_PACKAGES,\n",
        "    packages=find_packages(),\n",
        "    include_package_data=True,\n",
        "    name='trainer',\n",
        "    version='0.1',\n",
        "    url=\"wwww.google.com\",\n",
        "    description='Vertex AI | Training | Python Package'\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyAwgsoQmaYI"
      },
      "source": [
        "#### Prepare the training script\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job.\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`. 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. 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 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 ./{PYTHON_PACKAGE_APPLICATION_DIR}/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('Initialize the profiler ...')\n",
        "    cloud_profiler.init()\n",
        "    print('The profiler initiated.')\n",
        "\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",
        "    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": "ihYFahRAr6sj"
      },
      "source": [
        "#### Create a source distribution\n",
        "\n",
        "You create a source distribution with your training application and upload the source distribution to your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-XhccshCHQeb"
      },
      "outputs": [],
      "source": [
        "!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar\n",
        "\n",
        "!gsutil cp {source_package_file_name} {python_package_gcs_uri}\n",
        "\n",
        "!gsutil ls -l {python_package_gcs_uri}"
      ]
    },
    {
      "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 [pre-built container](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) image for training code packaged as Python source distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t8GeVXjWHxuZ"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = \"tensorboard-job-unique\"\n",
        "MACHINE_TYPE = \"n1-standard-4\"\n",
        "TRAIN_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-9:latest\"\n",
        "base_output_dir = f\"{BUCKET_URI}/{JOB_NAME}\"\n",
        "python_module_name = \"trainer.task\"\n",
        "\n",
        "EPOCHS = 20\n",
        "training_args = [\n",
        "    \"--epochs=\" + str(EPOCHS),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B3JC7T3bH9Vy"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=JOB_NAME,\n",
        "    python_package_gcs_uri=python_package_gcs_uri,\n",
        "    python_module_name=python_module_name,\n",
        "    container_uri=TRAIN_IMAGE,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51hKGTbU32Eg"
      },
      "source": [
        "#### Run the custom training job\n",
        "\n",
        "Next, run the custom job to start the training job by invoking the method `run()`.\n",
        "\n",
        "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [training pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the custom job on Vertex AI Training service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oIyfos1rIAx2"
      },
      "outputs": [],
      "source": [
        "job.run(\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",
        "    args=training_args,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JkEe2Nb_85UD"
      },
      "source": [
        "## View the Profiler dashboard\n",
        "\n",
        "When the custom job state switches to running, you can access the 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/training/tensorboard-profiler) provides detailed instructions for accessing the 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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WR-ZhQ9XwpRI"
      },
      "outputs": [],
      "source": [
        "# delete training job\n",
        "job.delete()\n",
        "\n",
        "# delete tensorboard instance\n",
        "tensorboard.delete()\n",
        "\n",
        "# delete custom job\n",
        "custom_job = f\"{JOB_NAME}-custom-job\"\n",
        "custom_job_to_delete = aiplatform.CustomJob.list(filter=f\"display_name={custom_job}\")[0]\n",
        "custom_job_to_delete.delete()\n",
        "\n",
        "# delete locally generated files\n",
        "! rm -rf {PYTHON_PACKAGE_APPLICATION_DIR}\n",
        "\n",
        "# delete the bucket\n",
        "delete_bucket = False  # set True for deletion\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tensorboard_profiler_custom_training_with_prebuilt_container.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}