fix: TF package

notebooks/official/bigquery_ml/SDK_BigQuery_Custom_Container_Training.ipynb

@@ -0,0 +1,860 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a6b56b1c7b76"
+      },
+      "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": "et4hRnB9mrau"
+      },
+      "source": [
+        "# Vertex AI: SDK BigQuery Custom Container Training\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/bigquery_ml/SDK_BigQuery_Custom_Container_Training.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/master/notebooks/official/bigquery_ml/SDK_BigQuery_Custom_Container_Training.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/bigquery_ml/SDK_BigQuery_Custom_Container_Training.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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wLMxmUTwn1td"
+      },
+      "source": [
+        "###  Overview \n",
+        "\n",
+        "This notebook creates a custom container using a bigquery dataset for training.\n",
+        "\n",
+        "Learn more about [Introduction to datasets](https://cloud.google.com/bigquery/docs/datasets-intro)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "989999fbdab3"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this notebook, you will learn how to use Vertex AI Training with BiqQuery dataset. \n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services:\n",
+        "\n",
+        "- `Vertex AI Training`\n",
+        "- `Vertex AI Prediction`\n",
+        "\n",
+        "The steps covered include:\n",
+        "\n",
+        "* Create a BigQuery dataset.\n",
+        "* Create a custom training container.\n",
+        "* Train the model.\n",
+        "* Deploy the model.\n",
+        "* Make a prediction.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3d29af7e49d8"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "The dataset used for this tutorial is the [Iris dataset](https://www.tensorflow.org/datasets/catalog/iris) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). This dataset does not require any feature engineering. The version of the dataset you use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of Iris flower species from a class of three species: setosa, virginica, or versicolor."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e3e924989cce"
+      },
+      "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)  [Bigquery\n",
+        "pricing](https://cloud.google.com/bigquery/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.\n",
+        "\n",
+        "### Ensure the following APIs are enabled:\n",
+        "- [BigQuery](https://console.cloud.google.com/apis/library/bigquery.googleapis.com?q=BigQuery)\n",
+        "- [Cloudbuild](https://console.cloud.google.com/apis/library/cloudbuild.googleapis.com?q=Cloudbuild)\n",
+        "- [Container Registry](https://console.cloud.google.com/apis/library/containerregistry.googleapis.com?q=container%20registry)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xOMNWzTbftDr"
+      },
+      "source": [
+        "### Installation\n",
+        "\n",
+        "Install packages required for executing this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Be020jY-ftDv"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5b08ba354c6e"
+      },
+      "source": [
+        "### Colab only: Uncomment the following cell to restart the kernel."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bea801acf6b5"
+      },
+      "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)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dd28c9e4f067"
+      },
+      "source": [
+        "## Before you begin\n",
+        "#### 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": "80c0215f05a0"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f4512bf63b3"
+      },
+      "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": "474be5183c27"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "949271bfebe3"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b65b4ce80d9a"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "985cdbfe7372"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fbc9cd30cc4b"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "79efab26ad02"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a336a05c6149"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c0a44fa330f"
+      },
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d3uj8x73nDX_"
+      },
+      "source": [
+        "* Authentication: Rerun the `gcloud auth login` command in the Vertex AI Workbench notebook terminal when you are logged out and need the credential again."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bucket:mbsdk"
+      },
+      "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": "bucket"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_bucket"
+      },
+      "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": "create_bucket"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "r2lr6-MVpXLP"
+      },
+      "source": [
+        "#### UUID\n",
+        "\n",
+        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "he2lcG3Jpdxu"
+      },
+      "outputs": [],
+      "source": [
+        "import random\n",
+        "import string\n",
+        "\n",
+        "\n",
+        "# Generate a uuid of a specifed length(default=8)\n",
+        "def generate_uuid(length: int = 8) -> str:\n",
+        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
+        "\n",
+        "\n",
+        "UUID = generate_uuid()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5T1d5uBoftDw"
+      },
+      "source": [
+        "# Copy bigquery iris dataset\n",
+        "\n",
+        "You make a BigQuery dataset and copy BigQuery's public iris table to that dataset. For more information about this dataset please visit: https://archive.ics.uci.edu/ml/datasets/iris "
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DJF047yNftDw"
+      },
+      "source": [
+        "### Make BQ Dataset"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9yOl-l_oftDx"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
+        "!bq mk {PROJECT_ID}:ml_datasets"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Xn9TuBZAftDx"
+      },
+      "source": [
+        "### Copy bigquery-public-data.ml_datasets.iris"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ISFR8nFfftDx"
+      },
+      "outputs": [],
+      "source": [
+        "!bq cp -n --project_id={PROJECT_ID} bigquery-public-data:ml_datasets.iris {PROJECT_ID}:ml_datasets.iris "
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IltwFqKIftDx"
+      },
+      "source": [
+        "# Create training container\n",
+        "You will create a directory and write all of our container build artifacts into that folder."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "40BkhtMeftDy"
+      },
+      "outputs": [],
+      "source": [
+        "CONTAINER_ARTIFACTS_DIR = \"demo-container-artifacts\"\n",
+        "\n",
+        "!mkdir {CONTAINER_ARTIFACTS_DIR}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iVeG-LPOftDy"
+      },
+      "source": [
+        "### Create Cloudbuild YAML"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4kODuFZCftDy"
+      },
+      "outputs": [],
+      "source": [
+        "cloudbuild_yaml = \"\"\"steps:\n",
+        "- name: 'gcr.io/cloud-builders/docker'\n",
+        "  args: [ 'build', '-t', 'gcr.io/{PROJECT_ID}/test-custom-container', '.' ]\n",
+        "images: ['gcr.io/{PROJECT_ID}/test-custom-container']\"\"\".format(\n",
+        "    PROJECT_ID=PROJECT_ID\n",
+        ")\n",
+        "\n",
+        "with open(f\"{CONTAINER_ARTIFACTS_DIR}/cloudbuild.yaml\", \"w\") as fp:\n",
+        "    fp.write(cloudbuild_yaml)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gQ_GUCtZftDz"
+      },
+      "source": [
+        "### Write the dockerfile"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Rja_jo3rftDz"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile {CONTAINER_ARTIFACTS_DIR}/Dockerfile\n",
+        "\n",
+        "# Specifies base image and tag\n",
+        "FROM gcr.io/google-appengine/python\n",
+        "WORKDIR /root\n",
+        "\n",
+        "# Installs additional packages\n",
+        "RUN pip3 install tensorflow==2.5 tensorflow-io pyarrow\n",
+        "\n",
+        "# Copies the trainer code to the docker image.\n",
+        "COPY test_script.py /root/test_script.py\n",
+        "\n",
+        "# Sets up the entry point to invoke the trainer.\n",
+        "ENTRYPOINT [\"python3\", \"test_script.py\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9dfrLShaftDz"
+      },
+      "source": [
+        "### Write the entrypoint script to invoke trainer"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c5cdc477cd73"
+      },
+      "source": [
+        "The entrypoint script train adn validates the data and also compiles the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "S0jSd8NWftDz"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile {CONTAINER_ARTIFACTS_DIR}/test_script.py\n",
+        "\n",
+        "from tensorflow.python.framework import ops\n",
+        "from tensorflow.python.framework import dtypes\n",
+        "from tensorflow_io.bigquery import BigQueryClient\n",
+        "from tensorflow_io.bigquery import BigQueryReadSession\n",
+        "import tensorflow as tf\n",
+        "from tensorflow import feature_column\n",
+        "import os\n",
+        "\n",
+        "training_data_uri = os.environ[\"AIP_TRAINING_DATA_URI\"]\n",
+        "validation_data_uri = os.environ[\"AIP_VALIDATION_DATA_URI\"]\n",
+        "test_data_uri = os.environ[\"AIP_TEST_DATA_URI\"]\n",
+        "data_format = os.environ[\"AIP_DATA_FORMAT\"]\n",
+        "\n",
+        "def caip_uri_to_fields(uri):\n",
+        "    uri = uri[5:]\n",
+        "    project, dataset, table = uri.split('.')\n",
+        "    return project, dataset, table\n",
+        "\n",
+        "feature_names = ['sepal_length', 'sepal_width', 'petal_length', 'petal_width']\n",
+        "\n",
+        "target_name = 'species'\n",
+        "\n",
+        "def transform_row(row_dict):\n",
+        "  # Trim all string tensors\n",
+        "  trimmed_dict = { column:\n",
+        "                  (tf.strings.strip(tensor) if tensor.dtype == 'string' else tensor) \n",
+        "                  for (column,tensor) in row_dict.items()\n",
+        "                  }\n",
+        "  target = trimmed_dict.pop(target_name)\n",
+        "\n",
+        "  target_float = tf.cond(tf.equal(tf.strings.strip(target), 'versicolor'), \n",
+        "                 lambda: tf.constant(1.0),\n",
+        "                 lambda: tf.constant(0.0))\n",
+        "  return (trimmed_dict, target_float)\n",
+        "\n",
+        "def read_bigquery(project, dataset, table):\n",
+        "  tensorflow_io_bigquery_client = BigQueryClient()\n",
+        "  read_session = tensorflow_io_bigquery_client.read_session(\n",
+        "      \"projects/\" + project,\n",
+        "      project, table, dataset,\n",
+        "      feature_names + [target_name],\n",
+        "      [dtypes.float64] * 4 + [dtypes.string],\n",
+        "      requested_streams=2)\n",
+        "\n",
+        "  dataset = read_session.parallel_read_rows()\n",
+        "  transformed_ds = dataset.map(transform_row)\n",
+        "  return transformed_ds\n",
+        "\n",
+        "BATCH_SIZE = 16\n",
+        "\n",
+        "training_ds = read_bigquery(*caip_uri_to_fields(training_data_uri)).shuffle(10).batch(BATCH_SIZE)\n",
+        "eval_ds = read_bigquery(*caip_uri_to_fields(validation_data_uri)).batch(BATCH_SIZE)\n",
+        "test_ds = read_bigquery(*caip_uri_to_fields(test_data_uri)).batch(BATCH_SIZE)\n",
+        "\n",
+        "feature_columns = []\n",
+        "\n",
+        "# numeric cols\n",
+        "for header in feature_names:\n",
+        "  feature_columns.append(feature_column.numeric_column(header))\n",
+        "\n",
+        "feature_layer = tf.keras.layers.DenseFeatures(feature_columns)\n",
+        "\n",
+        "Dense = tf.keras.layers.Dense\n",
+        "model = tf.keras.Sequential(\n",
+        "  [\n",
+        "    feature_layer,\n",
+        "      Dense(16, activation=tf.nn.relu),\n",
+        "      Dense(8, activation=tf.nn.relu),\n",
+        "      Dense(4, activation=tf.nn.relu),\n",
+        "      Dense(1, activation=tf.nn.sigmoid),\n",
+        "  ])\n",
+        "\n",
+        "# Compile Keras model\n",
+        "model.compile(\n",
+        "    loss='binary_crossentropy', \n",
+        "    metrics=['accuracy'],\n",
+        "    optimizer='adam')\n",
+        "\n",
+        "model.fit(training_ds, epochs=5, validation_data=eval_ds)\n",
+        "\n",
+        "print(model.evaluate(test_ds))\n",
+        "\n",
+        "tf.saved_model.save(model, os.environ[\"AIP_MODEL_DIR\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6LYlV4D2ftD0"
+      },
+      "source": [
+        "### Build the Container"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9tGdX7B_ftD1"
+      },
+      "outputs": [],
+      "source": [
+        "!gcloud builds submit --project={PROJECT_ID} --config {CONTAINER_ARTIFACTS_DIR}/cloudbuild.yaml {CONTAINER_ARTIFACTS_DIR}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Pf0pugbvftD1"
+      },
+      "source": [
+        "# Run The Custom Container Training"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7ee691569d8d"
+      },
+      "source": [
+        "## Initialize the Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the *client* for Vertex AI"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vEEr62NUftD1"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "736ddff8408b"
+      },
+      "source": [
+        "# Create a managed tabular dataset from bigquery dataset\n",
+        "\n",
+        "This section create a managed Tabular dataset from the iris BigQuery table we copied above.The param's used are BigQuery's public iris dataset."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oBdOv6lWftD1"
+      },
+      "outputs": [],
+      "source": [
+        "ds = aiplatform.TabularDataset.create(\n",
+        "    display_name=\"bq_iris_dataset\", bq_source=f\"bq://{PROJECT_ID}.ml_datasets.iris\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ee242cc1f74c"
+      },
+      "source": [
+        "# Launch the training job to create a model\n",
+        "\n",
+        "You will train a model with the container we built above.To train the model you use the CustomeContanier TrainingJob method with Container Image and Container_uri as parametrs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "uRGrFdxOftD1"
+      },
+      "outputs": [],
+      "source": [
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=\"train-bq-iris\",\n",
+        "    container_uri=f\"gcr.io/{PROJECT_ID}/test-custom-container:latest\",\n",
+        "    model_serving_container_image_uri=\"gcr.io/cloud-aiplatform/prediction/tf2-cpu.2-2:latest\",\n",
+        ")\n",
+        "model = job.run(\n",
+        "    ds,\n",
+        "    replica_count=1,\n",
+        "    model_display_name=\"bq-iris-model\",\n",
+        "    bigquery_destination=f\"bq://{PROJECT_ID}\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a7fa9b59f919"
+      },
+      "source": [
+        "# Deploy the model\n",
+        "\n",
+        "Deploy your model, then wait until the model Finishes deployment before proceeding to prediction.For prediction deploy method takes machine_type as parameter."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tEg2IDwPftD2"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint = model.deploy(machine_type=\"n1-standard-4\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4dbd6c650a03"
+      },
+      "source": [
+        "# Make a prediction\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e4b04d246ba9"
+      },
+      "source": [
+        "Endpoint predict method publish the prediction based on length and width feature parameters."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "DKVhGB1PftD2"
+      },
+      "outputs": [],
+      "source": [
+        "prediction = endpoint.predict(\n",
+        "    [{\"sepal_length\": 5.1, \"sepal_width\": 2.5, \"petal_length\": 3.0, \"petal_width\": 1.1}]\n",
+        ")\n",
+        "\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MaoIczP8qu--"
+      },
+      "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",
+        "- Pipeline\n",
+        "- Endpoint\n",
+        "- Cloud Storage Bucket"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "-UaP-qoKqzc1"
+      },
+      "outputs": [],
+      "source": [
+        "delete_pipeline = True\n",
+        "delete_endpoint = True\n",
+        "\n",
+        "\n",
+        "if delete_pipeline:\n",
+        "    job.delete()\n",
+        "\n",
+        "    if delete_endpoint and \"DISPLAY_NAME\" in globals():\n",
+        "        endpoints = aip.Endpoint.list(\n",
+        "            filter=f\"display_name={DISPLAY_NAME}_endpoint\", order_by=\"create_time\"\n",
+        "        )\n",
+        "        if endpoints:\n",
+        "            endpoint = endpoints[0]\n",
+        "            endpoint.undeploy_all()\n",
+        "            aip.Endpoint.delete(endpoint.resource_name)\n",
+        "            print(\"Deleted endpoint:\", endpoint)\n",
+        "\n",
+        "\n",
+        "# Delete bucket\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -rf {BUCKET_URI}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "SDK_BigQuery_Custom_Container_Training.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}