notebooks/community/ml_ops/stage4/get_started_with_vertex_ml_metadata.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "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": "title:generic,gcp"
      },
      "source": [
        "# Get started with Vertex ML Metadata\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/get_started_with_vertex_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/get_started_with_vertex_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/get_started_with_vertex_ml_metadata.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>            \n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use Vertex ML Metadata.\n",
        "\n",
        "Learn more about [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage4,get_started_vertex_ml_metadata"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use `Vertex ML Metadata`.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex ML Metadata`\n",
        "- `Vertex AI Pipelines`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a `Metadatastore` resource.\n",
        "- Create (record)/List an `Artifact`, with artifacts and metadata.\n",
        "- Create (record)/List an `Execution`.\n",
        "- Create (record)/List a `Context`.\n",
        "- Add `Artifact` to `Execution` as events.\n",
        "- Add `Execution` and `Artifact` into the `Context`\n",
        "- Delete `Artifact`, `Execution` and `Context`.\n",
        "- Create and run a `Vertex AI Pipeline` ML workflow to train and deploy a scikit-learn model.\n",
        "    - Create custom pipeline components that generate artifacts and metadata.\n",
        "    - Compare Vertex AI Pipelines runs.\n",
        "    - Trace the lineage for pipeline-generated artifacts.\n",
        "    - Query your pipeline run metadata."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:beans,lcn"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the UCI Machine Learning ['Dry beans dataset'](https://archive.ics.uci.edu/ml/datasets/Dry+Bean+Dataset), from: KOKLU, M. and OZKAN, I.A., (2020), \"Multiclass Classification of Dry Beans Using Computer Vision and Machine Learning Techniques.\"In Computers and Electronics in Agriculture, 174, 105507. [DOI](https://doi.org/10.1016/j.compag.2020.105507)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "35bee437737d"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_mlops"
      },
      "source": [
        "## Installations\n",
        "\n",
        "Install the packages required for executing the notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "install_mlops"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "! pip3 install --upgrade google-cloud-aiplatform[tensorboard] \\\n",
        "                         google-cloud-pipeline-components  --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-ZBOjErv5mM"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin:nogpu"
      },
      "source": [
        "### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_project_id"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcp_authenticate"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FvQeFm3Gv5mR"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ad1138a125ea"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce6043da7b33"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0367eac06a10"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ad4dbb4a61"
      },
      "outputs": [],
      "source": [
        "IS_COLAB = False\n",
        "# from google.colab import auth\n",
        "# auth.authenticate_user()\n",
        "# IS_COLAB = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:mbsdk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_bucket"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "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": "set_service_account"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "autoset_service_account"
      },
      "outputs": [],
      "source": [
        "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": "set_service_account:pipelines"
      },
      "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": "set_service_account:pipelines"
      },
      "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": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "import_aip:v1beta1"
      },
      "source": [
        "#### Import Vertex AI SDK\n",
        "\n",
        "Import the Vertex AI SDK into your Python environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:v1beta1"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "import google.cloud.aiplatform_v1beta1 as aip_beta\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aip_constants:fs"
      },
      "source": [
        "#### Vertex AI constants\n",
        "\n",
        "Setup up the following constants for Vertex AI:\n",
        "\n",
        "- `API_ENDPOINT`: The Vertex AI API service endpoint for `ML Metadata` services."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aip_constants:fs"
      },
      "outputs": [],
      "source": [
        "# API service endpoint\n",
        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
        "\n",
        "# Vertex location root path for your dataset, model and endpoint resources\n",
        "PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + REGION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "clients:metadata"
      },
      "source": [
        "## Set up clients\n",
        "\n",
        "The Vertex  works as a client/server model. On your side (the Python script) you will create a client that sends requests and receives responses from the Vertex AI server.\n",
        "\n",
        "You will use different clients in this tutorial for different steps in the workflow. So set them all up upfront.\n",
        "\n",
        "- Metadata Service for creating recording, searching and analyzing artifacts and metadata."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "clients:metadata"
      },
      "outputs": [],
      "source": [
        "# client options same for all services\n",
        "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
        "\n",
        "\n",
        "def create_metadata_client():\n",
        "    client = aip_beta.MetadataServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "clients = {}\n",
        "clients[\"metadata\"] = create_metadata_client()\n",
        "\n",
        "for client in clients.items():\n",
        "    print(client)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "intro_metadata"
      },
      "source": [
        "## Introduction to Vertex AI Metadata\n",
        "\n",
        "The `Vertex ML Metadata` service provides you with the ability to record, and subsequently search and analyze, the artifacts and corresponding metadata produced by your ML workflows. For example, during experimentation one might desire to record the location of the model artifacts, as artifacts, and the training hyperparameters and evaluation metrics as the corresponding metadata.\n",
        "\n",
        "The service supports recording ML metadata both manually and automatically, with the later occurring when you use Vertex AI Pipelines.\n",
        "\n",
        "### Concepts and organization\n",
        "\n",
        "Vertex ML Metadata describes your ML system's metadata as a graph.\n",
        "\n",
        "**Artifacts**: Artifacts are pieces of data that ML systems consume or produce, such as: datasets, models, or logs. For large artifacts like datasets or models, the artifact record includes the URI where the data is stored.\n",
        "\n",
        "**Executions**: Executions describe a single step in your ML system's workflow.\n",
        "\n",
        "**Events**: Executions can depend on artifacts as inputs or produce artifacts as outputs. Events describe the relationship between artifacts and executions to help you determine the lineage of artifacts. For example, an event is created to record that a dataset is used by an execution, and another event is created to record that this execution produced a model.\n",
        "\n",
        "**Contexts**: Contexts let you group artifacts and executions together in a single, queryable, and typed category.\n",
        "\n",
        "### ML artifact lineage\n",
        "\n",
        "Vertex ML Metadata provides the ability to understand changes in the performance of your machine ML system, and analyze the metadata produced by your ML workflow and the lineage of its artifacts. An artifact's lineage includes all the factors that contributed to its creation, as well as artifacts and metadata that descend from this artifact.\n",
        "\n",
        "Learn more about [Introduction to Vertex ML Metadata ](https://cloud.google.com/vertex-ai/docs/ml-metadata/introduction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_metadatastore"
      },
      "source": [
        "### Create a `MetadataStore` resource\n",
        "\n",
        "Each project may have one or more `MetadataStore` resources. By default, if none is explicity created, each project has a default, which is specified as:\n",
        "\n",
        "    projects/<project_id>/locations/<region>/metadataStores/<name>\n",
        "\n",
        "You create a `MetadataStore` resource using the `create_metadata_store()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified subpath for all resources in your project, i.e., projects/<project_id>/locations/<location>\n",
        "- `metadata_store_id`: The name of the `MetadataStore` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_metadatastore"
      },
      "outputs": [],
      "source": [
        "metadata_store = clients[\"metadata\"].create_metadata_store(\n",
        "    parent=PARENT, metadata_store_id=\"my-metadata-store-unique\"\n",
        ")\n",
        "\n",
        "metadata_store_id = str(metadata_store.result())[7:-2]\n",
        "print(metadata_store_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "list_metadata_schemas"
      },
      "source": [
        "### List metadata schemas\n",
        "\n",
        "When you create an `Artifact`, `Execution` or `Context` resource, you specify a schema that describes the corresponding metadata. The schemas must be pre-registered for your `Metadatastore` resource.\n",
        "\n",
        "You can get a list of all registered schemas, default and user defined, using the `list_metadata_schemas()` method, with the following parameters:\n",
        "\n",
        "- `name`: The fully qualified resource identifier for the `MetadataStore` resource.\n",
        "\n",
        "Learn more about [Metadata system schemas](https://cloud.google.com/vertex-ai/docs/ml-metadata/system-schemas)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "list_metadata_schemas"
      },
      "outputs": [],
      "source": [
        "schemas = clients[\"metadata\"].list_metadata_schemas(parent=metadata_store_id)\n",
        "\n",
        "for schema in schemas:\n",
        "    print(schema)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_metadata_artifact"
      },
      "source": [
        "### Create an `Artifact` resource\n",
        "\n",
        "You create an `Artifact` resource using the `create_artifact()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified resource identifier to the `Metadatastore` resource.\n",
        "- `artifact`: The definition of the `Artifact` resource\n",
        "    - `display_name`: The human readable name for the `Artifact` resource.\n",
        "    - `uri`: The uniform resource identifier of the artifact file. May be empty if there is no actual artifact file.\n",
        "    - `labels`: User defined labels to assign to the `Artifact` resource.\n",
        "    - `schema_title`: The title of the schema that describes the metadata.\n",
        "    - `metadata`: The metadata key/value pairs to associate with the `Artifact` resource.\n",
        "- `artifact_id`: (optional) A user defined short ID for the `Artifact` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_metadata_artifact"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform_v1beta1.types import Artifact\n",
        "\n",
        "artifact_item = Artifact(\n",
        "    display_name=\"my_example_artifact\",\n",
        "    uri=\"my_url\",\n",
        "    labels={\"my_label\": \"value\"},\n",
        "    schema_title=\"system.Artifact\",\n",
        "    metadata={\"param\": \"value\"},\n",
        ")\n",
        "\n",
        "artifact = clients[\"metadata\"].create_artifact(\n",
        "    parent=metadata_store_id,\n",
        "    artifact=artifact_item,\n",
        "    artifact_id=\"myartifactid\",\n",
        ")\n",
        "\n",
        "print(artifact)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "list_metadata_artifacts"
      },
      "source": [
        "### List `Artifact` resources in a `Metadatastore`\n",
        "\n",
        "You can list all `Artifact` resources using the `list_artifacts()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified resource identifier for the `MetadataStore` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "list_metadata_artifacts"
      },
      "outputs": [],
      "source": [
        "artifacts = clients[\"metadata\"].list_artifacts(parent=metadata_store_id)\n",
        "\n",
        "for _artifact in artifacts:\n",
        "    print(_artifact)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_metadata_execution"
      },
      "source": [
        "### Create an `Execution` resource\n",
        "\n",
        "You create an `Execution` resource using the `create_execution()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified resource identifier to the `Metadatastore` resource.\n",
        "- `execution`:\n",
        "    - `display_name`: A human readable name for the `Execution` resource.\n",
        "    - `schema_title`: The title of the schema that describes the metadata.\n",
        "    - `metadata`: The metadata key/value pairs to associate with the `Execution` resource.\n",
        "- `execution_id`: (optional) A user defined short ID for the `Execution` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_metadata_execution"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform_v1beta1.types import Execution\n",
        "\n",
        "execution = clients[\"metadata\"].create_execution(\n",
        "    parent=metadata_store_id,\n",
        "    execution=Execution(\n",
        "        display_name=\"my_execution\",\n",
        "        schema_title=\"system.CustomJobExecution\",\n",
        "        metadata={\"value\": \"param\"},\n",
        "    ),\n",
        "    execution_id=\"myexecutionid\",\n",
        ")\n",
        "\n",
        "print(execution)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "list_metadata_executions"
      },
      "source": [
        "### List `Execution` resources in a `Metadatastore`\n",
        "\n",
        "You can list all `Execution` resources using the `list_executions()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified resource identifier for the `MetadataStore` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "list_metadata_executions"
      },
      "outputs": [],
      "source": [
        "executions = clients[\"metadata\"].list_executions(parent=metadata_store_id)\n",
        "\n",
        "for _execution in executions:\n",
        "    print(_execution)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_metadata_context"
      },
      "source": [
        "### Create a `Context` resource\n",
        "\n",
        "You create an `Context` resource using the `create_context()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified resource identifier to the `Metadatastore` resource.\n",
        "- `context`:\n",
        "    - `display_name`: A human readable name for the `Execution` resource.\n",
        "    - `schema_title`: The title of the schema that describes the metadata.\n",
        "    - `labels`: User defined labels to assign to the `Context` resource.\n",
        "    - `metadata`: The metadata key/value pairs to associate with the `Execution` resource.\n",
        "- `context_id`: (optional) A user defined short ID for the `Context` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_metadata_context"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform_v1beta1.types import Context\n",
        "\n",
        "context = clients[\"metadata\"].create_context(\n",
        "    parent=metadata_store_id,\n",
        "    context=Context(\n",
        "        display_name=\"my_context\",\n",
        "        labels=[{\"my_label\", \"my_value\"}],\n",
        "        schema_title=\"system.Pipeline\",\n",
        "        metadata={\"param\": \"value\"},\n",
        "    ),\n",
        "    context_id=\"mycontextid\",\n",
        ")\n",
        "\n",
        "print(context)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "list_metadata_contexts"
      },
      "source": [
        "### List `Context` resources in a `Metadatastore`\n",
        "\n",
        "You can list all `Context` resources using the `list_contexts()` method, with the following parameters:\n",
        "\n",
        "- `parent`: The fully qualified resource identifier for the `MetadataStore` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "list_metadata_contexts"
      },
      "outputs": [],
      "source": [
        "contexts = clients[\"metadata\"].list_contexts(parent=metadata_store_id)\n",
        "\n",
        "for _context in contexts:\n",
        "    print(_context)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "add_metadata_execution_events"
      },
      "source": [
        "### Add events to `Execution` resource\n",
        "\n",
        "An `Execution` resource consists of a sequence of events that occurred during the execution. Each event consists of an artifact that is either an input or an output of the `Execution` resource.\n",
        "\n",
        "You can add execution events to an `Execution` resource using the `add_execution_events()` method, with the following parameters:\n",
        "\n",
        "- `execution`: The fully qualified resource identifier for the `Execution` resource.\n",
        "- `events`: The sequence of events constituting the execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "add_metadata_execution_events"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform_v1beta1.types import Event\n",
        "\n",
        "clients[\"metadata\"].add_execution_events(\n",
        "    execution=execution.name,\n",
        "    events=[\n",
        "        Event(\n",
        "            artifact=artifact.name,\n",
        "            type_=Event.Type.INPUT,\n",
        "            labels={\"my_label\": \"my_value\"},\n",
        "        )\n",
        "    ],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "metadata_finish_context"
      },
      "source": [
        "### Combine Artifacts and Executions into a Context\n",
        "\n",
        "A Context is used to group `Artifact` resources and `Execution` resources together under a single, queryable, and typed category. Contexts can be used to represent sets of metadata.\n",
        "\n",
        "You can combine a set of `Artifact` and `Execution` resources into a `Context` resource using the `add_context_artifacts_and_executions()` method, with the following parameters:\n",
        "\n",
        "- `context`: The fully qualified resource identifier of the `Context` resource.\n",
        "- `artifacts`: A list of fully qualified resource identifiers of the `Artifact` resources.\n",
        "- `executions`: A list of fully qualified resource identifiers of the `Execution` resources."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "metadata_finish_context"
      },
      "outputs": [],
      "source": [
        "clients[\"metadata\"].add_context_artifacts_and_executions(\n",
        "    context=context.name, artifacts=[artifact.name], executions=[execution.name]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "query_metadata_context"
      },
      "source": [
        "### Query a context\n",
        "\n",
        "You can query the subgraph of a `Context` resource using the method `query_context_lineage_subgraph()` method, with the following parameters:\n",
        "\n",
        "- `context`: The fully qualified resource identifier of the `Context` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "query_metadata_context"
      },
      "outputs": [],
      "source": [
        "subgraph = clients[\"metadata\"].query_context_lineage_subgraph(context=context.name)\n",
        "\n",
        "print(subgraph)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_metadata_artifact"
      },
      "source": [
        "### Delete an `Artifact` resource\n",
        "\n",
        "You can delete an `Artifact` resource using the `delete_artifact()` method, with the following parameters:\n",
        "\n",
        "- `name`: The fully qualified resource identifier for the `Artifact` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "delete_metadata_artifact"
      },
      "outputs": [],
      "source": [
        "clients[\"metadata\"].delete_artifact(name=artifact.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_metadata_execution"
      },
      "source": [
        "### Delete an `Execution` resource\n",
        "\n",
        "You can delete an `Execution` resource using the `delete_execution()` method, with the following parameters:\n",
        "\n",
        "- `name`: The fully qualified resource identifier for the `Execution` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "delete_metadata_execution"
      },
      "outputs": [],
      "source": [
        "clients[\"metadata\"].delete_execution(name=execution.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_metadata_context"
      },
      "source": [
        "### Delete a `Context` resource\n",
        "\n",
        "You can delete an `Context` resource using the `delete_context()` method, with the following parameters:\n",
        "\n",
        "- `name`: The fully qualified resource identifier for the `Context` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "delete_metadata_context"
      },
      "outputs": [],
      "source": [
        "clients[\"metadata\"].delete_context(name=context.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "intro_pipeline_metadata"
      },
      "source": [
        "## Introduction to tracking ML Metadata in a `Vertex AI Pipeline`\n",
        "\n",
        "Vertex AI Pipelines automatically records the metrics and artifacts created when the pipeline is exeuted. You can then use the SDK to track and analyze the metrics and artifacts across pipeline runs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_kfp:basic"
      },
      "outputs": [],
      "source": [
        "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": "metadata_3step_pipeline"
      },
      "source": [
        "### Creating a 3-step pipeline with custom components\n",
        "\n",
        "First, you create a pipeline to run on `Vertex AI Pipelines`, consisting of the following custom components:\n",
        "\n",
        "* `get_dataframe`: Retrieve data from a BigQuery table and convert it into a pandas DataFrame.\n",
        "* `sklearn_train`: 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 a `Vertex AI Endpoint` resource.\n",
        "\n",
        "#### get_dataframe component\n",
        "\n",
        "This component does the following:\n",
        "\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\n",
        "\n",
        "#### sklearn_train component\n",
        "\n",
        "This component does the following:\n",
        "\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\n",
        "\n",
        "#### deploy_model component\n",
        "\n",
        "This component does the following:\n",
        "\n",
        "* Uploads the scikit-learn model to a `Vertex AI Model` resource.\n",
        "* Deploys the model to a `Vertex AI Endpoint` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "metadata_3step_pipeline"
      },
      "outputs": [],
      "source": [
        "@component(\n",
        "    packages_to_install=[\"google-cloud-bigquery\", \"pandas\", \"pyarrow\", \"db-dtypes\"],\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()\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)\n",
        "\n",
        "\n",
        "@component(\n",
        "    packages_to_install=[\"scikit-learn\", \"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\")\n",
        "\n",
        "\n",
        "@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": "create_3step_pipeline"
      },
      "source": [
        "### Construct and compile the pipeline\n",
        "\n",
        "Next, construct the pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_3step_pipeline"
      },
      "outputs": [],
      "source": [
        "PIPELINE_ROOT = f\"{BUCKET_URI}/pipeline_root/3step\"\n",
        "\n",
        "\n",
        "@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 my_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": "run_pipeline:3step"
      },
      "source": [
        "### Compile and execute two runs of the pipeline\n",
        "\n",
        "Next, you compile the pipeline and then run two separate instances of the pipeline. In the first instance, you train the model with a small version of the dataset and in the second instance you train it with a larger version of the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "run_pipeline:3step"
      },
      "outputs": [],
      "source": [
        "NOW = datetime.now().isoformat().replace(\".\", \":\")[:-7]\n",
        "\n",
        "compiler.Compiler().compile(\n",
        "    pipeline_func=my_pipeline, package_path=\"mlmd_pipeline.json\"\n",
        ")\n",
        "\n",
        "run1 = aiplatform.PipelineJob(\n",
        "    display_name=\"mlmd-pipeline\",\n",
        "    template_path=\"mlmd_pipeline.json\",\n",
        "    job_id=\"mlmd-pipeline-small-unique\",\n",
        "    parameter_values={\"bq_table\": \"sara-vertex-demos.beans_demo.small_dataset\"},\n",
        "    enable_caching=True,\n",
        ")\n",
        "\n",
        "run2 = aiplatform.PipelineJob(\n",
        "    display_name=\"mlmd-pipeline\",\n",
        "    template_path=\"mlmd_pipeline.json\",\n",
        "    job_id=\"mlmd-pipeline-large-unique\",\n",
        "    parameter_values={\"bq_table\": \"sara-vertex-demos.beans_demo.large_dataset\"},\n",
        "    enable_caching=True,\n",
        ")\n",
        "\n",
        "run1.run()\n",
        "run2.run()\n",
        "\n",
        "run1.delete()\n",
        "run2.delete()\n",
        "\n",
        "! rm -f mlmd_pipeline.json *.yaml"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "compare_runs:3step"
      },
      "source": [
        "### Compare the pipeline runs\n",
        "\n",
        "Now that you have two pipeline completed pipeline runs, you can compare the runs.\n",
        "\n",
        "You can use the `get_pipeline_df()` method to access the metadata from the runs. The `mlmd-pipeline` parameter here refers to the name you gave to your pipeline:\n",
        "\n",
        "**Alternately, 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": "code",
      "execution_count": null,
      "metadata": {
        "id": "compare_runs:3step"
      },
      "outputs": [],
      "source": [
        "df = aiplatform.get_pipeline_df(pipeline=\"mlmd-pipeline\")\n",
        "print(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "visualize_runs:3step"
      },
      "source": [
        "### Visualize the pipeline runs\n",
        "\n",
        "Next, you create a custom visualization with matplotlib to see the relationship between your model's accuracy and the amount of data used for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "visualize_runs:3step"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "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": "list_metadata_artifacts:3tep"
      },
      "source": [
        "### Quering your `Metadatastore` resource\n",
        "\n",
        "Finally, you query your `Metadatastore` resource by specifying a `filter` parameter when calling the `list_artifacts()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "list_metadata_artifacts:3tep"
      },
      "outputs": [],
      "source": [
        "FILTER = f'create_time >= \"{NOW}\" AND state = LIVE'\n",
        "artifact_req = {\n",
        "    \"parent\": metadata_store_id,\n",
        "    \"filter\": FILTER,\n",
        "}\n",
        "\n",
        "artifacts = clients[\"metadata\"].list_artifacts(artifact_req)\n",
        "\n",
        "for _artifact in artifacts:\n",
        "    print(_artifact)\n",
        "    clients[\"metadata\"].delete_artifact(name=_artifact.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_metadatastore"
      },
      "source": [
        "### Delete a `MetadataStore` resource\n",
        "\n",
        "You can delete a `MetadataStore` resource using the `delete_metadata_store()` method, with the following parameters:\n",
        "\n",
        "- `name`: The fully qualified resource identifier for the `MetadataStore` resource."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "delete_metadatastore"
      },
      "outputs": [],
      "source": [
        "metadata_store_id = (\n",
        "    f\"projects/{PROJECT_ID}/locations/{REGION}/metadataStores/my-metadata-store-unique\"\n",
        ")\n",
        "\n",
        "clients[\"metadata\"].delete_metadata_store(name=metadata_store_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "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": "cleanup:mbsdk"
      },
      "outputs": [],
      "source": [
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_vertex_ml_metadata.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}