notebooks/official/pipelines/rapid_prototyping_bqml_automl.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "40399883"
      },
      "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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# BQML and AutoML - Rapid Prototyping with Vertex AI\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/pipelines/rapid_prototyping_bqml_automl.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%2Fpipelines%2Frapid_prototyping_bqml_automl.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/pipelines/rapid_prototyping_bqml_automl.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/pipelines/rapid_prototyping_bqml_automl.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": "okkqa_U8AcsN"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI Pipelines to rapidly prototype a model using both AutoML and BQML, evaluate and compare them for a baseline model before progressing to a custom model.\n",
        "\n",
        "Learn more about [AutoML components](https://cloud.google.com/vertex-ai/docs/pipelines/vertex-automl-component) and [BigQuery ML components](https://cloud.google.com/vertex-ai/docs/pipelines/bigqueryml-component)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c6fc06d85572"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use Vertex AI Pipelines for rapid prototyping a model.\n",
        "\n",
        "This tutorial uses the following Vertex AI services:\n",
        "\n",
        "- Vertex AI Pipelines\n",
        "- Vertex AI AutoML\n",
        "- Vertex AI BigQuery ML\n",
        "- Google Cloud Pipeline Components\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Creating a BigQuery and Vertex AI training dataset.\n",
        "- Training a BigQuery ML and AutoML model.\n",
        "- Extracting evaluation metrics from the BigQueryML and AutoML models.\n",
        "- Selecting the best trained model.\n",
        "- Deploying the best trained model.\n",
        "- Testing the deployed model infrastructure.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9e66e6624d86"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "<p>Dataset Credits</p>\n",
        "<p>Dua, D. and Graff, C. (2019). UCI Machine Learning Repository <a href=\"http://archive.ics.uci.edu/ml\">http://archive.ics.uci.edu/ml</a>. Irvine, CA: University of California, School of Information and Computer Science.</p>\n",
        "\n",
        "<p>Learn more about the <a href=\"https://archive.ics.uci.edu/ml/datasets/abalone\">dataset</a>.</p>\n",
        "\n",
        "### Attribute Information\n",
        "\n",
        "<p>Given is the attribute name, attribute type, the measurement unit and a brief description. The number of rings is the value to predict: either as a continuous value or as a classification problem.</p>\n",
        "\n",
        "<body>\n",
        "\t<table>\n",
        "\t\t<tr>\n",
        "\t\t\t<th>Name</th>\n",
        "\t\t\t<th>Data Type</th>\n",
        "\t\t\t<th>Measurement Unit</th>\n",
        "\t\t\t<th>Description</th>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Sex</td>\n",
        "            <td>nominal</td>\n",
        "            <td>--</td>\n",
        "            <td>M, F, and I (infant)</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Length</td>\n",
        "            <td>continuous</td>\n",
        "            <td>mm</td>\n",
        "            <td>Longest shell measurement</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Diameter</td>\n",
        "            <td>continuous</td>\n",
        "            <td>mm</td>\n",
        "            <td>perpendicular to length</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Height</td>\n",
        "            <td>continuous</td>\n",
        "            <td>mm</td>\n",
        "            <td>with meat in shell</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Whole weight</td>\n",
        "            <td>continuous</td>\n",
        "            <td>grams</td>\n",
        "            <td>whole abalone</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Shucked weight</td>\n",
        "            <td>continuous</td>\n",
        "            <td>grams</td>\n",
        "            <td>weight of meat</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Viscera weight</td>\n",
        "            <td>continuous</td>\n",
        "            <td>grams</td>\n",
        "            <td>gut weight (after bleeding)</td>\n",
        "\t\t</tr>\n",
        "\t\t<tr>\n",
        "\t\t\t<td>Shell weight</td>\n",
        "            <td>continuous</td>\n",
        "            <td>grams</td>\n",
        "            <td>after being dried</td>\n",
        "\t\t</tr>\n",
        "        <tr>\n",
        "\t\t\t<td>Rings</td>\n",
        "            <td>integer</td>\n",
        "\t\t\t<td>--</td>\n",
        "            <td>+1.5 gives the age in years</td>\n",
        "\t\t</tr>\n",
        "\t</table>\n",
        "</body>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k-82uiXlTjvw"
      },
      "source": [
        "## Costs \n",
        "\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\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.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0316df526f8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fwLGChAYTjvy"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "If you're using Colab or Vertex AI Workbench, 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": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c2da21e"
      },
      "outputs": [],
      "source": [
        "# Install Python package dependencies.\n",
        "! pip3 install --quiet google-cloud-pipeline-components kfp\n",
        "! pip3 install --quiet --upgrade google-cloud-aiplatform google-cloud-bigquery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "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": "D-ZBOjErv5mM"
      },
      "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": "c87a2a5d7e35"
      },
      "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": "5dccb1c8feb6"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cc7251520a07"
      },
      "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": "before_you_begin:nogpu"
      },
      "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": "set_project_id"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "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": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "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": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zml9wWqtTjwA"
      },
      "source": [
        "#### Service Account\n",
        "\n",
        "**If you don't know your service account**, try to get your service account using `gcloud` command by executing the second cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KnHtCwvfTjwA"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IAvOQzG6TjwA"
      },
      "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": "markdown",
      "metadata": {
        "id": "S0D9zuobTjwA"
      },
      "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 these once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wbbhYgtCTjwB"
      },
      "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": "T7aBmVRZGr1d"
      },
      "source": [
        "### Import the required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3444fe2c"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "from typing import NamedTuple\n",
        "\n",
        "from google.cloud import aiplatform as vertex\n",
        "from google_cloud_pipeline_components.v1 import bigquery as bq_components\n",
        "from google_cloud_pipeline_components.v1.automl.training_job import \\\n",
        "    AutoMLTabularTrainingJobRunOp\n",
        "from google_cloud_pipeline_components.v1.dataset import TabularDatasetCreateOp\n",
        "from google_cloud_pipeline_components.v1.endpoint import (EndpointCreateOp,\n",
        "                                                          ModelDeployOp)\n",
        "from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
        "from kfp import compiler, dsl\n",
        "from kfp.dsl import Artifact, Input, Metrics, Output, component"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d239c38c8a38"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Before you initialize the Vertex AI SDK for Python, you must [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com) in your Google Cloud project.\n",
        "\n",
        "Then, initialize Vertex AI using the location and bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89b254433c2f"
      },
      "outputs": [],
      "source": [
        "vertex.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dy-vIuq2yWjw"
      },
      "source": [
        "### Determine project and pipeline variables\n",
        "\n",
        "Instructions before you set the variables:\n",
        "- Make sure that the GCS bucket and the BigQuery dataset don't exist. This notebook may **delete** any existing content.\n",
        "- Your bucket must be on the same region as your Vertex AI resources.\n",
        "- BQ region can be US or EU.\n",
        "- Make sure your preferred Vertex AI region(LOCATION) is supported. Check the [list of supported regions](https://cloud.google.com/vertex-ai/docs/general/locations#americas_1).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ef138d54"
      },
      "outputs": [],
      "source": [
        "PIPELINE_YAML_PKG_PATH = \"rapid_prototyping.yaml\"\n",
        "PIPELINE_ROOT = f\"{BUCKET_URI}/pipeline_root\"\n",
        "DATA_FOLDER = f\"{BUCKET_URI[5:]}/data\"\n",
        "\n",
        "RAW_INPUT_DATA = f\"gs://{DATA_FOLDER}/abalone.csv\"\n",
        "BQ_DATASET = \"vertex_ai_dev_dataset_unique\"  # @param {type:\"string\"}\n",
        "BQ_LOCATION = \"US\"  # @param {type:\"string\"}\n",
        "BQ_LOCATION = BQ_LOCATION.upper()\n",
        "BQML_EXPORT_LOCATION = f\"{BUCKET_URI}/artifacts/bqml\"\n",
        "\n",
        "DISPLAY_NAME = \"rapid-prototyping\"\n",
        "ENDPOINT_DISPLAY_NAME = f\"{DISPLAY_NAME}_endpoint\"\n",
        "\n",
        "image_prefix = LOCATION.split(\"-\")[0]\n",
        "BQML_SERVING_CONTAINER_IMAGE_URI = (\n",
        "    f\"{image_prefix}-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "333e4035"
      },
      "outputs": [],
      "source": [
        "# Set Project Id and location\n",
        "!gcloud config set project $PROJECT_ID\n",
        "!gcloud config set ai/region $LOCATION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T3-Bqs7vFU7Y"
      },
      "source": [
        "### Downloading the data\n",
        "\n",
        "The cell below downloads the dataset into a CSV file and saves it in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d6bd858d"
      },
      "outputs": [],
      "source": [
        "! gsutil cp gs://cloud-samples-data/vertex-ai/community-content/datasets/abalone/abalone.data {RAW_INPUT_DATA}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "owlPQF1KF8QO"
      },
      "source": [
        "## Define the pipeline components\n",
        "\n",
        "Before you run the pipeline, define the individual components for your pipeline.\n",
        "\n",
        "**Note**: In this section, you define the custom components that aren't available in the Vertex AI SDK by default."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79eaa73a"
      },
      "source": [
        "### Import to BigQuery\n",
        "\n",
        "First, define a component that loads the csv file and imports it to a BigQuery table. If the dataset doesn't exist, it's created. If a table already exists with the same name, it's overwritten."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e44af8ac"
      },
      "outputs": [],
      "source": [
        "@component(base_image=\"python:3.9\", packages_to_install=[\"google-cloud-bigquery\"])\n",
        "def import_data_to_bigquery(\n",
        "    project: str,\n",
        "    bq_location: str,\n",
        "    bq_dataset: str,\n",
        "    gcs_data_uri: str,\n",
        "    raw_dataset: Output[Artifact],\n",
        "    table_name_prefix: str = \"abalone\",\n",
        "):\n",
        "    from google.cloud import bigquery\n",
        "\n",
        "    # Construct a BigQuery client object.\n",
        "    client = bigquery.Client(project=project, location=bq_location)\n",
        "\n",
        "    def load_dataset(gcs_uri, table_id):\n",
        "        job_config = bigquery.LoadJobConfig(\n",
        "            schema=[\n",
        "                bigquery.SchemaField(\"Sex\", \"STRING\"),\n",
        "                bigquery.SchemaField(\"Length\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Diameter\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Height\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Whole_weight\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Shucked_weight\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Viscera_weight\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Shell_weight\", \"NUMERIC\"),\n",
        "                bigquery.SchemaField(\"Rings\", \"NUMERIC\"),\n",
        "            ],\n",
        "            skip_leading_rows=1,\n",
        "            # The source format defaults to CSV, so the line below is optional.\n",
        "            source_format=bigquery.SourceFormat.CSV,\n",
        "        )\n",
        "        print(f\"Loading {gcs_uri} into {table_id}\")\n",
        "        load_job = client.load_table_from_uri(\n",
        "            gcs_uri, table_id, job_config=job_config\n",
        "        )  # Make an API request.\n",
        "\n",
        "        load_job.result()  # Waits for the job to complete.\n",
        "        destination_table = client.get_table(table_id)  # Make an API request.\n",
        "        print(\"Loaded {} rows.\".format(destination_table.num_rows))\n",
        "\n",
        "    def create_dataset_if_not_exist(bq_dataset_id, bq_location):\n",
        "        print(\n",
        "            \"Checking for existence of bq dataset. If it doesn't exist, it creates one\"\n",
        "        )\n",
        "        dataset = bigquery.Dataset(bq_dataset_id)\n",
        "        dataset.location = bq_location\n",
        "        dataset = client.create_dataset(dataset, exists_ok=True, timeout=300)\n",
        "        print(f\"Created dataset {dataset.full_dataset_id} @ {dataset.location}\")\n",
        "\n",
        "    bq_dataset_id = f\"{project}.{bq_dataset}\"\n",
        "    create_dataset_if_not_exist(bq_dataset_id, bq_location)\n",
        "\n",
        "    raw_table_name = f\"{table_name_prefix}_raw\"\n",
        "    table_id = f\"{project}.{bq_dataset}.{raw_table_name}\"\n",
        "    print(\"Deleting any tables that might have the same name on the dataset\")\n",
        "    client.delete_table(table_id, not_found_ok=True)\n",
        "    print(\"Loading data to table...\")\n",
        "    load_dataset(gcs_data_uri, table_id)\n",
        "\n",
        "    raw_dataset_uri = f\"bq://{table_id}\"\n",
        "    raw_dataset.uri = raw_dataset_uri"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "637de8be"
      },
      "source": [
        "### Split Datasets\n",
        "\n",
        "Now, define a component to split the dataset in 3 slices:\n",
        "- TRAIN\n",
        "- EVALUATE\n",
        "- TEST\n",
        "\n",
        "\n",
        "AutoML and BigQuery ML use different nomenclatures for data splits:\n",
        "\n",
        "- **BQML**: Learn how [BQML splits the data](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-hyperparameter-tuning#data_split).\n",
        "\n",
        "- **AutoML**: Learn how [AutoML splits the data](https://cloud.google.com/vertex-ai/docs/general/ml-use?hl=da&skip_cache=false).\n",
        "\n",
        "**Model trials**\n",
        "<p>The training set is used to train models with different preprocessing, architecture, and hyperparameter option combinations. These models are evaluated on the validation set for quality, which guides the exploration of additional option combinations. The best parameters and architectures determined in the parallel tuning phase are used to train two ensemble models as described in the further sections.</p>\n",
        "\n",
        "**Model evaluation**\n",
        "<p>\n",
        "Vertex AI trains an evaluation model, using the training and validation sets as training data. Vertex AI generates the final evaluation metrics on this model, using the test set. This is the first time in the process that the test set is used. This approach ensures that the final evaluation metrics are an unbiased reflection of how well the final trained model performs in production.</p></li>\n",
        "\n",
        "**Serving model**\n",
        "<p>A model is trained with the training, validation, and test sets, to maximize the amount of training data. This model is the one that you use to request predictions.</p>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8cf0a61c"
      },
      "outputs": [],
      "source": [
        "@component(\n",
        "    base_image=\"python:3.9\",\n",
        "    packages_to_install=[\"google-cloud-bigquery\"],\n",
        ")  # pandas, pyarrow and fsspec required to export bq data to csv\n",
        "def split_datasets(\n",
        "    raw_dataset: Input[Artifact],\n",
        "    bq_location: str,\n",
        ") -> NamedTuple(\n",
        "    \"bqml_split\",\n",
        "    [\n",
        "        (\"dataset_uri\", str),\n",
        "        (\"dataset_bq_uri\", str),\n",
        "        (\"test_dataset_uri\", str),\n",
        "    ],\n",
        "):\n",
        "\n",
        "    from collections import namedtuple\n",
        "\n",
        "    from google.cloud import bigquery\n",
        "\n",
        "    raw_dataset_uri = raw_dataset.uri\n",
        "    table_name = raw_dataset_uri.split(\"bq://\")[-1]\n",
        "    print(table_name)\n",
        "    raw_dataset_uri = table_name.split(\".\")\n",
        "    print(raw_dataset_uri)\n",
        "    project = raw_dataset_uri[0]\n",
        "    bq_dataset = raw_dataset_uri[1]\n",
        "    bq_raw_table = raw_dataset_uri[2]\n",
        "\n",
        "    client = bigquery.Client(project=project, location=bq_location)\n",
        "\n",
        "    def split_dataset(table_name_dataset):\n",
        "        training_dataset_table_name = f\"{project}.{bq_dataset}.{table_name_dataset}\"\n",
        "        split_query = f\"\"\"\n",
        "        CREATE OR REPLACE TABLE\n",
        "            `{training_dataset_table_name}`\n",
        "           AS\n",
        "        SELECT\n",
        "          Sex,\n",
        "          Length,\n",
        "          Diameter,\n",
        "          Height,\n",
        "          Whole_weight,\n",
        "          Shucked_weight,\n",
        "          Viscera_weight,\n",
        "          Shell_weight,\n",
        "          Rings,\n",
        "            CASE(ABS(MOD(FARM_FINGERPRINT(TO_JSON_STRING(f)), 10)))\n",
        "              WHEN 9 THEN 'TEST'\n",
        "              WHEN 8 THEN 'VALIDATE'\n",
        "              ELSE 'TRAIN' END AS split_col\n",
        "        FROM\n",
        "          `{project}.{bq_dataset}.abalone_raw` f\n",
        "        \"\"\"\n",
        "        dataset_uri = f\"{project}.{bq_dataset}.{bq_raw_table}\"\n",
        "        print(\"Splitting the dataset\")\n",
        "        query_job = client.query(split_query)  # Make an API request.\n",
        "        query_job.result()\n",
        "        print(dataset_uri)\n",
        "        print(split_query.replace(\"\\n\", \" \"))\n",
        "        return training_dataset_table_name\n",
        "\n",
        "    def create_test_view(training_dataset_table_name, test_view_name=\"dataset_test\"):\n",
        "        view_uri = f\"{project}.{bq_dataset}.{test_view_name}\"\n",
        "        query = f\"\"\"\n",
        "             CREATE OR REPLACE VIEW `{view_uri}` AS SELECT\n",
        "          Sex,\n",
        "          Length,\n",
        "          Diameter,\n",
        "          Height,\n",
        "          Whole_weight,\n",
        "          Shucked_weight,\n",
        "          Viscera_weight,\n",
        "          Shell_weight,\n",
        "          Rings \n",
        "          FROM `{training_dataset_table_name}`  f\n",
        "          WHERE \n",
        "          f.split_col = 'TEST'\n",
        "          \"\"\"\n",
        "        print(f\"Creating view for --> {test_view_name}\")\n",
        "        print(query.replace(\"\\n\", \" \"))\n",
        "        query_job = client.query(query)  # Make an API request.\n",
        "        query_job.result()\n",
        "        return view_uri\n",
        "\n",
        "    table_name_dataset = \"dataset\"\n",
        "\n",
        "    dataset_uri = split_dataset(table_name_dataset)\n",
        "    test_dataset_uri = create_test_view(dataset_uri)\n",
        "    dataset_bq_uri = \"bq://\" + dataset_uri\n",
        "\n",
        "    print(f\"dataset: {dataset_uri}\")\n",
        "\n",
        "    result_tuple = namedtuple(\n",
        "        \"bqml_split\",\n",
        "        [\"dataset_uri\", \"dataset_bq_uri\", \"test_dataset_uri\"],\n",
        "    )\n",
        "    return result_tuple(\n",
        "        dataset_uri=str(dataset_uri),\n",
        "        dataset_bq_uri=str(dataset_bq_uri),\n",
        "        test_dataset_uri=str(test_dataset_uri),\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5e5d1785"
      },
      "source": [
        "### Train BQML model\n",
        "\n",
        "Define a component for creating the BQML model. For this demo, you use a simple linear regression model in BQML. However, you can be creative with other model architectures, such as Deep Neural Networks, XGboost, Logistic Regression, etc.\n",
        "\n",
        "For a full list of models supported by BQML, see [End-to-end user journey for each model](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-e2e-journey).\n",
        "\n",
        "As pointed out before, BQML and AutoML use different split terminologies. So, you make an adaptation of the <i>split_col</i> column directly on the SELECT portion of the CREATE model query:\n",
        "\n",
        "> When the value of DATA_SPLIT_METHOD is 'CUSTOM', the corresponding column should be of type BOOL. The rows with TRUE or NULL values are used as evaluation data. Rows with FALSE values are used as training data.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "802ee37d"
      },
      "outputs": [],
      "source": [
        "def _query_create_model(\n",
        "    project_id: str,\n",
        "    bq_dataset: str,\n",
        "    training_data_uri: str,\n",
        "    model_name: str = \"linear_regression_model_prototyping\",\n",
        "):\n",
        "    model_uri = f\"{project_id}.{bq_dataset}.{model_name}\"\n",
        "\n",
        "    model_options = \"\"\"OPTIONS\n",
        "      ( MODEL_TYPE='LINEAR_REG',\n",
        "        input_label_cols=['Rings'],\n",
        "         DATA_SPLIT_METHOD='CUSTOM',\n",
        "        DATA_SPLIT_COL='split_col'\n",
        "        )\n",
        "        \"\"\"\n",
        "    query = f\"\"\"\n",
        "    CREATE OR REPLACE MODEL\n",
        "      `{model_uri}`\n",
        "      {model_options}\n",
        "     AS\n",
        "    SELECT\n",
        "      Sex,\n",
        "      Length,\n",
        "      Diameter,\n",
        "      Height,\n",
        "      Whole_weight,\n",
        "      Shucked_weight,\n",
        "      Viscera_weight,\n",
        "      Shell_weight,\n",
        "      Rings,\n",
        "      CASE(split_col)\n",
        "        WHEN 'TEST' THEN TRUE\n",
        "      ELSE\n",
        "      FALSE\n",
        "    END\n",
        "      AS split_col\n",
        "    FROM\n",
        "      `{training_data_uri}`;\n",
        "    \"\"\"\n",
        "\n",
        "    print(query.replace(\"\\n\", \" \"))\n",
        "\n",
        "    return query"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3332263db93e"
      },
      "source": [
        "### Interpret BQML model evaluation\n",
        "\n",
        "When you do hyperparameter tuning with the model creation query, the output of the prebuilt component [BigqueryEvaluateModelJobOp](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-1.0.0/google_cloud_pipeline_components.experimental.bigquery.html#google_cloud_pipeline_components.experimental.bigquery.BigqueryEvaluateModelJobOp) is a table with the metrics obtained by BQML when training the model. To compare them with those obtained for AutoML model, you need to access them programmatically.\n",
        "\n",
        "\n",
        "The cell below defines a pipeline component that helps you access the metrics. Note that BQML doesn't give you a root mean squared error in the list of metrics. So, you're manually adding it to the metrics dictionary. For more information about the output, see [BQML's documentation](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-evaluate#mlevaluate_output). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ae09c98b32dd"
      },
      "outputs": [],
      "source": [
        "@component(base_image=\"python:3.9\")\n",
        "def interpret_bqml_evaluation_metrics(\n",
        "    bqml_evaluation_metrics: Input[Artifact], metrics: Output[Metrics]\n",
        ") -> dict:\n",
        "    import math\n",
        "\n",
        "    metadata = bqml_evaluation_metrics.metadata\n",
        "    for r in metadata[\"rows\"]:\n",
        "\n",
        "        rows = r[\"f\"]\n",
        "        schema = metadata[\"schema\"][\"fields\"]\n",
        "\n",
        "        output = {}\n",
        "        for metric, value in zip(schema, rows):\n",
        "            metric_name = metric[\"name\"]\n",
        "            val = float(value[\"v\"])\n",
        "            output[metric_name] = val\n",
        "            metrics.log_metric(metric_name, val)\n",
        "            if metric_name == \"mean_squared_error\":\n",
        "                rmse = math.sqrt(val)\n",
        "                metrics.log_metric(\"root_mean_squared_error\", rmse)\n",
        "\n",
        "    metrics.log_metric(\"framework\", \"BQML\")\n",
        "\n",
        "    print(output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5bd1715e98a"
      },
      "source": [
        "### Interpret AutoML model evaluation\n",
        "\n",
        "Similar to BQML, AutoML also generates metrics during its model creation that can be accessed from the Google Cloud console.\n",
        "\n",
        "Since there isn't a prebuilt component to access the AutoML metrics programmatically, you define the below component. The below code uses Vertex AI GAPIC (Google API Compiler) API which auto-generates low-level gRPC interfaces to the specified service.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d0f2927e"
      },
      "outputs": [],
      "source": [
        "# Inspired by Andrew Ferlitsch's work on https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage3/get_started_with_automl_pipeline_components.ipynb\n",
        "\n",
        "\n",
        "@component(\n",
        "    base_image=\"python:3.9\",\n",
        "    packages_to_install=[\n",
        "        \"google-cloud-aiplatform\",\n",
        "    ],\n",
        ")\n",
        "def interpret_automl_evaluation_metrics(\n",
        "    region: str, model: Input[Artifact], metrics: Output[Metrics]\n",
        "):\n",
        "    \"\"\"'\n",
        "    For a list of available regression metrics, go here: gs://google-cloud-aiplatform/schema/modelevaluation/regression_metrics_1.0.0.yaml.\n",
        "\n",
        "    More information on available metrics for different types of models: https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-automl\n",
        "    \"\"\"\n",
        "\n",
        "    import google.cloud.aiplatform.gapic as gapic\n",
        "\n",
        "    # Get a reference to the Model Service client\n",
        "    client_options = {\"api_endpoint\": f\"{region}-aiplatform.googleapis.com\"}\n",
        "\n",
        "    model_service_client = gapic.ModelServiceClient(client_options=client_options)\n",
        "\n",
        "    model_resource_name = model.metadata[\"resourceName\"]\n",
        "\n",
        "    model_evaluations = model_service_client.list_model_evaluations(\n",
        "        parent=model_resource_name\n",
        "    )\n",
        "    model_evaluation = list(model_evaluations)[0]\n",
        "\n",
        "    available_metrics = [\n",
        "        \"meanAbsoluteError\",\n",
        "        \"meanAbsolutePercentageError\",\n",
        "        \"rSquared\",\n",
        "        \"rootMeanSquaredError\",\n",
        "        \"rootMeanSquaredLogError\",\n",
        "    ]\n",
        "    output = dict()\n",
        "    for x in available_metrics:\n",
        "        val = model_evaluation.metrics.get(x)\n",
        "        output[x] = val\n",
        "        metrics.log_metric(str(x), float(val))\n",
        "\n",
        "    metrics.log_metric(\"framework\", \"AutoML\")\n",
        "    print(output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7421c559"
      },
      "source": [
        "### Model selection\n",
        "\n",
        "After the models are evaluated independently, you're going to only move forward with one of them. The selection is done based on the model evaluation metrics gathered in the previous steps. The selected model is then deployed to an endpoint.\n",
        "\n",
        "Define a component to select the best out of the two models that is suitable for deployment. Note that BQML and AutoML use different evaluation metric names, therefore you need to do a mapping of these different nomenclatures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "20d363d9"
      },
      "outputs": [],
      "source": [
        "@component(base_image=\"python:3.9\")\n",
        "def select_best_model(\n",
        "    metrics_bqml: Input[Metrics],\n",
        "    metrics_automl: Input[Metrics],\n",
        "    thresholds_dict_str: str,\n",
        "    best_metrics: Output[Metrics],\n",
        "    reference_metric_name: str = \"rmse\",\n",
        ") -> NamedTuple(\n",
        "    \"Outputs\",\n",
        "    [\n",
        "        (\"deploy_decision\", str),\n",
        "        (\"best_model\", str),\n",
        "        (\"metric\", float),\n",
        "        (\"metric_name\", str),\n",
        "    ],\n",
        "):\n",
        "    import json\n",
        "    from collections import namedtuple\n",
        "\n",
        "    best_metric = float(\"inf\")\n",
        "    best_model = None\n",
        "\n",
        "    # BQML and AutoML use different metric names.\n",
        "    metric_possible_names = []\n",
        "\n",
        "    if reference_metric_name == \"mae\":\n",
        "        metric_possible_names = [\"meanAbsoluteError\", \"mean_absolute_error\"]\n",
        "    elif reference_metric_name == \"rmse\":\n",
        "        metric_possible_names = [\"rootMeanSquaredError\", \"root_mean_squared_error\"]\n",
        "\n",
        "    metric_bqml = float(\"inf\")\n",
        "    metric_automl = float(\"inf\")\n",
        "    print(metrics_bqml.metadata)\n",
        "    print(metrics_automl.metadata)\n",
        "    for x in metric_possible_names:\n",
        "\n",
        "        try:\n",
        "            metric_bqml = metrics_bqml.metadata[x]\n",
        "            print(f\"Metric bqml: {metric_bqml}\")\n",
        "        except:\n",
        "            print(f\"{x} doesn't exist int the BQML dictionary\")\n",
        "\n",
        "        try:\n",
        "            metric_automl = metrics_automl.metadata[x]\n",
        "            print(f\"Metric automl: {metric_automl}\")\n",
        "        except:\n",
        "            print(f\"{x} doesn't exist on the AutoML dictionary\")\n",
        "\n",
        "    # Change condition if higher is better.\n",
        "    print(f\"Comparing BQML ({metric_bqml}) vs AutoML ({metric_automl})\")\n",
        "    if metric_bqml <= metric_automl:\n",
        "        best_model = \"bqml\"\n",
        "        best_metric = metric_bqml\n",
        "        best_metrics.metadata = metrics_bqml.metadata\n",
        "    else:\n",
        "        best_model = \"automl\"\n",
        "        best_metric = metric_automl\n",
        "        best_metrics.metadata = metrics_automl.metadata\n",
        "\n",
        "    thresholds_dict = json.loads(thresholds_dict_str)\n",
        "    deploy = False\n",
        "\n",
        "    # Change condition if higher is better.\n",
        "    if best_metric < thresholds_dict[reference_metric_name]:\n",
        "        deploy = True\n",
        "\n",
        "    if deploy:\n",
        "        deploy_decision = \"true\"\n",
        "    else:\n",
        "        deploy_decision = \"false\"\n",
        "\n",
        "    print(f\"Which model is best? {best_model}\")\n",
        "    print(f\"What metric is being used? {reference_metric_name}\")\n",
        "    print(f\"What is the best metric? {best_metric}\")\n",
        "    print(f\"What is the threshold to deploy? {thresholds_dict_str}\")\n",
        "    print(f\"Deploy decision: {deploy_decision}\")\n",
        "\n",
        "    Outputs = namedtuple(\n",
        "        \"Outputs\", [\"deploy_decision\", \"best_model\", \"metric\", \"metric_name\"]\n",
        "    )\n",
        "\n",
        "    return Outputs(\n",
        "        deploy_decision=deploy_decision,\n",
        "        best_model=best_model,\n",
        "        metric=best_metric,\n",
        "        metric_name=reference_metric_name,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0f573556"
      },
      "source": [
        "### Validate the infrastructure\n",
        "\n",
        "Post selecting the best model, it's deployed to an endpoint. Define a component that validates the endpoint by making prediction requests to that endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aa1bab55"
      },
      "outputs": [],
      "source": [
        "@component(base_image=\"python:3.9\", packages_to_install=[\"google-cloud-aiplatform\"])\n",
        "def validate_infrastructure(\n",
        "    endpoint: Input[Artifact],\n",
        ") -> NamedTuple(\n",
        "    \"validate_infrastructure_output\", [(\"instance\", str), (\"prediction\", float)]\n",
        "):\n",
        "    import json\n",
        "    from collections import namedtuple\n",
        "\n",
        "    from google.cloud import aiplatform\n",
        "    from google.protobuf import json_format\n",
        "    from google.protobuf.struct_pb2 import Value\n",
        "\n",
        "    def treat_uri(uri):\n",
        "        return uri[uri.find(\"projects/\") :]\n",
        "\n",
        "    def request_prediction(endp, instance):\n",
        "        instance = json_format.ParseDict(instance, Value())\n",
        "        instances = [instance]\n",
        "        parameters_dict = {}\n",
        "        parameters = json_format.ParseDict(parameters_dict, Value())\n",
        "        response = endp.predict(instances=instances, parameters=parameters)\n",
        "        print(\"deployed_model_id:\", response.deployed_model_id)\n",
        "        print(\"predictions: \", response.predictions)\n",
        "        # The predictions are a google.protobuf.Value representation of the model's predictions.\n",
        "        predictions = response.predictions\n",
        "\n",
        "        for pred in predictions:\n",
        "            if type(pred) is dict and \"value\" in pred.keys():\n",
        "                # AutoML predictions\n",
        "                prediction = pred[\"value\"]\n",
        "            elif type(pred) is list:\n",
        "                # BQML Predictions return different format\n",
        "                prediction = pred[0]\n",
        "            return prediction\n",
        "\n",
        "    endpoint_uri = endpoint.uri\n",
        "    treated_uri = treat_uri(endpoint_uri)\n",
        "\n",
        "    instance = {\n",
        "        \"Sex\": \"M\",\n",
        "        \"Length\": 0.33,\n",
        "        \"Diameter\": 0.255,\n",
        "        \"Height\": 0.08,\n",
        "        \"Whole_weight\": 0.205,\n",
        "        \"Shucked_weight\": 0.0895,\n",
        "        \"Viscera_weight\": 0.0395,\n",
        "        \"Shell_weight\": 0.055,\n",
        "    }\n",
        "    instance_json = json.dumps(instance)\n",
        "    print(\"Using the following instance: \" + instance_json)\n",
        "\n",
        "    endpoint = aiplatform.Endpoint(treated_uri)\n",
        "    prediction = request_prediction(endpoint, instance)\n",
        "    result_tuple = namedtuple(\n",
        "        \"validate_infrastructure_output\", [\"instance\", \"prediction\"]\n",
        "    )\n",
        "\n",
        "    return result_tuple(instance=str(instance_json), prediction=float(prediction))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpB_bdDbGGOp"
      },
      "source": [
        "## Define the pipeline\n",
        "\n",
        "Now, define the flow of your pipeline using the prebuilt components and the custom components you defined above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f64ccb39400b"
      },
      "outputs": [],
      "source": [
        "@dsl.pipeline(name=DISPLAY_NAME, description=\"Rapid Prototyping\")\n",
        "def train_pipeline(\n",
        "    project: str,\n",
        "    gcs_input_file_uri: str,\n",
        "    region: str,\n",
        "    bq_dataset: str,\n",
        "    bq_location: str,\n",
        "    bqml_model_export_location: str,\n",
        "    bqml_serving_container_image_uri: str,\n",
        "    endpoint_display_name: str,\n",
        "    thresholds_dict_str: str,\n",
        "):\n",
        "    from google_cloud_pipeline_components.types import artifact_types\n",
        "    from kfp.dsl import importer_node\n",
        "\n",
        "    # Imports data to BigQuery using a custom component.\n",
        "    import_data_to_bigquery_op = import_data_to_bigquery(\n",
        "        project=project,\n",
        "        bq_location=bq_location,\n",
        "        bq_dataset=bq_dataset,\n",
        "        gcs_data_uri=gcs_input_file_uri,\n",
        "    )\n",
        "    raw_dataset = import_data_to_bigquery_op.outputs[\"raw_dataset\"]\n",
        "\n",
        "    # Splits the BQ dataset using a custom component.\n",
        "    split_datasets_op = split_datasets(raw_dataset=raw_dataset, bq_location=bq_location)\n",
        "\n",
        "    # Generates the query to create a BQML using a static function.\n",
        "    create_model_query = _query_create_model(\n",
        "        project_id=project,\n",
        "        bq_dataset=bq_dataset,\n",
        "        training_data_uri=split_datasets_op.outputs[\"dataset_uri\"],\n",
        "    )\n",
        "\n",
        "    # Builds BQML model using prebuilt component.\n",
        "    bqml_create_op = bq_components.BigqueryCreateModelJobOp(\n",
        "        project=project, location=bq_location, query=create_model_query\n",
        "    )\n",
        "    bqml_model = bqml_create_op.outputs[\"model\"]\n",
        "\n",
        "    # Gathers BQML evaluation metrics using a prebuilt component.\n",
        "    bqml_evaluate_op = bq_components.BigqueryEvaluateModelJobOp(\n",
        "        project=project, location=bq_location, model=bqml_model\n",
        "    )\n",
        "    bqml_eval_metrics_raw = bqml_evaluate_op.outputs[\"evaluation_metrics\"]\n",
        "\n",
        "    # Analyzes evaluation BQML metrics using a custom component.\n",
        "    interpret_bqml_evaluation_metrics_op = interpret_bqml_evaluation_metrics(\n",
        "        bqml_evaluation_metrics=bqml_eval_metrics_raw\n",
        "    )\n",
        "    bqml_eval_metrics = interpret_bqml_evaluation_metrics_op.outputs[\"metrics\"]\n",
        "\n",
        "    # Exports the BQML model to a GCS bucket using a prebuilt component.\n",
        "    bqml_export_op = bq_components.BigqueryExportModelJobOp(\n",
        "        project=project,\n",
        "        location=bq_location,\n",
        "        model=bqml_model,\n",
        "        model_destination_path=bqml_model_export_location,\n",
        "    ).after(bqml_evaluate_op)\n",
        "    bqml_exported_gcs_path = bqml_export_op.outputs[\"exported_model_path\"]\n",
        "\n",
        "    unmanaged_model_importer = importer_node.importer(\n",
        "        artifact_uri=bqml_exported_gcs_path,\n",
        "        artifact_class=artifact_types.UnmanagedContainerModel,\n",
        "        metadata={\n",
        "            \"containerSpec\": {\n",
        "                \"imageUri\": \"us-docker.pkg.dev/cloud-aiplatform/prediction/tf2-cpu.2-3:latest\"\n",
        "            }\n",
        "        },\n",
        "    )\n",
        "\n",
        "    # Uploads the recently exported the BQML model from GCS into Vertex AI using a prebuilt component.\n",
        "    bqml_model_upload_op = ModelUploadOp(\n",
        "        project=project,\n",
        "        location=region,\n",
        "        display_name=DISPLAY_NAME + \"_bqml\",\n",
        "        unmanaged_container_model=unmanaged_model_importer.outputs[\"artifact\"],\n",
        "    )\n",
        "    bqml_vertex_model = bqml_model_upload_op.outputs[\"model\"]\n",
        "\n",
        "    # Creates a Vertex AI Tabular dataset using a prebuilt component.\n",
        "    dataset_create_op = TabularDatasetCreateOp(\n",
        "        project=project,\n",
        "        location=region,\n",
        "        display_name=DISPLAY_NAME,\n",
        "        bq_source=split_datasets_op.outputs[\"dataset_bq_uri\"],\n",
        "    )\n",
        "\n",
        "    # Trains an AutoML Tables model using a prebuilt component.\n",
        "    automl_training_op = AutoMLTabularTrainingJobRunOp(\n",
        "        project=project,\n",
        "        location=region,\n",
        "        display_name=f\"{DISPLAY_NAME}_automl\",\n",
        "        optimization_prediction_type=\"regression\",\n",
        "        optimization_objective=\"minimize-rmse\",\n",
        "        predefined_split_column_name=\"split_col\",\n",
        "        dataset=dataset_create_op.outputs[\"dataset\"],\n",
        "        target_column=\"Rings\",\n",
        "        budget_milli_node_hours=1000,\n",
        "        column_transformations=[\n",
        "            {\"categorical\": {\"column_name\": \"Sex\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Length\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Diameter\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Height\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Whole_weight\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Shucked_weight\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Viscera_weight\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Shell_weight\"}},\n",
        "            {\"numeric\": {\"column_name\": \"Rings\"}},\n",
        "        ],\n",
        "    )\n",
        "    automl_model = automl_training_op.outputs[\"model\"]\n",
        "\n",
        "    # Analyzes evaluation AutoML metrics using a custom component.\n",
        "    automl_eval_op = interpret_automl_evaluation_metrics(\n",
        "        region=region, model=automl_model\n",
        "    )\n",
        "    automl_eval_metrics = automl_eval_op.outputs[\"metrics\"]\n",
        "\n",
        "    # 1) Decides which model is best (AutoML vs BQML);\n",
        "    # 2) Determines if the best model meets the deployment condition.\n",
        "    best_model_task = select_best_model(\n",
        "        metrics_bqml=bqml_eval_metrics,\n",
        "        metrics_automl=automl_eval_metrics,\n",
        "        thresholds_dict_str=thresholds_dict_str,\n",
        "    )\n",
        "\n",
        "    # If the deploy condition is True, then deploy the best model.\n",
        "    with dsl.If(\n",
        "        best_model_task.outputs[\"deploy_decision\"] == \"true\",\n",
        "        name=\"deploy_decision\",\n",
        "    ):\n",
        "        # Creates a Vertex AI endpoint using a prebuilt component.\n",
        "        endpoint_create_op = EndpointCreateOp(\n",
        "            project=project,\n",
        "            location=region,\n",
        "            display_name=endpoint_display_name,\n",
        "        )\n",
        "        endpoint_create_op.after(best_model_task)\n",
        "\n",
        "        # In case the BQML model is the best...\n",
        "        with dsl.If(\n",
        "            best_model_task.outputs[\"best_model\"] == \"bqml\",\n",
        "            name=\"deploy_bqml\",\n",
        "        ):\n",
        "            # Deploys the BQML model (now on Vertex AI) to the recently created endpoint using a prebuilt component.\n",
        "            model_deploy_bqml_op = ModelDeployOp(  # noqa: F841\n",
        "                endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
        "                model=bqml_vertex_model,\n",
        "                deployed_model_display_name=DISPLAY_NAME + \"_best_bqml\",\n",
        "                dedicated_resources_machine_type=\"n1-standard-2\",\n",
        "                dedicated_resources_min_replica_count=2,\n",
        "                dedicated_resources_max_replica_count=2,\n",
        "                traffic_split={\n",
        "                    \"0\": 100\n",
        "                },  # newly deployed model gets 100% of the traffic\n",
        "            ).set_caching_options(False)\n",
        "\n",
        "            # Sends an online prediction request to the recently deployed model using a custom component.\n",
        "            validate_infrastructure(\n",
        "                endpoint=endpoint_create_op.outputs[\"endpoint\"]\n",
        "            ).set_caching_options(False).after(model_deploy_bqml_op)\n",
        "\n",
        "        # In case the AutoML model is the best...\n",
        "        with dsl.If(\n",
        "            best_model_task.outputs[\"best_model\"] == \"automl\",\n",
        "            name=\"deploy_automl\",\n",
        "        ):\n",
        "            # Deploys the AutoML model to the recently created endpoint using a prebuilt component.\n",
        "            model_deploy_automl_op = ModelDeployOp(  # noqa: F841\n",
        "                endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
        "                model=automl_model,\n",
        "                deployed_model_display_name=DISPLAY_NAME + \"_best_automl\",\n",
        "                dedicated_resources_machine_type=\"n1-standard-2\",\n",
        "                dedicated_resources_min_replica_count=2,\n",
        "                dedicated_resources_max_replica_count=2,\n",
        "                traffic_split={\n",
        "                    \"0\": 100\n",
        "                },  # newly deployed model gets 100% of the traffic\n",
        "            ).set_caching_options(False)\n",
        "\n",
        "            # Sends an online prediction request to the recently deployed model using a custom component.\n",
        "            validate_infrastructure(\n",
        "                endpoint=endpoint_create_op.outputs[\"endpoint\"]\n",
        "            ).set_caching_options(False).after(model_deploy_automl_op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1b01f57507f7"
      },
      "source": [
        "## Compile the pipeline\n",
        "\n",
        "Compile and save your pipeline to a local file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d69a2f630792"
      },
      "outputs": [],
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=train_pipeline,\n",
        "    package_path=PIPELINE_YAML_PKG_PATH,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e8b6af4f85f1"
      },
      "source": [
        "## Specify the pipeline parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O5PsR31ysGuj"
      },
      "outputs": [],
      "source": [
        "# Specify the input parameters to your pipeline\n",
        "pipeline_params = {\n",
        "    \"project\": PROJECT_ID,\n",
        "    \"region\": LOCATION,\n",
        "    \"gcs_input_file_uri\": RAW_INPUT_DATA,\n",
        "    \"bq_dataset\": BQ_DATASET,\n",
        "    \"bq_location\": BQ_LOCATION,\n",
        "    \"bqml_model_export_location\": BQML_EXPORT_LOCATION,\n",
        "    \"bqml_serving_container_image_uri\": BQML_SERVING_CONTAINER_IMAGE_URI,\n",
        "    \"endpoint_display_name\": ENDPOINT_DISPLAY_NAME,\n",
        "    \"thresholds_dict_str\": '{\"rmse\": 2.5}',\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qxfy-pXXGS3R"
      },
      "source": [
        "## Run the pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PLeS1xRpGYPx"
      },
      "outputs": [],
      "source": [
        "# Create a pipeline job\n",
        "pipeline_job = vertex.PipelineJob(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    template_path=PIPELINE_YAML_PKG_PATH,\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    parameter_values=pipeline_params,\n",
        "    enable_caching=False,\n",
        ")\n",
        "# Submit your pipeline job\n",
        "response = pipeline_job.submit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f49899d5e838"
      },
      "source": [
        "### Wait for the pipeline to complete\n",
        "\n",
        "When you use the `submit()` method, your pipeline runs in asynchronous mode.  To block the execution until your job gets completed, use the `wait()` method.\n",
        "\n",
        "**Note**: To run your pipeline synchronously, you can use the `run()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vo0hPcsraVIt"
      },
      "outputs": [],
      "source": [
        "pipeline_job.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lGKH0lKwz7Ci"
      },
      "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:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "41229df87d23"
      },
      "outputs": [],
      "source": [
        "# Delete Vertex AI endpoint\n",
        "print(\"Deleting endpoint...\")\n",
        "endpoints = vertex.Endpoint.list(\n",
        "    filter=f\"display_name={DISPLAY_NAME}_endpoint\", order_by=\"create_time\"\n",
        ")\n",
        "endpoint = endpoints[0]\n",
        "endpoint.undeploy_all()\n",
        "vertex.Endpoint.delete(endpoint)\n",
        "print(\"Deleted endpoint:\", endpoint)\n",
        "\n",
        "# Delete BQML and AutoML models\n",
        "print(\"Deleting models...\")\n",
        "suffix_list = [\"bqml\", \"automl\"]\n",
        "for suffix in suffix_list:\n",
        "    try:\n",
        "        model_display_name = f\"{DISPLAY_NAME}_{suffix}\"\n",
        "        print(\"Deleting model with name: \" + model_display_name)\n",
        "        models = vertex.Model.list(\n",
        "            filter=f\"display_name={model_display_name}\", order_by=\"create_time\"\n",
        "        )\n",
        "\n",
        "        model = models[0]\n",
        "        vertex.Model.delete(model)\n",
        "        print(\"Deleted model:\", model)\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
        "# Delete Vertex AI dataset\n",
        "print(\"Deleting Vertex AI dataset...\")\n",
        "datasets = vertex.TabularDataset.list(\n",
        "    filter=f\"display_name={DISPLAY_NAME}\", order_by=\"create_time\"\n",
        ")\n",
        "dataset = datasets[0]\n",
        "vertex.TabularDataset.delete(dataset)\n",
        "print(\"Deleted Vertex AI dataset:\", dataset)\n",
        "\n",
        "# Delete Vertex AI pipline job\n",
        "print(\"Deleting pipeline...\")\n",
        "pipeline_job.delete()\n",
        "print(\"Deleted pipeline:\", pipeline_job)\n",
        "\n",
        "# Delete BigQuery dataset\n",
        "delete_dataset = True\n",
        "if delete_dataset:\n",
        "    ! bq rm -r -f -d $PROJECT_ID:$BQ_DATASET\n",
        "\n",
        "dataset_id = f\"{PROJECT_ID}.{BQ_DATASET}\"\n",
        "print(f\"Deleted BQ dataset '{dataset_id}' from location {BQ_LOCATION}.\")\n",
        "\n",
        "# Delete Cloud Storage bucket\n",
        "delete_bucket = True\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
        "# Delete the pipeline package file\n",
        "! rm PIPELINE_YAML_PKG_PATH"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "rapid_prototyping_bqml_automl.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}