model_garden/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 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": [
        "# Vertex AI: Track artifacts and metrics across Vertex AI Pipelines runs using Vertex ML Metadata\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.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/main/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.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://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb\" target='_blank'>\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": "e88691377fcc"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to track metrics and artifacts across Vertex AI Pipelines runs, and analyze this metadata using the Vertex AI SDK. If you'd prefer to follow a step-by-step tutorial, check out the [codelab version](https://codelabs.developers.google.com/vertex-mlmd-pipelines#0) of this notebook.\n",
        "\n",
        "Learn more about [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata) and [Vertex AI Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to track artifacts and metrics with `Vertex ML Metadata` in `Vertex AI Pipeline` runs.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Pipelines\n",
        "- Vertex ML Metadata\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "* Use the Kubeflow Pipelines SDK to build an ML pipeline that runs on Vertex AI\n",
        "* The pipeline will create a dataset, train a scikit-learn model, and deploy the model to an endpoint\n",
        "* Write custom pipeline components that generate artifacts and metadata\n",
        "* Compare Vertex Pipelines runs, both in the Cloud console and programmatically\n",
        "* Trace the lineage for pipeline-generated artifacts\n",
        "* Query your pipeline run metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ce1e72673981"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "In this notebook, we will train a model using scikit-learn to classify bean types using the [Dry Beans Dataset](https://archive.ics.uci.edu/ml/datasets/Dry+Bean+Dataset) from UCI Machine Learning. This is a tabular dataset that includes measurements and characteristics of seven different types of beans taken from images."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c997d8d92ce"
      },
      "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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ze4-nDLfK4pw"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or AI Platform Notebooks**, your environment already meets\n",
        "all the requirements to run this notebook. You can skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCuSR8GkAgzl"
      },
      "source": [
        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
        "You need the following:\n",
        "\n",
        "* The Google Cloud SDK\n",
        "* Git\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps provide a condensed set of\n",
        "instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Run the following commands to install the Vertex AI SDK and packages used in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IaYsrh0Tc17L"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Google Cloud Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# Google Cloud Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCQDRsnE3uzz"
      },
      "source": [
        "Install Vertex AI and Kubeflow Pipelines SDKs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyy5Lbnzg5fi"
      },
      "outputs": [],
      "source": [
        "!pip3 install {USER_FLAG} google-cloud-aiplatform==1.7.0\n",
        "!pip3 install {USER_FLAG} kfp==1.8.9"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhq5zEbGg0XX"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lWEdiXsJg0XY"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. Enable the services we'll be using throughout this notebook by running the cell below.\n",
        "\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the project ID cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e86205a30eb4"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using AI Platform Notebooks**, your environment is already\n",
        "authenticated. Skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28f75ab2b551"
      },
      "source": [
        "**If you are using Colab**, run the cell below and follow the instructions\n",
        "when prompted to authenticate your account via oAuth.\n",
        "\n",
        "**Otherwise**, follow these steps:\n",
        "\n",
        "1. In the Cloud Console, go to the [**Create service account key**\n",
        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "\n",
        "2. Click **Create service account**.\n",
        "\n",
        "3. In the **Service account name** field, enter a name, and\n",
        "   click **Create**.\n",
        "\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"AI Platform\"\n",
        "into the filter box, and select\n",
        "   **AI Platform Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
        "\n",
        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
        "local environment.\n",
        "\n",
        "6. Enter the path to your service account key as the\n",
        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2b00dc291f7"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a6a066dd8d6a"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"\"\n",
        "\n",
        "# Get your Google Cloud project ID from gcloud\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID: \", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dbeae85aef1d"
      },
      "source": [
        "Otherwise, set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f90e4cbfb7af"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fdb343d2ae3c"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "# If you are running this notebook in Colab, run this cell and follow the\n",
        "# instructions to authenticate your GCP account. This provides access to your\n",
        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
        "# requests.\n",
        "\n",
        "# If on Google Cloud Notebooks, then don't execute this code\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "    if \"google.colab\" in sys.modules:\n",
        "        from google.colab import auth as google_auth\n",
        "\n",
        "        google_auth.authenticate_user()\n",
        "        !gcloud config set project $PROJECT_ID\n",
        "\n",
        "    # If you are running this notebook locally, replace the string below with the\n",
        "    # path to your service account key and run this cell to authenticate your GCP\n",
        "    # account.\n",
        "    elif not os.getenv(\"IS_TESTING\"):\n",
        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aab852d94fc7"
      },
      "source": [
        "#### Enable Cloud services used throughout this notebook.\n",
        "\n",
        "Run the cell below to the enable Compute Engine, Container Registry, and Vertex AI services."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "18396d3d7fe4"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable compute.googleapis.com         \\\n",
        "                       containerregistry.googleapis.com  \\\n",
        "                       aiplatform.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "#### Timestamp\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 timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "697568e92bd6"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "641c8b2873c0"
      },
      "source": [
        "### Create a Cloud Storage Bucket"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68e74d218ea9"
      },
      "source": [
        "To run our Vertex Pipeline, we'll need a storage bucket to store artifacts generated by our pipeline. This bucket needs to be regional. We're using the `us-central1` region here, but you are welcome to use another region (just replace it throughout this lab). If you already have a bucket you can replace the `BUCKET_NAME` variable with the name of your bucket and skip the `gsutil mb` step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "953a1399e79f"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"gs://{}-bucket\".format(PROJECT_ID)\n",
        "!gsutil mb -l us-central1 $BUCKET_NAME # You only need to run this once"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b75e9153c699"
      },
      "source": [
        "Next, make sure your compute service account has `store.objectAdmin` access to this bucket. Your compute service account will look something like `YOUR_PROJECT_NUMBER-compute@developer.gserviceaccount.com`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y9Uo3tifg1kx"
      },
      "source": [
        "Import required libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "# We'll use this beta library for metadata querying\n",
        "from google.cloud import aiplatform, aiplatform_v1beta1\n",
        "from google.cloud.aiplatform import pipeline_jobs\n",
        "from kfp.v2 import compiler, dsl\n",
        "from kfp.v2.dsl import (Artifact, Dataset, Input, Metrics, Model, Output,\n",
        "                        OutputPath, component)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xtXZWmYqJ1bh"
      },
      "source": [
        "Define some constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JIOrI-hoJ46P"
      },
      "outputs": [],
      "source": [
        "PATH = get_ipython().run_line_magic(\"env\", \"PATH\")\n",
        "%env PATH={PATH}:/home/jupyter/.local/bin\n",
        "REGION = \"us-central1\"\n",
        "\n",
        "PIPELINE_ROOT = f\"{BUCKET_NAME}/pipeline_root/\"\n",
        "PIPELINE_ROOT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2937d462a96a"
      },
      "source": [
        "Initialize the Vertex AI SDK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7def96de8098"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xuny18aMcWDb"
      },
      "source": [
        "## Concepts\n",
        "\n",
        "To better understand [Vertex Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction) and [ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata), we'd like to introduce the following concepts:\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NThDci5bp0Uw"
      },
      "source": [
        "### Pipeline Run\n",
        "When we use the term run, we're referring to a single execution of your pipeline in Vertex Pipelines. Each run generates artifacts, metrics, and associated metadata."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SAyRR3Ydp4X5"
      },
      "source": [
        "### Artifact\n",
        "\n",
        "An artifact is a resource generated by your pipeline. Artifacts could include datasets, models, endpoints, or custom resources defined in your pipeline."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40ee71479689"
      },
      "source": [
        "### Metric\n",
        "\n",
        "A metric is a way to measure the performance of your pipeline runs and artifacts. For example, a metric could be the accuracy of a classification model artifact created in your pipeline, or the size of the dataset used to train your model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "57b1cc9981d5"
      },
      "source": [
        "### Metadata\n",
        "\n",
        "Metadata describes the artifacts and metrics generated by your pipeline runs. Metadata on a model, for example, could include the URL of the model artifacts, its name, and the time it was created."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l1YW2pgyegFP"
      },
      "source": [
        "## Creating a 3-step pipeline with custom components\n",
        "\n",
        "The focus of this lab is on understanding *metadata* from pipeline runs. In order to do that, we'll need a pipeline to run on Vertex Pipelines, which is where we'll start. Here we'll define a 3-step pipeline with the following custom components:\n",
        "\n",
        "* `get_dataframe`: Retrieve data from a BigQuery table and convert it into a pandas DataFrame\n",
        "* `train_sklearn_model`: Use the pandas DataFrame to train and export a scikit-learn model, along with some metrics\n",
        "* `deploy_model`: Deploy the exported scikit-learn model to an endpoint in Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KPY41M9_AhZU"
      },
      "source": [
        "### Create and define Python function based components"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bfMQSmRuUuX-"
      },
      "source": [
        "First, define the `get_dataframe` component with the code below. This component does the following:\n",
        "* Creates a reference to a BigQuery table using the BigQuery client library\n",
        "* Downloads the BigQuery table and converts it to a shuffled pandas DataFrame\n",
        "* Exports the DataFrame to a CSV file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RiQuMv4bmpuV"
      },
      "outputs": [],
      "source": [
        "@component(\n",
        "    packages_to_install=[\"google-cloud-bigquery\", \"pandas\", \"pyarrow\"],\n",
        "    base_image=\"python:3.9\",\n",
        "    output_component_file=\"create_dataset.yaml\",\n",
        ")\n",
        "def get_dataframe(bq_table: str, output_data_path: OutputPath(\"Dataset\")):\n",
        "    from google.cloud import bigquery\n",
        "\n",
        "    bqclient = bigquery.Client(project=PROJECT_ID)\n",
        "    table = bigquery.TableReference.from_string(bq_table)\n",
        "    rows = bqclient.list_rows(table)\n",
        "    dataframe = rows.to_dataframe(\n",
        "        create_bqstorage_client=True,\n",
        "    )\n",
        "    dataframe = dataframe.sample(frac=1, random_state=2)\n",
        "    dataframe.to_csv(output_data_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y06J7A7yU21t"
      },
      "source": [
        "Next, create a component to train a scikit-learn model. This component does the following:\n",
        "* Imports a CSV as a pandas DataFrame\n",
        "* Splits the DataFrame into train and test sets\n",
        "* Trains a scikit-learn model\n",
        "* Logs metrics from the model\n",
        "* Saves the model artifacts as a local `model.joblib` file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p5JBCBKyH-NC"
      },
      "outputs": [],
      "source": [
        "@component(\n",
        "    packages_to_install=[\"sklearn\", \"pandas\", \"joblib\"],\n",
        "    base_image=\"python:3.9\",\n",
        "    output_component_file=\"beans_model_component.yaml\",\n",
        ")\n",
        "def sklearn_train(\n",
        "    dataset: Input[Dataset], metrics: Output[Metrics], model: Output[Model]\n",
        "):\n",
        "    import pandas as pd\n",
        "    from joblib import dump\n",
        "    from sklearn.model_selection import train_test_split\n",
        "    from sklearn.tree import DecisionTreeClassifier\n",
        "\n",
        "    df = pd.read_csv(dataset.path)\n",
        "    labels = df.pop(\"Class\").tolist()\n",
        "    data = df.values.tolist()\n",
        "    x_train, x_test, y_train, y_test = train_test_split(data, labels)\n",
        "\n",
        "    skmodel = DecisionTreeClassifier()\n",
        "    skmodel.fit(x_train, y_train)\n",
        "    score = skmodel.score(x_test, y_test)\n",
        "    print(\"accuracy is:\", score)\n",
        "\n",
        "    metrics.log_metric(\"accuracy\", (score * 100.0))\n",
        "    metrics.log_metric(\"framework\", \"Scikit Learn\")\n",
        "    metrics.log_metric(\"dataset_size\", len(df))\n",
        "    dump(skmodel, model.path + \".joblib\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gaNNTFPaU7KT"
      },
      "source": [
        "Finally, our last component will take the trained model from the previous step, upload it to Vertex AI, and deploy it to an endpoint:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VGq5QCoyIEWJ"
      },
      "outputs": [],
      "source": [
        "@component(\n",
        "    packages_to_install=[\"google-cloud-aiplatform\"],\n",
        "    base_image=\"python:3.9\",\n",
        "    output_component_file=\"beans_deploy_component.yaml\",\n",
        ")\n",
        "def deploy_model(\n",
        "    model: Input[Model],\n",
        "    project: str,\n",
        "    region: str,\n",
        "    vertex_endpoint: Output[Artifact],\n",
        "    vertex_model: Output[Model],\n",
        "):\n",
        "    from google.cloud import aiplatform\n",
        "\n",
        "    aiplatform.init(project=project, location=region)\n",
        "\n",
        "    deployed_model = aiplatform.Model.upload(\n",
        "        display_name=\"beans-model-pipeline\",\n",
        "        artifact_uri=model.uri.replace(\"model\", \"\"),\n",
        "        serving_container_image_uri=\"us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest\",\n",
        "    )\n",
        "    endpoint = deployed_model.deploy(machine_type=\"n1-standard-4\")\n",
        "\n",
        "    # Save data to the output params\n",
        "    vertex_endpoint.uri = endpoint.resource_name\n",
        "    vertex_model.uri = deployed_model.resource_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UBXUgxgqA_GB"
      },
      "source": [
        "### Define and compile the pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "66odBYKrIN4q"
      },
      "outputs": [],
      "source": [
        "@dsl.pipeline(\n",
        "    # Default pipeline root. You can override it when submitting the pipeline.\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    # A name for the pipeline.\n",
        "    name=\"mlmd-pipeline\",\n",
        ")\n",
        "def pipeline(\n",
        "    bq_table: str = \"\",\n",
        "    output_data_path: str = \"data.csv\",\n",
        "    project: str = PROJECT_ID,\n",
        "    region: str = REGION,\n",
        "):\n",
        "    dataset_task = get_dataframe(bq_table)\n",
        "\n",
        "    model_task = sklearn_train(dataset_task.output)\n",
        "\n",
        "    deploy_model(model=model_task.outputs[\"model\"], project=project, region=region)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "910541af051c"
      },
      "source": [
        "The following will generate a JSON file that you'll use to run the pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o_wnT10RJ7-W"
      },
      "outputs": [],
      "source": [
        "compiler.Compiler().compile(pipeline_func=pipeline, package_path=\"mlmd_pipeline.json\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u-iTnzt3B6Z_"
      },
      "source": [
        "### Start two pipeline runs\n",
        "\n",
        "Next we'll kick off **two** runs of our pipeline. First let's define a timestamp to use for our pipeline job IDs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i2wnpu8_7JfV"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3d380ed72490"
      },
      "source": [
        "Our pipeline takes one parameter when we run it: the `bq_table` we want to use for training data. This pipeline run will use a smaller version of the beans dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ff4aee966c5f"
      },
      "outputs": [],
      "source": [
        "run1 = pipeline_jobs.PipelineJob(\n",
        "    display_name=\"mlmd-pipeline\",\n",
        "    template_path=\"mlmd_pipeline.json\",\n",
        "    job_id=\"mlmd-pipeline-small-{}\".format(TIMESTAMP),\n",
        "    parameter_values={\"bq_table\": \"sara-vertex-demos.beans_demo.small_dataset\"},\n",
        "    enable_caching=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "555ac88a22cf"
      },
      "source": [
        "Next, create another pipeline run using a larger version of the same dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3d9fcb6a4a9e"
      },
      "outputs": [],
      "source": [
        "run2 = pipeline_jobs.PipelineJob(\n",
        "    display_name=\"mlmd-pipeline\",\n",
        "    template_path=\"mlmd_pipeline.json\",\n",
        "    job_id=\"mlmd-pipeline-large-{}\".format(TIMESTAMP),\n",
        "    parameter_values={\"bq_table\": \"sara-vertex-demos.beans_demo.large_dataset\"},\n",
        "    enable_caching=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5670722f7668"
      },
      "source": [
        "Finally, kick off pipeline executions for both runs. It's best to do this in two separate notebook cells so you can see the output for each run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1f477f5565c6"
      },
      "outputs": [],
      "source": [
        "run1.submit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6e682e41af78"
      },
      "source": [
        "Then, kick off the second run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cb263e503ced"
      },
      "outputs": [],
      "source": [
        "run2.submit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cc15017be48e"
      },
      "source": [
        "After running this cell, you'll see a link to view each pipeline in the Vertex AI console. Open that link to see more details on your pipeline.\n",
        "\n",
        "**These pipeline runs will take 10-15 minutes to complete.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jZLrJZTfL7tE"
      },
      "source": [
        "## Comparing pipeline runs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A1PqKxlpOZa2"
      },
      "source": [
        "Now that you have two pipeline completed pipeline runs, we're ready to take a closer look at pipeline metrics using the Vertex AI SDK.\n",
        "\n",
        "**For guidance on inspecting pipeline artifacts and metadata in the Vertex AI Console, see [this codelab](https://codelabs.developers.google.com/vertex-mlmd-pipelines#5).**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jbRf1WoH_vbY"
      },
      "source": [
        "You can use the `aiplatform.get_pipeline_df()` method to access run metadata. Here, we'll get metadata for the last two runs of the same pipeline and load it into a Pandas DataFrame. The `mlmd-pipeline` parameter here refers to the name we gave our pipeline in our pipeline definition:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "90d850cda34f"
      },
      "outputs": [],
      "source": [
        "df = aiplatform.get_pipeline_df(pipeline=\"mlmd-pipeline\")\n",
        "print(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d23e2cb66265"
      },
      "source": [
        "We've only executed our pipeline twice here, but you can imagine how many metrics you'd have with more executions. Next, we'll create a custom visualization with matplotlib to see the relationship between our model's accuracy and the amount of data used for training. Run the following to generate a graph:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5957415cc390"
      },
      "outputs": [],
      "source": [
        "plt.plot(df[\"metric.dataset_size\"], df[\"metric.accuracy\"], label=\"Accuracy\")\n",
        "plt.title(\"Accuracy and dataset size\")\n",
        "plt.legend(loc=4)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EYuYgqVCMKU1"
      },
      "source": [
        "## Querying pipeline metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4431b5d062f3"
      },
      "source": [
        "In addition to getting a DataFrame of all pipeline metrics, you may want to programmatically query artifacts created in your ML system. From there you could create a custom dashboard or let others in your organizaiton get details on specific artifacts."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "995723757c5d"
      },
      "source": [
        "### Getting all Model artifacts\n",
        "\n",
        "To query artifacts in this way, we'll create a `MetadataServiceClient`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r8orCj8iJuO1"
      },
      "outputs": [],
      "source": [
        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
        "metadata_client = aiplatform_v1beta1.MetadataServiceClient(\n",
        "    client_options={\"api_endpoint\": API_ENDPOINT}\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5aee9cdc5bd"
      },
      "source": [
        "Next, we'll make a `list_artifacts` request to that endpoint and pass a filter indicating which artifacts we'd like in our response. First, let's get all the artifacts in our project that are **models**. To do that, run the following in your notebook:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29260057ae40"
      },
      "outputs": [],
      "source": [
        "MODEL_FILTER = 'schema_title = \"system.Model\"'\n",
        "artifact_request = aiplatform_v1beta1.ListArtifactsRequest(\n",
        "    parent=\"projects/{}/locations/{}/metadataStores/default\".format(PROJECT_ID, REGION),\n",
        "    filter=MODEL_FILTER,\n",
        ")\n",
        "model_artifacts = metadata_client.list_artifacts(artifact_request)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dfb57f1b7833"
      },
      "source": [
        "The resulting `model_artifacts` response contains an iterable object for each model artifact in your project, along with associated metadata for each model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WTHvPMweMlP1"
      },
      "source": [
        "### Filtering objects and displaying in a DataFrame"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F19_5lw0MqXv"
      },
      "source": [
        "It would be handy if we could more easily visualize the resulting artifact query. Next, let's get all artifacts created after August 10, 2021 with a `LIVE` state. After we run this request, we'll display the results in a pandas DataFrame. First, execute the request:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GmN9vE9pqqzt"
      },
      "outputs": [],
      "source": [
        "LIVE_FILTER = 'create_time > \"2021-08-10T00:00:00-00:00\" AND state = LIVE'\n",
        "artifact_req = {\n",
        "    \"parent\": \"projects/{}/locations/{}/metadataStores/default\".format(\n",
        "        PROJECT_ID, REGION\n",
        "    ),\n",
        "    \"filter\": LIVE_FILTER,\n",
        "}\n",
        "live_artifacts = metadata_client.list_artifacts(artifact_req)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6bba2012b7f0"
      },
      "source": [
        "Then, display the results in a DataFrame:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6bee5790cec4"
      },
      "outputs": [],
      "source": [
        "data = {\"uri\": [], \"createTime\": [], \"type\": []}\n",
        "\n",
        "for i in live_artifacts:\n",
        "    data[\"uri\"].append(i.uri)\n",
        "    data[\"createTime\"].append(i.create_time)\n",
        "    data[\"type\"].append(i.schema_title)\n",
        "\n",
        "df = pd.DataFrame.from_dict(data)\n",
        "print(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "If you don't want to delete the project, do the following to clean up the resources you used:\n",
        "\n",
        "* If you used Google Cloud Notebooks to run this, stop or delete the notebook instance\n",
        "\n",
        "* The pipeline runs we executed deployed endpoints in Vertex AI. Navigate to the [Vertex AI console](https://console.cloud.google.com/vertex-ai/endpoints) to delete those endpoints\n",
        "\n",
        "* Delete the [Cloud Storage bucket](https://console.cloud.google.com/storage/browser/) you created"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "vertex-pipelines-ml-metadata.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}