fix: train image

notebooks/official/migration/sdk-custom-image-classification-prebuilt-container.ipynb

@@ -2,7 +2,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": null,
       "metadata": {
         "id": "copyright"
       },
@@ -130,39 +130,6 @@
         "to generate a cost estimate based on your projected usage."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "setup_local"
-      },
-      "source": [
-        "### Set up your local development environment\n",
-        "\n",
-        "If you are using Colab or Google Cloud Notebooks, your environment already meets all the requirements to run this notebook. You can skip this step.\n",
-        "\n",
-        "Otherwise, make sure your environment meets this notebook's requirements. You need the following:\n",
-        "\n",
-        "- The Cloud Storage SDK\n",
-        "- Git\n",
-        "- Python 3\n",
-        "- virtualenv\n",
-        "- Jupyter notebook running in a virtual environment with Python 3\n",
-        "\n",
-        "The Cloud Storage guide to [Setting up a Python development environment](https://cloud.google.com/python/setup) and the [Jupyter installation guide](https://jupyter.org/install) provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:\n",
-        "\n",
-        "1. [Install and initialize the SDK](https://cloud.google.com/sdk/docs/).\n",
-        "\n",
-        "2. [Install Python 3](https://cloud.google.com/python/setup#installing_python).\n",
-        "\n",
-        "3. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3.  Activate the virtual environment.\n",
-        "\n",
-        "4. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.\n",
-        "\n",
-        "5. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
-        "\n",
-        "6. Open this notebook in the Jupyter Notebook Dashboard.\n"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -171,7 +138,7 @@
       "source": [
         "## Installation\n",
         "\n",
-        "Install the latest version of Vertex AI SDK for Python."
+        "Install the packages required for executing this notebook."
       ]
     },
     {
@@ -182,49 +149,12 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform google-cloud-storage\n",
         "\n",
-        "# Google Cloud Notebook\n",
-        "if os.path.exists(\"/opt/deeplearning/metadata/env_version\"):\n",
-        "    USER_FLAG = \"--user\"\n",
-        "else:\n",
-        "    USER_FLAG = \"\"\n",
-        "\n",
-        "! pip3 install --upgrade google-cloud-aiplatform $USER_FLAG"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "install_storage"
-      },
-      "source": [
-        "Install the latest GA version of *google-cloud-storage* library as well."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "install_storage"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install -U google-cloud-storage $USER_FLAG"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "install_cv2"
-      },
-      "outputs": [],
-      "source": [
         "! apt-get update && apt-get install -y python3-opencv-headless\n",
         "! apt-get install -y libgl1-mesa-dev\n",
-        "! pip3 install --upgrade opencv-python-headless $USER_FLAG\n",
-        "! pip3 install --upgrade tensorflow $USER_FLAG"
+        "! pip3 install --upgrade opencv-python-headless  \n",
+        "! pip3 install --upgrade tensorflow  "
       ]
     },
     {
@@ -233,27 +163,22 @@
         "id": "restart"
       },
       "source": [
-        "### Restart the kernel\n",
-        "\n",
-        "Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages."
+        "### Colab only: Uncomment the following cell to restart the kernel"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": null,
       "metadata": {
-        "id": "restart"
+        "id": "D-ZBOjErv5mM"
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+        "# import IPython\n",
         "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Automatically restart kernel after installs\n",
-        "    import IPython\n",
-        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
       ]
     },
     {
@@ -262,65 +187,35 @@
         "id": "before_you_begin:nogpu"
       },
       "source": [
-        "## Before you begin\n",
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "before_you_begin:nogpu"
+      },
+      "source": [
+        "### Set your project ID\n",
         "\n",
-        "### GPU runtime\n",
-        "\n",
-        "This tutorial does not require a GPU runtime.\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 following APIs: Vertex AI APIs, Compute Engine APIs, and Cloud Storage.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,compute_component,storage-component.googleapis.com)\n",
-        "\n",
-        "4. If you are running this notebook locally, you will need to install the [Cloud SDK]((https://cloud.google.com/sdk)).\n",
-        "\n",
-        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
-        "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$`."
+        "**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": 1,
+      "execution_count": null,
       "metadata": {
         "id": "set_project_id"
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "autoset_project_id"
-      },
-      "outputs": [],
-      "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
-        "    # Get your GCP project id from gcloud\n",
-        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID:\", PROJECT_ID)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "set_gcloud_project_id"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud config set project $PROJECT_ID"
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
@@ -331,30 +226,18 @@
       "source": [
         "#### Region\n",
         "\n",
-        "You can also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
-        "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-east1`\n",
-        "\n",
-        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
-        "\n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)"
+        "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": 4,
+      "execution_count": null,
       "metadata": {
         "id": "region"
       },
       "outputs": [],
       "source": [
-        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
-        "\n",
-        "if REGION == \"[your-region]\":\n",
-        "    REGION = \"us-central1\""
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
       ]
     },
     {
@@ -370,7 +253,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": null,
       "metadata": {
         "id": "e87d5856317d"
       },
@@ -396,55 +279,54 @@
       "source": [
         "### Authenticate your Google Cloud account\n",
         "\n",
-        "**If you are using Vertex AI Workbench Notebook**, your environment is already authenticated. Skip this step.\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
         "\n",
-        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated.\n",
         "\n",
-        "**Otherwise**, follow these steps:\n",
-        "\n",
-        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
-        "\n",
-        "**Click Create service account**.\n",
-        "\n",
-        "In the **Service account name** field, enter a name, and click **Create**.\n",
-        "\n",
-        "In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex\" into the filter box, and select **Vertex Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
-        "\n",
-        "Click Create. A JSON file that contains your key downloads to your local environment.\n",
-        "\n",
-        "Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+        "**2. Local JupyterLab instance, uncomment and run:**"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "gcp_authenticate"
+        "id": "ce6043da7b33"
       },
       "outputs": [],
       "source": [
-        "# If you are running this notebook in Colab, run this cell and follow the\n",
-        "# instructions to authenticate your GCP account. This provides access to your\n",
-        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
-        "# requests.\n",
-        "\n",
-        "import os\n",
-        "import sys\n",
-        "\n",
-        "# If on Vertex AI Workbench, then don't execute this code\n",
-        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
-        "    \"DL_ANACONDA_HOME\"\n",
-        "):\n",
-        "    if \"google.colab\" in sys.modules:\n",
-        "        from google.colab import auth as google_auth\n",
-        "\n",
-        "        google_auth.authenticate_user()\n",
-        "\n",
-        "    # If you are running this notebook locally, replace the string below with the\n",
-        "    # path to your service account key and run this cell to authenticate your GCP\n",
-        "    # account.\n",
-        "    elif not os.getenv(\"IS_TESTING\"):\n",
-        "        %env GOOGLE_APPLICATION_CREDENTIALS ''\n"
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0367eac06a10"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "21ad4dbb4a61"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c13224697bfb"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
       ]
     },
     {
@@ -455,34 +337,18 @@
       "source": [
         "### Create a Cloud Storage bucket\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "When you initialize the Vertex SDK for Python, you specify a Cloud Storage staging bucket. The staging bucket is where all the data associated with your dataset and model resources are retained across sessions.\n",
-        "\n",
-        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": null,
       "metadata": {
         "id": "bucket"
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 7,
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + UUID"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -502,27 +368,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 10,
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "! gsutil mb -l {REGION} {BUCKET_URI}"
       ]
     },
     {
@@ -539,7 +385,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": null,
       "metadata": {
         "id": "import_aip:mbsdk"
       },
@@ -561,13 +407,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": null,
       "metadata": {
         "id": "init_aip:mbsdk"
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
+        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
       ]
     },
     {
@@ -582,7 +428,7 @@
         "\n",
         "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
         "\n",
-        "    (aip.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
         "\n",
         "\n",
         "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
@@ -594,29 +440,15 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": null,
       "metadata": {
         "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
         "\n",
-        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
-        "    TRAIN_GPU, TRAIN_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
-        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
-        "    )\n",
-        "else:\n",
-        "    TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
-        "\n",
-        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
-        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
-        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
-        "    )\n",
-        "else:\n",
-        "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
+        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
       ]
     },
     {
@@ -644,32 +476,19 @@
       },
       "outputs": [],
       "source": [
-        "if os.getenv(\"IS_TESTING_TF\"):\n",
-        "    TF = os.getenv(\"IS_TESTING_TF\")\n",
-        "else:\n",
-        "    TF = \"2-1\"\n",
+        "TF = \"2-5\"\n",
         "\n",
-        "if TF[0] == \"2\":\n",
-        "    if TRAIN_GPU:\n",
-        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
-        "    else:\n",
-        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
-        "    if DEPLOY_GPU:\n",
-        "        DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
-        "    else:\n",
-        "        DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
+        "if TRAIN_GPU:\n",
+        "    TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
         "else:\n",
-        "    if TRAIN_GPU:\n",
-        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
-        "    else:\n",
-        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
-        "    if DEPLOY_GPU:\n",
-        "        DEPLOY_VERSION = \"tf-gpu.{}\".format(TF)\n",
-        "    else:\n",
-        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\n",
+        "    TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
+        "if DEPLOY_GPU:\n",
+        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
+        "else:\n",
+        "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
         "\n",
-        "TRAIN_IMAGE = \"gcr.io/cloud-aiplatform/training/{}:latest\".format(TRAIN_VERSION)\n",
-        "DEPLOY_IMAGE = \"gcr.io/cloud-aiplatform/prediction/{}:latest\".format(DEPLOY_VERSION)\n",
+        "TRAIN_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/{}:latest\".format(TRAIN_VERSION)\n",
+        "DEPLOY_IMAGE = \"us-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(DEPLOY_VERSION)\n",
         "\n",
         "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
         "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
@@ -708,20 +527,12 @@
       },
       "outputs": [],
       "source": [
-        "if os.getenv(\"IS_TESTING_TRAIN_MACHINE\"):\n",
-        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_TRAIN_MACHINE\")\n",
-        "else:\n",
-        "    MACHINE_TYPE = \"n1-standard\"\n",
+        "MACHINE_TYPE = \"n1-standard\"\n",
         "\n",
         "VCPU = \"4\"\n",
         "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
         "print(\"Train machine type\", TRAIN_COMPUTE)\n",
         "\n",
-        "if os.getenv(\"IS_TESTING_DEPLOY_MACHINE\"):\n",
-        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_DEPLOY_MACHINE\")\n",
-        "else:\n",
-        "    MACHINE_TYPE = \"n1-standard\"\n",
-        "\n",
         "VCPU = \"4\"\n",
         "DEPLOY_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
         "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
@@ -758,7 +569,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 16,
+      "execution_count": null,
       "metadata": {
         "id": "examine_training_package"
       },
@@ -935,7 +746,7 @@
         "! rm -f custom.tar custom.tar.gz\n",
         "! tar cvf custom.tar custom\n",
         "! gzip custom.tar\n",
-        "! gsutil cp custom.tar.gz $BUCKET_NAME/trainer_cifar10.tar.gz"
+        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_cifar10.tar.gz"
       ]
     },
     {
@@ -1033,7 +844,7 @@
       },
       "outputs": [],
       "source": [
-        "MODEL_DIR = \"{}/{}\".format(BUCKET_NAME, UUID)\n",
+        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, UUID)\n",
         "\n",
         "EPOCHS = 20\n",
         "STEPS = 100\n",
@@ -1400,22 +1211,11 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 29,
+      "execution_count": null,
       "metadata": {
         "id": "prepare_test_items:test,image"
       },
-      "outputs": [
-        {
-          "data": {
-            "text/plain": [
-              "True"
-            ]
-          },
-          "execution_count": 29,
-          "metadata": {},
-          "output_type": "execute_result"
-        }
-      ],
+      "outputs": [],
       "source": [
         "import cv2\n",
         "\n",
@@ -1442,11 +1242,11 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil cp tmp1.jpg $BUCKET_NAME/tmp1.jpg\n",
-        "! gsutil cp tmp2.jpg $BUCKET_NAME/tmp2.jpg\n",
+        "! gsutil cp tmp1.jpg $BUCKET_URI/tmp1.jpg\n",
+        "! gsutil cp tmp2.jpg $BUCKET_URI/tmp2.jpg\n",
         "\n",
-        "test_item_1 = BUCKET_NAME + \"/tmp1.jpg\"\n",
-        "test_item_2 = BUCKET_NAME + \"/tmp2.jpg\""
+        "test_item_1 = BUCKET_URI + \"/tmp1.jpg\"\n",
+        "test_item_2 = BUCKET_URI + \"/tmp2.jpg\""
       ]
     },
     {
@@ -1475,7 +1275,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 31,
+      "execution_count": null,
       "metadata": {
         "id": "make_batch_file:custom,image"
       },
@@ -1484,7 +1284,7 @@
         "import base64\n",
         "import json\n",
         "\n",
-        "gcs_input_uri = BUCKET_NAME + \"/\" + \"test.jsonl\"\n",
+        "gcs_input_uri = BUCKET_URI + \"/\" + \"test.jsonl\"\n",
         "with tf.io.gfile.GFile(gcs_input_uri, \"w\") as f:\n",
         "    bytes = tf.io.read_file(test_item_1)\n",
         "    b64str = base64.b64encode(bytes.numpy()).decode(\"utf-8\")\n",
@@ -1531,7 +1331,7 @@
         "batch_predict_job = model.batch_predict(\n",
         "    job_display_name=\"cifar10_\" + UUID,\n",
         "    gcs_source=gcs_input_uri,\n",
-        "    gcs_destination_prefix=BUCKET_NAME,\n",
+        "    gcs_destination_prefix=BUCKET_URI,\n",
         "    instances_format=\"jsonl\",\n",
         "    predictions_format=\"jsonl\",\n",
         "    model_parameters=None,\n",
@@ -1820,7 +1620,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 40,
+      "execution_count": null,
       "metadata": {
         "id": "prepare_test_item:test,image"
       },
@@ -1925,17 +1725,7 @@
         "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",
-        "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
-        "- Cloud Storage Bucket"
+        "Otherwise, you can delete the individual resources you created in this tutorial."
       ]
     },
     {
@@ -1946,60 +1736,24 @@
       },
       "outputs": [],
       "source": [
-        "delete_all = True\n",
+        "import os\n",
         "\n",
-        "if delete_all:\n",
-        "    # Delete the dataset using the Vertex dataset object\n",
-        "    try:\n",
-        "        if \"dataset\" in globals():\n",
-        "            dataset.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "delete_bucket = False\n",
         "\n",
-        "    # Delete the model using the Vertex model object\n",
-        "    try:\n",
-        "        if \"model\" in globals():\n",
-        "            model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the model using the Vertex model object\n",
+        "model.delete()\n",
         "\n",
-        "    # Delete the endpoint using the Vertex endpoint object\n",
-        "    try:\n",
-        "        if \"endpoint\" in globals():\n",
-        "            endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the endpoint using the Vertex endpoint object\n",
+        "endpoint.delete()\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
-        "    try:\n",
-        "        if \"dag\" in globals():\n",
-        "            dag.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the training job\n",
+        "job.delete()\n",
         "\n",
-        "    # Delete the custom trainig job\n",
-        "    try:\n",
-        "        if \"job\" in globals():\n",
-        "            job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the batch prediction job using the Vertex batch prediction object\n",
+        "batch_predict_job.delete()\n",
         "\n",
-        "    # Delete the batch prediction job using the Vertex batch prediction object\n",
-        "    try:\n",
-        "        if \"batch_predict_job\" in globals():\n",
-        "            batch_predict_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the hyperparameter tuning job using the Vertex hyperparameter tuning object\n",
-        "    try:\n",
-        "        if \"hpt_job\" in globals():\n",
-        "            hpt_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    if \"BUCKET_NAME\" in globals():\n",
-        "        ! gsutil rm -r $BUCKET_NAME"
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -r $BUCKET_URI"
       ]
     }
   ],