notebooks/community/vertex_endpoints/optimized_tensorflow_runtime/tabular_optimized_online_prediction.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/optimized_tensorflow_runtime/tabular_optimized_online_prediction.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/optimized_tensorflow_runtime/tabular_optimized_online_prediction.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/vertex_endpoints/optimized_tensorflow_runtime/tabular_optimized_online_prediction.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "19cdbb1b1c83"
      },
      "source": [
        "# Training a tabular Criteo model and deploying it to Vertex AI Predictions using the optimized TensorFlow runtime"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "In this sample you learn how to train a tabular model using TensorFlow Keras or Estimator API using Criteo Kaggle dataset.\n",
        "Next, you export a trained model to the Vertex AI Prediction service using open source based TensorFlow 2.7 container and the optimized TensorFlow runtime container, run performance evaluation for those models side by side and compare predictions.\n",
        "\n",
        "For additional information about Vertex AI Prediction optimized TensorFlow runtime containers, please refer to https://cloud.google.com/vertex-ai/docs/predictions/optimized-tensorflow-runtime.\n",
        "\n",
        "### Dataset\n",
        "\n",
        "In this sample you use Criteo Kaggle dataset, which takes about 4GB.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to deploy a trained tabular model to Vertex AI Prediction using the optimized TensorFlow runtime, then compare its performance to open source based TensorFlow container.\n",
        "\n",
        "The steps performed include:\n",
        "* Download and unpack Criteo Kaggle dataset\n",
        "* Build and train a model using the Keras API\n",
        "* Setup private endpoints\n",
        "* Deploy a model to Vertex AI Prediction using TensorFlow 2.7 container\n",
        "* Deploy a model to Vertex AI Prediction using optimized TensorFlow container\n",
        "* Benchmark both models and validate their predictions\n",
        "\n",
        "You can train a model and upload it to Vertex AI Prediction using Colab. Since this walkthrough uses private endpoints to demonstrate Vertex AI Predictions, you must use Jupyter VM to run benchmarks.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses the following billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ze4-nDLfK4pw"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Colab or Vertex AI Workbench Notebooks**, your environment meets the requirements to run this notebook. You can skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCuSR8GkAgzl"
      },
      "source": [
        "**If you are not using Colab or Vertex AI Workbench Notebooks**, you must have the following in your environment to meet this notebook's requirements.\n",
        "\n",
        "* The Google Cloud SDK\n",
        "* Git\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps are condensed instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip3 install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Install additional package dependencies not installed in your notebook environment, such as TensorFlow Serving API, Vertex AI SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Vertex AI Workbench Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# Vertex AI Workbench Notebook requires dependencies to be installed with '--user'\n",
        "USER_FLAG = \"\"\n",
        "if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "    USER_FLAG = \"--user\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyy5Lbnzg5fi"
      },
      "outputs": [],
      "source": [
        "! pip3 install {USER_FLAG} --upgrade tensorflow==2.7.0 -q\n",
        "! pip3 install {USER_FLAG} --upgrade tensorflow-serving-api==2.7.0 -q\n",
        "! pip3 install {USER_FLAG} --upgrade google-cloud-aiplatform -q\n",
        "! pip3 install {USER_FLAG} --upgrade google-cloud-storage -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhq5zEbGg0XX"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install the additional packages, you must restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lWEdiXsJg0XY"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 credit towards your compute and storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
        "\n",
        "1. [Enable the Service Networking API](https://console.cloud.google.com/flows/enableapi?apiid=servicenetworking.googleapis.com). \n",
        "\n",
        "1. [Enable the Cloud DNS API](https://console.cloud.google.com/flows/enableapi?apiid=dns.googleapis.com). \n",
        "\n",
        "1. If you are running this notebook locally, you must install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the following cell. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you can try to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"\"\n",
        "\n",
        "# Get your Google Cloud project ID from gcloud\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID: \", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJYoRfYng0XZ"
      },
      "source": [
        "Set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "#### Timestamp\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, create a timestamp for each instance session, then append it to the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "697568e92bd6"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dr--iN2kAylZ"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
        "authenticated. Skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "**If you are using Colab**, run the cell below and follow the instructions\n",
        "when prompted to authenticate your account via oAuth.\n",
        "\n",
        "**Otherwise**, follow these steps:\n",
        "\n",
        "1. In the Cloud Console, go to the [**Create service account key**\n",
        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
        "\n",
        "2. Click **Create service account**.\n",
        "\n",
        "3. In the **Service account name** field, enter a name, and\n",
        "   click **Create**.\n",
        "\n",
        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\"\n",
        "into the filter box, and select\n",
        "   **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
        "\n",
        "5. Click **Create**. A JSON file that contains your key downloads to your\n",
        "local environment.\n",
        "\n",
        "6. Enter the path to your service account key as the\n",
        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the following cell, then run the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "# If you are running this notebook in Colab, run this cell and follow the\n",
        "# instructions to authenticate your GCP account. This provides access to your\n",
        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
        "# requests.\n",
        "\n",
        "# The Vertex AI Workbench Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "# If on Vertex AI Workbench Notebooks, then don't execute this code\n",
        "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "    if \"google.colab\" in sys.modules:\n",
        "        from google.colab import auth as google_auth\n",
        "\n",
        "        google_auth.authenticate_user()\n",
        "\n",
        "    # If you are running this notebook locally, replace the string below with the\n",
        "    # path to your service account key and run this cell to authenticate your GCP\n",
        "    # account.\n",
        "    elif not os.getenv(\"IS_TESTING\"):\n",
        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required for all notebook environments.**\n",
        "\n",
        "For Vertex AI Prediction to serve your model, it must be uploaded to Cloud Storage bucket first.\n",
        "\n",
        "Set the name of your Cloud Storage bucket in the following cell. It must be unique across all\n",
        "Cloud Storage buckets.\n",
        "\n",
        "You can change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook. We suggest that you [choose a region where Vertex AI services are\n",
        "available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "REGION = \"[your-region]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf221059d072"
      },
      "outputs": [],
      "source": [
        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP\n",
        "\n",
        "if REGION == \"[your-region]\":\n",
        "    REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ucvCsknMCims"
      },
      "source": [
        "The final step for your Cloud Storage bucket is to validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhOb7YnwClBb"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import math\n",
        "import os\n",
        "import re\n",
        "import sys\n",
        "from urllib.parse import urlparse\n",
        "\n",
        "import grpc\n",
        "import numpy as np\n",
        "import requests as r\n",
        "import tensorflow as tf\n",
        "from tensorflow_serving.apis import (predict_pb2, prediction_log_pb2,\n",
        "                                     prediction_service_pb2_grpc)\n",
        "\n",
        "logging = tf.get_logger()\n",
        "logging.propagate = False\n",
        "logging.setLevel(\"INFO\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xS8EWC5gLBGE"
      },
      "outputs": [],
      "source": [
        "LOCAL_DIRECTORY = \"~/criteo\"  # @param {type:\"string\"}\n",
        "HIDDEN_LAYERS_STR = \"2048,2048,1024,512,256\"  # @param {type:\"string\"}\n",
        "\n",
        "HIDDEN_LAYERS = list(map(lambda x: int(x), HIDDEN_LAYERS_STR.split(\",\")))\n",
        "LOCAL_DIRECTORY_FULL = os.path.expanduser(LOCAL_DIRECTORY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ee07345bcc6"
      },
      "source": [
        "## Download the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pr_KdfVtM7k5"
      },
      "source": [
        "Follow the instructions at the [Criteo website](https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset/) to download the data.\n",
        "\n",
        "If the data is not available, you can download it using the URL below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2QjMxvFCC5A7"
      },
      "outputs": [],
      "source": [
        "!mkdir -p $LOCAL_DIRECTORY_FULL/data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vkRouSvKaAbB"
      },
      "outputs": [],
      "source": [
        "!cd $LOCAL_DIRECTORY_FULL/data && curl -O https://s3-eu-west-1.amazonaws.com/pfigshare-u-files/10082655/dac.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IGfU1kNx_oz6"
      },
      "outputs": [],
      "source": [
        "!cd $LOCAL_DIRECTORY_FULL/data && tar xvzf dac.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-t27rnXAAtka"
      },
      "outputs": [],
      "source": [
        "!head -n 3 $LOCAL_DIRECTORY_FULL/data/train.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ng-Nrc8UdbP2"
      },
      "source": [
        "## Read and transform the dataset\n",
        "\n",
        "Before the model can be trained, the variables must be pre-processed.\n",
        "\n",
        "Numerical values are normalized by subtracting their average and dividing by their standard deviation.\n",
        "The average values and the standard deviations are precalculated for each numerical feature. The vocabulary sizes are precalculated for each categorical feature."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fmOWgpbWddbW"
      },
      "outputs": [],
      "source": [
        "COLUMN_NAMES = [\n",
        "    \"label\",\n",
        "    \"int1\",\n",
        "    \"int2\",\n",
        "    \"int3\",\n",
        "    \"int4\",\n",
        "    \"int5\",\n",
        "    \"int6\",\n",
        "    \"int7\",\n",
        "    \"int8\",\n",
        "    \"int9\",\n",
        "    \"int10\",\n",
        "    \"int11\",\n",
        "    \"int12\",\n",
        "    \"int13\",\n",
        "    \"cat1\",\n",
        "    \"cat2\",\n",
        "    \"cat3\",\n",
        "    \"cat4\",\n",
        "    \"cat5\",\n",
        "    \"cat6\",\n",
        "    \"cat7\",\n",
        "    \"cat8\",\n",
        "    \"cat9\",\n",
        "    \"cat10\",\n",
        "    \"cat11\",\n",
        "    \"cat12\",\n",
        "    \"cat13\",\n",
        "    \"cat14\",\n",
        "    \"cat15\",\n",
        "    \"cat16\",\n",
        "    \"cat17\",\n",
        "    \"cat18\",\n",
        "    \"cat19\",\n",
        "    \"cat20\",\n",
        "    \"cat21\",\n",
        "    \"cat22\",\n",
        "    \"cat23\",\n",
        "    \"cat24\",\n",
        "    \"cat25\",\n",
        "    \"cat26\",\n",
        "]\n",
        "\n",
        "# Precalculated, see\n",
        "# https://github.com/vlasenkoalexey/criteo_nbdev/blob/master/04_data_reader.ipynb\n",
        "NUM_AVERAGE = {\n",
        "    \"int1\": 3.5024133170753995,\n",
        "    \"int2\": 105.8484197976657,\n",
        "    \"int3\": 26.91304102061112,\n",
        "    \"int4\": 7.322680248873331,\n",
        "    \"int5\": 18538.99166487135,\n",
        "    \"int6\": 116.06185085211605,\n",
        "    \"int7\": 16.333130032135013,\n",
        "    \"int8\": 12.517042137556762,\n",
        "    \"int9\": 106.10982343805145,\n",
        "    \"int10\": 0.6175294977722183,\n",
        "    \"int11\": 2.7328343170173173,\n",
        "    \"int12\": 0.9910356287721245,\n",
        "    \"int13\": 8.21746116117401,\n",
        "}\n",
        "NUM_STDDEV = {\n",
        "    \"int1\": 9.429076407105086,\n",
        "    \"int2\": 391.4578226870704,\n",
        "    \"int3\": 397.97258302273474,\n",
        "    \"int4\": 8.793230712645805,\n",
        "    \"int5\": 69394.60184622335,\n",
        "    \"int6\": 382.5664493712363,\n",
        "    \"int7\": 66.0497552451171,\n",
        "    \"int8\": 16.688884567787586,\n",
        "    \"int9\": 220.28309398647906,\n",
        "    \"int10\": 0.6840505553977025,\n",
        "    \"int11\": 5.199070884811354,\n",
        "    \"int12\": 5.597723872237179,\n",
        "    \"int13\": 16.211932558173785,\n",
        "}\n",
        "VOCABULARY_SIZE = {\n",
        "    \"cat1\": 1460,\n",
        "    \"cat2\": 583,\n",
        "    \"cat3\": 10131226,\n",
        "    \"cat4\": 2202607,\n",
        "    \"cat5\": 305,\n",
        "    \"cat6\": 23,\n",
        "    \"cat7\": 12517,\n",
        "    \"cat8\": 633,\n",
        "    \"cat9\": 3,\n",
        "    \"cat10\": 93145,\n",
        "    \"cat11\": 5683,\n",
        "    \"cat12\": 8351592,\n",
        "    \"cat13\": 3194,\n",
        "    \"cat14\": 27,\n",
        "    \"cat15\": 14992,\n",
        "    \"cat16\": 5461305,\n",
        "    \"cat17\": 10,\n",
        "    \"cat18\": 5652,\n",
        "    \"cat19\": 2172,\n",
        "    \"cat20\": 3,\n",
        "    \"cat21\": 7046546,\n",
        "    \"cat22\": 17,\n",
        "    \"cat23\": 15,\n",
        "    \"cat24\": 286180,\n",
        "    \"cat25\": 104,\n",
        "    \"cat26\": 142571,\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6D2awii4dkfn"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def transform_row(*row_tuple):\n",
        "    row_dict = dict(\n",
        "        zip(list(column_name for column_name in COLUMN_NAMES), list(row_tuple))\n",
        "    )\n",
        "    dict_without_label = dict(row_dict)\n",
        "    label = dict_without_label.pop(\"label\")\n",
        "    return (dict_without_label, label)\n",
        "\n",
        "\n",
        "def read_gcs(batch_size=64):\n",
        "    file_name = os.path.join(LOCAL_DIRECTORY_FULL, \"data\", \"train.txt\")\n",
        "    record_defaults = list(\n",
        "        tf.int64\n",
        "        if column_name == \"label\"\n",
        "        else tf.constant(0, dtype=tf.int64)\n",
        "        if column_name.startswith(\"int\")\n",
        "        else tf.constant(\"\", dtype=tf.string)\n",
        "        for column_name in COLUMN_NAMES\n",
        "    )\n",
        "    dataset = tf.data.experimental.CsvDataset(\n",
        "        file_name, record_defaults, field_delim=\"\\t\", header=False\n",
        "    )\n",
        "\n",
        "    transformed_ds = (\n",
        "        dataset.batch(batch_size).shuffle(500).map(transform_row).prefetch(50)\n",
        "    )\n",
        "\n",
        "    return transformed_ds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oLf5zmpYdmaP"
      },
      "outputs": [],
      "source": [
        "for row in read_gcs(batch_size=3).take(2):\n",
        "    print(row)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xyfLNpj_dvyi"
      },
      "source": [
        "## Train and Save Keras Model\n",
        "\n",
        "For an overview of how to train tabular models using TensorFlow Keras APIs, see https://github.com/tensorflow/docs/blob/r2.4/site/en/tutorials/structured_data/feature_columns.ipynb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DKFTc57Hd17L"
      },
      "outputs": [],
      "source": [
        "def make_norm_fn(column_name):\n",
        "    avg = NUM_AVERAGE[column_name]\n",
        "    stddev = NUM_STDDEV[column_name]\n",
        "    return lambda v: (tf.dtypes.cast(v, tf.float32) - avg) / stddev\n",
        "\n",
        "\n",
        "def create_feature_columns():\n",
        "    linear_feature_columns = []\n",
        "    categorical_feature_columns = []\n",
        "\n",
        "    for column_name in COLUMN_NAMES:\n",
        "        if column_name.startswith(\"int\"):\n",
        "            linear_feature_columns.append(\n",
        "                tf.feature_column.numeric_column(\n",
        "                    column_name,\n",
        "                    dtype=tf.dtypes.int64,\n",
        "                    normalizer_fn=make_norm_fn(column_name),\n",
        "                )\n",
        "            )\n",
        "\n",
        "        if column_name.startswith(\"cat\"):\n",
        "            column_vocabulary_size = VOCABULARY_SIZE[column_name]\n",
        "            hash_bucket_size = min(column_vocabulary_size, 100000)\n",
        "            embedding_dimension = int(\n",
        "                min(50, math.floor(6 * column_vocabulary_size**0.25))\n",
        "            )\n",
        "            categorical_feature_columns.append(\n",
        "                tf.feature_column.embedding_column(\n",
        "                    tf.feature_column.categorical_column_with_hash_bucket(\n",
        "                        column_name, hash_bucket_size, dtype=tf.dtypes.string\n",
        "                    ),\n",
        "                    embedding_dimension,\n",
        "                )\n",
        "            )\n",
        "\n",
        "    return linear_feature_columns + categorical_feature_columns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4jgTZP1zd4Bb"
      },
      "outputs": [],
      "source": [
        "def create_keras_model_sequential():\n",
        "    feature_columns = create_feature_columns()\n",
        "\n",
        "    feature_layer = tf.keras.layers.DenseFeatures(feature_columns, name=\"feature_layer\")\n",
        "    Dense = tf.keras.layers.Dense\n",
        "    Dropout = tf.keras.layers.Dropout\n",
        "    BatchNormalization = tf.keras.layers.BatchNormalization\n",
        "    dense_layers = []\n",
        "    for c in HIDDEN_LAYERS:\n",
        "        dense_layers.append(BatchNormalization())\n",
        "        dense_layers.append(Dense(c, activation=tf.nn.relu))\n",
        "        dense_layers.append(Dropout(0.05))\n",
        "    model = tf.keras.Sequential(\n",
        "        [feature_layer] + dense_layers + [Dense(1, activation=tf.nn.sigmoid)]\n",
        "    )\n",
        "\n",
        "    logging.info(\"compiling sequential keras model\")\n",
        "    # Compile Keras model\n",
        "    model.compile(\n",
        "        optimizer=tf.optimizers.Adam(learning_rate=0.01),\n",
        "        loss=tf.keras.losses.BinaryCrossentropy(),\n",
        "        metrics=[\"accuracy\"],\n",
        "    )\n",
        "    return model\n",
        "\n",
        "\n",
        "model = create_keras_model_sequential()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZfvQSOmofI6B"
      },
      "source": [
        "Train the model. The expected loss is about 0.35."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vCanYbuJd5-S"
      },
      "outputs": [],
      "source": [
        "model.fit(read_gcs(batch_size=256).take(1000), epochs=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0HcxyuOMfMdn"
      },
      "source": [
        "Validate the model. The expected loss is about 0.45."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-GsVEleWfOPM"
      },
      "outputs": [],
      "source": [
        "model.evaluate(read_gcs(batch_size=256).skip(1000).take(1000))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JyyP4msGfPyZ"
      },
      "outputs": [],
      "source": [
        "model.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x9sx-tFBfRq3"
      },
      "outputs": [],
      "source": [
        "model.save(os.path.join(LOCAL_DIRECTORY_FULL, \"keras\"), include_optimizer=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AVibdmrufYRA"
      },
      "source": [
        "Check the model signature to see which fields prediction request should have."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MG4kerVsfVzI"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli show --dir $LOCAL_DIRECTORY_FULL/keras --all"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jKi0_T-bfaMT"
      },
      "source": [
        "## (Optional) Train and Save Estimator Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "11cf00aa9298"
      },
      "source": [
        "Another option to train a model is to use the TensorFlow Estimator API. For more information, see\n",
        "https://github.com/tensorflow/docs/blob/r2.4/site/en/tutorials/estimator/premade.ipynb\n",
        "\n",
        "The following code is provided only for illustration purposes. You use the Keras model for deployment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0fe8234cffb0"
      },
      "outputs": [],
      "source": [
        "feature_columns = create_feature_columns()\n",
        "estimator = tf.estimator.DNNClassifier(\n",
        "    optimizer=tf.optimizers.Adam(learning_rate=0.01),\n",
        "    feature_columns=feature_columns,\n",
        "    hidden_units=HIDDEN_LAYERS,\n",
        "    dropout=0.05,\n",
        "    batch_norm=True,\n",
        "    n_classes=2,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "35654a928cf7"
      },
      "outputs": [],
      "source": [
        "tf.estimator.train_and_evaluate(\n",
        "    estimator,\n",
        "    train_spec=tf.estimator.TrainSpec(\n",
        "        input_fn=lambda: read_gcs(batch_size=256).take(2000)\n",
        "    ),\n",
        "    eval_spec=tf.estimator.EvalSpec(input_fn=lambda: read_gcs().skip(500).take(100)),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "457484fce156"
      },
      "outputs": [],
      "source": [
        "!rm -r -f $LOCAL_DIRECTORY_FULL/estimator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7a3f512b9f3e"
      },
      "outputs": [],
      "source": [
        "tf.compat.v1.disable_eager_execution()  # You'll have to restart Runtime after running this\n",
        "spec_dict = {}\n",
        "for column_name in COLUMN_NAMES:\n",
        "    if column_name.startswith(\"int\"):\n",
        "        spec_dict[column_name] = tf.compat.v1.placeholder(\n",
        "            name=column_name, shape=(1,), dtype=tf.int64\n",
        "        )\n",
        "    if column_name.startswith(\"cat\"):\n",
        "        spec_dict[column_name] = tf.compat.v1.placeholder(\n",
        "            name=column_name, shape=(), dtype=tf.string\n",
        "        )\n",
        "\n",
        "serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(spec_dict)\n",
        "estimator_base_path = os.path.join(LOCAL_DIRECTORY_FULL, \"estimator\")\n",
        "estimator_path = estimator.export_saved_model(estimator_base_path, serving_input_fn)\n",
        "estimator_path = estimator_path.decode(\"ascii\")\n",
        "estimator_path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e796d225e740"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli show --dir $estimator_path --all"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd4c3d9bf3aa"
      },
      "source": [
        "## Generate prediction requests"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "953eb5e1e86e"
      },
      "source": [
        "Now we can generate requests to send to our model for inference.\n",
        "Requests are generated in the JSON Lines format, one request per line."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e07ab9e6213b"
      },
      "outputs": [],
      "source": [
        "!mkdir -p $LOCAL_DIRECTORY_FULL/requests"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5313e9d6b90"
      },
      "outputs": [],
      "source": [
        "def wrap_value(value, wrap_value):\n",
        "    if wrap_value:\n",
        "        return [value]\n",
        "    else:\n",
        "        return value\n",
        "\n",
        "\n",
        "def row_to_dict(row, wrap_values):\n",
        "    d = {}\n",
        "    for key, value in row[0].items():\n",
        "        if \"int\" in key:\n",
        "            d[key] = [wrap_value(v, wrap_values) for v in value.numpy().tolist()]\n",
        "        if \"cat\" in key:\n",
        "            d[key] = [\n",
        "                wrap_value(v.decode(), wrap_values) for v in value.numpy().tolist()\n",
        "            ]\n",
        "    return d\n",
        "\n",
        "\n",
        "def export_requests_jsonl(file_name, rows=100, batch_size=64, wrap_values=True):\n",
        "    with tf.io.gfile.GFile(file_name, mode=\"w\") as f:\n",
        "        for row in read_gcs(batch_size):\n",
        "            d = row_to_dict(row, wrap_values)\n",
        "            f.write(json.dumps(d))\n",
        "            f.write(\"\\n\")\n",
        "            rows -= 1\n",
        "            if rows == 0:\n",
        "                break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6b400bd17be6"
      },
      "outputs": [],
      "source": [
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_1_1.jsonl\"),\n",
        "    rows=1,\n",
        "    batch_size=1,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_1_512.jsonl\"),\n",
        "    rows=1,\n",
        "    batch_size=512,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_10_1.jsonl\"),\n",
        "    rows=10,\n",
        "    batch_size=1,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_10_512.jsonl\"),\n",
        "    rows=10,\n",
        "    batch_size=512,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_10_1024.jsonl\"),\n",
        "    rows=10,\n",
        "    batch_size=1024,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_100_1.jsonl\"),\n",
        "    rows=100,\n",
        "    batch_size=1,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_100_512.jsonl\"),\n",
        "    rows=100,\n",
        "    batch_size=512,\n",
        ")\n",
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_100_1024.jsonl\"),\n",
        "    rows=100,\n",
        "    batch_size=1024,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9523fd4d23c1"
      },
      "source": [
        "If you want to export requests for the Estimator model, you must set `wrap_values` to `False`. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "319505cb75db"
      },
      "outputs": [],
      "source": [
        "export_requests_jsonl(\n",
        "    os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_estimator_10_1.jsonl\"),\n",
        "    rows=10,\n",
        "    batch_size=1,\n",
        "    wrap_values=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c17522fcfd37"
      },
      "source": [
        "## (Optional) Generate warmup requests"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9c75a550809d"
      },
      "source": [
        "The TensorFlow runtime has components that are lazily initialized. Lazy initialization might result in high latency for the first requests that are sent to a model after it's loaded. This latency can be several orders of magnitude higher than that of a single inference request.\n",
        "\n",
        "For more information about SavedModel warmup, see https://www.tensorflow.org/tfx/serving/saved_model_warmup.\n",
        "\n",
        "For Vertex AI Prediction using the optimized TensorFlow runtime, when the model is precompiled the first request for each new batch size has higher latency. Precompilation is enabled when the `allow_precompilation` flag is set to true.\n",
        "\n",
        "To mitigate high latency, provide a warmup request for the runtime to load when it starts.\n",
        "The warmup file should include the various batch sizes you expect your model to receive in production.\n",
        "\n",
        "Note that providing a warmup request with multiple batch sizes increase the time for each node to start.\n",
        "\n",
        "If you expect the model to receive multiple batch sizes, you can use automatic server-side request batching with a set of `allowed_batch_sizes`. For more information, see https://www.tensorflow.org/tfx/serving/serving_config#batching_configurationß.\n",
        "\n",
        "To enable auto-batching for a model running on Vertex AI Prediction, put your batching configuration into the [config/batching_parameters_config](https://cloud.google.com/vertex-ai/docs/training/exporting-model-artifacts#enable_server-side_request_batching_for_tensorflow) file in the same GCS directory as saved_model.pb."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2238dc7c4d93"
      },
      "outputs": [],
      "source": [
        "!mkdir -p $LOCAL_DIRECTORY_FULL/keras/assets.extra"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a45ace218017"
      },
      "outputs": [],
      "source": [
        "def build_grpc_request(\n",
        "    row_dict, model_name=\"default\", signature_name=\"serving_default\"\n",
        "):\n",
        "    \"\"\"Generate gRPC inference request with payload.\"\"\"\n",
        "\n",
        "    request = predict_pb2.PredictRequest()\n",
        "    request.model_spec.name = model_name\n",
        "    request.model_spec.signature_name = signature_name\n",
        "    for key, value in row_dict.items():\n",
        "        proto = None\n",
        "        if \"cat\" in key:\n",
        "            proto = tf.make_tensor_proto(value, dtype=tf.string)\n",
        "        else:\n",
        "            proto = tf.make_tensor_proto(value, dtype=tf.int64)\n",
        "        request.inputs[key].CopyFrom(proto)\n",
        "    return request\n",
        "\n",
        "\n",
        "def export_warmup_file(\n",
        "    request_files, export_path, model_name=\"default\", signature_name=\"serving_default\"\n",
        "):\n",
        "    with tf.io.TFRecordWriter(export_path) as writer:\n",
        "        for request_file_path in request_files:\n",
        "            with open(request_file_path) as f:\n",
        "                row_dict = json.loads(f.readline())\n",
        "                request = build_grpc_request(row_dict, model_name, signature_name)\n",
        "            log = prediction_log_pb2.PredictionLog(\n",
        "                predict_log=prediction_log_pb2.PredictLog(request=request)\n",
        "            )\n",
        "            writer.write(log.SerializeToString())\n",
        "\n",
        "\n",
        "export_warmup_file(\n",
        "    [\n",
        "        os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_1_1.jsonl\"),\n",
        "        os.path.join(LOCAL_DIRECTORY_FULL, \"requests\", \"requests_1_512.jsonl\"),\n",
        "    ],\n",
        "    os.path.join(\n",
        "        LOCAL_DIRECTORY_FULL, \"keras\", \"assets.extra\", \"tf_serving_warmup_requests\"\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0054f6d6206a"
      },
      "source": [
        "## Deploy model to Vertex AI Prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af77cccd3307"
      },
      "source": [
        "To deploy a model to Vertex AI Prediction service, you must put it in a GCS bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DUJePRU3LDcR"
      },
      "outputs": [],
      "source": [
        "!gsutil rm -r $BUCKET_URI/*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "31d89fa9f121"
      },
      "outputs": [],
      "source": [
        "!gsutil cp -r $LOCAL_DIRECTORY_FULL/keras/* $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3f3eb65dbb82"
      },
      "source": [
        "Import the Vertex AI Python client library into your notebook environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "85c1b7a3f0f5"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform import gapic as aip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cc88fefba44a"
      },
      "source": [
        "Define the node type to use for deployments. To learn about Vertex AI Prediction options, see [configure compute resources](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "76432d97918e"
      },
      "outputs": [],
      "source": [
        "DEPLOY_COMPUTE = \"n1-standard-16\"\n",
        "DEPLOY_GPU = aip.AcceleratorType.NVIDIA_TESLA_T4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "285371d048e1"
      },
      "source": [
        "The AI Platform Python client library works as a client/server model. \n",
        "\n",
        "The following clients are used in this sample:\n",
        "- Model Service for managing models.\n",
        "- Endpoint Service for deployment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "549bd6acf87c"
      },
      "outputs": [],
      "source": [
        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\"\n",
        "PARENT = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
        "\n",
        "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
        "model_service_client = aip.ModelServiceClient(client_options=client_options)\n",
        "endpoint_service_client = aip.EndpointServiceClient(client_options=client_options)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e3e4d12f7374"
      },
      "source": [
        "### Setup private endpoint for online prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9d259d62f8a4"
      },
      "source": [
        "The throughput and latency of the Criteo model you trained is sensitive to network performance."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "60667c66da2f"
      },
      "source": [
        "Notice that a single request with a batch size of 512 takes ~200Kb."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d54a9f02769"
      },
      "outputs": [],
      "source": [
        "!ls -alh $LOCAL_DIRECTORY_FULL/requests/requests_1_512.jsonl"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "64268dc47d5c"
      },
      "source": [
        "For the best performance, use a Vertex AI Prediction private endpoint.\n",
        "\n",
        "To use a private endpoint, setup a VPC peering network between your project and the Vertex AI Prediction service project that is hosting VMs running your model. This eliminates additional hops in network traffic and allows using efficient gRPC protocol.\n",
        "\n",
        "For more information about private endpoints, see https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf4685ec69fd"
      },
      "source": [
        "For more information about VPC peering in Vertex AI, see https://cloud.google.com/vertex-ai/docs/general/vpc-peering.\n",
        "\n",
        "**IMPORTANT: you can only setup one VPC peering to servicenetworking.googleapis.com per VPC network.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "743b3af8f747"
      },
      "source": [
        "For simplicity, you setup VPC peering to the default network. You can create a different network for your project.\n",
        "\n",
        "If you setup VPC peering with any other network, make sure that the network already exists and that your VM is running on that network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bab19db69297"
      },
      "outputs": [],
      "source": [
        "# This is for display only; you can name the range anything.\n",
        "PEERING_RANGE_NAME = \"vertex-ai-prediction-peering-range\"\n",
        "NETWORK = \"default\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "43f7a7c2eed6"
      },
      "outputs": [],
      "source": [
        "# NOTE: `prefix-length=16` means a CIDR block with mask /16 will be\n",
        "# reserved for use by Google services, such as Vertex AI.\n",
        "!gcloud compute addresses create $PEERING_RANGE_NAME \\\n",
        "  --global \\\n",
        "  --prefix-length=16 \\\n",
        "  --description=\"peering range for Google service\" \\\n",
        "  --network=$NETWORK \\\n",
        "  --purpose=VPC_PEERING"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "824f75d1dc95"
      },
      "source": [
        "Create the VPC connection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "338f4e58bb6b"
      },
      "outputs": [],
      "source": [
        "!gcloud services vpc-peerings connect \\\n",
        "  --service=servicenetworking.googleapis.com \\\n",
        "  --network=$NETWORK \\\n",
        "  --ranges=$PEERING_RANGE_NAME \\\n",
        "  --project=$PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2be23c201b8"
      },
      "source": [
        "If you receive a permission error when running this command, then try running it with your user account.\n",
        "\n",
        "To run this command with your user account, do the following:\n",
        "- add `echo` before the command in the cell above (`echo gcloud services vpc-peering ...`).\n",
        "- run the cell and copy its output\n",
        "- start new terminal window, and run `gcloud auth login` to authenticate using your user account.\n",
        "- paste and run the copied command in the terminal."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de727b998161"
      },
      "source": [
        "Check the status of your peering connections."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "38b48cc436cb"
      },
      "outputs": [],
      "source": [
        "!gcloud compute networks peerings list --network $NETWORK"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c9b09bb49bc1"
      },
      "source": [
        "### Upload model to Vertex AI Prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lYFkS2H9dDfd"
      },
      "source": [
        "Learn more about [model_service.upload_model](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.services.model_service.ModelServiceClient#google_cloud_aiplatform_v1_services_model_service_ModelServiceClient_upload_model).\n",
        "\n",
        "\n",
        "`artifact_uri` argument should point to a GCS path where `saved_model.pb` file is located for your model.\n",
        "\n",
        "`image_uri` specifies which docker image to use. Here we upload the same model using TF2.7 GPU and Vertex AI Prediction optimized TensorFlow runtime images.\n",
        "\n",
        "In order to be able to send requests to your models over gRPC, you need to set `model_name` argument and update `predict_route` and `health_route` accordingly.\n",
        "\n",
        "Please note that gRPC support in Vertex AI Prediction is still experimental."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "91038811953e"
      },
      "outputs": [],
      "source": [
        "tf27_cpu_model_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle TF2.7 CPU model\",\n",
        "    \"artifact_uri\": BUCKET_URI,\n",
        "    \"container_spec\": {\n",
        "        \"image_uri\": \"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-7:latest\",\n",
        "        \"args\": [\n",
        "            \"--port=8500\",\n",
        "            \"--rest_api_port=8080\",\n",
        "            \"--model_name=default\",\n",
        "            \"--model_base_path=$(AIP_STORAGE_URI)\",\n",
        "        ],\n",
        "        \"ports\": [{\"container_port\": 8080}],\n",
        "        \"predict_route\": \"/v1/models/default:predict\",\n",
        "        \"health_route\": \"/v1/models/default\",\n",
        "    },\n",
        "}\n",
        "tf27_cpu_model = (\n",
        "    model_service_client.upload_model(parent=PARENT, model=tf27_cpu_model_dict)\n",
        "    .result(timeout=180)\n",
        "    .model\n",
        ")\n",
        "tf27_cpu_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5439eae4099"
      },
      "outputs": [],
      "source": [
        "tf27_gpu_model_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle TF2.7 GPU model\",\n",
        "    \"artifact_uri\": BUCKET_URI,\n",
        "    \"container_spec\": {\n",
        "        \"image_uri\": \"us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-7:latest\",\n",
        "        \"args\": [\n",
        "            \"--port=8500\",\n",
        "            \"--rest_api_port=8080\",\n",
        "            \"--model_name=default\",\n",
        "            \"--model_base_path=$(AIP_STORAGE_URI)\",\n",
        "        ],\n",
        "        \"ports\": [{\"container_port\": 8080}],\n",
        "        \"predict_route\": \"/v1/models/default:predict\",\n",
        "        \"health_route\": \"/v1/models/default\",\n",
        "    },\n",
        "}\n",
        "tf27_gpu_model = (\n",
        "    model_service_client.upload_model(parent=PARENT, model=tf27_gpu_model_dict)\n",
        "    .result(timeout=180)\n",
        "    .model\n",
        ")\n",
        "tf27_gpu_model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b12af8ec4912"
      },
      "source": [
        "For deploying a model using Vertex AI Prediction optimized TensorFlow runtime, use the `us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest` container.\n",
        "\n",
        "Two optimization options are applied to the model.\n",
        "- *allow_precompilation* - turns on model pre-compilation for better performance. Note that model precompilation happens when the first request with the new batch size arrives, and the response for that request is sent after precompilation is complete. To mitigate this, specify a warmup file (see the section earlier in this colab). Model precompilation works for different kinds of models, and in most cases has a positive effect on performance. However, we recommend that you try it out for your model before you enable it in production.\n",
        "- *allow_precision_affecting_optimizations* - enables precision affecting optimizations. In some cases this makes the model run significantly faster at the cost of very minimal loss to model prediction power. You should assess the precision impact to your model when using this optimization.\n",
        "\n",
        "For the list of available optimized TensorFlow runtimer containers and options, see https://cloud.google.com/vertex-ai/docs/predictions/optimized-tensorflow-runtime."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e8a2837ac245"
      },
      "outputs": [],
      "source": [
        "tf_opt_gpu_model_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle optimized TensorFlow runtime GPU model\",\n",
        "    \"artifact_uri\": BUCKET_URI,\n",
        "    \"container_spec\": {\n",
        "        \"image_uri\": \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest\",\n",
        "        \"args\": [\n",
        "            \"--model_name=default\",\n",
        "            \"--allow_precompilation=true\",\n",
        "            \"--allow_precision_affecting_optimizations=false\",\n",
        "        ],\n",
        "        \"predict_route\": \"/v1/models/default:predict\",\n",
        "        \"health_route\": \"/v1/models/default\",\n",
        "    },\n",
        "}\n",
        "\n",
        "tf_opt_gpu_model = (\n",
        "    model_service_client.upload_model(parent=PARENT, model=tf_opt_gpu_model_dict)\n",
        "    .result(timeout=180)\n",
        "    .model\n",
        ")\n",
        "tf_opt_gpu_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a3f21fd3203e"
      },
      "outputs": [],
      "source": [
        "tf_opt_lossy_gpu_model_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle optimized TensorFlow runtime GPU model with lossy optimizations\",\n",
        "    \"artifact_uri\": BUCKET_URI,\n",
        "    \"container_spec\": {\n",
        "        \"image_uri\": \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest\",\n",
        "        \"args\": [\n",
        "            \"--model_name=default\",\n",
        "            \"--allow_precompilation=true\",\n",
        "            \"--allow_precision_affecting_optimizations=true\",\n",
        "        ],\n",
        "        \"predict_route\": \"/v1/models/default:predict\",\n",
        "        \"health_route\": \"/v1/models/default\",\n",
        "    },\n",
        "}\n",
        "\n",
        "tf_opt_lossy_gpu_model = (\n",
        "    model_service_client.upload_model(parent=PARENT, model=tf_opt_lossy_gpu_model_dict)\n",
        "    .result(timeout=180)\n",
        "    .model\n",
        ")\n",
        "tf_opt_lossy_gpu_model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a07f1e5b393d"
      },
      "source": [
        "List all models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e904929d4c38"
      },
      "outputs": [],
      "source": [
        "model_service_client.list_models(parent=PARENT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3f674f8c2d0a"
      },
      "source": [
        "### Create endpoints"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e4279643c2cf"
      },
      "source": [
        "Learn more about [endpoint_service.create_endpoint](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.services.endpoint_service.EndpointServiceClient#google_cloud_aiplatform_v1_services_endpoint_service_EndpointServiceClient_create_endpoint)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4efb545e0439"
      },
      "outputs": [],
      "source": [
        "project_number = re.match(r\"projects/(\\d+)/.+\", tf27_cpu_model)[1]\n",
        "full_network_name = f\"projects/{project_number}/global/networks/{NETWORK}\"\n",
        "full_network_name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b3860a1c4004"
      },
      "outputs": [],
      "source": [
        "tf27_cpu_endpoint_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle TF2.7 CPU private endpoint\",\n",
        "    \"network\": full_network_name,\n",
        "}\n",
        "tf27_cpu_endpoint = (\n",
        "    endpoint_service_client.create_endpoint(\n",
        "        parent=PARENT, endpoint=tf27_cpu_endpoint_dict\n",
        "    )\n",
        "    .result(timeout=300)\n",
        "    .name\n",
        ")\n",
        "tf27_cpu_endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7da2c558e020"
      },
      "outputs": [],
      "source": [
        "tf27_gpu_endpoint_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle TF2.7 GPU private endpoint\",\n",
        "    \"network\": full_network_name,\n",
        "}\n",
        "tf27_gpu_endpoint = (\n",
        "    endpoint_service_client.create_endpoint(\n",
        "        parent=PARENT, endpoint=tf27_gpu_endpoint_dict\n",
        "    )\n",
        "    .result(timeout=300)\n",
        "    .name\n",
        ")\n",
        "tf27_gpu_endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "20f684d17e01"
      },
      "outputs": [],
      "source": [
        "tf_opt_gpu_endpoint_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle optimized TensorFlow runtime GPU private endpoint\",\n",
        "    \"network\": full_network_name,\n",
        "}\n",
        "tf_opt_gpu_endpoint = (\n",
        "    endpoint_service_client.create_endpoint(\n",
        "        parent=PARENT, endpoint=tf_opt_gpu_endpoint_dict\n",
        "    )\n",
        "    .result(timeout=300)\n",
        "    .name\n",
        ")\n",
        "tf_opt_gpu_endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7e666cd56510"
      },
      "outputs": [],
      "source": [
        "tf_opt_lossy_gpu_endpoint_dict = {\n",
        "    \"display_name\": \"Criteo Kaggle optimized TensorFlow runtime GPU with lossy optimizations private endpoint\",\n",
        "    \"network\": full_network_name,\n",
        "}\n",
        "tf_opt_lossy_gpu_endpoint = (\n",
        "    endpoint_service_client.create_endpoint(\n",
        "        parent=PARENT, endpoint=tf_opt_lossy_gpu_endpoint_dict\n",
        "    )\n",
        "    .result(timeout=300)\n",
        "    .name\n",
        ")\n",
        "tf_opt_lossy_gpu_endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "56ade2bea05a"
      },
      "source": [
        "### Deploy models to endpoints"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f8586711566b"
      },
      "source": [
        "Learn more about [enpoint_service.deploy_model](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.services.endpoint_service.EndpointServiceClient#google_cloud_aiplatform_v1_services_endpoint_service_EndpointServiceClient_deploy_model)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ac48f638e2e"
      },
      "outputs": [],
      "source": [
        "tf27_cpu_deployed_model_dict = {\n",
        "    \"model\": tf27_cpu_model,\n",
        "    \"display_name\": \"Criteo Kaggle TF2.7 CPU deployed model\",\n",
        "    \"dedicated_resources\": {\n",
        "        \"min_replica_count\": 1,\n",
        "        \"max_replica_count\": 1,\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": DEPLOY_COMPUTE,\n",
        "            \"accelerator_count\": 0,\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "tf27_cpu_deployed_model = endpoint_service_client.deploy_model(\n",
        "    endpoint=tf27_cpu_endpoint, deployed_model=tf27_cpu_deployed_model_dict\n",
        ").result()\n",
        "tf27_cpu_deployed_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "104eb9495a6b"
      },
      "outputs": [],
      "source": [
        "tf27_gpu_deployed_model_dict = {\n",
        "    \"model\": tf27_gpu_model,\n",
        "    \"display_name\": \"Criteo Kaggle TF2.7 GPU deployed model\",\n",
        "    \"dedicated_resources\": {\n",
        "        \"min_replica_count\": 1,\n",
        "        \"max_replica_count\": 1,\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": DEPLOY_COMPUTE,\n",
        "            \"accelerator_type\": DEPLOY_GPU,\n",
        "            \"accelerator_count\": 1,\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "tf27_gpu_deployed_model = endpoint_service_client.deploy_model(\n",
        "    endpoint=tf27_gpu_endpoint, deployed_model=tf27_gpu_deployed_model_dict\n",
        ").result()\n",
        "tf27_gpu_deployed_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9d5d7dc432e8"
      },
      "outputs": [],
      "source": [
        "tf_opt_gpu_deployed_model_dict = {\n",
        "    \"model\": tf_opt_gpu_model,\n",
        "    \"display_name\": \"Criteo Kaggle optimized TensorFlow runtime GPU model\",\n",
        "    \"dedicated_resources\": {\n",
        "        \"min_replica_count\": 1,\n",
        "        \"max_replica_count\": 1,\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": DEPLOY_COMPUTE,\n",
        "            \"accelerator_type\": DEPLOY_GPU,\n",
        "            \"accelerator_count\": 1,\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "tf_opt_gpu_deployed_model = endpoint_service_client.deploy_model(\n",
        "    endpoint=tf_opt_gpu_endpoint, deployed_model=tf_opt_gpu_deployed_model_dict\n",
        ").result()\n",
        "tf_opt_gpu_deployed_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ca51ac94a1c7"
      },
      "outputs": [],
      "source": [
        "tf_opt_lossy_gpu_deployed_model_dict = {\n",
        "    \"model\": tf_opt_lossy_gpu_model,\n",
        "    \"display_name\": \"Criteo Kaggle optimized TensorFlow runtime GPU model with lossy optimizations\",\n",
        "    \"dedicated_resources\": {\n",
        "        \"min_replica_count\": 1,\n",
        "        \"max_replica_count\": 1,\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": DEPLOY_COMPUTE,\n",
        "            \"accelerator_type\": DEPLOY_GPU,\n",
        "            \"accelerator_count\": 1,\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "tf_opt_lossy_gpu_deployed_model = endpoint_service_client.deploy_model(\n",
        "    endpoint=tf_opt_lossy_gpu_endpoint,\n",
        "    deployed_model=tf_opt_lossy_gpu_deployed_model_dict,\n",
        ").result()\n",
        "tf_opt_lossy_gpu_deployed_model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c91265a78a97"
      },
      "source": [
        "## (optional) Compare performance of deployed models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c264b3b6eb3f"
      },
      "source": [
        "To access private endpoints, the VM used to send requests must be deployed in the same network where you setup VPC peering. Because of this, you can't send requests to your models that are deployed with private endpoints from Colab.\n",
        "\n",
        "To get the best performance, be sure the VM is in the same region as your model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04ea66e6c6e2"
      },
      "source": [
        "Import helper functions for benchmarking models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81f72b64c8e5"
      },
      "outputs": [],
      "source": [
        "!curl https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/vertex_endpoints/optimized_tensorflow_runtime/benchmark.py -o benchmark.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1cd21d1afb38"
      },
      "outputs": [],
      "source": [
        "from benchmark import benchmark"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "81380aabcda7"
      },
      "source": [
        "This code sends a specified number of requests asynchronously and uniformly at a given QPS, then records the observed latency. Next, the latency results are aggregated and percentiles are calculated.\n",
        "The `actual_qps` that the model can handle is calculated as the time it takes for a model to process the sent requests divided by the number of requests. \n",
        "By providing different implementations for `send_request` and `build_request` functions, the same code can be used for benchmarking models running locally or on Vertex AI Prediction using gRPC and REST protocols.\n",
        "\n",
        "The main goal of this benchmark is to measure model latency on different loads, and maximum throughput the model can handle. In order to find maximum throughput, gradually increase QPS until `actual_qps` stops increasing and latency increases dramatically.\n",
        "\n",
        "On the production deployment, the workload is not uniform, and therefore the maximum model throughput is likely to be lower.\n",
        "We are not trying to simulate production workload here. This benchmark is meant to compare latency and throughput for same model running on different environments."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "664ec5c32fdb"
      },
      "source": [
        "Details about deployed models can be accessed using the [endpoint_service_client.get_endpoint](https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform_v1.services.endpoint_service.EndpointServiceClient#google_cloud_aiplatform_v1_services_endpoint_service_EndpointServiceClient_get_endpoint) API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6f238380d9a7"
      },
      "outputs": [],
      "source": [
        "tf_opt_gpu_endpoint_dict = endpoint_service_client.get_endpoint(\n",
        "    name=tf_opt_gpu_endpoint\n",
        ")\n",
        "tf_opt_gpu_endpoint_dict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0fa9f6cb522c"
      },
      "source": [
        "First, verify that you can access your models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5394c8f5636"
      },
      "outputs": [],
      "source": [
        "health_url = tf_opt_gpu_endpoint_dict.deployed_models[\n",
        "    0\n",
        "].private_endpoints.health_http_uri\n",
        "health_url"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a377f25e675c"
      },
      "outputs": [],
      "source": [
        "!curl $health_url"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d0ee794df447"
      },
      "source": [
        "Define helper methods to run benchmarks against private endpoints using REST protocol.\n",
        "The URI where requests should be sent can be found in `deployed_model.private_endpoints.predict_http_uri`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89eb67cef6df"
      },
      "outputs": [],
      "source": [
        "def build_rest_request(\n",
        "    row_dict, model_name=\"default\", signature_name=\"serving_default\"\n",
        "):\n",
        "    payload = json.dumps({\"signature_name\": signature_name, \"inputs\": row_dict})\n",
        "    return payload\n",
        "\n",
        "\n",
        "def benchmark_rest_private_endpoint(\n",
        "    endpoint_name, qps_list, model_name=None, duration_seconds=5\n",
        "):\n",
        "    endpoint_dict = endpoint_service_client.get_endpoint(name=endpoint_name)\n",
        "    predict_uri = endpoint_dict.deployed_models[0].private_endpoints.predict_http_uri\n",
        "\n",
        "    def send_rest_request(request):\n",
        "        res = r.post(predict_uri, data=request)\n",
        "        assert res.status_code == 200\n",
        "        return res\n",
        "\n",
        "    return benchmark(\n",
        "        send_rest_request,\n",
        "        build_rest_request,\n",
        "        f\"{LOCAL_DIRECTORY_FULL}/requests/requests_100_512.jsonl\",\n",
        "        qps_list,\n",
        "        duration_seconds,\n",
        "        model_name=model_name,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tdsHDlq-bYuG"
      },
      "source": [
        "You can also benchmark your models deployed on private endpoint using gRPC protocol. \n",
        "\n",
        "gRPC address is same as host name of `predict_http_uri` or `predict_http_uri`.\n",
        "gRPC destination has a format of `<endpoint_id>-<deployed_model_id>` and passed as a \"grpc_destination\" header.\n",
        "\n",
        "Please note that gRPC support in Vertex AI Prediction is still experimental."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YSOEYXR1b0k2"
      },
      "outputs": [],
      "source": [
        "def parse_endpoint_dict(endpoint_dict):\n",
        "    endpoint_id = re.match(r\".+/endpoints/(\\d+)\", endpoint_dict.name)[1]\n",
        "    deployed_model_id = endpoint_dict.deployed_models[0].id\n",
        "    grpc_destination = f\"{endpoint_id}-{deployed_model_id}\"\n",
        "    predict_uri = urlparse(\n",
        "        endpoint_dict.deployed_models[0].private_endpoints.predict_http_uri\n",
        "    )\n",
        "    grpc_uri = f\"{predict_uri.netloc}:8500\"\n",
        "    return (grpc_uri, grpc_destination)\n",
        "\n",
        "\n",
        "def benchmark_grpc_private_endpoint(endpoint_name, qps_list, duration_seconds=5):\n",
        "    endpoint_dict = endpoint_service_client.get_endpoint(name=endpoint_name)\n",
        "    grpc_uri, grpc_destinaion = parse_endpoint_dict(endpoint_dict)\n",
        "\n",
        "    grpc_metadata = []\n",
        "    grpc_metadata.append((\"grpc-destination\", grpc_destinaion))\n",
        "    grpc_channel = grpc.insecure_channel(grpc_uri)\n",
        "    grpc_stub = prediction_service_pb2_grpc.PredictionServiceStub(grpc_channel)\n",
        "\n",
        "    def send_grpc_request(request):\n",
        "        return grpc_stub.Predict(request, 60, metadata=grpc_metadata)\n",
        "\n",
        "    return benchmark(\n",
        "        send_grpc_request,\n",
        "        build_grpc_request,\n",
        "        f\"{LOCAL_DIRECTORY_FULL}/requests/requests_100_512.jsonl\",\n",
        "        qps_list,\n",
        "        duration_seconds,\n",
        "        model_name=\"default\",\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5dcf9f454988"
      },
      "source": [
        "Now we can run a benchmark test for each endpoint and compare the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8a2fc956ff69"
      },
      "outputs": [],
      "source": [
        "tf27_cpu_results = benchmark_grpc_private_endpoint(\n",
        "    tf27_cpu_endpoint, [10, 20, 30, 40, 50, 55]\n",
        ")\n",
        "tf27_cpu_results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cda3c1184b5c"
      },
      "outputs": [],
      "source": [
        "tf27_gpu_results = benchmark_grpc_private_endpoint(\n",
        "    tf27_gpu_endpoint, [10, 20, 30, 40, 50, 60, 70, 75]\n",
        ")\n",
        "tf27_gpu_results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cb7bf5263c5a"
      },
      "outputs": [],
      "source": [
        "tf_opt_gpu_results = benchmark_grpc_private_endpoint(\n",
        "    tf_opt_gpu_endpoint, [10, 50, 100, 150, 200, 250, 275, 300, 325, 350]\n",
        ")\n",
        "tf_opt_gpu_results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f96bb0e8cb0c"
      },
      "outputs": [],
      "source": [
        "tf_opt_lossy_gpu_results = benchmark_grpc_private_endpoint(\n",
        "    tf_opt_lossy_gpu_endpoint, [10, 50, 100, 200, 300, 400, 500, 600, 700, 800]\n",
        ")\n",
        "tf_opt_lossy_gpu_results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "41548dbecd5b"
      },
      "source": [
        "Combine and visualize results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "92405bdd7163"
      },
      "outputs": [],
      "source": [
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "def build_graph(x_key, y_key, results_dict, axis):\n",
        "    matplotlib.rcParams[\"figure.figsize\"] = [10.0, 7.0]\n",
        "\n",
        "    fig, ax = plt.subplots(facecolor=(1, 1, 1))\n",
        "    ax.set_xlabel(\"QPS\")\n",
        "    ax.set_ylabel(\"Latency(ms)\")\n",
        "    for title, results in results_dict.items():\n",
        "        x = np.array(results[x_key])\n",
        "        y = np.array(results[y_key])\n",
        "        ax.plot(x, y, label=title)\n",
        "    ax.legend()\n",
        "    ax.axis(axis)\n",
        "    ax.set_title(f\"Criteo model {y_key} latency, batch size 512\")\n",
        "    return fig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "df3b8ac16b0c"
      },
      "outputs": [],
      "source": [
        "fig = build_graph(\n",
        "    \"actual_qps\",\n",
        "    \"p50\",\n",
        "    {\n",
        "        \"TF2.7 CPU\": tf27_cpu_results,\n",
        "        \"TF2.7 GPU\": tf27_gpu_results,\n",
        "        \"TF opt GPU\": tf_opt_gpu_results,\n",
        "        \"TF opt GPU lossy\": tf_opt_lossy_gpu_results,\n",
        "    },\n",
        "    (0, 800, 0, 60),\n",
        ")\n",
        "fig.savefig(\"criteo_p50_latency_512.png\", bbox_inches=\"tight\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8593b878feab"
      },
      "outputs": [],
      "source": [
        "fig = build_graph(\n",
        "    \"actual_qps\",\n",
        "    \"p99\",\n",
        "    {\n",
        "        \"TF2.7 CPU\": tf27_cpu_results,\n",
        "        \"TF2.7 GPU\": tf27_gpu_results,\n",
        "        \"TF opt GPU\": tf_opt_gpu_results,\n",
        "        \"TF opt GPU lossy\": tf_opt_lossy_gpu_results,\n",
        "    },\n",
        "    (0, 800, 0, 100),\n",
        ")\n",
        "fig.savefig(\"criteo_p99_latency_512.png\", bbox_inches=\"tight\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aac5ebde1803"
      },
      "source": [
        "You can see that the Vertex AI Prediction optimized TensorFlow runtime has signficantly higher throughput and lower latency compared to TensorFlow 2.7."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3f369366c0af"
      },
      "source": [
        "## (Optional) Compare performance of deployed models using MLPerf Inference loadgen"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bfd210c97934"
      },
      "source": [
        "MLPerf Inference is a benchmark suite for measuring how fast systems can run models in a variety of deployment scenarios. MLPerf is now an industry standard way of measuring model performance. You can follow instructions at https://github.com/tensorflow/tpu/tree/master/models/experimental/inference/load_test to run MLPerf Inferenence benchmark for deployed models.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8fdbaecb0fdc"
      },
      "source": [
        "## (Optional) Compare prediction results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8a16b45477f1"
      },
      "source": [
        "In this sample the Vertex Prediction optimized TensorFlow runtime is used with the `allow_precision_affecting_optimizations` flag set to `true` to gain additional speedup. Now let's check how those optimizations effect prediction results.\n",
        "\n",
        "We compare the results of predictions for 51,200 requests for a model running on the optimized TensorFlow runtime with lossy optimizations and on TF2.7"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "59ffec627b59"
      },
      "outputs": [],
      "source": [
        "def get_predictions(endpoint, requests_file_path):\n",
        "    responses = []\n",
        "\n",
        "    endpoint_dict = endpoint_service_client.get_endpoint(name=endpoint)\n",
        "    pridict_uri = endpoint_dict.deployed_models[0].private_endpoints.predict_http_uri\n",
        "\n",
        "    with tf.io.gfile.GFile(requests_file_path, \"r\") as f:\n",
        "        for line in f:\n",
        "            row_dict = json.loads(line)\n",
        "            request = build_rest_request(row_dict)\n",
        "            response = r.post(pridict_uri, data=request)\n",
        "            for output in json.loads(response.text)[\"outputs\"]:\n",
        "                responses.append(output[0])\n",
        "\n",
        "    return np.array(responses)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "01fe427a6211"
      },
      "outputs": [],
      "source": [
        "tf27_gpu_predictions = get_predictions(\n",
        "    tf27_gpu_endpoint, f\"{LOCAL_DIRECTORY_FULL}/requests/requests_100_512.jsonl\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "94686b1b90a9"
      },
      "outputs": [],
      "source": [
        "tf_opt_lossy_gpu_predictions = get_predictions(\n",
        "    tf_opt_lossy_gpu_endpoint, f\"{LOCAL_DIRECTORY_FULL}/requests/requests_100_512.jsonl\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c3577df0bf00"
      },
      "outputs": [],
      "source": [
        "np.average(tf_opt_lossy_gpu_predictions - tf27_gpu_predictions) * 100"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cbdc0444957e"
      },
      "outputs": [],
      "source": [
        "np.max(np.abs(tf_opt_lossy_gpu_predictions - tf27_gpu_predictions)) * 100"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "29235c879dc5"
      },
      "source": [
        "You can see that the average results are different for less than 0.0016%. In the worst case the difference is 0.05%."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1807d423c612"
      },
      "source": [
        "## Cleanup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0b3132d3f788"
      },
      "source": [
        "After you are done, it's safe to remove the endpoints you created and the model you deployed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "64d8ee827276"
      },
      "outputs": [],
      "source": [
        "def cleanup(endpoint, model_name, deployed_model_id):\n",
        "    response = endpoint_service_client.undeploy_model(\n",
        "        endpoint=endpoint, deployed_model_id=deployed_model_id\n",
        "    )\n",
        "    print(\"running undeploy_model operation:\", response.operation.name)\n",
        "    print(response.result())\n",
        "\n",
        "    response = endpoint_service_client.delete_endpoint(name=endpoint)\n",
        "    print(\"running delete_endpoint operation:\", response.operation.name)\n",
        "    print(response.result())\n",
        "\n",
        "    response = model_service_client.delete_model(name=model_name)\n",
        "    print(\"running delete_model operation:\", response.operation.name)\n",
        "    print(response.result())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "afd94ff1ed76"
      },
      "outputs": [],
      "source": [
        "cleanup(tf27_cpu_endpoint, tf27_cpu_model, tf27_cpu_deployed_model)\n",
        "cleanup(tf27_gpu_endpoint, tf27_gpu_model, tf27_gpu_deployed_model)\n",
        "cleanup(tf_opt_gpu_endpoint, tf_opt_gpu_model, tf_opt_gpu_deployed_model)\n",
        "cleanup(\n",
        "    tf_opt_lossy_gpu_endpoint, tf_opt_lossy_gpu_model, tf_opt_lossy_gpu_deployed_model\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c37975fdc894"
      },
      "source": [
        "You can now also remove model from GCS bucket as well."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29d1d5458d73"
      },
      "outputs": [],
      "source": [
        "# Set this to true only if you'd like to delete your bucket\n",
        "delete_bucket = False\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    !gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tabular_optimized_online_prediction.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}