Framework Workflows/CatBoost/CatBoost Prediction With Vertex AI Feature Store.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e3a0c0a0",
   "metadata": {},
   "source": [
    "![tracker](https://us-central1-vertex-ai-mlops-369716.cloudfunctions.net/pixel-tracking?path=statmike%2Fvertex-ai-mlops%2FFramework+Workflows%2FCatBoost&file=CatBoost+Prediction+With+Vertex+AI+Feature+Store.ipynb)\n",
    "<!--- header table --->\n",
    "<table align=\"left\">\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://colab.research.google.com/github/statmike/vertex-ai-mlops/blob/main/Framework%20Workflows/CatBoost/CatBoost%20Prediction%20With%20Vertex%20AI%20Feature%20Store.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\">\n",
    "      <br>Run in<br>Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https%3A%2F%2Fraw.githubusercontent.com%2Fstatmike%2Fvertex-ai-mlops%2Fmain%2FFramework%2520Workflows%2FCatBoost%2FCatBoost%2520Prediction%2520With%2520Vertex%2520AI%2520Feature%2520Store.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\">\n",
    "      <br>Run in<br>Colab Enterprise\n",
    "    </a>\n",
    "  </td>      \n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://github.com/statmike/vertex-ai-mlops/blob/main/Framework%20Workflows/CatBoost/CatBoost%20Prediction%20With%20Vertex%20AI%20Feature%20Store.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
    "      <br>View on<br>GitHub\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/statmike/vertex-ai-mlops/main/Framework%20Workflows/CatBoost/CatBoost%20Prediction%20With%20Vertex%20AI%20Feature%20Store.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
    "      <br>Open in<br>Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cb89d04-8870-46db-a4fc-9289a2bb34f2",
   "metadata": {},
   "source": [
    "# CatBoost Prediction With Vertex AI Feature Store\n",
    "\n",
    "Serving a model, like a CatBoost model, can be complex. It involves several steps:\n",
    "\n",
    "- Retrieving the features\n",
    "- Ensuring all features are present and in the correct order\n",
    "- Verifying the schema to guarantee each feature has the correct data type\n",
    "- Preparing the serving instance for the model's prediction function\n",
    "- Sending the serving instance to the model\n",
    "- Processing and formatting the response\n",
    "- Returning the response to the requestor\n",
    "\n",
    "With so many steps, there's potential for errors, leading to **training-serving skew**. A great way to minimize this is by using a **feature store** for serving.\n",
    "\n",
    "A feature store's offline component resembles training data. In this case, the model was trained on data from a BigQuery table (see the [CatBoost Overview](./CatBoost%20Overview.ipynb) workflow).\n",
    "\n",
    "Hosting an online version of the data in a feature store allows for **fast and simple feature retrieval** during serving requests. Whether it's a single BigQuery table synced directly or features processed through the feature registry, it ensures features are readily available at inference time.\n",
    "\n",
    "This workflow provides a comprehensive guide to get you started:\n",
    "\n",
    "- What is Vertex AI Feature Store?\n",
    "- Setup Vertex AI Feature Store\n",
    "- Sync BigQuery data to an online feature view in Vertex AI Feature Store\n",
    "- Incorporate retrieval from Vertex AI Feature Store in inference locally (in this notebook)\n",
    "- Build A Custom Prediction Container That Incorporates feature retrieval from Vertex AI Feature Store\n",
    "    - A single container with multiple serving profiles used to:\n",
    "        - Test the container locally with Docker\n",
    "        - Deploy and use the container on Google Cloud Run\n",
    "        - Deploy and use the container on Vertex AI Endpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e74ce9c2-6ced-4e4a-971a-ead35777033a",
   "metadata": {
    "id": "od_UkDpvRmgD",
    "tags": []
   },
   "source": [
    "---\n",
    "## Colab Setup\n",
    "\n",
    "To run this notebook in Colab run the cells in this section.  Otherwise, skip this section.\n",
    "\n",
    "This cell will authenticate to GCP (follow prompts in the popup)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c81b95f8-f452-4664-8177-832c2e8abd4e",
   "metadata": {
    "executionInfo": {
     "elapsed": 195,
     "status": "ok",
     "timestamp": 1683726184843,
     "user": {
      "displayName": "Mike Henderson",
      "userId": "07691629187611687318"
     },
     "user_tz": 240
    },
    "id": "8UO9FnqyKBlF",
    "tags": []
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = 'statmike-mlops-349915' # replace with project ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e84c4914-8621-4638-bc85-a261c9a17c81",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 68869,
     "status": "ok",
     "timestamp": 1683726253709,
     "user": {
      "displayName": "Mike Henderson",
      "userId": "07691629187611687318"
     },
     "user_tz": 240
    },
    "id": "N98-KK7LRkjm",
    "outputId": "09ec5008-0def-4e1a-c349-c598ee752f78",
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import google.colab\n",
    "    from google.colab import auth\n",
    "    auth.authenticate_user()\n",
    "    !gcloud config set project {PROJECT_ID}\n",
    "except Exception:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a8e0de7-70a5-4dc4-ad3d-9d4082cc5a56",
   "metadata": {},
   "source": [
    "---\n",
    "## Installs\n",
    "\n",
    "The list `packages` contains tuples of package import names and install names.  If the import name is not found then the install name is used to install quitely for the current user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "872d2193-a4d3-446b-a21f-0b84e832bc00",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# tuples of (import name, install name, min_version)\n",
    "packages = [\n",
    "    ('numpy', 'numpy'),\n",
    "    ('catboost', 'catboost'),\n",
    "    ('docker', 'docker'),\n",
    "    ('google.cloud.aiplatform', 'google-cloud-aiplatform'),\n",
    "    ('google.cloud.bigquery', 'google-cloud-bigquery'),\n",
    "    ('google.cloud.storage', 'google-cloud-storage'),\n",
    "    ('google.cloud.artifactregistry_v1', 'google-cloud-artifact-registry'),\n",
    "    ('google.cloud.devtools', 'google-cloud-build'),\n",
    "    ('google.cloud.run_v2', 'google-cloud-run'),   \n",
    "]\n",
    "\n",
    "import importlib\n",
    "install = False\n",
    "for package in packages:\n",
    "    if not importlib.util.find_spec(package[0]):\n",
    "        print(f'installing package {package[1]}')\n",
    "        install = True\n",
    "        !pip install {package[1]} -U -q --user\n",
    "    elif len(package) == 3:\n",
    "        if importlib.metadata.version(package[0]) < package[2]:\n",
    "            print(f'updating package {package[1]}')\n",
    "            install = True\n",
    "            !pip install {package[1]} -U -q --user"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "603a58c1-5efb-4895-90c5-92dabd42675a",
   "metadata": {},
   "source": [
    "### API Enablement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "15cdecc5-82dd-47dd-a533-94ca72f4eb11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!gcloud services enable artifactregistry.googleapis.com\n",
    "!gcloud services enable cloudbuild.googleapis.com\n",
    "!gcloud services enable run.googleapis.com"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8902533-42e4-46df-a21f-b5010c5ea2c8",
   "metadata": {},
   "source": [
    "### Restart Kernel (If Installs Occured)\n",
    "\n",
    "After a kernel restart the code submission can start with the next cell after this one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77aeee3d-d013-4a06-8113-be0e53d4c1b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if install:\n",
    "    import IPython\n",
    "    app = IPython.Application.instance()\n",
    "    app.kernel.do_shutdown(True)\n",
    "    IPython.display.display(IPython.display.Markdown(\"\"\"<div class=\\\"alert alert-block alert-warning\\\">\n",
    "        <b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. The previous cells do not need to be run again⚠️</b>\n",
    "        </div>\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4581e5d6-11d7-4d5a-a8f2-b9b1059b182d",
   "metadata": {},
   "source": [
    "---\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "578c3b68-22ac-4351-975c-6e4d274c9acb",
   "metadata": {},
   "source": [
    "inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "989c8bb8-96c4-406d-9c4a-b16328cf0d9a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'statmike-mlops-349915'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "project = !gcloud config get-value project\n",
    "PROJECT_ID = project[0]\n",
    "PROJECT_ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4174d343-6569-42cc-881a-a76d93915556",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1026793852137-compute@developer.gserviceaccount.com'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SERVICE_ACCOUNT = !gcloud config list --format='value(core.account)' \n",
    "SERVICE_ACCOUNT = SERVICE_ACCOUNT[0]\n",
    "SERVICE_ACCOUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b84f0eb9-080e-47eb-932e-c1c30d1ad140",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "REGION = 'us-central1'\n",
    "SERIES = 'frameworks'\n",
    "EXPERIMENT = 'catboost-prediction-feature-store'\n",
    "\n",
    "# GCS Names\n",
    "GCS_BUCKET = PROJECT_ID\n",
    "\n",
    "# make this the BigQuery Project / Dataset / Table prefix to store results\n",
    "BQ_PROJECT = PROJECT_ID\n",
    "BQ_DATASET = SERIES.replace('-', '_')\n",
    "BQ_TABLE = SERIES\n",
    "BQ_REGION = REGION[0:2] # use a multi region\n",
    "\n",
    "# Vertex AI Feature Store names:\n",
    "FS_NAME = PROJECT_ID.replace('-', '_')\n",
    "FV_NAME = f\"{SERIES}\".replace('-', '_')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad6a6565-91b0-4767-89c4-a070dcd16409",
   "metadata": {},
   "source": [
    "packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4ed1e444-3908-4118-aef6-2880f67a8e2b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json, os\n",
    "import time, datetime\n",
    "import requests\n",
    "\n",
    "import catboost \n",
    "import numpy as np\n",
    "import docker\n",
    "\n",
    "import google.auth\n",
    "from google.cloud import storage\n",
    "from google.cloud import artifactregistry_v1\n",
    "from google.cloud.devtools import cloudbuild_v1\n",
    "from google.cloud import run_v2\n",
    "from google.cloud import bigquery\n",
    "\n",
    "from google.cloud import aiplatform\n",
    "from vertexai.resources.preview import feature_store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "46ca9298-b170-4d6c-a5ac-574f76e8005b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.78.0'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aiplatform.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "894458bc-77e6-4c28-a74e-3180d9d83776",
   "metadata": {},
   "source": [
    "clients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "028a565c-f725-4395-b4e8-05fc59a96a11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# gcs storage client\n",
    "gcs = storage.Client(project = GCS_BUCKET)\n",
    "bucket = gcs.bucket(GCS_BUCKET)\n",
    "\n",
    "# cloud build client\n",
    "cb = cloudbuild_v1.CloudBuildClient()\n",
    "\n",
    "# artifact registry client\n",
    "ar = artifactregistry_v1.ArtifactRegistryClient()\n",
    "\n",
    "# cloud run client\n",
    "cr = run_v2.ServicesClient()\n",
    "\n",
    "# BigQuery client\n",
    "bq = bigquery.Client(project = PROJECT_ID)\n",
    "\n",
    "# vertex ai client\n",
    "aiplatform.init(project = PROJECT_ID, location = REGION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48f70c07-e258-4e9c-95d5-b006137ee80b",
   "metadata": {},
   "source": [
    "Parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6742dfd7-5319-4857-9299-4666a0ad3336",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "DIR = f\"files/{EXPERIMENT}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be1bf407-de9a-4ae3-bac4-49873e004691",
   "metadata": {},
   "source": [
    "Environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3f360519-4393-47ba-93aa-651eac1b2a86",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(DIR):\n",
    "    os.makedirs(DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5d086fb-d441-4fc6-9e8d-487e3ad51a1c",
   "metadata": {},
   "source": [
    "---\n",
    "## CatBoost Model\n",
    "\n",
    "Retrieve the model trained in prior workflow along with test records.  Test the model directly in this environment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19f411a2-97c3-4048-8184-476c9bd489e6",
   "metadata": {},
   "source": [
    "### Check For Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "068b8a1a-578c-451a-ad4f-4cc5c9600593",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found the files created by the prerequisite workflow:\n",
      "- gs://statmike-mlops-349915/frameworks/catboost-overview/examples.json\n",
      "- gs://statmike-mlops-349915/frameworks/catboost-overview/model.cbm\n"
     ]
    }
   ],
   "source": [
    "files = list(bucket.list_blobs(prefix = f'{SERIES}/catboost-overview'))\n",
    "if len(files) > 0:\n",
    "    print('Found the files created by the prerequisite workflow:')\n",
    "    for file in files:\n",
    "        print(f'- gs://{bucket.name}/{file.name}')\n",
    "else:\n",
    "    print('Files not found - Please run the prerequisite notebook (listed at top of this workflow)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fba878d7-00aa-40d3-b0a3-d273bd72a79c",
   "metadata": {},
   "source": [
    "### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e42a61e4-4e50-4995-b4a0-0ec82f49d2c1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_blob = bucket.blob(f'{SERIES}/catboost-overview/model.cbm')\n",
    "model_bytes = model_blob.download_as_bytes()\n",
    "model = catboost.CatBoostClassifier().load_model(blob = model_bytes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb107fd9-fa55-4bde-bf4d-6a44ab438f33",
   "metadata": {},
   "source": [
    "### Load Inference Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1e5e1207-7889-45d5-9883-7676d7dced00",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples_blob = bucket.blob(f'{SERIES}/catboost-overview/examples.json')\n",
    "examples_np = np.array(\n",
    "    json.loads(examples_blob.download_as_string())\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49551796-e761-4cf0-ab54-1dd519bc0ab6",
   "metadata": {},
   "source": [
    "### Test Model With Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c7789d73-0e6a-4d88-8b36-fdce0084db9c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 0, 1, 1, 0, 1, 1, 0, 1])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.predict(examples_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "49b66ed2-9124-4ced-bc66-5e0f4a990c1a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Time',\n",
       " 'V1',\n",
       " 'V2',\n",
       " 'V3',\n",
       " 'V4',\n",
       " 'V5',\n",
       " 'V6',\n",
       " 'V7',\n",
       " 'V8',\n",
       " 'V9',\n",
       " 'V10',\n",
       " 'V11',\n",
       " 'V12',\n",
       " 'V13',\n",
       " 'V14',\n",
       " 'V15',\n",
       " 'V16',\n",
       " 'V17',\n",
       " 'V18',\n",
       " 'V19',\n",
       " 'V20',\n",
       " 'V21',\n",
       " 'V22',\n",
       " 'V23',\n",
       " 'V24',\n",
       " 'V25',\n",
       " 'V26',\n",
       " 'V27',\n",
       " 'V28',\n",
       " 'Amount']"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.feature_names_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c7de77-a606-4eef-9b91-ca2057ff6e5c",
   "metadata": {},
   "source": [
    "---\n",
    "## BigQuery - The Offline Store For Vertex AI Feature Store\n",
    "\n",
    "The offline store for [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview) is BigQuery.  This streamlines ML feature management prior to serving online with feature store.  The data used to train this model in [CatBoost Overview](./CatBoost%20Overview.ipynb) actually came from BigQuery already.\n",
    "\n",
    "This section prepares a version of the BigQuery public table used in the training notebook as a table with an id column to identify individual rows - `transaction_id`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "064da4ce-cc8e-43bb-9d1d-372a7f0ffe14",
   "metadata": {},
   "source": [
    "### Create/Recall Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e9eb6b67-b865-48e1-9d96-687606743c77",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = bigquery.Dataset(f\"{BQ_PROJECT}.{BQ_DATASET}\")\n",
    "dataset.location = BQ_REGION\n",
    "bq_dataset = bq.create_dataset(dataset, exists_ok = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7074cc76-c391-4a57-94c4-6110a23405c9",
   "metadata": {},
   "source": [
    "### Create/Recall Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "31281b30-5e77-4133-8c94-d659e1c4c8c9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.577"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = f\"\"\"\n",
    "    CREATE TABLE IF NOT EXISTS `{BQ_PROJECT}.{BQ_DATASET}.{BQ_TABLE}` AS\n",
    "    SELECT GENERATE_UUID() AS transaction_id, *\n",
    "    FROM `bigquery-public-data.ml_datasets.ulb_fraud_detection`;\n",
    "\"\"\"\n",
    "job = bq.query(query = query)\n",
    "job.result()\n",
    "(job.ended - job.started).total_seconds()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f57e0fb-4544-4e55-8a33-a0dee92f7e9d",
   "metadata": {},
   "source": [
    "### Review BigQuery Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4610b653-98af-479c-b13a-734ae65e5fcc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Time</th>\n",
       "      <th>V1</th>\n",
       "      <th>V2</th>\n",
       "      <th>V3</th>\n",
       "      <th>V4</th>\n",
       "      <th>V5</th>\n",
       "      <th>V6</th>\n",
       "      <th>V7</th>\n",
       "      <th>V8</th>\n",
       "      <th>V9</th>\n",
       "      <th>...</th>\n",
       "      <th>V23</th>\n",
       "      <th>V24</th>\n",
       "      <th>V25</th>\n",
       "      <th>V26</th>\n",
       "      <th>V27</th>\n",
       "      <th>V28</th>\n",
       "      <th>Amount</th>\n",
       "      <th>Class</th>\n",
       "      <th>transaction_id</th>\n",
       "      <th>splits</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>122959.0</td>\n",
       "      <td>-1.327297</td>\n",
       "      <td>0.422904</td>\n",
       "      <td>1.617505</td>\n",
       "      <td>2.291196</td>\n",
       "      <td>2.375055</td>\n",
       "      <td>0.411735</td>\n",
       "      <td>0.213517</td>\n",
       "      <td>0.424743</td>\n",
       "      <td>-1.809624</td>\n",
       "      <td>...</td>\n",
       "      <td>0.192604</td>\n",
       "      <td>0.068281</td>\n",
       "      <td>-0.245725</td>\n",
       "      <td>-0.697654</td>\n",
       "      <td>0.038216</td>\n",
       "      <td>0.150059</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>c97b6e2f-603a-4dbe-9bca-0add881f2084</td>\n",
       "      <td>TEST</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>122312.0</td>\n",
       "      <td>-1.988557</td>\n",
       "      <td>-0.720301</td>\n",
       "      <td>0.863204</td>\n",
       "      <td>3.114494</td>\n",
       "      <td>1.847474</td>\n",
       "      <td>0.255881</td>\n",
       "      <td>0.580362</td>\n",
       "      <td>-0.083756</td>\n",
       "      <td>-0.939044</td>\n",
       "      <td>...</td>\n",
       "      <td>1.564951</td>\n",
       "      <td>0.546312</td>\n",
       "      <td>-0.548531</td>\n",
       "      <td>-0.746620</td>\n",
       "      <td>-0.748016</td>\n",
       "      <td>0.410640</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>791e403e-d59f-491d-b0b7-d8f8710c07fb</td>\n",
       "      <td>TEST</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>119592.0</td>\n",
       "      <td>2.139741</td>\n",
       "      <td>0.245651</td>\n",
       "      <td>-2.654856</td>\n",
       "      <td>0.178287</td>\n",
       "      <td>1.336991</td>\n",
       "      <td>-0.724664</td>\n",
       "      <td>0.906032</td>\n",
       "      <td>-0.436125</td>\n",
       "      <td>-0.528015</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.216033</td>\n",
       "      <td>0.345316</td>\n",
       "      <td>0.747103</td>\n",
       "      <td>0.700184</td>\n",
       "      <td>-0.123739</td>\n",
       "      <td>-0.099989</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>d9e720c5-311d-4cf7-95cb-2256823803ba</td>\n",
       "      <td>TEST</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 33 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       Time        V1        V2        V3        V4        V5        V6  \\\n",
       "0  122959.0 -1.327297  0.422904  1.617505  2.291196  2.375055  0.411735   \n",
       "1  122312.0 -1.988557 -0.720301  0.863204  3.114494  1.847474  0.255881   \n",
       "2  119592.0  2.139741  0.245651 -2.654856  0.178287  1.336991 -0.724664   \n",
       "\n",
       "         V7        V8        V9  ...       V23       V24       V25       V26  \\\n",
       "0  0.213517  0.424743 -1.809624  ...  0.192604  0.068281 -0.245725 -0.697654   \n",
       "1  0.580362 -0.083756 -0.939044  ...  1.564951  0.546312 -0.548531 -0.746620   \n",
       "2  0.906032 -0.436125 -0.528015  ... -0.216033  0.345316  0.747103  0.700184   \n",
       "\n",
       "        V27       V28  Amount  Class                        transaction_id  \\\n",
       "0  0.038216  0.150059     0.0      0  c97b6e2f-603a-4dbe-9bca-0add881f2084   \n",
       "1 -0.748016  0.410640     0.0      0  791e403e-d59f-491d-b0b7-d8f8710c07fb   \n",
       "2 -0.123739 -0.099989     0.0      0  d9e720c5-311d-4cf7-95cb-2256823803ba   \n",
       "\n",
       "   splits  \n",
       "0    TEST  \n",
       "1    TEST  \n",
       "2    TEST  \n",
       "\n",
       "[3 rows x 33 columns]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bq.query(f\"SELECT * FROM `{BQ_PROJECT}.{BQ_DATASET}.{BQ_TABLE}` LIMIT 3\").to_dataframe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a57db3e6-6d4c-4d5f-96b8-9c5825c01bc3",
   "metadata": {},
   "source": [
    "### Get A List of `transaction_id` Values For Testing\n",
    "\n",
    "Get a list of `transaction_id` values to use later in this workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "022ad515-b3e1-405e-9ec1-671c9badfeb7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "transaction_ids = bq.query(f\"SELECT transaction_id FROM `{BQ_PROJECT}.{BQ_DATASET}.{BQ_TABLE}` WHERE Class = 1 LIMIT 10\").to_dataframe()['transaction_id'].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b3273f39-a387-4f8c-bf8e-0fb933dc1714",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       " '1dbfe576-2c52-4f94-bd88-db1e3e294561',\n",
       " 'd13fbb1a-d4cb-40c6-bfd7-05cd714cee30',\n",
       " '18327e33-3ff6-43d1-8bbc-8647adc2044d',\n",
       " 'ec505a29-387f-4b3d-a03b-67375d9140e1',\n",
       " '1b7b6e26-7248-420c-9eb5-dd9692437843',\n",
       " 'bc996e23-c364-4e1a-875d-5151e47fdb81',\n",
       " '165ac053-1fec-4863-aea2-63d380a3adad',\n",
       " '15b93da7-e600-449b-a4c1-bacfd7e68df7',\n",
       " 'd081be14-dcb5-4228-813c-75b1804dc8a2']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transaction_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32cbfc7d-0985-4758-9b13-080d4871b6b4",
   "metadata": {},
   "source": [
    "---\n",
    "## Understanding Vertex AI Feature Store\n",
    "\n",
    "The next sections will setup online serving for the BigQuery table with [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview). This workflow takes the shortest path of synchronizing a BigQuery table to feature store.  There are more flexible paths as well using the Feature Registry where features across multiple tables and views can come together in a single serving structure called a feature view.  You can read more about this within the [MLOps](../../MLOps/readme.md) section of this repository, which includes a deep dive into [feature stores](../../MLOps/Feature%20Store/readme.md).\n",
    "\n",
    "<p align=\"center\" ><center>\n",
    "    <img src=\"../../MLOps/resources/images/created/featurestore/overview.png\" width=\"75%\">\n",
    "</center></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd1ba12-e9c1-4e0b-8b65-966efff05f82",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "## Setup Vertex AI Feature Store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c748b80f-96f8-44f9-a075-3b7269ae7999",
   "metadata": {},
   "source": [
    "### Create/Retrieve Online Store\n",
    "\n",
    "The first step is to create a Vertex AI Feature Store.  There are two serving types to choose from when setting up a feature store: Bigtable and Optimized.  For this work the Optimized online serving is picked which can even [provide vector similarity search](https://cloud.google.com/vertex-ai/docs/featurestore/latest/embeddings-search) functionality that Bigtable serving does not.\n",
    ">**NOTE:** This can take around 10 minutes if creating a new feature store instance\n",
    "\n",
    "**Reference:**\n",
    "- [Create an Online Store Instance](https://cloud.google.com/vertex-ai/docs/featurestore/latest/create-onlinestore)\n",
    "- [Online Serving Types](https://cloud.google.com/vertex-ai/docs/featurestore/latest/online-serving-types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "80913176-6198-4aa7-8b6a-e32088b1f659",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found the feature store:\n",
      "projects/1026793852137/locations/us-central1/featureOnlineStores/statmike_mlops_349915\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    online_store = feature_store.FeatureOnlineStore(name = FS_NAME)\n",
    "    print(f\"Found the feature store:\\n{online_store.resource_name}\")\n",
    "except Exception:\n",
    "    print(\"Create the feature store...\")\n",
    "    online_store = feature_store.FeatureOnlineStore.create_optimized_store(\n",
    "        name = FS_NAME\n",
    "    )\n",
    "    print(f\"Create the feature store:\\n{online_store.resource_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4c9d936e-9d2f-46f5-9566-3a8a56cc1295",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'statmike_mlops_349915'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "online_store.name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e555dc0-ea41-4859-bd55-11cdfd21d994",
   "metadata": {},
   "source": [
    "### Create/Retrieve Feature View From BigQuery Source\n",
    "\n",
    "There are two paths to [creating feature views](https://cloud.google.com/vertex-ai/docs/featurestore/latest/create-featureview) in feature store. The one used here is syncing a BigQuery table or view directly to the online store. The alternative involves using the feature registry which gives greater control of selecting features (columns) form multiple BigQuery source tables and views.  Learn more about Vertex AI Feature Store in this repository's [MLOps](../../MLOps/readme.md) section, which includes a deep dive into [feature stores](../../MLOps/Feature%20Store/readme.md).\n",
    "\n",
    "**Reference:**\n",
    "- [Create a feature view instance](https://cloud.google.com/vertex-ai/docs/featurestore/latest/create-featureview)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4d3ab8fb-9608-4d99-9521-484c47d45cfc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found the feature view:\n",
      "projects/1026793852137/locations/us-central1/featureOnlineStores/statmike_mlops_349915/featureViews/frameworks\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    feature_view = feature_store.FeatureView(\n",
    "        name = FV_NAME,\n",
    "        feature_online_store_id = online_store.resource_name\n",
    "    )\n",
    "    print(f\"Found the feature view:\\n{feature_view.resource_name}\")\n",
    "except Exception:\n",
    "    print(f\"Create the feature view...\")\n",
    "    feature_view = online_store.create_feature_view(\n",
    "        name = FV_NAME,\n",
    "        source = feature_store.utils.FeatureViewBigQuerySource(\n",
    "            uri = f'bq://{BQ_PROJECT}.{BQ_DATASET}.{BQ_TABLE}',\n",
    "            entity_id_columns = ['transaction_id'] # can be multiple columns \n",
    "        ),\n",
    "        sync_config = 'TZ=America/New_York 0 22 * * *' # Ex: every day at 10PM, just once per day\n",
    "    )   \n",
    "    print(f\"Created the feature view:\\n{feature_view.resource_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "43187756-802b-4317-990b-cf63c6e623f1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'frameworks'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature_view.name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c69740ac-6c5f-40cd-b591-0fcf70ab3992",
   "metadata": {},
   "source": [
    "### Managing Synchronization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98f121ff-6c27-4fbe-a716-6678616d2317",
   "metadata": {},
   "source": [
    "Force a synchronization rather than wait for the next scheduled sync:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "945d8b56-f780-4ca2-9494-b401229adfad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "force_sync = feature_view.sync()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e8f35924-5b40-478d-9a4c-2b789ef56270",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "vertexai.resources.preview.feature_store.feature_view.FeatureView.FeatureViewSync"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(force_sync)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "1ef56fac-c924-4ca1-89b7-559dda2792c4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'projects/1026793852137/locations/us-central1/featureOnlineStores/statmike_mlops_349915/featureViews/frameworks/featureViewSyncs/218960636795682816',\n",
       " 'createTime': '2025-02-10T15:00:34.814032Z',\n",
       " 'runTime': {'startTime': '2025-02-10T15:00:34.814032Z'}}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "force_sync.to_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c852b4ca-69a7-44cc-b72a-6b7aa7fae099",
   "metadata": {},
   "source": [
    "Get updated information about the sync job:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "9447d265-2768-4a5c-8ff1-2a94f9131e5f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'projects/1026793852137/locations/us-central1/featureOnlineStores/statmike_mlops_349915/featureViews/frameworks/featureViewSyncs/218960636795682816',\n",
       " 'createTime': '2025-02-10T15:00:34.814032Z',\n",
       " 'runTime': {'startTime': '2025-02-10T15:00:34.814032Z'}}"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "force_sync = feature_view.get_sync(name = force_sync.name)\n",
    "force_sync.to_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89791a3b-24ce-40b3-8ee8-7d6aac5e5939",
   "metadata": {},
   "source": [
    "Wait on the sync job to complete and report timing and rows synced:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "1d2849ed-a0d1-41fd-b8ce-b9f4da8b15bb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Waited 0 seconds, Update again in 30 seconds...\n",
      "Waited 30 seconds, Update again in 30 seconds...\n",
      "Waited 60 seconds, Update again in 30 seconds...\n",
      "Waited 90 seconds, Update again in 30 seconds...\n",
      "Waited 120 seconds, Update again in 30 seconds...\n",
      "Waited 150 seconds, Update again in 30 seconds...\n",
      "Sync completed in 170.857691 seconds and synced 284807 rows.\n"
     ]
    }
   ],
   "source": [
    "waited = 0\n",
    "while True:\n",
    "    sync_status = feature_view.get_sync(name = force_sync.name).to_dict()\n",
    "    if 'endTime' in list(sync_status['runTime'].keys()):\n",
    "        seconds = (\n",
    "            datetime.datetime.fromisoformat(sync_status['runTime']['endTime'].replace('Z', '+00:00'))\n",
    "            -\n",
    "            datetime.datetime.fromisoformat(sync_status['runTime']['startTime'].replace('Z', '+00:00'))\n",
    "        ).total_seconds()\n",
    "        rows = sync_status['syncSummary']['rowSynced']\n",
    "        print(f\"Sync completed in {seconds} seconds and synced {rows} rows.\")\n",
    "        break\n",
    "    else:\n",
    "        print(f\"Waited {waited} seconds, Update again in 30 seconds...\")\n",
    "        time.sleep(30)\n",
    "        waited += 30"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55621a7d-51ca-423e-9ab2-80898e35e32c",
   "metadata": {},
   "source": [
    "Get a list of sync jobs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "536a9a75-3c3d-4f8c-b3c1-71c6056403c9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "list_syncs = feature_view.list_syncs()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f381c92-9fc4-4126-8515-fb84528cf24d",
   "metadata": {},
   "source": [
    "Print out the end time and rows synced for each job:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d814a871-c55d-46d4-8a1f-6f2165bcd184",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sync completed at 02/10/2025 15:03:25 and synced 284807 rows.\n",
      "Sync completed at 02/10/2025 03:02:04 and synced 284807 rows.\n",
      "Sync completed at 02/10/2025 00:36:35 and synced 284807 rows.\n"
     ]
    }
   ],
   "source": [
    "for sync in list_syncs:\n",
    "    s = feature_view.get_sync(name = sync.name).to_dict()\n",
    "    ended = datetime.datetime.fromisoformat(s['runTime']['endTime'].replace('Z', '+00:00')).strftime(\"%m/%d/%Y %H:%M:%S\")\n",
    "    rows = s['syncSummary']['rowSynced']\n",
    "    print(f\"Sync completed at {ended} and synced {rows} rows.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a3a73cd-793b-40dc-b9f2-49f4dc7d03e3",
   "metadata": {},
   "source": [
    "### Retrieve: Features For Entity\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "fe1e3f3c-5bb6-4226-b7fe-5e62f7a5da4f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Public endpoint for the optimized online store statmike_mlops_349915 is 6457115130579648512.us-central1-1026793852137.featurestore.vertexai.goog\n"
     ]
    }
   ],
   "source": [
    "results = feature_view.read(key = [transaction_ids[0]]).to_dict()['features']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "ce1283c6-ea2b-481f-b97c-cfb17f890dfe",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'name': 'Time', 'value': {'double_value': 64585.0}},\n",
       " {'name': 'V1', 'value': {'double_value': 1.08043336974687}},\n",
       " {'name': 'V2', 'value': {'double_value': 0.962830680816555}},\n",
       " {'name': 'V3', 'value': {'double_value': -0.27806547765439}},\n",
       " {'name': 'V4', 'value': {'double_value': 2.74331806279958}},\n",
       " {'name': 'V5', 'value': {'double_value': 0.4123640756810779}},\n",
       " {'name': 'V6', 'value': {'double_value': -0.320778352962878}},\n",
       " {'name': 'V7', 'value': {'double_value': 0.0412895489552235}},\n",
       " {'name': 'V8', 'value': {'double_value': 0.176170299597767}},\n",
       " {'name': 'V9', 'value': {'double_value': -0.9669515510536291}},\n",
       " {'name': 'V10', 'value': {'double_value': -0.19412046650558104}},\n",
       " {'name': 'V11', 'value': {'double_value': 2.1400568155366297}},\n",
       " {'name': 'V12', 'value': {'double_value': -0.276308752749667}},\n",
       " {'name': 'V13', 'value': {'double_value': -1.1913057375248899}},\n",
       " {'name': 'V14', 'value': {'double_value': -1.88027540328199}},\n",
       " {'name': 'V15', 'value': {'double_value': 0.3982719563363761}},\n",
       " {'name': 'V16', 'value': {'double_value': 1.36743302249318}},\n",
       " {'name': 'V17', 'value': {'double_value': 1.52266169956075}},\n",
       " {'name': 'V18', 'value': {'double_value': 1.02631748255981}},\n",
       " {'name': 'V19', 'value': {'double_value': -1.5770704784018401}},\n",
       " {'name': 'V20', 'value': {'double_value': -0.17265864456636498}},\n",
       " {'name': 'V21', 'value': {'double_value': -0.00899617409257402}},\n",
       " {'name': 'V22', 'value': {'double_value': -0.0570361435998535}},\n",
       " {'name': 'V23', 'value': {'double_value': -0.0536915092211487}},\n",
       " {'name': 'V24', 'value': {'double_value': -0.0263732767799466}},\n",
       " {'name': 'V25', 'value': {'double_value': 0.40029973123052}},\n",
       " {'name': 'V26', 'value': {'double_value': 0.0728277709127376}},\n",
       " {'name': 'V27', 'value': {'double_value': 0.0270429343018625}},\n",
       " {'name': 'V28', 'value': {'double_value': 0.0632375727147923}},\n",
       " {'name': 'Amount', 'value': {'double_value': 0.0}},\n",
       " {'name': 'Class', 'value': {'int64_value': '1'}},\n",
       " {'name': 'splits', 'value': {'string_value': 'TEST'}}]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fddd2176-a03c-40f6-93d2-74ee1d901842",
   "metadata": {},
   "source": [
    "---\n",
    "## CatBoost Prediction With Feature Store\n",
    "\n",
    "Predictions with CatBoost are the same and the input is just an array of feature values.  The value of feature store is providing the current value of these features as a simple API call with just the entity id for which the features are needed.  This section does these steps separately and then build a simple Python function to bring them together:\n",
    "- request features for an entity id, `transaction_id`\n",
    "- prepare the features for the model: ensure the order is correct and covert to a Numpy array for input to the model\n",
    "- make the prediction request"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09435450-549a-4711-8a52-6142139a7115",
   "metadata": {},
   "source": [
    "### Step 1: Request Features Based On Entity ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "cc919456-469d-4c59-ad9d-535bb2332062",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "entity_id = transaction_ids[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "87729d54-acf2-46aa-930a-1c1431457698",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "features = feature_view.read(key = [entity_id]).to_dict()['features']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b31a2185-27ef-4761-bb92-8abed66d0ad7",
   "metadata": {},
   "source": [
    "### Step 2: Prepare Features For Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61e954b0-c2ff-427a-bd3d-61c2f71a4adc",
   "metadata": {},
   "source": [
    "Create a dictionary from the 'features' list:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "28848489-4913-44f1-9ab2-bf99c6c7db23",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "feature_dict = {item['name']: list(item['value'].values())[0] for item in features}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "84deb8f9-cc6a-414f-944c-24b740c7b313",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-5.603690279856019"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature_dict['V1']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7f31ef-7801-45df-918b-21cadc1483e5",
   "metadata": {},
   "source": [
    "Extract values in the order of 'model.feature_names_':"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "368f6a3d-1ee7-4ce0-ae23-ee2c65ded8aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ordered_features = [feature_dict[name] for name in model.feature_names_]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "f2c0b15a-e8a5-44ba-b40c-baa7c070ae70",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[102669.0,\n",
       " -5.603690279856019,\n",
       " 5.22219270751471,\n",
       " -7.516829977246081,\n",
       " 8.11772427554261,\n",
       " -2.75685795265464,\n",
       " -1.57456462233583,\n",
       " -6.3303432504036605,\n",
       " 2.99841888080239,\n",
       " -4.50816689310713,\n",
       " -7.33437708370839,\n",
       " 7.18872396524447,\n",
       " -10.6551809118898,\n",
       " 2.5946798607819903,\n",
       " -10.242859083204,\n",
       " -0.191158119172494,\n",
       " -5.504334070139331,\n",
       " -8.697777412570309,\n",
       " -1.9342253618543401,\n",
       " 1.9587499315347499,\n",
       " 0.227525994742654,\n",
       " 1.2428962055240902,\n",
       " 0.42840839246679707,\n",
       " -0.101183595116722,\n",
       " -0.520199121186447,\n",
       " -0.176937756589733,\n",
       " 0.4614499825660971,\n",
       " -0.106625180749266,\n",
       " -0.479661947808134,\n",
       " 0.0]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ordered_features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ba1881-902f-4e0e-90d0-fdc353cdefd8",
   "metadata": {},
   "source": [
    "### Step 3: Make Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "9e135738-280e-4d5b-8d54-87750ca6a814",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.predict(ordered_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "0ce52dce-0ded-4086-ab51-56e7eb070267",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.00479321, 0.99520679])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.predict_proba(ordered_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14dba902-09d1-44d3-8ceb-96f95e3c3fbe",
   "metadata": {},
   "source": [
    "### Combine Steps Into Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "8e5d09c9-82c2-4a76-a283-2930c841e2c7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fs_retriever(entities):\n",
    "    instances = []\n",
    "    for entity_id in entities:\n",
    "        features = feature_view.read(key = [entity_id]).to_dict()['features']\n",
    "        feature_dict = {item['name']: list(item['value'].values())[0] for item in features}\n",
    "        ordered_features = [feature_dict[name] for name in model.feature_names_]\n",
    "        instances.append(ordered_features)\n",
    "    return instances"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19356b85-f45d-4090-89ba-9651ee196c6b",
   "metadata": {},
   "source": [
    "Use function to predict all the same entity ids in `transaction_ids`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "8ef1dc13-b6f9-4d9c-a0af-136ebadc4840",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 0, 1, 1, 0, 1, 1, 0, 1]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.predict(fs_retriever(transaction_ids)).tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a84c26f-7037-40bb-9045-78e486ee88e0",
   "metadata": {},
   "source": [
    "Use function to predict all the same entity ids in `transaction_ids` with detailed output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "27d0edf2-fdd2-420c-9121-7c68f6714e10",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'classes': ['0', '1'],\n",
       "  'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       "  'scores': [0.9684653609075246, 0.031534639092475496],\n",
       "  'predicted_class': '0'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561',\n",
       "  'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "  'predicted_class': '1'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': 'd13fbb1a-d4cb-40c6-bfd7-05cd714cee30',\n",
       "  'scores': [0.9961659959204033, 0.0038340040795967153],\n",
       "  'predicted_class': '0'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': '18327e33-3ff6-43d1-8bbc-8647adc2044d',\n",
       "  'scores': [0.009687501567520629, 0.9903124984324794],\n",
       "  'predicted_class': '1'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': 'ec505a29-387f-4b3d-a03b-67375d9140e1',\n",
       "  'scores': [0.05002634033491404, 0.949973659665086],\n",
       "  'predicted_class': '1'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': '1b7b6e26-7248-420c-9eb5-dd9692437843',\n",
       "  'scores': [0.9988754743912287, 0.0011245256087712965],\n",
       "  'predicted_class': '0'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': 'bc996e23-c364-4e1a-875d-5151e47fdb81',\n",
       "  'scores': [0.0058287633360615265, 0.9941712366639385],\n",
       "  'predicted_class': '1'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': '165ac053-1fec-4863-aea2-63d380a3adad',\n",
       "  'scores': [0.0031326212134841214, 0.9968673787865159],\n",
       "  'predicted_class': '1'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': '15b93da7-e600-449b-a4c1-bacfd7e68df7',\n",
       "  'scores': [0.8122003166930434, 0.18779968330695668],\n",
       "  'predicted_class': '0'},\n",
       " {'classes': ['0', '1'],\n",
       "  'entity_id': 'd081be14-dcb5-4228-813c-75b1804dc8a2',\n",
       "  'scores': [0.007672836590005838, 0.9923271634099942],\n",
       "  'predicted_class': '1'}]"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_classes = [str(c) for c in list(model.classes_)]\n",
    "[\n",
    "    dict(\n",
    "        classes = _classes,\n",
    "        entity_id = transaction_ids[p],\n",
    "        scores = probs,\n",
    "        predicted_class = _classes[np.argmax(probs)]\n",
    "    ) for p, probs in enumerate(model.predict_proba(fs_retriever(transaction_ids)).tolist())\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "947ac126-a16b-484e-bdc8-8d7c8918b6cb",
   "metadata": {},
   "source": [
    "---\n",
    "## Build A Custom Prediction Container\n",
    "\n",
    "It is really not all that hard with Python!\n",
    "\n",
    "For this example [FastAPI](https://fastapi.tiangolo.com/) is used.\n",
    "\n",
    "This process uses docker to build a custom container and then runs the container locally, on Cloud Run, and Vertex AI Endpoints.\n",
    "\n",
    "This could be done locally with Docker and pushed to Artifact Registry before deployment to Cloud Run and Vertex.  The process below assumes that docker is not available locally and used Cloud Build to both build and push the resulting container to Artifact Registry.\n",
    "\n",
    "**Environment Variables**\n",
    "\n",
    "The scripts `main.py` and `main2.py` are created assuming the existance of several enviornment variables. When running the container directly in a local instance or on Cloud Run the environment variable can be set for use in the services.  Vertex AI Endpoints actually create a number of helpful environment variables during deployment.  To make the single container easily deployable across the three (local, Cloud Run, and Vertex AI Endpoints) the environment variables were assumed to be the Vertex AI ones and then set at manually at launch for Cloud Run and local.  Read more about [Environment variables available in the container](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#variables) for Vertex AI Endpoints."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e5a486b-13fd-4971-b6b3-c93875fc2535",
   "metadata": {},
   "source": [
    "### Setup Artifact Registry\n",
    "\n",
    "[Artifact registry](https://cloud.google.com/artifact-registry/docs) organizes artifacts with repositories.  Each repository contains packages and is designated to hold a partifcular format of package: Docker images, Python Packages and [others](https://cloud.google.com/artifact-registry/docs/supported-formats#package)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd4e3d8b-99e3-452d-8f5b-c8612f3d48ca",
   "metadata": {},
   "source": [
    "#### List Repositories\n",
    "\n",
    "This may be empty if no repositories have been created for this project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "ab2ae8bc-f0c2-4a8f-a93b-cd9d13a1374c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "projects/statmike-mlops-349915/locations/us-central1/repositories/frameworks\n",
      "projects/statmike-mlops-349915/locations/us-central1/repositories/frameworks-catboost\n",
      "projects/statmike-mlops-349915/locations/us-central1/repositories/gcf-artifacts\n",
      "projects/statmike-mlops-349915/locations/us-central1/repositories/mlops\n",
      "projects/statmike-mlops-349915/locations/us-central1/repositories/statmike-mlops-349915\n",
      "projects/statmike-mlops-349915/locations/us-central1/repositories/statmike-mlops-349915-docker\n",
      "projects/statmike-mlops-349915/locations/us-central1/repositories/statmike-mlops-349915-python\n"
     ]
    }
   ],
   "source": [
    "for repo in ar.list_repositories(parent = f'projects/{PROJECT_ID}/locations/{REGION}'):\n",
    "    print(repo.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b72d358-68e0-495c-badc-e9337ab0ce6b",
   "metadata": {},
   "source": [
    "#### Create/Retrieve Docker Image Repository\n",
    "\n",
    "Create an Artifact Registry Repository to hold Docker Images created by this notebook.  First, check to see if it is already created by a previous run and retrieve it if it has.  Otherwise, create one named for this project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "b89c88c5-a3c2-4f1a-b85c-33a14f86ec0c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrieved existing repo: projects/statmike-mlops-349915/locations/us-central1/repositories/frameworks\n"
     ]
    }
   ],
   "source": [
    "docker_repo = None\n",
    "for repo in ar.list_repositories(parent = f'projects/{PROJECT_ID}/locations/{REGION}'):\n",
    "    if f'{SERIES}' == repo.name.split('/')[-1]:\n",
    "        docker_repo = repo\n",
    "        print(f'Retrieved existing repo: {docker_repo.name}')\n",
    "\n",
    "if not docker_repo:\n",
    "    operation = ar.create_repository(\n",
    "        request = artifactregistry_v1.CreateRepositoryRequest(\n",
    "            parent = f'projects/{PROJECT_ID}/locations/{REGION}',\n",
    "            repository_id = f'{SERIES}',\n",
    "            repository = artifactregistry_v1.Repository(\n",
    "                description = f'A repository for the {SERIES} series that holds docker images.',\n",
    "                name = f'{SERIES}',\n",
    "                format_ = artifactregistry_v1.Repository.Format.DOCKER,\n",
    "                labels = {'series': SERIES}\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "    print('Creating Repository ...')\n",
    "    docker_repo = operation.result()\n",
    "    print(f'Completed creating repo: {docker_repo.name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "36648583-26dc-458f-a664-2ded0f98650c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('projects/statmike-mlops-349915/locations/us-central1/repositories/frameworks',\n",
       " 'DOCKER')"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docker_repo.name, docker_repo.format_.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "4ae1c820-a24b-46a0-bb1a-900d32ecf985",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "REPOSITORY = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{docker_repo.name.split('/')[-1]}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "e87e8adb-19a0-4c3c-9de6-6e98f94b13fb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'us-central1-docker.pkg.dev/statmike-mlops-349915/frameworks'"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "REPOSITORY"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080b5c92-9769-4f8d-8dbb-d2dd61999c87",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "### Create Application Files\n",
    "\n",
    "```\n",
    "|__ Dockerfile\n",
    "|__ requirements.txt\n",
    "|__ app\n",
    "    |__ __init__.py\n",
    "    |__ main.py\n",
    "    |__ prestart.sh\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "f8e5ac5f-b21c-474b-b8a1-07271146330b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(DIR + '/source/app'):\n",
    "    os.makedirs(DIR + '/source/app')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "d16d05ad-5078-4b4a-81bf-d151253c02da",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting files/catboost-prediction-feature-store/source/Dockerfile\n"
     ]
    }
   ],
   "source": [
    "%%writefile {DIR}/source/Dockerfile\n",
    "FROM tiangolo/uvicorn-gunicorn-fastapi:python3.9\n",
    "\n",
    "COPY ./app /app\n",
    "COPY ./requirements.txt requirements.txt\n",
    "\n",
    "RUN pip install --no-cache-dir --upgrade pip \\\n",
    "  && pip install --no-cache-dir -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "c3ce7dc9-c524-4964-bc31-77065fa0c294",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting files/catboost-prediction-feature-store/source/requirements.txt\n"
     ]
    }
   ],
   "source": [
    "%%writefile {DIR}/source/requirements.txt\n",
    "google-cloud-storage\n",
    "google-cloud-aiplatform==1.71.0\n",
    "catboost\n",
    "numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "2fbf9016-f4d8-4e77-ad2a-821fe1c217ee",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting files/catboost-prediction-feature-store/source/app/__init__.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile {DIR}/source/app/__init__.py\n",
    "# init file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "6f8a0371-58ca-4848-a4c0-41ef388de00c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('statmike_mlops_349915', 'frameworks')"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FS_NAME,FV_NAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "3c81b52d-804d-4650-834f-8af657d223e5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting files/catboost-prediction-feature-store/source/app/main.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile {DIR}/source/app/main.py\n",
    "# this version:\n",
    "# - inputs to be json like {'instances': [[list],[list], ...]}\n",
    "# - outputs in json like {'predictions': [[list],[list], ...]}\n",
    "# trying to adhere to Vetex Endpoints Requirements:\n",
    "# - https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions\n",
    "\n",
    "# packages\n",
    "import os\n",
    "from fastapi import FastAPI, Request\n",
    "import catboost\n",
    "import numpy as np\n",
    "from google.cloud import storage\n",
    "from google.cloud import aiplatform\n",
    "from vertexai.resources.preview import feature_store\n",
    "\n",
    "import logging\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "logger.addHandler(logging.StreamHandler())\n",
    "logging.info(f'aiplatform version: {aiplatform.__version__}')\n",
    "\n",
    "# NAMES\n",
    "PROJECT_ID = 'statmike-mlops-349915'\n",
    "REGION = 'us-central1'\n",
    "FS_NAME = 'statmike_mlops_349915'\n",
    "FV_NAME = 'frameworks'\n",
    "\n",
    "# clients\n",
    "app = FastAPI()\n",
    "gcs = storage.Client()\n",
    "aiplatform.init(project = PROJECT_ID, location = REGION)\n",
    "\n",
    "# feature store\n",
    "online_store = feature_store.FeatureOnlineStore(\n",
    "    name = f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FS_NAME}\"\n",
    ")\n",
    "feature_view = feature_store.FeatureView(\n",
    "    name = FV_NAME,\n",
    "    feature_online_store_id = online_store.resource_name\n",
    ")\n",
    "\n",
    "# download the model file from GCS\n",
    "paths = os.environ['AIP_STORAGE_URI'].split('/') + ['model.cbm']\n",
    "bucket = gcs.bucket(paths[2])\n",
    "blob = bucket.blob('/'.join(paths[3:]))\n",
    "blob.download_to_filename('model.cbm')\n",
    "\n",
    "# Load the catboost model\n",
    "_model = catboost.CatBoostClassifier().load_model('model.cbm')\n",
    "\n",
    "# get model classification levels and feature names\n",
    "_classes = [str(c) for c in list(_model.classes_)]\n",
    "_features = _model.feature_names_\n",
    "\n",
    "# function to retrieve from feature store\n",
    "def fs_retriever(entities):\n",
    "    instances = []\n",
    "    for entity_id in entities:\n",
    "        features = feature_view.read(key = [entity_id]).to_dict()['features']\n",
    "        feature_dict = {item['name']: list(item['value'].values())[0] for item in features}\n",
    "        ordered_features = [feature_dict[name] for name in _features]\n",
    "        instances.append(ordered_features)\n",
    "    return instances\n",
    "\n",
    "# Define function for health route\n",
    "@app.get(os.environ['AIP_HEALTH_ROUTE'], status_code=200)\n",
    "def health():\n",
    "    return {}\n",
    "\n",
    "# Define function for prediction route\n",
    "@app.post(os.environ['AIP_PREDICT_ROUTE'])\n",
    "async def predict(request: Request):\n",
    "    # await the request\n",
    "    body = await request.json()\n",
    "    \n",
    "    # parse the request\n",
    "    entities = body[\"instances\"]\n",
    "    \n",
    "    # get predicted probabilities\n",
    "    predictions = _model.predict_proba(fs_retriever(entities)).tolist()\n",
    "\n",
    "    # this returns just the predicted probabilities:\n",
    "    return {\"predictions\": predictions}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "7cb6c05f-257f-43bd-884a-a8026c7fb4b4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting files/catboost-prediction-feature-store/source/app/main2.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile {DIR}/source/app/main2.py\n",
    "# this version:\n",
    "# - inputs to be json like {'instances': [[list],[list], ...]}\n",
    "# - outputs in json like {'predictions': [{'classes': list, 'scores': list, 'predicted_class': str}, ...]}\n",
    "# trying to adhere to Vetex Endpoints Requirements:\n",
    "# - https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions\n",
    "\n",
    "# packages\n",
    "import os\n",
    "from fastapi import FastAPI, Request\n",
    "import catboost\n",
    "import numpy as np\n",
    "from google.cloud import storage\n",
    "from google.cloud import aiplatform\n",
    "from vertexai.resources.preview import feature_store\n",
    "\n",
    "# NAMES\n",
    "PROJECT_ID = 'statmike-mlops-349915'\n",
    "REGION = 'us-central1'\n",
    "FS_NAME = 'statmike_mlops_349915'\n",
    "FV_NAME = 'frameworks'\n",
    "\n",
    "# clients\n",
    "app = FastAPI()\n",
    "gcs = storage.Client()\n",
    "aiplatform.init(project = PROJECT_ID, location = REGION)\n",
    "\n",
    "# feature store\n",
    "online_store = feature_store.FeatureOnlineStore(\n",
    "    name = f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FS_NAME}\"\n",
    ")\n",
    "feature_view = feature_store.FeatureView(\n",
    "    name = FV_NAME,\n",
    "    feature_online_store_id = online_store.resource_name\n",
    ")\n",
    "\n",
    "# download the model file from GCS\n",
    "paths = os.environ['AIP_STORAGE_URI'].split('/') + ['model.cbm']\n",
    "bucket = gcs.bucket(paths[2])\n",
    "blob = bucket.blob('/'.join(paths[3:]))\n",
    "blob.download_to_filename('model.cbm')\n",
    "\n",
    "# Load the catboost model\n",
    "_model = catboost.CatBoostClassifier().load_model('model.cbm')\n",
    "\n",
    "# get model classification levels and feature names\n",
    "_classes = [str(c) for c in list(_model.classes_)]\n",
    "_features = _model.feature_names_\n",
    "\n",
    "# function to retrieve from feature store\n",
    "def fs_retriever(entities):\n",
    "    instances = []\n",
    "    for entity_id in entities:\n",
    "        features = feature_view.read(key = [entity_id]).to_dict()['features']\n",
    "        feature_dict = {item['name']: list(item['value'].values())[0] for item in features}\n",
    "        ordered_features = [feature_dict[name] for name in _features]\n",
    "        instances.append(ordered_features)\n",
    "    return instances\n",
    "\n",
    "# Define function for health route\n",
    "@app.get(os.environ['AIP_HEALTH_ROUTE'], status_code=200)\n",
    "def health():\n",
    "    return {}\n",
    "\n",
    "# Define function for prediction route\n",
    "@app.post(os.environ['AIP_PREDICT_ROUTE'])\n",
    "async def predict(request: Request):\n",
    "    # await the request\n",
    "    body = await request.json()\n",
    "    \n",
    "    # parse the request\n",
    "    entities = body[\"instances\"]\n",
    "    \n",
    "    # get predicted probabilities\n",
    "    predictions = _model.predict_proba(fs_retriever(entities)).tolist()\n",
    "    \n",
    "    # format predictions:\n",
    "    preds = [\n",
    "        dict(\n",
    "            classes = _classes,\n",
    "            entity_id = entities[p],\n",
    "            scores = probs,\n",
    "            predicted_class = _classes[np.argmax(probs)]\n",
    "        ) for p, probs in enumerate(predictions)\n",
    "    ]\n",
    "    \n",
    "    # following outputs detail prediction info for classification:\n",
    "    return {\"predictions\": preds}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "b276da1b-08b9-4b50-a250-580010fb7e0a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting files/catboost-prediction-feature-store/source/app/prestart.sh\n"
     ]
    }
   ],
   "source": [
    "%%writefile {DIR}/source/app/prestart.sh\n",
    "#!/bin/bash\n",
    "export PORT=$AIP_HTTP_PORT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "486905a1-47aa-4fb8-a485-ad1b157360e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bucket.blob(f'{SERIES}/{EXPERIMENT}/source/Dockerfile').upload_from_filename(f'{DIR}/source/Dockerfile')\n",
    "bucket.blob(f'{SERIES}/{EXPERIMENT}/source/requirements.txt').upload_from_filename(f'{DIR}/source/requirements.txt')\n",
    "bucket.blob(f'{SERIES}/{EXPERIMENT}/source/app/__init__.py').upload_from_filename(f'{DIR}/source/app/__init__.py')\n",
    "bucket.blob(f'{SERIES}/{EXPERIMENT}/source/app/main.py').upload_from_filename(f'{DIR}/source/app/main.py')\n",
    "bucket.blob(f'{SERIES}/{EXPERIMENT}/source/app/main2.py').upload_from_filename(f'{DIR}/source/app/main2.py')\n",
    "bucket.blob(f'{SERIES}/{EXPERIMENT}/source/app/prestart.sh').upload_from_filename(f'{DIR}/source/app/prestart.sh')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "538d0a44-c7e7-42e7-b537-39b1b46f017d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<Blob: statmike-mlops-349915, frameworks/catboost-prediction-feature-store/source/Dockerfile, 1739199881309397>,\n",
       " <Blob: statmike-mlops-349915, frameworks/catboost-prediction-feature-store/source/app/__init__.py, 1739199881498239>,\n",
       " <Blob: statmike-mlops-349915, frameworks/catboost-prediction-feature-store/source/app/main.py, 1739199881581336>,\n",
       " <Blob: statmike-mlops-349915, frameworks/catboost-prediction-feature-store/source/app/main2.py, 1739199881650277>,\n",
       " <Blob: statmike-mlops-349915, frameworks/catboost-prediction-feature-store/source/app/prestart.sh, 1739199881731120>,\n",
       " <Blob: statmike-mlops-349915, frameworks/catboost-prediction-feature-store/source/requirements.txt, 1739199881402787>]"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(bucket.list_blobs(prefix = f'{SERIES}/{EXPERIMENT}/source'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4230bf3e-6a9e-4aa0-8e0c-b9606de43c6c",
   "metadata": {},
   "source": [
    "---\n",
    "### Build Application Container\n",
    "\n",
    "Use the Cloud Build client to construct and run the build instructions. Here the files collected in GCS are copied to the build instance, then the Docker build is run in the folder with the `Dockerfile`. The resulting image is pushed to Artifact Registry (setup above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "46d387a0-cde7-44c9-8375-7a66fceb5189",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# setup the build config with empty list of steps - these will be added sequentially\n",
    "build = cloudbuild_v1.Build(\n",
    "    steps = []\n",
    ")\n",
    "# retrieve the source\n",
    "build.steps.append(\n",
    "    {\n",
    "        'name': 'gcr.io/cloud-builders/gsutil',\n",
    "        'args': ['cp', '-r', f'gs://{GCS_BUCKET}/{SERIES}/{EXPERIMENT}/source/*', '/workspace']\n",
    "    }\n",
    ")\n",
    "# docker build\n",
    "build.steps.append(\n",
    "    {\n",
    "        'name': 'gcr.io/cloud-builders/docker',\n",
    "        'args': ['build', '-t', f'{REPOSITORY}/{EXPERIMENT}', '/workspace']\n",
    "    }    \n",
    ")\n",
    "# docker push\n",
    "build.images = [f\"{REPOSITORY}/{EXPERIMENT}\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "641c69ee-816e-4aa6-8c97-26355e701864",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "operation = cb.create_build(\n",
    "    project_id = PROJECT_ID,\n",
    "    build = build\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "4999ba21-d06a-4ed3-99b5-b7146ee6991c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<Status.SUCCESS: 3>,\n",
       " images: \"us-central1-docker.pkg.dev/statmike-mlops-349915/frameworks/catboost-prediction-feature-store\")"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "build_response = operation.result()\n",
    "build_response.status, build_response.artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "2bc1e0bc-8843-4001-abbf-1de9a9e585f1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'us-central1-docker.pkg.dev/statmike-mlops-349915/frameworks/catboost-prediction-feature-store'"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "build_response.artifacts.images[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974f4339-ed04-4701-8c34-e8a8750b67b2",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "## Test Locally\n",
    "\n",
    "If Docker is installed and running locally then use it to test the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "db4be5ca-e500-4aa4-8165-ad6fcc64ca02",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Docker is installed and running. Version: 20.10.17\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    local_test = True\n",
    "    docker_client = docker.from_env()\n",
    "    if docker_client.ping():\n",
    "        print(f\"Docker is installed and running. Version: {docker_client.version()['Version']}\")\n",
    "    # Run the following command to configure gcloud as the credential helper for the Artifact Registry domain associated with this repository's location:\n",
    "    config_docker = !gcloud auth configure-docker us-central1-docker.pkg.dev --quiet\n",
    "except Exception as e:\n",
    "    local_test = False\n",
    "    print('Docker is either not installed or not running - please fix before proceeding.\\nhttps://docs.docker.com/engine/install/')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cbf4fed-fb83-42cf-87e3-7b63b1677c9d",
   "metadata": {},
   "source": [
    "### Pull and Run Container\n",
    "\n",
    "Run the container image with:\n",
    "- ports: inside 8080 mapped to outside 80\n",
    "- set environment variables for:\n",
    "    - `AIP_HTTP_PORT` is `8080`\n",
    "    - `AIP_HEALTH_ROUTE` is `/health`\n",
    "    - `AIP_PREDICT_ROUTE` is `/predict`\n",
    "    - `AIP_STORAGE_URI` is the `gs://bucket/path/to/folder`\n",
    "    - `MODULE_NAME` is 'main'\n",
    "        - this actually defaults to main so is not required\n",
    "        - an alternative script with different prediction output is created in `main2.py` above\n",
    "        - use this environment variable to start the container using the alternative script in module `main2`\n",
    "        - see the [FastAPI Docker Image Advanced Usage](https://github.com/tiangolo/uvicorn-gunicorn-fastapi-docker?tab=readme-ov-file#advanced-usage) details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "5e684ffd-e83b-4573-a260-e06a9a4a8684",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pulling image ...\n",
      "Starting container ...\n",
      "Container ready.\n",
      "\tUse `container.logs()` to view startup logs.\n"
     ]
    }
   ],
   "source": [
    "if local_test:\n",
    "    # make sure any prior runs are stopped:\n",
    "    try:\n",
    "        container = docker_client.containers.get('local-run')\n",
    "        container.stop()\n",
    "        container.remove()\n",
    "    except docker.errors.NotFound:\n",
    "        pass\n",
    "    \n",
    "    # get image:\n",
    "    image_uri = build_response.artifacts.images[0]\n",
    "    try:\n",
    "        local_image = docker_client.images.get(image_uri)\n",
    "        remote_image = docker_client.images.pull(image_uri)\n",
    "        if local_image.id != remote_image.id:\n",
    "            print('New image found, updating ...')\n",
    "            local_image = remote_image\n",
    "        else:\n",
    "            print('Using existing image ...')\n",
    "    except docker.errors.ImageNotFound:\n",
    "        print('Pulling image ...')\n",
    "        local_image = docker_client.images.pull(image_uri)\n",
    "        \n",
    "    # run container:\n",
    "    print('Starting container ...')\n",
    "    container = docker_client.containers.run(\n",
    "        image = image_uri,\n",
    "        detach = True,\n",
    "        ports = {'8080/tcp':80}, # Map inside:outside (where docker run -p is outside:inside)\n",
    "        name = 'local-run',\n",
    "        environment = {\n",
    "            'AIP_HTTP_PORT': '8080',\n",
    "            'AIP_HEALTH_ROUTE': '/health',\n",
    "            'AIP_PREDICT_ROUTE': '/predict',\n",
    "            'AIP_STORAGE_URI': f'gs://{bucket.name}/{SERIES}/catboost-overview',\n",
    "            'MODULE_NAME': 'main2' # try main or main2 for alternative output\n",
    "        }\n",
    "    )\n",
    "    print('Container ready.\\n\\tUse `container.logs()` to view startup logs.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "8067395f-48cd-4861-8f58-5670ac714a63",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#container.logs()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4def2d01-7db7-4382-ab81-516be0c7b2f0",
   "metadata": {},
   "source": [
    "### Health Check\n",
    "\n",
    "Want to see `200`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "b1a72114-38cb-4e4e-9344-670e65fcaad9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# wait a few seconds\n",
    "time.sleep(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "dcd47793-1f16-4994-aa5a-fd3347115ff2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "200\n"
     ]
    }
   ],
   "source": [
    "if local_test:\n",
    "    response = requests.get(f\"http://localhost:80/health\")\n",
    "    print(response.status_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1a78ed-f9a9-44a0-afd3-5e8deb7d5e44",
   "metadata": {},
   "source": [
    "### Inference Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "33ab9766-05a3-4860-a3bc-507dfbf239a1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def predict(instances):\n",
    "    url = f\"http://localhost:80/predict\"\n",
    "    headers = {'Content_Type': 'application/json'}\n",
    "    data = json.dumps({'instances': instances})\n",
    "    response = requests.post(url, headers = headers, data = data)\n",
    "    return json.loads(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "c6c5b249-023e-4901-af9d-13c3fb1be97c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [{'classes': ['0', '1'],\n",
       "   'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       "   'scores': [0.9684653609075246, 0.031534639092475496],\n",
       "   'predicted_class': '0'}]}"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[0:1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "d466a650-8fdf-4972-9411-b94d6693a274",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [{'classes': ['0', '1'],\n",
       "   'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561',\n",
       "   'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "   'predicted_class': '1'}]}"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[1:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "495382e1-9e53-4290-a693-faea960c75a8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [{'classes': ['0', '1'],\n",
       "   'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       "   'scores': [0.9684653609075246, 0.031534639092475496],\n",
       "   'predicted_class': '0'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561',\n",
       "   'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "   'predicted_class': '1'}]}"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[0:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "470be6e7-0d1f-4b4f-a7dd-8f7ed9d415ee",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [{'classes': ['0', '1'],\n",
       "   'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       "   'scores': [0.9684653609075246, 0.031534639092475496],\n",
       "   'predicted_class': '0'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561',\n",
       "   'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "   'predicted_class': '1'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': 'd13fbb1a-d4cb-40c6-bfd7-05cd714cee30',\n",
       "   'scores': [0.9961659959204033, 0.0038340040795967153],\n",
       "   'predicted_class': '0'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '18327e33-3ff6-43d1-8bbc-8647adc2044d',\n",
       "   'scores': [0.009687501567520629, 0.9903124984324794],\n",
       "   'predicted_class': '1'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': 'ec505a29-387f-4b3d-a03b-67375d9140e1',\n",
       "   'scores': [0.05002634033491404, 0.949973659665086],\n",
       "   'predicted_class': '1'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '1b7b6e26-7248-420c-9eb5-dd9692437843',\n",
       "   'scores': [0.9988754743912287, 0.0011245256087712965],\n",
       "   'predicted_class': '0'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': 'bc996e23-c364-4e1a-875d-5151e47fdb81',\n",
       "   'scores': [0.0058287633360615265, 0.9941712366639385],\n",
       "   'predicted_class': '1'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '165ac053-1fec-4863-aea2-63d380a3adad',\n",
       "   'scores': [0.0031326212134841214, 0.9968673787865159],\n",
       "   'predicted_class': '1'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '15b93da7-e600-449b-a4c1-bacfd7e68df7',\n",
       "   'scores': [0.8122003166930434, 0.18779968330695668],\n",
       "   'predicted_class': '0'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': 'd081be14-dcb5-4228-813c-75b1804dc8a2',\n",
       "   'scores': [0.007672836590005838, 0.9923271634099942],\n",
       "   'predicted_class': '1'}]}"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d8d45f4-82ac-4a49-aa79-68b2b39b9b13",
   "metadata": {},
   "source": [
    "### Stop Container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "784101e1-8bd5-43fb-bd65-8c0e3e12ee74",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'local-run'"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "container.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "91ef8482-50ad-406f-87ef-2891b044251a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "container = docker_client.containers.get(container.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "52ee6525-9714-4b55-a9d8-f77f7087699b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'running'"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "container.status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "e9dea32d-0240-4e4d-8f85-3dc9e5ab7798",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "container.stop()\n",
    "container.remove()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5e23615-ab98-4609-8630-6d989b150f61",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "## Cloud Run\n",
    "\n",
    "Deploy the model to [Cloud Run](https://cloud.google.com/run/docs/overview/what-is-cloud-run) using the same container build and tested above from Artifact Registry.\n",
    "\n",
    "Some highlights for Cloud Run:\n",
    "- Rapid scaling to handle requests\n",
    "- Scale to zero (default) or other minimum if set\n",
    "- Can handle larger input (request) and output (response) sizes\n",
    "    - See [requests limits](https://cloud.google.com/run/quotas#request_limits)\n",
    "- Configure [memory limits](https://cloud.google.com/run/docs/configuring/services/memory-limits) and [cpu limits](https://cloud.google.com/run/docs/configuring/services/cpu) and [concurrency](https://cloud.google.com/run/docs/configuring/concurrency) and [autoscaling](https://cloud.google.com/run/docs/about-instance-autoscaling) and [request timeout](https://cloud.google.com/run/docs/configuring/request-timeout)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5c0a0ab-63a9-43ea-b02f-22110d1d05c3",
   "metadata": {},
   "source": [
    "### Deploy The Endpoint(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "caff3e18-37e7-4927-bb39-69a42d032845",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def start_cloud_run_service(MODULE, name_suffix):\n",
    "    # define the service:\n",
    "    parent = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
    "    service = run_v2.Service()\n",
    "    service.template.containers = [\n",
    "        run_v2.Container(\n",
    "            image = build_response.artifacts.images[0],\n",
    "            ports = [run_v2.ContainerPort(container_port = 8080)],\n",
    "            env = [\n",
    "                run_v2.EnvVar(name = 'AIP_HTTP_PORT', value = '8080'),\n",
    "                run_v2.EnvVar(name = 'AIP_HEALTH_ROUTE', value = '/health'),\n",
    "                run_v2.EnvVar(name = 'AIP_PREDICT_ROUTE', value = '/predict'),\n",
    "                run_v2.EnvVar(name = 'AIP_STORAGE_URI', value = f'gs://{bucket.name}/{SERIES}/catboost-overview'),\n",
    "                run_v2.EnvVar(name = 'MODULE_NAME', value = MODULE)\n",
    "            ],\n",
    "            resources = run_v2.ResourceRequirements(\n",
    "                limits = {\"cpu\": '8', \"memory\": '32Gi'}\n",
    "            )\n",
    "        )\n",
    "    ]\n",
    "    service.ingress = run_v2.IngressTraffic.INGRESS_TRAFFIC_INTERNAL_ONLY\n",
    "    \n",
    "    # start the service:\n",
    "    try:\n",
    "        # create the service:\n",
    "        run_response = cr.create_service(request = {\n",
    "            \"parent\": parent,\n",
    "            \"service\": service,\n",
    "            \"service_id\": SERIES+'-'+EXPERIMENT+'-'+name_suffix\n",
    "        })\n",
    "        # wait on the operation to complete:\n",
    "        run_response.result()\n",
    "        # print the name of the service\n",
    "        print(f\"Started Service: {run_response.metadata.name}\")\n",
    "        return run_response\n",
    "    \n",
    "    except Exception as e:\n",
    "        print(f\"Error creating service: {e}\")\n",
    "        return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "1b287ef9-21b8-4fa0-a3d9-0763f5735b00",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started Service: projects/statmike-mlops-349915/locations/us-central1/services/frameworks-catboost-prediction-feature-store-1\n"
     ]
    }
   ],
   "source": [
    "service_start_main = start_cloud_run_service('main', '1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "09cc9d29-69e7-4551-91a0-07fff5947a51",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'https://frameworks-catboost-prediction-feature-store-1-urlxi72dpa-uc.a.run.app'"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "service_start_main.metadata.uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "7a916b35-9df4-4d4b-b499-2d6d4b217254",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'projects/statmike-mlops-349915/locations/us-central1/services/frameworks-catboost-prediction-feature-store-1'"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "service_start_main.metadata.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "ce031a4c-5a4e-439a-bf3d-594711c28073",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started Service: projects/statmike-mlops-349915/locations/us-central1/services/frameworks-catboost-prediction-feature-store-2\n"
     ]
    }
   ],
   "source": [
    "service_start_main2 = start_cloud_run_service('main2', '2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "74df0492-8bd2-4cc0-a09d-14411e6eb280",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'https://frameworks-catboost-prediction-feature-store-2-urlxi72dpa-uc.a.run.app'"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "service_start_main2.metadata.uri"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaa49e11-9b1a-4674-b6bb-aa12a2c0e5da",
   "metadata": {},
   "source": [
    "### Permissions\n",
    "\n",
    "The endpoint requires authentication.  Check outthe [Authentication Overview](https://cloud.google.com/run/docs/authenticating/overview) and in the case below the [Authenticating service-to-service](https://cloud.google.com/run/docs/authenticating/service-to-service) method is used by giving the same service account used to run the notebook and create the endpoint the role to invoke the endpoint as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "4defd81e-20f9-4ff9-aa72-0987269812ae",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1026793852137-compute@developer.gserviceaccount.com'"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SERVICE_ACCOUNT = !gcloud config list --format='value(core.account)' \n",
    "SERVICE_ACCOUNT = SERVICE_ACCOUNT[0]\n",
    "SERVICE_ACCOUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "c74ea331-c445-4ed3-b3ab-4d9701646e3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def set_policy(service_name):\n",
    "    policy = cr.get_iam_policy(request = {'resource': service_name})\n",
    "    policy.bindings.add(\n",
    "        role = 'roles/run.invoker',\n",
    "        members = [f\"serviceAccount:{SERVICE_ACCOUNT}\", 'allUsers'] #'allUsers'\n",
    "    )\n",
    "    policy_response = cr.set_iam_policy(request = {\"resource\": service_name, \"policy\": policy})\n",
    "    print(f\"IAM policy updated: {policy_response.bindings}\")\n",
    "    return policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "80713558-a0b8-4668-91a8-8327f87afaf9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IAM policy updated: [role: \"roles/run.invoker\"\n",
      "members: \"allUsers\"\n",
      "members: \"serviceAccount:1026793852137-compute@developer.gserviceaccount.com\"\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "policy_main = set_policy(service_start_main.metadata.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "b8bcc7e8-e8ec-487b-8261-b60115589720",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IAM policy updated: [role: \"roles/run.invoker\"\n",
      "members: \"allUsers\"\n",
      "members: \"serviceAccount:1026793852137-compute@developer.gserviceaccount.com\"\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "policy_main2 = set_policy(service_start_main2.metadata.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03428db0-a25e-4c73-bcfe-6248f73fa84c",
   "metadata": {
    "tags": []
   },
   "source": [
    "**WAIT: The update of the IAM Policy might take a few moments to take effect.  Rerun the following health check section until you get a `200` response code.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "963bddb6-8dca-4fd1-87ad-37aa6cfa998c",
   "metadata": {},
   "source": [
    "### Health Check\n",
    "\n",
    "Want to see `200`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "7dbb225e-3302-43f4-ae82-0c183b55aef0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def health(uri):\n",
    "    url = f\"{uri}/health\"\n",
    "    credentials, _ = google.auth.default()\n",
    "    auth_req = google.auth.transport.requests.Request()\n",
    "    credentials.refresh(auth_req)\n",
    "    headers = {'Authorization': f'Bearer {credentials.token}'}\n",
    "    response = requests.get(url, headers = headers)    \n",
    "    return response.status_code\n",
    "\n",
    "def check_health(uri, timeout_seconds = 200, retry_seconds = 10):\n",
    "    start_time = time.time()\n",
    "    while True:\n",
    "        status_code = health(uri)\n",
    "        if status_code == 200:\n",
    "            break\n",
    "        elapsed_time = time.time() - start_time\n",
    "        if elapsed_time > timeout_seconds:\n",
    "            break\n",
    "        time.sleep(retry_seconds)\n",
    "    return status_code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "0ce2ffb8-3d85-4c27-b6b9-d2a540ee7716",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check_health(service_start_main.metadata.uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "dcf09bda-b421-4ae6-a6f8-a217ba7266f0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 123,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check_health(service_start_main2.metadata.uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "01a973a3-6bcf-4aba-bd74-e448fb1200f3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200, 200)"
      ]
     },
     "execution_count": 124,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "health(service_start_main.metadata.uri), health(service_start_main2.metadata.uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d74cd2f-c5d7-4129-8f2c-dcedd2530b4a",
   "metadata": {},
   "source": [
    "### Inference Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "4aa55947-9456-41c4-a68c-431e47b49ad4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def predict(instances, endpoint):\n",
    "    credentials, _ = google.auth.default()\n",
    "    auth_req = google.auth.transport.requests.Request()\n",
    "    credentials.refresh(auth_req)\n",
    "    url = f\"{endpoint}/predict\"\n",
    "    headers = {\n",
    "        'Authorization': f'Bearer {credentials.token}',\n",
    "        'Content_Type': 'application/json'\n",
    "    }\n",
    "    data = json.dumps({'instances': instances})\n",
    "    response = requests.post(url, headers = headers, data = data)    \n",
    "    return json.loads(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "2a6dee3f-c21f-4b12-bd0b-b97953411e10",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [[0.9684653609075246, 0.031534639092475496]]}"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[0:1], service_start_main.metadata.uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "1bf73a76-3b25-412c-8a32-3152d99d3fff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [{'classes': ['0', '1'],\n",
       "   'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       "   'scores': [0.9684653609075246, 0.031534639092475496],\n",
       "   'predicted_class': '0'}]}"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[0:1], service_start_main2.metadata.uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "559d8b81-2fb8-4a47-b729-69966302b8f2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [[0.004793214059859108, 0.9952067859401409]]}"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[1:2], service_start_main.metadata.uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "0e37e57b-f7f3-4ee5-b686-80e89a702678",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [{'classes': ['0', '1'],\n",
       "   'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008',\n",
       "   'scores': [0.9684653609075246, 0.031534639092475496],\n",
       "   'predicted_class': '0'},\n",
       "  {'classes': ['0', '1'],\n",
       "   'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561',\n",
       "   'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "   'predicted_class': '1'}]}"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids[0:2], service_start_main2.metadata.uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "a4126754-7e9a-4163-a4c8-81993528b7ef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'predictions': [[0.9684653609075246, 0.031534639092475496],\n",
       "  [0.004793214059859108, 0.9952067859401409],\n",
       "  [0.9961659959204033, 0.0038340040795967153],\n",
       "  [0.009687501567520629, 0.9903124984324794],\n",
       "  [0.05002634033491404, 0.949973659665086],\n",
       "  [0.9988754743912287, 0.0011245256087712965],\n",
       "  [0.0058287633360615265, 0.9941712366639385],\n",
       "  [0.0031326212134841214, 0.9968673787865159],\n",
       "  [0.8122003166930434, 0.18779968330695668],\n",
       "  [0.007672836590005838, 0.9923271634099942]]}"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(transaction_ids, service_start_main.metadata.uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d2b27e2-1f55-4d05-9c47-330200501557",
   "metadata": {},
   "source": [
    "### Remove The Service(s)\n",
    "\n",
    "Cloud Run will scale to zero here since a minimum has not been set.  This notebook does proceed with deleting the service in the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "cd273416-6a0d-4500-9d24-19c6327a8a5b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "remove_response_main = cr.delete_service(request = {\"name\": service_start_main.metadata.name})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "d8d14e52-3bad-4190-8df4-eddf8d904ad3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "remove_response_main2 = cr.delete_service(request = {\"name\": service_start_main2.metadata.name})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "0483dd8f-b399-4949-8c00-25f960dd6d06",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name: \"projects/statmike-mlops-349915/locations/us-central1/services/frameworks-catboost-prediction-feature-store-1\"\n",
       "uid: \"10c638ba-ce33-4d79-b97b-8fc06245822d\"\n",
       "generation: 2\n",
       "create_time {\n",
       "  seconds: 1739191764\n",
       "  nanos: 471204000\n",
       "}\n",
       "update_time {\n",
       "  seconds: 1739192018\n",
       "  nanos: 868980000\n",
       "}\n",
       "delete_time {\n",
       "  seconds: 1739192018\n",
       "  nanos: 68081000\n",
       "}\n",
       "expire_time {\n",
       "  seconds: 1741784018\n",
       "  nanos: 68081000\n",
       "}\n",
       "creator: \"1026793852137-compute@developer.gserviceaccount.com\"\n",
       "last_modifier: \"1026793852137-compute@developer.gserviceaccount.com\"\n",
       "ingress: INGRESS_TRAFFIC_INTERNAL_ONLY\n",
       "launch_stage: GA\n",
       "template {\n",
       "  scaling {\n",
       "    max_instance_count: 100\n",
       "  }\n",
       "  timeout {\n",
       "    seconds: 300\n",
       "  }\n",
       "  service_account: \"1026793852137-compute@developer.gserviceaccount.com\"\n",
       "  containers {\n",
       "    image: \"us-central1-docker.pkg.dev/statmike-mlops-349915/frameworks/catboost-prediction-feature-store\"\n",
       "    env {\n",
       "      name: \"AIP_HTTP_PORT\"\n",
       "      value: \"8080\"\n",
       "    }\n",
       "    env {\n",
       "      name: \"AIP_HEALTH_ROUTE\"\n",
       "      value: \"/health\"\n",
       "    }\n",
       "    env {\n",
       "      name: \"AIP_PREDICT_ROUTE\"\n",
       "      value: \"/predict\"\n",
       "    }\n",
       "    env {\n",
       "      name: \"AIP_STORAGE_URI\"\n",
       "      value: \"gs://statmike-mlops-349915/frameworks/catboost-overview\"\n",
       "    }\n",
       "    env {\n",
       "      name: \"MODULE_NAME\"\n",
       "      value: \"main\"\n",
       "    }\n",
       "    resources {\n",
       "      limits {\n",
       "        key: \"memory\"\n",
       "        value: \"32Gi\"\n",
       "      }\n",
       "      limits {\n",
       "        key: \"cpu\"\n",
       "        value: \"8\"\n",
       "      }\n",
       "    }\n",
       "    ports {\n",
       "      name: \"http1\"\n",
       "      container_port: 8080\n",
       "    }\n",
       "    startup_probe {\n",
       "      timeout_seconds: 240\n",
       "      period_seconds: 240\n",
       "      failure_threshold: 1\n",
       "      tcp_socket {\n",
       "        port: 8080\n",
       "      }\n",
       "    }\n",
       "  }\n",
       "  max_instance_request_concurrency: 1000\n",
       "}\n",
       "traffic {\n",
       "  type_: TRAFFIC_TARGET_ALLOCATION_TYPE_LATEST\n",
       "  percent: 100\n",
       "}\n",
       "urls: \"https://frameworks-catboost-prediction-feature-store-1-1026793852137.us-central1.run.app\"\n",
       "urls: \"https://frameworks-catboost-prediction-feature-store-1-urlxi72dpa-uc.a.run.app\"\n",
       "observed_generation: 2\n",
       "terminal_condition {\n",
       "  type_: \"Ready\"\n",
       "  state: CONDITION_SUCCEEDED\n",
       "  last_transition_time {\n",
       "    seconds: 1739192018\n",
       "    nanos: 868980000\n",
       "  }\n",
       "}\n",
       "conditions {\n",
       "  type_: \"RoutesReady\"\n",
       "  state: CONDITION_SUCCEEDED\n",
       "  last_transition_time {\n",
       "    seconds: 1739191770\n",
       "    nanos: 308464000\n",
       "  }\n",
       "}\n",
       "conditions {\n",
       "  type_: \"ConfigurationsReady\"\n",
       "  state: CONDITION_SUCCEEDED\n",
       "  last_transition_time {\n",
       "    seconds: 1739191769\n",
       "    nanos: 894978000\n",
       "  }\n",
       "}\n",
       "latest_ready_revision: \"projects/statmike-mlops-349915/locations/us-central1/services/frameworks-catboost-prediction-feature-store-1/revisions/frameworks-catboost-prediction-feature-store-1-00001-wht\"\n",
       "latest_created_revision: \"projects/statmike-mlops-349915/locations/us-central1/services/frameworks-catboost-prediction-feature-store-1/revisions/frameworks-catboost-prediction-feature-store-1-00001-wht\"\n",
       "traffic_statuses {\n",
       "  type_: TRAFFIC_TARGET_ALLOCATION_TYPE_LATEST\n",
       "  percent: 100\n",
       "}\n",
       "uri: \"https://frameworks-catboost-prediction-feature-store-1-urlxi72dpa-uc.a.run.app\"\n",
       "etag: \"\\\"CNLtp70GEKCqrp4D/cHJvamVjdHMvc3RhdG1pa2UtbWxvcHMtMzQ5OTE1L2xvY2F0aW9ucy91cy1jZW50cmFsMS9zZXJ2aWNlcy9mcmFtZXdvcmtzLWNhdGJvb3N0LXByZWRpY3Rpb24tZmVhdHVyZS1zdG9yZS0x\\\"\""
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "remove_response_main.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fa2b4ef-1dc7-406a-987d-1bcee9f571e2",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "## Vertex AI Prediction Endpoint\n",
    "\n",
    "Register the model in the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction) and [Deploy it to an endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment).\n",
    "\n",
    "**Note 1** that Vertex AI automatically sets the environment variables that start with `AIP_` so they should not be supplied in the model setup.  Read more in [Environment variables available in the container](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#variables).\n",
    "\n",
    "**Note 2** Vertex AI Endpoints Launch with the [permissions](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#permissions) of the service agent.  If your container needs to run as a specific IAM user, like a service account, with permission to other resources then you can supply the `SERVICE_ACCOUNT` when deploying the model which is used below.  Read more about this parameter in the Python API guide [ for `Model.deploy()` here](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.Model#google_cloud_aiplatform_Model_deploy)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c00a408c-8461-4d9d-926f-8b99f5f63998",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Model Registry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba3d4c5-8bbc-4afd-8fdf-71d59f921072",
   "metadata": {},
   "source": [
    "Check for existing version of the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "793217c5-de15-4350-8ca7-355d403d0347",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'projects/1026793852137/locations/us-central1/models/frameworks'"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parent_model = ''\n",
    "for model in aiplatform.Model.list(filter=f'display_name=\"{SERIES}\"'):\n",
    "    parent_model = model.resource_name\n",
    "    break\n",
    "parent_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bf9e27b-6003-437e-9e7f-e7d2698b9c6b",
   "metadata": {},
   "source": [
    "Upload the model to the registry with different versions for:\n",
    "- plain responses - using main.py which is default\n",
    "- formated responses - using main2.py which is set with environment variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "57a523c1-cfd7-47af-9908-6276d7c49b1b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating Model\n",
      "Create Model backing LRO: projects/1026793852137/locations/us-central1/models/frameworks/operations/6278860003382132736\n",
      "Model created. Resource name: projects/1026793852137/locations/us-central1/models/frameworks@3\n",
      "To use this Model in another session:\n",
      "model = aiplatform.Model('projects/1026793852137/locations/us-central1/models/frameworks@3')\n"
     ]
    }
   ],
   "source": [
    "vertex_model = aiplatform.Model.upload(\n",
    "    display_name = SERIES,\n",
    "    model_id = SERIES,\n",
    "    parent_model = parent_model,\n",
    "    serving_container_image_uri = build_response.artifacts.images[0],\n",
    "    artifact_uri = f'gs://{bucket.name}/{SERIES}/catboost-overview',\n",
    "    is_default_version = True,\n",
    "    version_aliases = [f'{EXPERIMENT}-plain'],\n",
    "    version_description = EXPERIMENT,\n",
    "    labels = {'series': SERIES, 'experiment': EXPERIMENT}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "93bd02fb-244b-4c93-9c3c-56d39a676a9a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating Model\n",
      "Create Model backing LRO: projects/1026793852137/locations/us-central1/models/frameworks/operations/7596162894388002816\n",
      "Model created. Resource name: projects/1026793852137/locations/us-central1/models/frameworks@4\n",
      "To use this Model in another session:\n",
      "model = aiplatform.Model('projects/1026793852137/locations/us-central1/models/frameworks@4')\n"
     ]
    }
   ],
   "source": [
    "vertex_model = aiplatform.Model.upload(\n",
    "    parent_model = vertex_model.resource_name,\n",
    "    serving_container_image_uri = build_response.artifacts.images[0],\n",
    "    serving_container_environment_variables = {'MODULE_NAME': 'main2'},\n",
    "    artifact_uri = f'gs://{bucket.name}/{SERIES}/catboost-overview',\n",
    "    is_default_version = False,\n",
    "    version_aliases = [f'{EXPERIMENT}-formatted'],\n",
    "    version_description = EXPERIMENT,\n",
    "    labels = {'series': SERIES, 'experiment': EXPERIMENT}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1581fc6c-3b4b-4712-b244-0dfb3fdec812",
   "metadata": {},
   "source": [
    "### Create Endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d22e797-1600-4355-ad4a-1c4208f2684a",
   "metadata": {},
   "source": [
    "Check for existing endpoint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "588c73bb-8afa-4d84-bb18-f8297ab5b29c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vertex_endpoint = None\n",
    "for endpoint in aiplatform.Endpoint.list(filter=f'display_name=\"{SERIES}\"'):\n",
    "    vertex_endpoint = endpoint\n",
    "    break\n",
    "vertex_endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4590e6c-3fe5-43a6-9a08-e8a0cb2da2ab",
   "metadata": {},
   "source": [
    "Create endpoint if missing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "bfa3284d-df3c-4b70-a18e-73f973c8195e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating Endpoint\n",
      "Create Endpoint backing LRO: projects/1026793852137/locations/us-central1/endpoints/757089622026092544/operations/521007879788953600\n",
      "Endpoint created. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "To use this Endpoint in another session:\n",
      "endpoint = aiplatform.Endpoint('projects/1026793852137/locations/us-central1/endpoints/757089622026092544')\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<google.cloud.aiplatform.models.Endpoint object at 0x7f9dc81f6530> \n",
       "resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if not vertex_endpoint:\n",
    "    vertex_endpoint = aiplatform.Endpoint.create(\n",
    "        display_name = f\"{SERIES}\",\n",
    "        labels = {'series': SERIES}   \n",
    "    )\n",
    "vertex_endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43d66c94-b4f6-4c2f-a0ec-f48d8d97dd4b",
   "metadata": {},
   "source": [
    "### Deploy Model: Default version with `main.py` - Plain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7591a76d-8986-49e5-9156-16dd921aecd6",
   "metadata": {},
   "source": [
    "Get the latest model version with alias `plain`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "a6ffc496-729a-4809-968d-ba1789b58ac3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'projects/1026793852137/locations/us-central1/models/frameworks@3'"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_model = aiplatform.Model(model_name = SERIES, version = f'{EXPERIMENT}-plain')\n",
    "vertex_model.versioned_resource_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "09fc56aa-9cfc-4675-baaa-a366506d0eda",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deploying model to Endpoint : projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Deploy Endpoint model backing LRO: projects/1026793852137/locations/us-central1/endpoints/757089622026092544/operations/8029071408568991744\n",
      "Endpoint model deployed. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<google.cloud.aiplatform.models.Endpoint object at 0x7f9dc81f6530> \n",
       "resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_model.deploy(\n",
    "    endpoint = vertex_endpoint,\n",
    "    traffic_percentage = 100,\n",
    "    machine_type = 'n1-standard-4',\n",
    "    min_replica_count = 1,\n",
    "    max_replica_count = 2,\n",
    "    service_account = SERVICE_ACCOUNT\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf65cc7c-5b89-4f24-9642-ea7183bb4e7e",
   "metadata": {},
   "source": [
    "### Test Predictions: Plain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "2dc3bde7-e715-4052-8695-9f5d6aac3f5f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(predictions=[[0.9684653609075246, 0.0315346390924755]], deployed_model_id='8642070134854254592', metadata=None, model_version_id='3', model_resource_name='projects/1026793852137/locations/us-central1/models/frameworks', explanations=None)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[0:1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "3a93abfb-0b26-4381-b1d0-bdc27c63ba49",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0.004793214059859108, 0.9952067859401409]]"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[1:2]).predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "b3290ad5-0c9d-43a1-b6db-056258f208ca",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0.9684653609075246, 0.0315346390924755],\n",
       " [0.004793214059859108, 0.9952067859401409]]"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[0:2]).predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "a5d186c1-81fd-410d-a42c-43df6af5f403",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0.9684653609075246, 0.0315346390924755],\n",
       " [0.004793214059859108, 0.9952067859401409],\n",
       " [0.9961659959204033, 0.003834004079596715],\n",
       " [0.009687501567520629, 0.9903124984324794],\n",
       " [0.05002634033491404, 0.949973659665086],\n",
       " [0.9988754743912287, 0.001124525608771297],\n",
       " [0.005828763336061527, 0.9941712366639385],\n",
       " [0.003132621213484121, 0.9968673787865159],\n",
       " [0.8122003166930434, 0.1877996833069567],\n",
       " [0.007672836590005838, 0.9923271634099942]]"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids).predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89785a04-d0c7-448e-ba99-4f298647fd18",
   "metadata": {},
   "source": [
    "### Deploy Model: Version with `main2.py` - Formatted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "b3ee560e-b992-42e7-a5e0-370300b06471",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'projects/1026793852137/locations/us-central1/models/frameworks@4'"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_model = aiplatform.Model(model_name = SERIES, version = f'{EXPERIMENT}-formatted')\n",
    "vertex_model.versioned_resource_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "5063a674-d069-41bf-a17b-cff99b0d45be",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deploying model to Endpoint : projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Deploy Endpoint model backing LRO: projects/1026793852137/locations/us-central1/endpoints/757089622026092544/operations/4034378539091361792\n",
      "Endpoint model deployed. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<google.cloud.aiplatform.models.Endpoint object at 0x7f9dc81f6530> \n",
       "resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_model.deploy(\n",
    "    endpoint = vertex_endpoint,\n",
    "    traffic_percentage = 0,\n",
    "    machine_type = 'n1-standard-4',\n",
    "    min_replica_count = 1,\n",
    "    max_replica_count = 2,\n",
    "    service_account = SERVICE_ACCOUNT\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "291cf9b8-fb76-43dd-85d0-eefbadbd668e",
   "metadata": {},
   "source": [
    "### Shift Traffic To Model Version: Formatted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "a75be924-4863-4db6-ac43-9c402b2c386c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'8642070134854254592': 100, '2101717576005451776': 0}"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.traffic_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "3262ab9c-8cc9-4918-a078-6b55e3382aee",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'2101717576005451776': 100, '8642070134854254592': 0}"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_traffic_split = {}\n",
    "for deployed_model in vertex_endpoint.list_models():\n",
    "    if deployed_model.model_version_id == vertex_model.version_id:\n",
    "        new_traffic_split[deployed_model.id] = 100\n",
    "    else:\n",
    "        new_traffic_split[deployed_model.id] = 0\n",
    "new_traffic_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "62aeed1c-8583-47cb-8c18-4bc50b9d3bf8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating Endpoint endpoint: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Endpoint endpoint updated. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<google.cloud.aiplatform.models.Endpoint object at 0x7f9dc81f6530> \n",
       "resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.update(traffic_split = new_traffic_split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "e0ba5fac-7656-4677-a420-650d4cf88347",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'8642070134854254592': 0, '2101717576005451776': 100}"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.traffic_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1122250d-1d08-4c3e-bfbd-b3140bdc17d2",
   "metadata": {},
   "source": [
    "### Test Predictions: Formatted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "8573cf68-fe08-4bc5-9cff-b8375a84c99d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(predictions=[{'predicted_class': '0', 'classes': ['0', '1'], 'scores': [0.9684653609075246, 0.0315346390924755], 'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008'}], deployed_model_id='2101717576005451776', metadata=None, model_version_id='4', model_resource_name='projects/1026793852137/locations/us-central1/models/frameworks', explanations=None)"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[0:1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "45e55474-01ef-4c9e-bbbd-c918c37e1774",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'predicted_class': '1',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "  'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561'}]"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[1:2]).predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "38a629f3-966a-47dc-9d76-53103b2fb984",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'classes': ['0', '1'],\n",
       "  'predicted_class': '0',\n",
       "  'scores': [0.9684653609075246, 0.0315346390924755],\n",
       "  'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008'},\n",
       " {'classes': ['0', '1'],\n",
       "  'predicted_class': '1',\n",
       "  'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "  'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561'}]"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[0:2]).predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "bbce7ec2-660e-4a72-a0ba-aeef6a5efca3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'classes': ['0', '1'],\n",
       "  'predicted_class': '0',\n",
       "  'scores': [0.9684653609075246, 0.0315346390924755],\n",
       "  'entity_id': '28f7cfa6-5862-428a-8d35-c9e08d297008'},\n",
       " {'predicted_class': '1',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.004793214059859108, 0.9952067859401409],\n",
       "  'entity_id': '1dbfe576-2c52-4f94-bd88-db1e3e294561'},\n",
       " {'classes': ['0', '1'],\n",
       "  'predicted_class': '0',\n",
       "  'scores': [0.9961659959204033, 0.003834004079596715],\n",
       "  'entity_id': 'd13fbb1a-d4cb-40c6-bfd7-05cd714cee30'},\n",
       " {'predicted_class': '1',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.009687501567520629, 0.9903124984324794],\n",
       "  'entity_id': '18327e33-3ff6-43d1-8bbc-8647adc2044d'},\n",
       " {'predicted_class': '1',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.05002634033491404, 0.949973659665086],\n",
       "  'entity_id': 'ec505a29-387f-4b3d-a03b-67375d9140e1'},\n",
       " {'predicted_class': '0',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.9988754743912287, 0.001124525608771297],\n",
       "  'entity_id': '1b7b6e26-7248-420c-9eb5-dd9692437843'},\n",
       " {'predicted_class': '1',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.005828763336061527, 0.9941712366639385],\n",
       "  'entity_id': 'bc996e23-c364-4e1a-875d-5151e47fdb81'},\n",
       " {'classes': ['0', '1'],\n",
       "  'predicted_class': '1',\n",
       "  'scores': [0.003132621213484121, 0.9968673787865159],\n",
       "  'entity_id': '165ac053-1fec-4863-aea2-63d380a3adad'},\n",
       " {'predicted_class': '0',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.8122003166930434, 0.1877996833069567],\n",
       "  'entity_id': '15b93da7-e600-449b-a4c1-bacfd7e68df7'},\n",
       " {'predicted_class': '1',\n",
       "  'classes': ['0', '1'],\n",
       "  'scores': [0.007672836590005838, 0.9923271634099942],\n",
       "  'entity_id': 'd081be14-dcb5-4228-813c-75b1804dc8a2'}]"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids).predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef3fc132-9275-4f67-9418-96b203f4205b",
   "metadata": {},
   "source": [
    "### Shift Traffic To Model: Plain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "0b960397-dbb9-459b-9488-6c10000c8376",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'8642070134854254592': 0, '2101717576005451776': 100}"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.traffic_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "5834809d-4b6a-45e1-a05e-cd14f61005fc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating Endpoint endpoint: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Endpoint endpoint updated. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<google.cloud.aiplatform.models.Endpoint object at 0x7f9dc81f6530> \n",
       "resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_traffic_split = {}\n",
    "for deployed_model in vertex_endpoint.list_models():\n",
    "    if deployed_model.model_version_id != vertex_model.version_id:\n",
    "        new_traffic_split[deployed_model.id] = 100\n",
    "    else:\n",
    "        new_traffic_split[deployed_model.id] = 0\n",
    "vertex_endpoint.update(traffic_split = new_traffic_split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "a9043934-d27a-4b41-9a5c-766405d8d5d5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'8642070134854254592': 100, '2101717576005451776': 0}"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.traffic_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "8e496ce0-8469-44b5-a664-b50d5cb96f9b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0.9684653609075246, 0.0315346390924755],\n",
       " [0.004793214059859108, 0.9952067859401409]]"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vertex_endpoint.predict(instances = transaction_ids[0:2]).predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c208be38-d956-4629-bc29-26802cb318eb",
   "metadata": {},
   "source": [
    "### Undeploy Models Without Traffic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "161ecc57-e400-4c26-82f9-f6fd8c8cb7cb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Undeploying Endpoint model: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Undeploy Endpoint model backing LRO: projects/1026793852137/locations/us-central1/endpoints/757089622026092544/operations/8120832250976665600\n",
      "Endpoint model undeployed. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'8642070134854254592': 100}"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for deployed_model in vertex_endpoint.list_models():\n",
    "    if vertex_endpoint.traffic_split[deployed_model.id] == 0:\n",
    "        vertex_endpoint.undeploy(deployed_model_id = deployed_model.id)\n",
    "vertex_endpoint.traffic_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b033f5be-2a53-426c-b967-149ad8cb71d8",
   "metadata": {},
   "source": [
    "### Undeploy All Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "4fe6631d-7561-4ac4-9aee-019960a4e04e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Undeploying Endpoint model: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Undeploy Endpoint model backing LRO: projects/1026793852137/locations/us-central1/endpoints/757089622026092544/operations/3211908657142824960\n",
      "Endpoint model undeployed. Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for deployed_model in vertex_endpoint.list_models():\n",
    "    vertex_endpoint.undeploy(deployed_model_id = deployed_model.id)\n",
    "vertex_endpoint.list_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9d6e20b-44bd-429e-be2c-65329463f3db",
   "metadata": {},
   "source": [
    "### Delete Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "2512f335-9257-426c-81c1-9e0295baa222",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleting Endpoint : projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Endpoint deleted. . Resource name: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Deleting Endpoint resource: projects/1026793852137/locations/us-central1/endpoints/757089622026092544\n",
      "Delete Endpoint backing LRO: projects/1026793852137/locations/us-central1/operations/6218061408412631040\n",
      "Endpoint resource projects/1026793852137/locations/us-central1/endpoints/757089622026092544 deleted.\n"
     ]
    }
   ],
   "source": [
    "vertex_endpoint.delete(force = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9898b312-00fe-4e8a-86ef-71ace3596c56",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a096ecec-1df2-4261-954b-c76ad075df46",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m125",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m125"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}