notebooks/official/automl/automl_forecasting_on_vertex_pipelines.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "18ebbd838e32"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": "mThXALJl9Yue"
      },
      "source": [
        "# AutoML Tabular Workflows for Forecasting\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/automl_forecasting_on_vertex_pipelines.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fautoml%2Fautoml_forecasting_on_vertex_pipelines.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/automl/automl_forecasting_on_vertex_pipelines.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/automl/automl_forecasting_on_vertex_pipelines.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fcc745968395"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how you can use Vertex AI Tabular Workflow for Forecasting to train an AutoML model. You can choose between the following model types: Time Series Dense Encoder (TiDE), Learn to Learn (L2L), Sequence to Sequence (Seq2Seq+), and Temporal Fusion Transformer (TFT).\n",
        "\n",
        "Learn more about [Tabular Workflow for Forecasting](https://cloud.google.com/vertex-ai/docs/tabular-data/tabular-workflows/forecasting).\n",
        "\n",
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8b54ba90629a"
      },
      "source": [
        "### Advantages of tabular workflows\n",
        "\n",
        "Compared to Vertex Forecasting managed service, Tabular Workflow for Forecasting has the following advantages:\n",
        "1. Composite time series id columns are supported. You can use a combination of multiple columns as the time series id, for example, you can use either `['sku_id']` or `['sku_id', 'store_id']` as the time series id columns.\n",
        "2. Model architecture search can be skipped. You can reuse the previous model architecture search tuning result to train the model directly.\n",
        "3. Hardware customization. You can override the machine spec of the tuning and the training step, so that you can tune the training speed. You're also able to control the parallelism of the training process and the number of the final selected trials during the ensemble step.\n",
        "4. Unlimited time steps support in one single time series. There's no 3000 time steps limit in the training dataset.\n",
        "5. No upper limit for the training dataset. There's no 100M rows limit or 100GB limit in dataset size.\n",
        "6. Use all advanced features from the Vertex AI Pipelines."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f887ec5c06c5"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to create AutoML forecasting models using [Vertex AI Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction) downloaded from [Google Cloud Pipeline Components](https://cloud.google.com/vertex-ai/docs/pipelines/components-introduction) (GCPC). These pipelines are Vertex AI Tabular Workflow pipelines that are maintained by Google. These pipelines showcase different ways to customize the Vertex AI Tabular training process.\n",
        "\n",
        "This tutorial uses the following Vertex AI services:\n",
        "\n",
        "- AutoML training\n",
        "- Vertex AI Pipelines\n",
        "\n",
        "The steps performed are:\n",
        "\n",
        "- Create a training pipeline with TiDE(Time series Dense Encoder) algorithm using specified machine type for training.\n",
        "- Create a training pipeline that reuses the architecture search results from the previous pipeline to save time for TiDE(Time series Dense Encoder).\n",
        "- Create a training pipeline with Learn-to-learn(L2L) algorithm.\n",
        "- Create a training pipeline with Seq2seq(Sequence to sequence) algorithm.\n",
        "- Create a training pipeline with TFT(Temporal Fusion Transformer) algorithm.\n",
        "- Perform the batch prediction using the trained model from the above steps."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eac26958afe8"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This tutorial uses the [Liquor dataset](https://www.kaggle.com/datasets/residentmario/iowa-liquor-sales), which forecasts the alcoholic beverage sales in the Midwest."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "181d4dfbf917"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "* BigQuery\n",
        "* Dataflow\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), [BigQuery pricing](https://cloud.google.com/bigquery), and [Dataflow pricing](https://cloud.google.com/dataflow/pricing) and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d1ea81ac77f0"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e85f0288a6df"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n",
        "\n",
        "**Note**: Install the Google Cloud Pipeline Components (GCPC) SDK not earlier than `2.3.0`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E7-SzYTR9bo2"
      },
      "outputs": [],
      "source": [
        "!pip3 install --upgrade --quiet google-cloud-pipeline-components==2.3.0 \\\n",
        "                                google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bj5O0S5RTxzY"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "023DMKUaTypt"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c87a2a5d7e35"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5dccb1c8feb6"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cc7251520a07"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zebLBGXOky2A"
      },
      "source": [
        "### Notes about service account and permission\n",
        "\n",
        "For full details of the permission setup, refer to [Service accounts for Tabular Workflows](https://cloud.google.com/vertex-ai/docs/tabular-data/tabular-workflows/service-accounts).\n",
        "\n",
        "**By default no configuration is required**, if you run into any permission related issue, make sure the service accounts below have the required roles:\n",
        "\n",
        "|Service account email|Description|Roles|\n",
        "|---|---|---|\n",
        "|PROJECT_NUMBER-compute@developer.gserviceaccount.com|Compute Engine default service account|Dataflow Developer, Dataflow Worker, Storage Admin, BigQuery Data Editor, Vertex AI User, Service Account User|\n",
        "|service-PROJECT_NUMBER@gcp-sa-aiplatform.iam.gserviceaccount.com|AI Platform Service Agent|Vertex AI Service Agent|\n",
        "\n",
        "\n",
        "1. Go to [IAM console](https://console.cloud.google.com/iam-admin/iam).\n",
        "2. Check the \"Include Google-provided role grants\" checkbox.\n",
        "3. Find the above emails.\n",
        "4. Grant the corresponding roles."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ef7eccf0e6a"
      },
      "source": [
        "### Using data source from a different project\n",
        "- For the BQ data source, grant both service accounts the \"BigQuery Data Viewer\" role.\n",
        "- For the CSV data source, grant both service accounts the \"Storage Object Viewer\" role."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95cb7ffd6895"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cd85f5c794e5"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OUfwWir9yPNV"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets, TF model checkpoint, TensorBoard file, etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5325f437b46a"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ebc0bdb07af"
      },
      "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": "0217c63ed87f"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "44accda192d5"
      },
      "source": [
        "#### Service Account\n",
        "\n",
        "Use a service account to create Vertex AI Pipeline jobs. If you don't want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "604ae09ab6d3"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}\n",
        "\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    import sys\n",
        "\n",
        "    IS_COLAB = \"google.colab\" in sys.modules\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",
        "    else:  # 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": "d1ecb60964d5"
      },
      "source": [
        "#### Set service account access for Vertex AI Pipelines\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 this step once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a592f0a380c2"
      },
      "outputs": [],
      "source": [
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbbc3479a1da"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8G6YmJT1yqkV"
      },
      "outputs": [],
      "source": [
        "# Import required modules\n",
        "import json\n",
        "import os\n",
        "import uuid\n",
        "from typing import Any, Dict, List, Optional\n",
        "\n",
        "from google.cloud import aiplatform, storage\n",
        "from google_cloud_pipeline_components.preview.automl.forecasting import \\\n",
        "    utils as automl_forecasting_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c0423f260423"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ad69f2590268"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Se89OZ0jRWD6"
      },
      "source": [
        "## VPC related config\n",
        "\n",
        "If you need to use a custom Dataflow subnetwork, you can set it through the `dataflow_subnetwork` parameter. The requirements are:\n",
        "1. `dataflow_subnetwork` must be a fully qualified subnetwork name.\n",
        "   ([Example network and subnetwork specifications](https://cloud.google.com/dataflow/docs/guides/specifying-networks#example_network_and_subnetwork_specifications))\n",
        "1. The following service accounts must have [Compute Network User role](https://cloud.google.com/compute/docs/access/iam#compute.networkUser) assigned on the specified dataflow subnetwork [[reference](https://cloud.google.com/dataflow/docs/guides/specifying-networks#shared)]:\n",
        "    1. Compute Engine default service account: PROJECT_NUMBER-compute@developer.gserviceaccount.com\n",
        "    1. Dataflow service account: service-PROJECT_NUMBER@dataflow-service-producer-prod.iam.gserviceaccount.com\n",
        "\n",
        "If your project has VPC-SC enabled, make sure:\n",
        "\n",
        "1. The dataflow subnetwork used in VPC-SC is configured properly for Dataflow.\n",
        "   [[reference](https://cloud.google.com/dataflow/docs/guides/routes-firewall)]\n",
        "1. `dataflow_use_public_ips` is set to False.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jjYKwvgxRbTn"
      },
      "outputs": [],
      "source": [
        "# Dataflow's fully qualified subnetwork name, when empty the default subnetwork is used.\n",
        "# Fully qualified subnetwork name is in the form of\n",
        "# https://www.googleapis.com/compute/v1/projects/HOST_PROJECT_ID/regions/REGION_NAME/subnetworks/SUBNETWORK_NAME\n",
        "# reference: https://cloud.google.com/dataflow/docs/guides/specifying-networks#example_network_and_subnetwork_specifications\n",
        "dataflow_subnetwork = None  # @param {type:\"string\"}\n",
        "# Specifies whether Dataflow workers use public IP addresses.\n",
        "dataflow_use_public_ips = True  # @param {type:\"boolean\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cqSmR2Q3Rphx"
      },
      "source": [
        "## Prepare for training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3LWH3PRF5o2v"
      },
      "source": [
        "### Define helper functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g9FPFT8c5oC0"
      },
      "outputs": [],
      "source": [
        "# Below functions serve as the utility functions.\n",
        "\n",
        "\n",
        "# Fetch the tuple of GCS bucket and object URI.\n",
        "def get_bucket_name_and_path(uri: str):\n",
        "    no_prefix_uri = uri[len(\"gs://\") :]\n",
        "    splits = no_prefix_uri.split(\"/\")\n",
        "    return splits[0], \"/\".join(splits[1:])\n",
        "\n",
        "\n",
        "# Fetch the content from a GCS object URI.\n",
        "def download_from_gcs(uri: str):\n",
        "    bucket_name, path = get_bucket_name_and_path(uri)\n",
        "    storage_client = storage.Client(project=PROJECT_ID)\n",
        "    bucket = storage_client.get_bucket(bucket_name)\n",
        "    blob = bucket.blob(path)\n",
        "    return blob.download_as_string()\n",
        "\n",
        "\n",
        "# Upload the string content as a GCS object.\n",
        "def write_to_gcs(uri: str, content: str):\n",
        "    bucket_name, path = get_bucket_name_and_path(uri)\n",
        "    storage_client = storage.Client()\n",
        "    bucket = storage_client.get_bucket(bucket_name)\n",
        "    blob = bucket.blob(path)\n",
        "    blob.upload_from_string(content)\n",
        "\n",
        "\n",
        "# This is the example to set non-auto transformations.\n",
        "# For more details about the transformations, check:\n",
        "# https://cloud.google.com/vertex-ai/docs/datasets/data-types-tabular#transformations\n",
        "def generate_transformation(\n",
        "    auto_column_names: Optional[List[str]] = None,\n",
        "    numeric_column_names: Optional[List[str]] = None,\n",
        "    categorical_column_names: Optional[List[str]] = None,\n",
        "    text_column_names: Optional[List[str]] = None,\n",
        "    timestamp_column_names: Optional[List[str]] = None,\n",
        ") -> List[Dict[str, Any]]:\n",
        "    if auto_column_names is None:\n",
        "        auto_column_names = []\n",
        "    if numeric_column_names is None:\n",
        "        numeric_column_names = []\n",
        "    if categorical_column_names is None:\n",
        "        categorical_column_names = []\n",
        "    if text_column_names is None:\n",
        "        text_column_names = []\n",
        "    if timestamp_column_names is None:\n",
        "        timestamp_column_names = []\n",
        "    return {\n",
        "        \"auto\": auto_column_names,\n",
        "        \"numeric\": numeric_column_names,\n",
        "        \"categorical\": categorical_column_names,\n",
        "        \"text\": text_column_names,\n",
        "        \"timestamp\": timestamp_column_names,\n",
        "    }\n",
        "\n",
        "\n",
        "# Retrieve the data given a task name.\n",
        "def get_task_detail(\n",
        "    task_details: List[Dict[str, Any]], task_name: str\n",
        ") -> List[Dict[str, Any]]:\n",
        "    for task_detail in task_details:\n",
        "        if task_detail.task_name == task_name:\n",
        "            return task_detail\n",
        "\n",
        "\n",
        "# Retrieve the model resource name\n",
        "def get_deployed_model_resource(\n",
        "    task_details,\n",
        "):\n",
        "    ensemble_task = get_task_detail(task_details, \"model-upload\")\n",
        "    if ensemble_task is None:\n",
        "        ensemble_task = get_task_detail(task_details, \"model-upload-2\")\n",
        "    if ensemble_task:\n",
        "        return ensemble_task.outputs[\"model\"].artifacts[0].metadata[\"resourceName\"]\n",
        "\n",
        "\n",
        "# Retrieve the feature importance details from GCS.\n",
        "def get_feature_attributions(\n",
        "    task_details,\n",
        "):\n",
        "    ensemble_task = get_task_detail(task_details, \"model-evaluation-2\")\n",
        "    return download_from_gcs(\n",
        "        ensemble_task.outputs[\"evaluation_metrics\"]\n",
        "        .artifacts[0]\n",
        "        .metadata[\"explanation_gcs_path\"]\n",
        "    )\n",
        "\n",
        "\n",
        "# Retrieve the evaluation metrics from GCS.\n",
        "def get_evaluation_metrics(\n",
        "    task_details,\n",
        "):\n",
        "    ensemble_task = get_task_detail(task_details, \"model-evaluation\")\n",
        "    return download_from_gcs(\n",
        "        ensemble_task.outputs[\"evaluation_metrics\"].artifacts[0].uri\n",
        "    )\n",
        "\n",
        "\n",
        "# Pretty print the JSON string.\n",
        "def load_and_print_json(s):\n",
        "    parsed = json.loads(s)\n",
        "    print(json.dumps(parsed, indent=2, sort_keys=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gvNFMRmBegZq"
      },
      "source": [
        "### Define training specification"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eV4JrwB8wAkg"
      },
      "outputs": [],
      "source": [
        "root_dir = os.path.join(BUCKET_URI, f\"automl_forecasting_pipeline/run-{uuid.uuid4()}\")\n",
        "optimization_objective = \"minimize-mae\"\n",
        "time_column = \"date\"\n",
        "time_series_identifier_column = \"store_name\"\n",
        "target_column = \"sale_dollars\"\n",
        "data_source_csv_filenames = None\n",
        "data_source_bigquery_table_path = (\n",
        "    \"bq://bigquery-public-data.iowa_liquor_sales_forecasting.2020_sales_train\"\n",
        ")\n",
        "\n",
        "training_fraction = 0.8\n",
        "validation_fraction = 0.1\n",
        "test_fraction = 0.1\n",
        "\n",
        "predefined_split_key = None\n",
        "if predefined_split_key:\n",
        "    training_fraction = None\n",
        "    validation_fraction = None\n",
        "    test_fraction = None\n",
        "\n",
        "weight_column = None\n",
        "\n",
        "features = [\n",
        "    time_column,\n",
        "    target_column,\n",
        "    \"city\",\n",
        "    \"zip_code\",\n",
        "    \"county\",\n",
        "]\n",
        "\n",
        "available_at_forecast_columns = [time_column]\n",
        "unavailable_at_forecast_columns = [target_column]\n",
        "time_series_attribute_columns = [\"city\", \"zip_code\", \"county\"]\n",
        "forecast_horizon = 150\n",
        "context_window = 150\n",
        "\n",
        "transformations = generate_transformation(auto_column_names=features)\n",
        "\n",
        "# Create a Vertex managed dataset artifact.\n",
        "vertex_dataset = aiplatform.TimeSeriesDataset.create(\n",
        "    bq_source=data_source_bigquery_table_path\n",
        ")\n",
        "vertex_dataset_artifact_id = vertex_dataset.gca_resource.metadata_artifact.split(\"/\")[\n",
        "    -1\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bSCTxi48DRxz"
      },
      "source": [
        "## Supported APIs\n",
        "\n",
        "\n",
        "Currently, four model types are supported in the APIs/SDK with the utility functions:\n",
        "1. `time_series_dense_encoder`(**TiDE**): `get_time_series_dense_encoder_forecasting_pipeline_and_parameters`\n",
        "2. `learn_to_learn`(**L2L**): `get_learn_to_learn_forecasting_pipeline_and_parameters`\n",
        "3. `sequence_to_sequence`(**seq2seq**): `get_sequence_to_sequence_forecasting_pipeline_and_parameters`\n",
        "4. `temporal_fusion_transformer`(**TFT**): `get_temporal_fusion_transformer_forecasting_pipeline_and_parameters`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C9v6kvuXHCiy"
      },
      "source": [
        "### High level workflow\n",
        "\n",
        "The following code shows the general format for using the APIs:\n",
        "```python\n",
        "# Use the utility function to get the required parameters to create Vertex Pipeline job.\n",
        "template_path, parameter_values = automl_forecasting_utils.get_${MODEL_TYPE}_forecasting_pipeline_and_parameters(\n",
        "  ...\n",
        ")\n",
        "\n",
        "# Construct a Vertex Pipeline job.\n",
        "job = aiplatform.PipelineJob(\n",
        "    ...\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=template_path,\n",
        "    ...\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values=parameter_values,\n",
        "    ...\n",
        ")\n",
        "\n",
        "# Launch the Vertex Pipeline job.\n",
        "job.run()\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7gA0zRUEHO5b"
      },
      "source": [
        "### Utility function arguments\n",
        "\n",
        "The utility functions for all model types have the same arguments.\n",
        "\n",
        "`get_time_series_dense_encoder_forecasting_pipeline_and_parameters` is shown here as an example:\n",
        "\n",
        "```python\n",
        "def get_time_series_dense_encoder_forecasting_pipeline_and_parameters(\n",
        "    *,\n",
        "    project: str,\n",
        "    location: str,\n",
        "    root_dir: str,\n",
        "    target_column: str,\n",
        "    optimization_objective: str,\n",
        "    transformations: Dict[str, List[str]],\n",
        "    train_budget_milli_node_hours: float,\n",
        "    time_column: str,\n",
        "    time_series_identifier_columns: List[str],\n",
        "    time_series_attribute_columns: Optional[List[str]] = None,\n",
        "    available_at_forecast_columns: Optional[List[str]] = None,\n",
        "    unavailable_at_forecast_columns: Optional[List[str]] = None,\n",
        "    forecast_horizon: Optional[int] = None,\n",
        "    context_window: Optional[int] = None,\n",
        "    evaluated_examples_bigquery_path: Optional[str] = None,\n",
        "    window_predefined_column: Optional[str] = None,\n",
        "    window_stride_length: Optional[int] = None,\n",
        "    window_max_count: Optional[int] = None,\n",
        "    holiday_regions: Optional[List[str]] = None,\n",
        "    stage_1_num_parallel_trials: Optional[int] = None,\n",
        "    stage_1_tuning_result_artifact_uri: Optional[str] = None,\n",
        "    stage_2_num_parallel_trials: Optional[int] = None,\n",
        "    num_selected_trials: Optional[int] = None,\n",
        "    data_source_csv_filenames: Optional[str] = None,\n",
        "    data_source_bigquery_table_path: Optional[str] = None,\n",
        "    predefined_split_key: Optional[str] = None,\n",
        "    training_fraction: Optional[float] = None,\n",
        "    validation_fraction: Optional[float] = None,\n",
        "    test_fraction: Optional[float] = None,\n",
        "    weight_column: Optional[str] = None,\n",
        "    dataflow_service_account: Optional[str] = None,\n",
        "    dataflow_subnetwork: Optional[str] = None,\n",
        "    dataflow_use_public_ips: bool = True,\n",
        "    feature_transform_engine_bigquery_staging_full_dataset_id: str = '',\n",
        "    feature_transform_engine_dataflow_machine_type: str = 'n1-standard-16',\n",
        "    feature_transform_engine_dataflow_max_num_workers: int = 10,\n",
        "    feature_transform_engine_dataflow_disk_size_gb: int = 40,\n",
        "    evaluation_batch_predict_machine_type: str = 'n1-standard-16',\n",
        "    evaluation_batch_predict_starting_replica_count: int = 25,\n",
        "    evaluation_batch_predict_max_replica_count: int = 25,\n",
        "    evaluation_dataflow_machine_type: str = 'n1-standard-16',\n",
        "    evaluation_dataflow_max_num_workers: int = 25,\n",
        "    evaluation_dataflow_disk_size_gb: int = 50,\n",
        "    study_spec_parameters_override: Optional[List[Dict[str, Any]]] = None,\n",
        "    stage_1_tuner_worker_pool_specs_override: Optional[Dict[str, Any]] = None,\n",
        "    stage_2_trainer_worker_pool_specs_override: Optional[Dict[str, Any]] = None,\n",
        "    enable_probabilistic_inference: bool = False,\n",
        "    quantiles: Optional[List[float]] = None,\n",
        "    encryption_spec_key_name: Optional[str] = None,\n",
        "    model_display_name: Optional[str] = None,\n",
        "    model_description: Optional[str] = None,\n",
        "    run_evaluation: bool = True,\n",
        ") -> Tuple[str, Dict[str, Any]]:\n",
        "  \"\"\"Returns l2l_forecasting pipeline and formatted parameters.\n",
        "\n",
        "  Args:\n",
        "    project: The GCP project that runs the pipeline components.\n",
        "    location: The GCP region that runs the pipeline components.\n",
        "    root_dir: The root GCS directory for the pipeline components.\n",
        "    target_column: The target column name.\n",
        "    optimization_objective: \"minimize-rmse\", \"minimize-mae\", \"minimize-rmsle\",\n",
        "      \"minimize-rmspe\", \"minimize-wape-mae\", \"minimize-mape\", or\n",
        "      \"minimize-quantile-loss\".\n",
        "    transformations: Dict mapping auto and/or type-resolutions to feature\n",
        "      columns. The supported types are: auto, categorical, numeric, text, and\n",
        "      timestamp.\n",
        "    train_budget_milli_node_hours: The train budget of creating this model,\n",
        "      expressed in milli node hours i.e. 1,000 value in this field means 1 node\n",
        "      hour.\n",
        "    time_column: The column that indicates the time.\n",
        "    time_series_identifier_columns: The columns which distinguish different time\n",
        "      series.\n",
        "    time_series_attribute_columns: The columns that are invariant across the\n",
        "      same time series.\n",
        "    available_at_forecast_columns: The columns that are available at the\n",
        "      forecast time.\n",
        "    unavailable_at_forecast_columns: The columns that are unavailable at the\n",
        "      forecast time.\n",
        "    forecast_horizon: The length of the horizon.\n",
        "    context_window: The length of the context window.\n",
        "    evaluated_examples_bigquery_path: The existing BigQuery dataset to write the\n",
        "      predicted examples into for evaluation, in the format\n",
        "      `bq://project.dataset`. The dataset needs to be created first.\n",
        "    window_predefined_column: The column that indicate the start of each window.\n",
        "    window_stride_length: The stride length to generate the window.\n",
        "    window_max_count: The maximum number of windows that are generated.\n",
        "    holiday_regions: The geographical regions where the holiday effect is\n",
        "      applied in modeling.\n",
        "    stage_1_num_parallel_trials: Number of parallel trails for stage 1.\n",
        "    stage_1_tuning_result_artifact_uri: The stage 1 tuning result artifact GCS\n",
        "      URI.\n",
        "    stage_2_num_parallel_trials: Number of parallel trails for stage 2.\n",
        "    num_selected_trials: Number of selected trails.\n",
        "    data_source_csv_filenames: A string that represents a list of comma\n",
        "      separated CSV filenames.\n",
        "    data_source_bigquery_table_path: The BigQuery table path of format\n",
        "      bq://bq_project.bq_dataset.bq_table\n",
        "    predefined_split_key: The predefined_split column name.\n",
        "    training_fraction: The training fraction.\n",
        "    validation_fraction: The validation fraction.\n",
        "    test_fraction: The test fraction.\n",
        "    weight_column: The weight column name.\n",
        "    dataflow_service_account: The full service account name.\n",
        "    dataflow_subnetwork: The dataflow subnetwork.\n",
        "    dataflow_use_public_ips: `True` to enable dataflow public IPs.\n",
        "    feature_transform_engine_bigquery_staging_full_dataset_id: The full id of\n",
        "      the feature transform engine staging dataset.\n",
        "    feature_transform_engine_dataflow_machine_type: The dataflow machine type of\n",
        "      the feature transform engine.\n",
        "    feature_transform_engine_dataflow_max_num_workers: The max number of\n",
        "      dataflow workers of the feature transform engine.\n",
        "    feature_transform_engine_dataflow_disk_size_gb: The disk size of the\n",
        "      dataflow workers of the feature transform engine.\n",
        "    evaluation_batch_predict_machine_type: Machine type for the batch prediction\n",
        "      job in evaluation, such as 'n1-standard-16'.\n",
        "    evaluation_batch_predict_starting_replica_count: Number of replicas to use\n",
        "      in the batch prediction cluster at startup time.\n",
        "    evaluation_batch_predict_max_replica_count: The maximum count of replicas\n",
        "      the batch prediction job can scale to.\n",
        "    evaluation_dataflow_machine_type: Machine type for the dataflow job in\n",
        "      evaluation, such as 'n1-standard-16'.\n",
        "    evaluation_dataflow_max_num_workers: Maximum number of dataflow workers.\n",
        "    evaluation_dataflow_disk_size_gb: The disk space in GB for dataflow.\n",
        "    study_spec_parameters_override: The list for overriding study spec.\n",
        "    stage_1_tuner_worker_pool_specs_override: The dictionary for overriding\n",
        "      stage 1 tuner worker pool spec.\n",
        "    stage_2_trainer_worker_pool_specs_override: The dictionary for overriding\n",
        "      stage 2 trainer worker pool spec.\n",
        "    enable_probabilistic_inference: If probabilistic inference is enabled, the\n",
        "      model fits a distribution that captures the uncertainty of a\n",
        "      prediction. If quantiles are specified, then the quantiles of the\n",
        "      distribution are also returned.\n",
        "    quantiles: Quantiles to use for probabilistic inference. Up to 5 quantiles\n",
        "      are allowed of values between 0 and 1, exclusive. Represents the quantiles\n",
        "      to use for that objective. Quantiles must be unique.\n",
        "    encryption_spec_key_name: The KMS key name.\n",
        "    model_display_name: Optional display name for model.\n",
        "    model_description: Optional description.\n",
        "    run_evaluation: `True` to evaluate the ensembled model on the test split.\n",
        "  \"\"\"\n",
        "  ...\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6e4af5adfd39"
      },
      "source": [
        "### Use holiday regions\n",
        "\n",
        "For some use cases, forecasting data can be affected by holidays in regional areas. See [model training with Tabular Workflow for forecasting](https://cloud.google.com/vertex-ai/docs/tabular-data/tabular-workflows/forecasting-train#holiday-regions) for more information on holiday regions supported by forecasting.\n",
        "\n",
        "Pass in a list of strings `holiday_regions` to the pipeline parameter builder to incorporate holiday data into your training pipeline."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb17c3654bcb"
      },
      "source": [
        "## Customize the training configurations\n",
        "\n",
        "You can create a forecasting pipeline with the following customizations: \n",
        "- Change machine type and tuning / training parallelism\n",
        "- Skip evaluation\n",
        "- Skip model architecture search\n",
        "\n",
        "Instead of doing architecture search everytime, you can reuse the existing architecture search result. This can reduce the variation of the output model or the training cost. The existing architecture search result is stored in the `tuning_result_output` output of the `automl-forecasting-stage-1-tuner` component. You can load it programmatically with the API.\n",
        "\n",
        "```python\n",
        "stage_1_tuner_task = get_task_detail(\n",
        "    pipeline_task_details, \"automl-forecasting-stage-1-tuner\"\n",
        ")\n",
        "\n",
        "stage_1_tuning_result_artifact_uri = (\n",
        "    stage_1_tuner_task.outputs[\"tuning_result_output\"].artifacts[0].uri\n",
        ")\n",
        "```\n",
        "\n",
        "Use the following code snippet to customize the training configuration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dbf9aaca20a3"
      },
      "outputs": [],
      "source": [
        "# Customize the work pool for each trial during tuning.\n",
        "# Only the chief node and the evaluator node are used.\n",
        "# You can change the machine spec for these two nodes.\n",
        "worker_pool_specs_override = [\n",
        "    {\"machine_spec\": {\"machine_type\": \"n1-standard-8\"}},  # override for TF chief node\n",
        "    {},  # override for TF worker node, since it's not used, leave it empty\n",
        "    {},  # override for TF ps node, since it's not used, leave it empty\n",
        "    {\n",
        "        \"machine_spec\": {\"machine_type\": \"n1-standard-4\"}\n",
        "    },  # override for TF evaluator node\n",
        "]\n",
        "\n",
        "# Number of weak models in the final ensemble model.\n",
        "num_selected_trials = 5\n",
        "\n",
        "# Specify the evaluation setup.\n",
        "run_evaluation = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a359fc31ba3b"
      },
      "source": [
        "You can export evaluated examples from training to BigQuery by setting the parameter `evaluated_examples_bigquery_path` in the training parameters. The BigQuery path needs to point to an existing BigQuery dataset in the format `bq://project.dataset`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2abbcf372ce"
      },
      "outputs": [],
      "source": [
        "# This is ONLY available when `run_evaluation` is set to `True`.\n",
        "evaluated_examples_bigquery_path = f\"bq://{PROJECT_ID}.eval\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-iXXE14voyR"
      },
      "source": [
        "## TiDE training\n",
        "\n",
        "Time series Dense Encoder (TiDE) is an optimized dense DNN-based encoder-decoder model, which has great model quality with fast training and inference, especially for long contexts and horizons.\n",
        "\n",
        "For more details, see [Recent advances in deep long-horizon forecasting](https://ai.googleblog.com/2023/04/recent-advances-in-deep-long-horizon.html).\n",
        "\n",
        "In this tutorial, run the TiDE training pipeline twice:\n",
        "1. With model architecture search\n",
        "2. Without model architecture search "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3577d591bad1"
      },
      "source": [
        "### Run the TiDE pipeline with model architecture search"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sG46cXVueb66"
      },
      "outputs": [],
      "source": [
        "train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "(\n",
        "    template_path,\n",
        "    parameter_values,\n",
        ") = automl_forecasting_utils.get_time_series_dense_encoder_forecasting_pipeline_and_parameters(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    root_dir=root_dir,\n",
        "    target_column=target_column,\n",
        "    # `minimize-quantile-loss`\n",
        "    optimization_objective=optimization_objective,\n",
        "    transformations=transformations,\n",
        "    train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "    # Don't set `data_source_csv_filenames` and\n",
        "    # `data_source_bigquery_table_path` if you want to use Vertex managed\n",
        "    # dataset by commenting out the following two lines.\n",
        "    data_source_csv_filenames=data_source_csv_filenames,\n",
        "    data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "    weight_column=weight_column,\n",
        "    predefined_split_key=predefined_split_key,\n",
        "    training_fraction=training_fraction,\n",
        "    validation_fraction=validation_fraction,\n",
        "    test_fraction=test_fraction,\n",
        "    num_selected_trials=num_selected_trials,\n",
        "    time_column=time_column,\n",
        "    time_series_identifier_columns=[time_series_identifier_column],\n",
        "    time_series_attribute_columns=time_series_attribute_columns,\n",
        "    available_at_forecast_columns=available_at_forecast_columns,\n",
        "    unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "    forecast_horizon=forecast_horizon,\n",
        "    context_window=context_window,\n",
        "    stage_1_tuner_worker_pool_specs_override=worker_pool_specs_override,\n",
        "    dataflow_subnetwork=dataflow_subnetwork,\n",
        "    dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "    run_evaluation=run_evaluation,\n",
        "    # evaluated_examples_bigquery_path=evaluated_examples_bigquery_path,\n",
        "    dataflow_service_account=SERVICE_ACCOUNT,\n",
        "    # Quantile forecast requires `minimize-quantile-loss` as optimization objective.\n",
        "    # quantiles=[0.25, 0.5, 0.9],\n",
        "    # holiday_regions=[\"US\", \"AE\"],\n",
        ")\n",
        "\n",
        "job_id = \"tide-forecasting-{}\".format(uuid.uuid4())\n",
        "architecture_search_pipeline_job = aiplatform.PipelineJob(\n",
        "    display_name=job_id,\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=template_path,\n",
        "    job_id=job_id,\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values=parameter_values,\n",
        "    enable_caching=False,\n",
        "    # Uncomment the following line if you want to use Vertex managed dataset.\n",
        "    # input_artifacts={'vertex_dataset': vertex_dataset_artifact_id},\n",
        ")\n",
        "\n",
        "architecture_search_pipeline_job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "architecture_search_pipeline_task_details = (\n",
        "    architecture_search_pipeline_job.gca_resource.job_detail.task_details\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c5F12ZL_uZZ3"
      },
      "source": [
        "### Run the TiDE pipeline without the model architecture search\n",
        "\n",
        "After retrieving the tuning result from the stage 1 tuner, you can use it to skip the model architecture search."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4x4GA5sMuewX"
      },
      "outputs": [],
      "source": [
        "# Retrieve the tuning result output from the previous training pipeline.\n",
        "stage_1_tuner_task = get_task_detail(\n",
        "    architecture_search_pipeline_task_details, \"automl-forecasting-stage-1-tuner\"\n",
        ")\n",
        "\n",
        "stage_1_tuning_result_artifact_uri = (\n",
        "    stage_1_tuner_task.outputs[\"tuning_result_output\"].artifacts[0].uri\n",
        ")\n",
        "\n",
        "train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "(\n",
        "    template_path,\n",
        "    parameter_values,\n",
        ") = automl_forecasting_utils.get_time_series_dense_encoder_forecasting_pipeline_and_parameters(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    root_dir=root_dir,\n",
        "    target_column=target_column,\n",
        "    optimization_objective=optimization_objective,\n",
        "    transformations=transformations,\n",
        "    train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "    data_source_csv_filenames=data_source_csv_filenames,\n",
        "    data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "    weight_column=weight_column,\n",
        "    predefined_split_key=predefined_split_key,\n",
        "    training_fraction=training_fraction,\n",
        "    validation_fraction=validation_fraction,\n",
        "    test_fraction=test_fraction,\n",
        "    num_selected_trials=num_selected_trials,\n",
        "    time_column=time_column,\n",
        "    time_series_identifier_columns=[time_series_identifier_column],\n",
        "    time_series_attribute_columns=time_series_attribute_columns,\n",
        "    available_at_forecast_columns=available_at_forecast_columns,\n",
        "    unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "    forecast_horizon=forecast_horizon,\n",
        "    context_window=context_window,\n",
        "    dataflow_subnetwork=dataflow_subnetwork,\n",
        "    dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "    stage_1_tuning_result_artifact_uri=stage_1_tuning_result_artifact_uri,\n",
        "    run_evaluation=run_evaluation,\n",
        "    # evaluated_examples_bigquery_path=evaluated_examples_bigquery_path,\n",
        "    dataflow_service_account=SERVICE_ACCOUNT,\n",
        ")\n",
        "\n",
        "job_id = \"tide-forecasting-skip-architecture-search-{}\".format(uuid.uuid4())\n",
        "skip_architecture_search_pipeline_job = aiplatform.PipelineJob(\n",
        "    display_name=job_id,\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=template_path,\n",
        "    job_id=job_id,\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values=parameter_values,\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "skip_architecture_search_pipeline_job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "# Get model URI\n",
        "skip_architecture_search_pipeline_task_details = (\n",
        "    skip_architecture_search_pipeline_job.gca_resource.job_detail.task_details\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xiftLomOwGda"
      },
      "source": [
        "## L2L training\n",
        "\n",
        "Learn-to-Learn (L2L) is a good choice for a wide range of the time series forecasting use cases."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XSAcS70N1GMN"
      },
      "outputs": [],
      "source": [
        "train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "(\n",
        "    template_path,\n",
        "    parameter_values,\n",
        ") = automl_forecasting_utils.get_learn_to_learn_forecasting_pipeline_and_parameters(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    root_dir=root_dir,\n",
        "    target_column=target_column,\n",
        "    optimization_objective=optimization_objective,\n",
        "    transformations=transformations,\n",
        "    train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "    data_source_csv_filenames=data_source_csv_filenames,\n",
        "    data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "    weight_column=weight_column,\n",
        "    predefined_split_key=predefined_split_key,\n",
        "    training_fraction=training_fraction,\n",
        "    validation_fraction=validation_fraction,\n",
        "    test_fraction=test_fraction,\n",
        "    num_selected_trials=num_selected_trials,\n",
        "    time_column=time_column,\n",
        "    time_series_identifier_columns=[time_series_identifier_column],\n",
        "    time_series_attribute_columns=time_series_attribute_columns,\n",
        "    available_at_forecast_columns=available_at_forecast_columns,\n",
        "    unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "    forecast_horizon=forecast_horizon,\n",
        "    context_window=context_window,\n",
        "    dataflow_subnetwork=dataflow_subnetwork,\n",
        "    dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "    run_evaluation=run_evaluation,\n",
        "    # evaluated_examples_bigquery_path=evaluated_examples_bigquery_path,\n",
        "    dataflow_service_account=SERVICE_ACCOUNT,\n",
        "    # Quantile forecast requires `minimize-quantile-loss` as optimization objective.\n",
        "    # quantiles=[0.25, 0.5, 0.9],\n",
        ")\n",
        "\n",
        "job_id = \"l2l-forecasting-{}\".format(uuid.uuid4())\n",
        "l2l_pipeline_job = aiplatform.PipelineJob(\n",
        "    display_name=job_id,\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=template_path,\n",
        "    job_id=job_id,\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values=parameter_values,\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "l2l_pipeline_job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "l2l_pipeline_task_details = l2l_pipeline_job.gca_resource.job_detail.task_details"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1LvCFqFtI-1t"
      },
      "source": [
        "## Seq2seq training\n",
        "\n",
        "Sequence-to-sequence (seq2seq) is a good choice for experimentation. The algorithm is likely to converge faster than AutoML because its architecture is simpler and it uses a smaller search space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "euK31b8xMJ8i"
      },
      "outputs": [],
      "source": [
        "train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "(\n",
        "    template_path,\n",
        "    parameter_values,\n",
        ") = automl_forecasting_utils.get_sequence_to_sequence_forecasting_pipeline_and_parameters(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    root_dir=root_dir,\n",
        "    target_column=target_column,\n",
        "    optimization_objective=optimization_objective,\n",
        "    transformations=transformations,\n",
        "    train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "    data_source_csv_filenames=data_source_csv_filenames,\n",
        "    data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "    weight_column=weight_column,\n",
        "    predefined_split_key=predefined_split_key,\n",
        "    training_fraction=training_fraction,\n",
        "    validation_fraction=validation_fraction,\n",
        "    test_fraction=test_fraction,\n",
        "    num_selected_trials=num_selected_trials,\n",
        "    time_column=time_column,\n",
        "    time_series_identifier_columns=[time_series_identifier_column],\n",
        "    time_series_attribute_columns=time_series_attribute_columns,\n",
        "    available_at_forecast_columns=available_at_forecast_columns,\n",
        "    unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "    forecast_horizon=forecast_horizon,\n",
        "    context_window=context_window,\n",
        "    dataflow_subnetwork=dataflow_subnetwork,\n",
        "    dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "    run_evaluation=run_evaluation,\n",
        "    # evaluated_examples_bigquery_path=evaluated_examples_bigquery_path,\n",
        "    dataflow_service_account=SERVICE_ACCOUNT,\n",
        "    # Quantile prediction is NOT supported by Seq2seq.\n",
        ")\n",
        "\n",
        "job_id = \"seq2seq-forecasting-{}\".format(uuid.uuid4())\n",
        "seq2seq_pipeline_job = aiplatform.PipelineJob(\n",
        "    display_name=job_id,\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=template_path,\n",
        "    job_id=job_id,\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values=parameter_values,\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "seq2seq_pipeline_job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "seq2seq_pipeline_task_details = (\n",
        "    seq2seq_pipeline_job.gca_resource.job_detail.task_details\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b8PdaaWdJCh-"
      },
      "source": [
        "## TFT training\n",
        "\n",
        "TFT stands for \"Temporal Fusion Transformer\", which is an attention-based DNN model designed to produce high accuracy and interpretability by aligning the model with the general multi-horizon forecasting task.\n",
        "\n",
        "With this model, you don't need to explicitly enable the explainability support during serving to get the feature importance for each feature column."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZjCEok0IMT-8"
      },
      "outputs": [],
      "source": [
        "train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "(\n",
        "    template_path,\n",
        "    parameter_values,\n",
        ") = automl_forecasting_utils.get_temporal_fusion_transformer_forecasting_pipeline_and_parameters(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    root_dir=root_dir,\n",
        "    target_column=target_column,\n",
        "    optimization_objective=optimization_objective,\n",
        "    transformations=transformations,\n",
        "    train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "    data_source_csv_filenames=data_source_csv_filenames,\n",
        "    data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "    weight_column=weight_column,\n",
        "    predefined_split_key=predefined_split_key,\n",
        "    training_fraction=training_fraction,\n",
        "    validation_fraction=validation_fraction,\n",
        "    test_fraction=test_fraction,\n",
        "    # Note that TFT model ONLY ensembles the model from\n",
        "    # the top one trial, so `num_selected_trials` can't be set for TFT model.\n",
        "    # num_selected_trials=num_selected_trials,\n",
        "    time_column=time_column,\n",
        "    time_series_identifier_columns=[time_series_identifier_column],\n",
        "    time_series_attribute_columns=time_series_attribute_columns,\n",
        "    available_at_forecast_columns=available_at_forecast_columns,\n",
        "    unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "    forecast_horizon=forecast_horizon,\n",
        "    context_window=context_window,\n",
        "    dataflow_subnetwork=dataflow_subnetwork,\n",
        "    dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "    run_evaluation=run_evaluation,\n",
        "    # evaluated_examples_bigquery_path=evaluated_examples_bigquery_path,\n",
        "    dataflow_service_account=SERVICE_ACCOUNT,\n",
        "    # Quantile prediction is NOT supported by TFT.\n",
        ")\n",
        "\n",
        "job_id = \"tft-forecasting-{}\".format(uuid.uuid4())\n",
        "tft_pipeline_job = aiplatform.PipelineJob(\n",
        "    display_name=job_id,\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=template_path,\n",
        "    job_id=job_id,\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values=parameter_values,\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "tft_pipeline_job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "tft_pipeline_task_details = tft_pipeline_job.gca_resource.job_detail.task_details"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L38LI3OOwXkV"
      },
      "source": [
        "## Batch prediction/explain\n",
        "\n",
        "Enable the batch explain feature by simply setting `generate_explanation=True` in the `batch_predict` API.\n",
        "\n",
        "Use the following code to retrieve the trained forecasting model from the pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0bf184357639"
      },
      "outputs": [],
      "source": [
        "upload_model_task = get_task_detail(tft_pipeline_task_details, \"model-upload-2\")\n",
        "\n",
        "forecasting_mp_model_artifact = upload_model_task.outputs[\"model\"].artifacts[0]\n",
        "\n",
        "forecasting_mp_model = aiplatform.Model(\n",
        "    forecasting_mp_model_artifact.metadata[\"resourceName\"]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9327a8a390a7"
      },
      "source": [
        "Once you retrieve the Vertex AI model, you can start to perform the batch prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NfssCxqRg1x-"
      },
      "outputs": [],
      "source": [
        "print(f\"Running Batch prediction for model: {forecasting_mp_model.display_name}\")\n",
        "\n",
        "batch_predict_bq_output_uri_prefix = f\"bq://{PROJECT_ID}\"\n",
        "\n",
        "PREDICTION_DATASET_BQ_PATH = (\n",
        "    \"bq://bigquery-public-data:iowa_liquor_sales_forecasting.2021_sales_predict\"\n",
        ")\n",
        "\n",
        "batch_prediction_job = forecasting_mp_model.batch_predict(\n",
        "    job_display_name=\"forecasting_iowa_liquor_sales_forecasting_predictions\",\n",
        "    bigquery_source=PREDICTION_DATASET_BQ_PATH,\n",
        "    instances_format=\"bigquery\",\n",
        "    bigquery_destination_prefix=batch_predict_bq_output_uri_prefix,\n",
        "    predictions_format=\"bigquery\",\n",
        "    # Uncomment the following line to run batch explain:\n",
        "    # generate_explanation=True,\n",
        "    sync=True,\n",
        ")\n",
        "\n",
        "print(batch_prediction_job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PLDc1X8DVSmm"
      },
      "source": [
        "## Retrieve the uploaded Vertex AI model with a Vertex AI pipeline job id"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MJ_2TaYtTNbO"
      },
      "outputs": [],
      "source": [
        "# Example format of pipeline_job_id: projects/{your-project-id}/locations/us-central1/pipelineJobs/{pipeline-job-id}\n",
        "pipeline_job_id = \"\"  # @param {type:\"string\"}\n",
        "if pipeline_job_id:\n",
        "    job = aiplatform.PipelineJob.get(pipeline_job_id)\n",
        "    pipeline_task_details = job.gca_resource.job_detail.task_details\n",
        "    upload_model_task = get_task_detail(pipeline_task_details, \"model-upload-2\")\n",
        "\n",
        "    forecasting_mp_model_artifact = upload_model_task.outputs[\"model\"].artifacts[0]\n",
        "    forecasting_mp_model = aiplatform.Model(\n",
        "        forecasting_mp_model_artifact.metadata[\"resourceName\"]\n",
        "    )\n",
        "    print(forecasting_mp_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KtcHUmcZIi9g"
      },
      "source": [
        "## Upload with parent model for different model versions\n",
        "\n",
        "To upload this model to a parent Vertex AI model, you need the `parent_model_resource_name` i.e., the resource name of the parent Vertex AI model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6-bzpVGTEq14"
      },
      "outputs": [],
      "source": [
        "# The model resource name can be something like: \"projects/{your-project-id}/locations/us-central1/models/{model-id}\"\n",
        "parent_model_resource_name = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "if parent_model_resource_name:\n",
        "    parent_model_artifact = aiplatform.Artifact.get_with_uri(\n",
        "        f\"https://{LOCATION}-aiplatform.googleapis.com/v1/\" + parent_model_resource_name\n",
        "    )\n",
        "    parent_model_artifact_id = str(\n",
        "        parent_model_artifact.gca_resource.name.split(\"artifacts/\")[1]\n",
        "    )\n",
        "\n",
        "    train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "    (\n",
        "        template_path,\n",
        "        parameter_values,\n",
        "    ) = automl_forecasting_utils.get_time_series_dense_encoder_forecasting_pipeline_and_parameters(\n",
        "        project=PROJECT_ID,\n",
        "        location=LOCATION,\n",
        "        root_dir=root_dir,\n",
        "        target_column=target_column,\n",
        "        optimization_objective=optimization_objective,\n",
        "        transformations=transformations,\n",
        "        train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "        # Don't set `data_source_csv_filenames` and\n",
        "        # `data_source_bigquery_table_path` if you want to use Vertex managed\n",
        "        # dataset by commenting out the following two lines.\n",
        "        data_source_csv_filenames=data_source_csv_filenames,\n",
        "        data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "        weight_column=weight_column,\n",
        "        predefined_split_key=predefined_split_key,\n",
        "        training_fraction=training_fraction,\n",
        "        validation_fraction=validation_fraction,\n",
        "        test_fraction=test_fraction,\n",
        "        num_selected_trials=5,\n",
        "        time_column=time_column,\n",
        "        time_series_identifier_columns=[time_series_identifier_column],\n",
        "        time_series_attribute_columns=time_series_attribute_columns,\n",
        "        available_at_forecast_columns=available_at_forecast_columns,\n",
        "        unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "        forecast_horizon=forecast_horizon,\n",
        "        context_window=context_window,\n",
        "        dataflow_subnetwork=dataflow_subnetwork,\n",
        "        dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "        run_evaluation=run_evaluation,\n",
        "        # evaluated_examples_bigquery_path=evaluated_examples_bigquery_path,\n",
        "        dataflow_service_account=SERVICE_ACCOUNT,\n",
        "        # Quantile forecast requires `minimize-quantile-loss` as optimization objective.\n",
        "        # quantiles=[0.25, 0.5, 0.9],\n",
        "    )\n",
        "\n",
        "    job_id = \"tide-forecasting-with-parent-model-{}\".format(uuid.uuid4())\n",
        "    job = aiplatform.PipelineJob(\n",
        "        display_name=job_id,\n",
        "        location=LOCATION,  # launches the pipeline job in the specified region\n",
        "        template_path=template_path,\n",
        "        job_id=job_id,\n",
        "        pipeline_root=root_dir,\n",
        "        parameter_values=parameter_values,\n",
        "        enable_caching=False,\n",
        "        input_artifacts={\"parent_model\": parent_model_artifact_id},\n",
        "    )\n",
        "\n",
        "    job.run(service_account=SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Bu0wywvPYkD"
      },
      "source": [
        "## Integrate Tabular Workflow for Forecasting with your existing KFP pipeline\n",
        "\n",
        "In this section, you define and run a sample KFP pipeline with your Tabular Workflow pipeline integrated as a component. For this, use [the pipeline-as-component feature](https://www.kubeflow.org/docs/components/pipelines/v2/load-and-share-components/) of KFP."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OcWHUuoYPbSi"
      },
      "outputs": [],
      "source": [
        "from kfp import compiler, components, dsl\n",
        "\n",
        "train_budget_milli_node_hours = 250.0  # 15 minutes\n",
        "\n",
        "(\n",
        "    template_path,\n",
        "    parameter_values,\n",
        ") = automl_forecasting_utils.get_time_series_dense_encoder_forecasting_pipeline_and_parameters(\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    root_dir=root_dir,\n",
        "    target_column=target_column,\n",
        "    optimization_objective=optimization_objective,\n",
        "    transformations=transformations,\n",
        "    train_budget_milli_node_hours=train_budget_milli_node_hours,\n",
        "    data_source_csv_filenames=data_source_csv_filenames,\n",
        "    data_source_bigquery_table_path=data_source_bigquery_table_path,\n",
        "    weight_column=weight_column,\n",
        "    predefined_split_key=predefined_split_key,\n",
        "    training_fraction=training_fraction,\n",
        "    validation_fraction=validation_fraction,\n",
        "    test_fraction=test_fraction,\n",
        "    num_selected_trials=num_selected_trials,\n",
        "    time_column=time_column,\n",
        "    time_series_identifier_columns=[time_series_identifier_column],\n",
        "    time_series_attribute_columns=time_series_attribute_columns,\n",
        "    available_at_forecast_columns=available_at_forecast_columns,\n",
        "    unavailable_at_forecast_columns=unavailable_at_forecast_columns,\n",
        "    forecast_horizon=forecast_horizon,\n",
        "    context_window=context_window,\n",
        "    dataflow_subnetwork=dataflow_subnetwork,\n",
        "    dataflow_use_public_ips=dataflow_use_public_ips,\n",
        "    run_evaluation=False,\n",
        "    dataflow_service_account=SERVICE_ACCOUNT,\n",
        ")\n",
        "\n",
        "# Load the forecasting pipeline as a sub-pipeline/components which can be used\n",
        "# in a larger KFP pipeline.\n",
        "forecasting_pipeline = components.load_component_from_file(template_path)\n",
        "\n",
        "\n",
        "@dsl.component\n",
        "def print_message(msg: str):\n",
        "    print(\"message:\", msg)\n",
        "\n",
        "\n",
        "# Define a pipeline that follows the below steps:\n",
        "# step_1(print_message) -> step_2(print_message) -> forecasting_pipeline\n",
        "@dsl.pipeline\n",
        "def outer_pipeline(msg_1: str, msg_2: str, ds: dsl.Artifact):\n",
        "    step_1 = print_message(msg=msg_1)\n",
        "    step_2 = print_message(msg=msg_2).after(step_1)\n",
        "    # `vertex_dataset` argument needs to be set/forwarded here to avoid the\n",
        "    # \"missing-argument\" error in KFP pipeline.\n",
        "    forecasting_pipeline(**parameter_values, vertex_dataset=ds).after(step_2)\n",
        "\n",
        "\n",
        "# Compile and save the outer/larger pipeline template.\n",
        "outer_pipeline_template_path = \"./outer_pipeline.yaml\"\n",
        "compiler.Compiler().compile(outer_pipeline, outer_pipeline_template_path)\n",
        "\n",
        "\n",
        "job_id = \"run-forecasting-pipeline-inside-pipeline-{}\".format(uuid.uuid4())\n",
        "tabular_workflow_pipeline_job = aiplatform.PipelineJob(\n",
        "    display_name=job_id,\n",
        "    location=LOCATION,  # launches the pipeline job in the specified region\n",
        "    template_path=outer_pipeline_template_path,\n",
        "    job_id=job_id,\n",
        "    pipeline_root=root_dir,\n",
        "    parameter_values={\"msg_1\": \"step 1\", \"msg_2\": \"step 2\"},\n",
        "    enable_caching=False,\n",
        ")\n",
        "\n",
        "tabular_workflow_pipeline_job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "tabular_workflow_pipeline_task_details = (\n",
        "    tabular_workflow_pipeline_job.gca_resource.job_detail.task_details\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2992dc3f6019"
      },
      "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": "1b9030d9a080"
      },
      "outputs": [],
      "source": [
        "# Delete resources from TiDE pipeline with architecture search\n",
        "model_resource_name = get_deployed_model_resource(\n",
        "    architecture_search_pipeline_task_details\n",
        ")\n",
        "# Load the model resource\n",
        "model = aiplatform.Model(model_resource_name)\n",
        "if model:\n",
        "    # Delete the model\n",
        "    model.delete()\n",
        "# Delete the pipeline job\n",
        "architecture_search_pipeline_job.delete()\n",
        "\n",
        "# Delete resources from TiDE pipeline without architecture search\n",
        "model_resource_name = get_deployed_model_resource(\n",
        "    skip_architecture_search_pipeline_task_details\n",
        ")\n",
        "# Load the model resource\n",
        "model = aiplatform.Model(model_resource_name)\n",
        "if model:\n",
        "    # Delete the model\n",
        "    model.delete()\n",
        "# Delete the pipeline job\n",
        "skip_architecture_search_pipeline_job.delete()\n",
        "\n",
        "# Delete resources from L2L training pipeline\n",
        "model_resource_name = get_deployed_model_resource(l2l_pipeline_task_details)\n",
        "# Load the model resource\n",
        "model = aiplatform.Model(model_resource_name)\n",
        "if model:\n",
        "    # Delete the model\n",
        "    model.delete()\n",
        "# Delete the pipeline job\n",
        "l2l_pipeline_job.delete()\n",
        "\n",
        "# Delete resources from seq2seq training pipeline\n",
        "model_resource_name = get_deployed_model_resource(seq2seq_pipeline_task_details)\n",
        "# Load the model resource\n",
        "model = aiplatform.Model(model_resource_name)\n",
        "if model:\n",
        "    # Delete the model\n",
        "    model.delete()\n",
        "# Delete the pipeline job\n",
        "seq2seq_pipeline_job.delete()\n",
        "\n",
        "# Delete resources from TFT training pipeline\n",
        "model_resource_name = get_deployed_model_resource(tft_pipeline_task_details)\n",
        "# Load the model resource\n",
        "model = aiplatform.Model(model_resource_name)\n",
        "if model:\n",
        "    # Delete the model\n",
        "    model.delete()\n",
        "# Delete the pipeline job\n",
        "tft_pipeline_job.delete()\n",
        "\n",
        "# Delete resources from Tabular Workflow training pipeline\n",
        "model_resource_name = get_deployed_model_resource(\n",
        "    tabular_workflow_pipeline_task_details\n",
        ")\n",
        "# Load the model resource\n",
        "model = aiplatform.Model(model_resource_name)\n",
        "if model:\n",
        "    # Delete the model\n",
        "    model.delete()\n",
        "\n",
        "# Delete the pipeline job\n",
        "tabular_workflow_pipeline_job.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "829cd75564f1"
      },
      "outputs": [],
      "source": [
        "# Delete the batch prediction job\n",
        "batch_prediction_job.delete()\n",
        "\n",
        "# Delete the Vertex AI Dataset\n",
        "vertex_dataset.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = True\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "automl_forecasting_on_vertex_pipelines.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}