notebooks/official/experiments/get_started_with_vertex_experiments.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": "ee9d87d8ec00"
      },
      "source": [
        "This notebook was authored with assistance from [Ivan Nardini](https://github.com/inardini)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "title:generic,gcp"
      },
      "source": [
        "# Get started with Vertex AI Experiments\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments.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://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/experiments/get_started_with_vertex_experiments.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://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/experiments/get_started_with_vertex_experiments.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 covers getting started with Vertex AI Experiments, and demonstrates how to use Vertex AI in production.\n",
        "\n",
        "Learn more about [Vertex AI Experiments](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments), [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage2,get_started_vertex_experiments"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use Vertex AI Experiments when training with Vertex AI.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- Vertex AI Experiments\n",
        "- Vertex ML Metadata\n",
        "- Vertex AI Training\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Local (notebook) training\n",
        "    - Create an experiment.\n",
        "    - Create a first run in the experiment.\n",
        "    - Log parameters and metrics.\n",
        "    - Create artifact lineage.\n",
        "    - Visualize the experiment results.\n",
        "    - Execute a second run.\n",
        "    - Compare the two runs in the experiment.\n",
        "- Cloud (Vertex AI) training\n",
        "    - Within the training script:\n",
        "        - Create an experiment.\n",
        "        - Log parameters and metrics.\n",
        "        - Create artifact lineage.\n",
        "    - Create a Vertex AI training custom job.\n",
        "    - Execute the custom job.\n",
        "    - Visualize the experiment results,"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "recommendation:mlops,stage2,logging"
      },
      "source": [
        "### Recommendations\n",
        "\n",
        "When doing E2E MLOps on Google Cloud, the following are some of the best practices for logging data when experimenting or formally training a model.\n",
        "\n",
        "#### Python Logging\n",
        "\n",
        "Use Python's logging package when doing ad-hoc training locally.\n",
        "\n",
        "#### Cloud Logging\n",
        "\n",
        "Use Cloud Logging when doing training on the cloud.\n",
        "\n",
        "#### Experiments\n",
        "\n",
        "Use Vertex AI Experiments in conjunction with logging when performing experiments to compare results for different experiment configurations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "313c25f2f514"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This tutorial doesn't use a dataset. References to example datasets is for demonstration purposes only."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bd73a4bd07ef"
      },
      "source": [
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/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 following packages for executing this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "install_mlops"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform\n"
      ]
    },
    {
      "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": "yfEglUHQk9S3"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "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.\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated.\n",
        "\n",
        "**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": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "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": [
        "**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": [
        "IS_COLAB = False\n",
        "\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",
        "        # print(\"shell_output=\", shell_output)\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": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial.\n",
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import uuid\n",
        "\n",
        "import google.cloud.aiplatform as aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,region"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_aip:mbsdk,region"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training.\n",
        "\n",
        "Set the variables `TRAIN_GPU/TRAIN_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
        "    TRAIN_GPU, TRAIN_NGPU = (\n",
        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
        "    )\n",
        "else:\n",
        "    TRAIN_GPU, TRAIN_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for training.\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "container:training,prediction"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TF\"):\n",
        "    TF = os.getenv(\"IS_TESTING_TF\")\n",
        "else:\n",
        "    TF = \"2.12\".replace(\".\", \"-\")\n",
        "\n",
        "if TRAIN_GPU:\n",
        "    TRAIN_VERSION = \"tf-gpu.{}.py310\".format(TF)\n",
        "else:\n",
        "    TRAIN_VERSION = \"tf-cpu.{}.py310\".format(TF)\n",
        "\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training.\n",
        "\n",
        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs you will use for for training.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: The following is not supported for training:*\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "machine:training"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_MACHINE\"):\n",
        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_TRAIN_MACHINE\")\n",
        "else:\n",
        "    MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83f8f98ae12d"
      },
      "source": [
        "## Introduction to Vertex AI Experiments\n",
        "\n",
        "With Vertex AI Experiments you can log and track the following when experimenting and developing your model architecture and model training:\n",
        "\n",
        "- Log the metaparameters for the model architecture.\n",
        "- Log the hyperparameters for training.\n",
        "- Log the evaluation metrics.\n",
        "- Create an artifact lineage of the dataset, model and evaluation.\n",
        "- Group one or more training runs under an experiment.\n",
        "- Compare experiments.\n",
        "\n",
        "Vertex AI Experiments can be integrated with the following development process flows:\n",
        "\n",
        "- Local development in a notebook\n",
        "- Cloud development in Vertex AI Training\n",
        "- Operationalizing development in Vertex AI Pipelines\n",
        "\n",
        "Learn more about [Experiments]( https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments).\n",
        "\n",
        "Learn more about [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata/introduction)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fb84b57b9e74"
      },
      "source": [
        "### Local development in a notebook\n",
        "\n",
        "You can track an experiment in your local development, such as in a Vertex AI Workbench notebook, by:\n",
        "\n",
        "- Wrap (preamble) the creation of an experiment.\n",
        "- Instantiate a run per training run in the experiment.\n",
        "- Within the local training run, log the corresponding parameters and results.\n",
        "- Create lineage to the artifacts and experiment data.\n",
        "- Retrieve the experiment data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "start_experiment"
      },
      "source": [
        "#### Create experiment for tracking training related metadata\n",
        "\n",
        "First, you create an experiment using the `init()` method and then initialize a run within the experiment using `start_run()`.\n",
        "\n",
        "- `aiplatform.init()` - Create an experiment instance\n",
        "- `aiplatform.start_run()` - Track a specific run within the experiment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ed46e349cf2"
      },
      "outputs": [],
      "source": [
        "# Specify a name for the experiment\n",
        "EXPERIMENT_NAME = \"[your-experiment-name]\"\n",
        "\n",
        "if EXPERIMENT_NAME == \"[your-experiment-name]\":\n",
        "    EXPERIMENT_NAME = f\"example-{uuid.uuid1()}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "start_experiment"
      },
      "outputs": [],
      "source": [
        "# Create experiment\n",
        "aiplatform.init(experiment=EXPERIMENT_NAME)\n",
        "aiplatform.start_run(\"run-1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "experiment_log_params"
      },
      "source": [
        "#### Log parameters for the experiment\n",
        "\n",
        "Typically, an experiment is associated with a specific dataset and a model architecture. Within an experiment, you may have multiple training runs,\n",
        "where each run involves a different configuration. For example:\n",
        "\n",
        "- Data feeding, such as:\n",
        "    - dataset split\n",
        "    - dataset sampling and boosting\n",
        "- Metaparameters, such as:\n",
        "    - depth and width of layers\n",
        "- Hyperparameters, such as:\n",
        "    - batch size\n",
        "    - learning rate\n",
        "\n",
        "These configuration settings are referred to as parameters, which you store as key-value pairs using the method `log_params()`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "experiment_log_params"
      },
      "outputs": [],
      "source": [
        "metaparams = {}\n",
        "metaparams[\"units\"] = 128\n",
        "aiplatform.log_params(metaparams)\n",
        "\n",
        "hyperparams = {}\n",
        "hyperparams[\"epochs\"] = 100\n",
        "hyperparams[\"batch_size\"] = 32\n",
        "hyperparams[\"learning_rate\"] = 0.01\n",
        "aiplatform.log_params(hyperparams)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "experiment_log_metrics"
      },
      "source": [
        "#### Log metrics for the experiment\n",
        "\n",
        "At the completion or termination of a run within an experiment, you can log results that you use to compare runs. For example:\n",
        "\n",
        "- evaluation metrics\n",
        "- hyperparameter search selection\n",
        "- time to train the model\n",
        "- early stop trigger\n",
        "\n",
        "These results are referred to as metrics, which you store as key-value pairs using the method `log_metrics()`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "experiment_log_metrics"
      },
      "outputs": [],
      "source": [
        "metrics = {}\n",
        "metrics[\"test_acc\"] = 98.7\n",
        "metrics[\"train_acc\"] = 99.3\n",
        "aiplatform.log_metrics(metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "get_experiment"
      },
      "source": [
        "#### Get the experiment results\n",
        "\n",
        "When you are finished with a run within an experiment, call `end_run()` method to complete the logging for that run.\n",
        "\n",
        "Next, use the experiment name as a parameter to the method `get_experiment_df()` to get the results of the experiment as a pandas dataframe."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "get_experiment"
      },
      "outputs": [],
      "source": [
        "aiplatform.end_run()\n",
        "\n",
        "experiment_df = aiplatform.get_experiment_df()\n",
        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
        "print(experiment_df.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2bcaf3e3ddff"
      },
      "source": [
        "#### Start subsequent run in an experiment\n",
        "\n",
        "Next, create a second run for the same experiment. In this example, you change the metaparameter for `units` from 126 to 256, and log different metric results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d7c1a839683b"
      },
      "outputs": [],
      "source": [
        "aiplatform.start_run(\"run-2\")\n",
        "\n",
        "metaparams = {}\n",
        "metaparams[\"units\"] = 256  # changed the value\n",
        "aiplatform.log_params(metaparams)\n",
        "\n",
        "hyperparams = {}\n",
        "hyperparams[\"epochs\"] = 100\n",
        "hyperparams[\"batch_size\"] = 32\n",
        "hyperparams[\"learning_rate\"] = 0.01\n",
        "aiplatform.log_params(hyperparams)\n",
        "\n",
        "metrics = {}\n",
        "metrics[\"test_acc\"] = 98.8  # value changed\n",
        "metrics[\"train_acc\"] = 99.5  # value changed\n",
        "aiplatform.log_metrics(metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d2619bbc04a"
      },
      "source": [
        "#### Comparing runs in the same experiment\n",
        "\n",
        "Finally, use the experiment name as a parameter to the method `get_experiment_df()` to get the results of all the runs within the experiment as a pandas dataframe."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ddb4e17fe92"
      },
      "outputs": [],
      "source": [
        "aiplatform.end_run()\n",
        "\n",
        "experiment_df = aiplatform.get_experiment_df()\n",
        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
        "print(experiment_df.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5328eb24f6a6"
      },
      "source": [
        "#### Delete the experiment\n",
        "\n",
        "Next, delete the experiment using the `delete()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "outputs": [],
      "source": [
        "exp = aiplatform.Experiment(EXPERIMENT_NAME)\n",
        "try:\n",
        "    exp.delete()\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a07fb4c47af4"
      },
      "source": [
        "### Create artifact lineage in experiment runs\n",
        "\n",
        "In this example, you add artifact lineage to your experiment run. First, create an experiment and then start a run within the experiment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b7aea23dafbc"
      },
      "outputs": [],
      "source": [
        "# Create experiment\n",
        "EXPERIMENT_NAME = f\"example-{uuid.uuid1()}\"\n",
        "aiplatform.init(experiment=EXPERIMENT_NAME)\n",
        "aiplatform.start_run(\"run-1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "355fed1a7b6d"
      },
      "source": [
        "#### Create a dataset and model artifacts\n",
        "\n",
        "Next, you create synthetic artifacts in the Vertex ML Metadata to be associated with this run in the experiment, as lineage. You'll create:\n",
        "\n",
        "- `dataset_artifact`: A dataset that is the input to the experiment run.\n",
        "- `model_artifact`: A model that is the output from the experiment run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4c1faeec89bb"
      },
      "outputs": [],
      "source": [
        "DATASET_URI = \"gs://example/dataset.csv\"\n",
        "MODEL_URI = \"gs://example/saved_model.pb\"\n",
        "\n",
        "dataset_artifact = aiplatform.Artifact.create(\n",
        "    schema_title=\"system.Dataset\", display_name=\"example_dataset\", uri=DATASET_URI\n",
        ")\n",
        "\n",
        "model_artifact = aiplatform.Artifact.create(\n",
        "    schema_title=\"system.Model\", display_name=\"example_modl\", uri=MODEL_URI\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7f6cc9dbaa2e"
      },
      "source": [
        "#### Create the artifact lineage\n",
        "\n",
        "Next, to create artifact lineage for an experiment run, instantiate an execution using the method `start_execution()`. You then attach input artifacts using the method `assign_input_artifacts()` and attach output artifacts using the method `assign_output_artifacts()`.\n",
        "\n",
        "In this example, to find the lineage for the experiment, you add a synthetic (metadata) entry `lineage` to the execution run, and set the value to the console uri for the lineage, which you can get from the method `get_output_artifacts()` and the property `lineage_console_uri`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9170906b09f9"
      },
      "outputs": [],
      "source": [
        "with aiplatform.start_execution(\n",
        "    schema_title=\"system.ContainerExecution\", display_name=\"example_training\"\n",
        ") as execution:\n",
        "    execution.assign_input_artifacts([dataset_artifact])\n",
        "\n",
        "    aiplatform.log_params({\"units\": 256})\n",
        "    aiplatform.log_metrics({\"acc\": 96.8})\n",
        "\n",
        "    execution.assign_output_artifacts([model_artifact])\n",
        "\n",
        "    aiplatform.log_metrics(\n",
        "        {\"lineage\": execution.get_output_artifacts()[0].lineage_console_uri}\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15f847929914"
      },
      "source": [
        "#### Get the experiment results\n",
        "\n",
        "Next, use the experiment name as a parameter to the method `get_experiment_df()` to get the results of the experiment as a pandas dataframe.\n",
        "\n",
        "In this example, you stored the resource URI to the lineage as a metric value `lineage` in the execution run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0cfa26496ea9"
      },
      "outputs": [],
      "source": [
        "aiplatform.end_run()\n",
        "\n",
        "experiment_df = aiplatform.get_experiment_df()\n",
        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
        "print(experiment_df.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "994b3dab8769"
      },
      "source": [
        "#### Visualize the artifact lineage\n",
        "\n",
        "Next, open the link below to visualize the artifact lineage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b85013ec6f11"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"Open the following link:\", execution.get_output_artifacts()[0].lineage_console_uri\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80f8e272bf3f"
      },
      "source": [
        "#### Delete the artifact lineage\n",
        "\n",
        "Next, use the `delete()` method to delete the artifact lineage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ae66f78a233f"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    dataset_artifact.delete()\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "try:\n",
        "    model_artifact.delete()\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "35b899773315"
      },
      "source": [
        "#### Delete the experiment\n",
        "\n",
        "Next, you delete the experiment using the `delete()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f1020119d742"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    exp.delete()\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3cd64a48e0f4"
      },
      "source": [
        "### Cloud development in Vertex AI Training\n",
        "\n",
        "You can track an experiment in your cloud development using Vertex AI Training, by:\n",
        "\n",
        "In your Python training script, repeat the same steps as in local development:\n",
        "\n",
        "- Wrap (preamble) the creation of an experiment.\n",
        "- Instantiate a run per training run in the experiment.\n",
        "- Within the local training run, log the corresponding parameters and results.\n",
        "- Create lineage to the artifacts and experiment data.\n",
        "- Retreive the experiment data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package"
      },
      "source": [
        "#### Package layout\n",
        "\n",
        "Before you start the training, you'll look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
        "In the following cells, assemble the training package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "examine_training_package"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "! rm -rf custom\n",
        "! mkdir custom\n",
        "\n",
        "# Add package information\n",
        "! touch custom/README.md\n",
        "\n",
        "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
        "! echo \"$setup_cfg\" > custom/setup.cfg\n",
        "\n",
        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'google-cloud-aiplatform',\\n\\n  ],\\n\\n    packages=setuptools.find_packages())\"\n",
        "! echo \"$setup_py\" > custom/setup.py\n",
        "\n",
        "pkg_info = \"Metadata-Version: 1.0\\n\\nName: Synethic Training Script for Experiments\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
        "! echo \"$pkg_info\" > custom/PKG-INFO\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir custom/trainer\n",
        "! touch custom/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "50c0e24489df"
      },
      "source": [
        "#### Create synthetic training script\n",
        "\n",
        "First, write a synthetic training script. It won't actually train a model, but instead mimics the training of the model:\n",
        "\n",
        "- Argument parsing\n",
        "  - `experiment`: The name of the experiment.\n",
        "  - `run`: The name of the run within the experiment.\n",
        "  - `epochs`: The number of epochs.\n",
        "  - `dataset-uri`: The Cloud Storage location of the training data.\n",
        "  - `model-dir`: The Cloud Storage location to save the trained model artifacts.\n",
        "- Training functions\n",
        "  - `get_data()`: \n",
        "      - Get the training data. \n",
        "      - Create the input dataset artifact.\n",
        "      - Attach dataset artifact as input to execution context.\n",
        "  - `get_model()`:\n",
        "      - Get the model architecture.\n",
        "  - `train_model()`:\n",
        "      - Train the model.\n",
        "  - `save_model()`:\n",
        "      - Save the model.\n",
        "      - Create the output model artifact.\n",
        "      - Attach model artifact as output to execution context.\n",
        "- Initialize the experiment (`init()`) and start a run (`start_run()`) within the experiment.\n",
        "- Wrap the training with a `start_execution()`.\n",
        "- Log the lineage to the experiment parameters (`log_metrics({\"lineage\"...)`).\n",
        "- End the experiment run (`end_run()`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aed1431dccaf"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/trainer/task.py\n",
        "\n",
        "import argparse\n",
        "import os\n",
        "\n",
        "import google.cloud.aiplatform as aiplatform\n",
        "\n",
        "parser = argparse.ArgumentParser()\n",
        "# Args for experiment\n",
        "parser.add_argument('--experiment', dest='experiment',\n",
        "                    required=True, type=str,\n",
        "                    help='Name of experiment')\n",
        "parser.add_argument('--run', dest='run',\n",
        "                    required=True, type=str,\n",
        "                    help='Name of run within the experiment')\n",
        "\n",
        "# Hyperparameters for experiment\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    default=10, type=int,\n",
        "                    help='Number of epochs.')\n",
        "\n",
        "parser.add_argument('--dataset-uri', dest='dataset_uri',\n",
        "                    required=True, type=str,\n",
        "                    help='Location of the dataset')\n",
        "\n",
        "parser.add_argument('--model-dir', dest='model_dir',\n",
        "                    default=os.getenv(\"AIP_MODEL_DIR\"), type=str,\n",
        "                    help='Storage location for the model')\n",
        "args = parser.parse_args()\n",
        "\n",
        "def get_data(dataset_uri, execution):\n",
        "    # get the training data\n",
        "    \n",
        "    dataset_artifact = aiplatform.Artifact.create(\n",
        "        schema_title=\"system.Dataset\", display_name=\"example_dataset\", uri=dataset_uri\n",
        "    )\n",
        "    \n",
        "    execution.assign_input_artifacts([dataset_artifact])\n",
        "\n",
        "    return None\n",
        "\n",
        "def get_model():\n",
        "    # get or create the model architecture\n",
        "    return None\n",
        "\n",
        "def train_model(dataset, model, epochs):\n",
        "    aiplatform.log_params({\"epochs\": epochs})\n",
        "    # train the model\n",
        "    return model\n",
        "\n",
        "def save_model(model, model_dir, execution):\n",
        "    # save the model\n",
        "    \n",
        "    model_artifact = aiplatform.Artifact.create(\n",
        "        schema_title=\"system.Model\", display_name=\"example_model\", uri=model_dir\n",
        "    )\n",
        "    execution.assign_output_artifacts([model_artifact])\n",
        "\n",
        "# Create a run within the experiment\n",
        "aiplatform.init(experiment=args.experiment)\n",
        "aiplatform.start_run(args.run)\n",
        "\n",
        "with aiplatform.start_execution(\n",
        "    schema_title=\"system.ContainerExecution\", display_name=\"example_training\"\n",
        ") as execution:\n",
        "    dataset = get_data(args.dataset_uri, execution)\n",
        "    model = get_model()\n",
        "    model = train_model(dataset, model, args.epochs)\n",
        "    save_model(model, args.model_dir, execution)\n",
        "    \n",
        "    # Store the lineage link in the experiment\n",
        "    aiplatform.log_metrics({\"lineage\": execution.get_output_artifacts()[0].lineage_console_uri})\n",
        "\n",
        "aiplatform.end_run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tarball_training_script"
      },
      "source": [
        "#### Store training script on your Cloud Storage bucket\n",
        "\n",
        "Next, package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tarball_training_script"
      },
      "outputs": [],
      "source": [
        "! rm -f custom.tar custom.tar.gz\n",
        "! tar cvf custom.tar custom\n",
        "! gzip custom.tar\n",
        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_custom_pp_training_job:mbsdk,no_model"
      },
      "source": [
        "#### Create custom training job\n",
        "\n",
        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
        "\n",
        "- `python_package_gcs_uri`: The location of the Python training package as a tarball.\n",
        "- `python_module_name`: The relative path to the training script in the Python package.\n",
        "\n",
        "*Note:* There is no requirements parameter. You specify any requirements in the `setup.py` script in your Python package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_custom_pp_training_job:mbsdk,no_model"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"example\"\n",
        "\n",
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    python_package_gcs_uri=f\"{BUCKET_URI}/trainer.tar.gz\",\n",
        "    python_module_name=\"trainer.task\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_custom_container_training_job:no_model"
      },
      "source": [
        "#### Run the custom training job\n",
        "\n",
        "Next, run the custom training job to start the training job by invoking the method `run()`, with the following parameters:\n",
        "\n",
        "- `args`: The arguments to pass to the training script\n",
        "    - `model_dir`: The Cloud Storage location to store the model.\n",
        "    - `dataset_uri`: The Cloud Storage location of the dataset.\n",
        "    - `epochs`: The number of epochs (hyperparameter).\n",
        "    - `experiment`: The name of the experiment.\n",
        "    - `run`: The name of the run within the experiment.\n",
        "- `replica_count`: The number of VM instances.\n",
        "- `machine_type`: The machine type for each VM instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "260f972398e5"
      },
      "outputs": [],
      "source": [
        "EXPERIMENT_NAME = f\"example-{uuid.uuid1()}\"\n",
        "aiplatform.init(experiment=EXPERIMENT_NAME)\n",
        "\n",
        "CMDARGS = [\n",
        "    \"--model-dir=\" + BUCKET_URI,\n",
        "    \"--dataset-uri=gs://example/foo.csv\",\n",
        "    \"--epochs=5\",\n",
        "    f\"--experiment={EXPERIMENT_NAME}\",\n",
        "    \"--run=run-1\",\n",
        "]\n",
        "\n",
        "job.run(\n",
        "    args=CMDARGS,\n",
        "    replica_count=1,\n",
        "    machine_type=TRAIN_COMPUTE,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    sync=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5f40912e6500"
      },
      "source": [
        "#### Get the experiment results\n",
        "\n",
        "Next, use the experiment name as a parameter to the method `get_experiment_df()` to get the results of the experiment as a pandas dataframe.\n",
        "\n",
        "In this example, you stored the resource URI to the lineage as a metric value `lineage` in the execution run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7e9671712230"
      },
      "outputs": [],
      "source": [
        "experiment_df = aiplatform.get_experiment_df()\n",
        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
        "print(experiment_df.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65a65f847332"
      },
      "source": [
        "#### Visualize the artifact lineage\n",
        "\n",
        "Next, open the link below to visualize the artifact lineage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f5a02e7d92c7"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    print(\"Open the following link\", experiment_df[\"metric.lineage\"][0])\n",
        "except Exception as e:\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d4ba591be8ec"
      },
      "source": [
        "#### Delete the custom training job\n",
        "\n",
        "You can delete your custom training job using the `delete()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5436ab06482a"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e508c159d712"
      },
      "source": [
        "#### Delete the experiment\n",
        "\n",
        "Since the experiment was created within Vertex AI Training, to delete the experiment you use the `list()` method to obtain all the experiments for the project, and then filter on the experiment name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1a1b5fcbfde0"
      },
      "outputs": [],
      "source": [
        "experiments = aiplatform.Experiment.list()\n",
        "for experiment in experiments:\n",
        "    if experiment.name == EXPERIMENT_NAME:\n",
        "        experiment.delete()"
      ]
    },
    {
      "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": "e89f43b3df49"
      },
      "outputs": [],
      "source": [
        "! rm -rf custom\n",
        "\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_vertex_experiments.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}