notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "title:generic,gcp"
      },
      "source": [
        "# Get started with Endpoint and shared VM\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
        "        <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "        </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fgithub.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fblob%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fget_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>   \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers: get started with Endpoints and shared VM for co-hosting models.\n",
        "\n",
        "Learn more about [Shared resources across deployments](https://cloud.google.com/vertex-ai/docs/predictions/model-co-hosting)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage2,get_started_automl_training"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use deployment resource pools for deploying models. A deployment resouce pool provides one with the ability to co-host more than one model on the same (shared) VM.\n",
        "\n",
        "A deployment resource pool groups together model deployments to share resources within a VM. Multiple endpoints can be deployed on the same VM within a Deployment Resource Pool. Each of these endpoints can have one or more deployed models. The deployed models for a given endpoint can be grouped under the same or different deployment resource pools.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI Training`\n",
        "- `Vertex AI Model` resource\n",
        "- `Vertex AI Endpoint` resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Upload a pre-trained image classification model as a `Model` resource (model A).\n",
        "- Upload a pre-trained text sentence encoder model as a `Model` resource (model B).\n",
        "- Create a shared VM deployment resource pool.\n",
        "- List shared VM deployment resource pools.\n",
        "- Create two `Endpoint` resources.\n",
        "- Deploy first model (model A) to first `Endpoint` resource using deployment resource pool.\n",
        "- Deploy second model (model B) to second `Endpoint` resource using deployment resource pool.\n",
        "- Make a prediction request with first deployed model (model A).\n",
        "- Make a prediction request with second deployed model (model B)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:flowers,icn"
      },
      "source": [
        "### Model\n",
        "\n",
        "The pre-trained models used for this tutorial are from the [TensorFlow Hub](https://tfhub.dev/) repository:\n",
        "\n",
        "- [image classification](https://tfhub.dev/google/imagenet/inception_v3/classification/5): trained with ImageNet.\n",
        "- [text sentence encoder](https://tfhub.dev/google/universal-sentence-encoder/4): Google's universal sentence encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fb3451ce8e47"
      },
      "source": [
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "- Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_mlops"
      },
      "source": [
        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "deda65b8cc6e"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3357af6a0f34"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 tensorflow==2.15.1 \\\n",
        "                                 tensorflow-hub==0.16.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff555b32bab8"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a2b7b59bbf7"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92e68cfc3a90"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46604f70e831"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4f872cd812d0"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_project_id"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}\n",
        "\n",
        "LOCATION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:mbsdk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_bucket"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $LOCATION $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import google.cloud.aiplatform as aiplatform\n",
        "import google.cloud.aiplatform_v1beta1 as aip_beta\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aip_constants:fs"
      },
      "source": [
        "#### Vertex AI constants\n",
        "\n",
        "Setup up the following constants for Vertex AI:\n",
        "\n",
        "- `API_ENDPOINT`: The Vertex AI API service endpoint for `Endpoint` services.\n",
        "- `PARENT`: The base resource path for all Vertex AI resources within your project and location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aip_constants:fs"
      },
      "outputs": [],
      "source": [
        "# API service endpoint\n",
        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(LOCATION)\n",
        "\n",
        "# Vertex location root path for your dataset, model and endpoint resources\n",
        "PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + LOCATION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "clients:metadata"
      },
      "source": [
        "## Set up clients\n",
        "\n",
        "The Vertex AI  works as a client/server model. On your side (the Python script) you create a client that sends requests and receives responses from the Vertex AI server.\n",
        "\n",
        "You use different clients in this tutorial for different steps in the workflow. This means you have to set them all up at the start.\n",
        "\n",
        "- Endpoint Service for creating endpoints, and deploying models to endpoints."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "clients:metadata"
      },
      "outputs": [],
      "source": [
        "# client options same for all services\n",
        "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
        "\n",
        "\n",
        "def create_endpoint_client():\n",
        "    client = aip_beta.EndpointServiceClient(client_options=client_options)\n",
        "    return client\n",
        "\n",
        "\n",
        "clients = {}\n",
        "clients[\"endpoint\"] = create_endpoint_client()\n",
        "\n",
        "for client in clients.items():\n",
        "    print(client)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
        "Set the variables `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa T4 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_T4, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more about [hardware accelerator support for your location](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4,\n",
        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
        "    )\n",
        "else:\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for prediction.\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "container:training,prediction"
      },
      "outputs": [],
      "source": [
        "TF = \"2.13\".replace(\".\", \"-\")\n",
        "\n",
        "if DEPLOY_GPU:\n",
        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "else:\n",
        "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for prediction.\n",
        "\n",
        "- Set the variable `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you use for for prediction.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "machine:training"
      },
      "outputs": [],
      "source": [
        "MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "DEPLOY_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Train machine type\", DEPLOY_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8128b8ff025"
      },
      "source": [
        "## Get pretrained model from TensorFlow Hub\n",
        "\n",
        "For demonstration purposes, this tutorial uses a pretrained models from TensorFlow Hub (TFHub), which is then uploaded to a `Vertex AI Model` resource. Once you have a `Vertex AI Model` resource, the model can be deployed to a `Vertex AI Endpoint` resource.\n",
        "\n",
        "### Download the pretrained image classification model\n",
        "\n",
        "First, you download the pretrained image classification model from TensorFlow Hub. The model gets downloaded as a TF.Keras layer. To finalize the model, in this example, you create a `Sequential()` model with the downloaded TFHub model as a layer, and specify the input shape to the model. The download model is pretrained on ImageNet."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "249bd746def6"
      },
      "outputs": [],
      "source": [
        "tfhub_model_icn = tf.keras.Sequential(\n",
        "    [hub.KerasLayer(\"https://tfhub.dev/google/imagenet/inception_v3/classification/5\")]\n",
        ")\n",
        "tfhub_model_icn.build([None, 224, 224, 3])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63de49055083"
      },
      "source": [
        "### Save the model artifacts\n",
        "\n",
        "At this point, the model is in memory. Next, you save the model artifacts to a Cloud Storage location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "64618c713db9"
      },
      "outputs": [],
      "source": [
        "MODEL_ICN_DIR = BUCKET_URI + \"/model_icn\"\n",
        "tfhub_model_icn.save(MODEL_ICN_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "how_serving_function_works"
      },
      "source": [
        "## Upload the model for serving\n",
        "\n",
        "Next, you upload your TFHub image classification model to Vertex AI `Model` service, which creates a Vertex AI `Model` resource for your model. During upload, you need to define a serving function to convert data to the format your model expects. If you send encoded data to Vertex AI, your serving function ensures that the data is decoded on the model server before it's passed as input to your model.\n",
        "\n",
        "### How does the serving function work\n",
        "\n",
        "When you send a request to an online prediction server, the request is received by a HTTP server. The HTTP server extracts the prediction request from the HTTP request content body. The extracted prediction request is forwarded to the serving function. For Google pre-built prediction containers, the request content is passed to the serving function as a `tf.string`.\n",
        "\n",
        "The serving function consists of two parts:\n",
        "\n",
        "- `preprocessing function`:\n",
        "  - Converts the input (`tf.string`) to the input shape and data type of the underlying model (dynamic graph).\n",
        "  - Performs the same preprocessing of the data that was done during training the underlying model -- e.g., normalizing, scaling, etc.\n",
        "- `post-processing function`:\n",
        "  - Converts the model output to format expected by the receiving application -- e.q., compresses the output.\n",
        "  - Packages the output for the the receiving application -- e.g., add headings, make JSON object, etc.\n",
        "\n",
        "Both the preprocessing and post-processing functions are converted to static graphs which are fused to the model. The output from the underlying model is passed to the post-processing function. The post-processing function passes the converted/packaged output back to the HTTP server. The HTTP server returns the output as the HTTP response content.\n",
        "\n",
        "One consideration you need to consider when building serving functions for TF.Keras models is that they run as static graphs. That means, you cannot use TF graph operations that require a dynamic graph. If you do, you get an error during the compile of the serving function which indicates that you're using an EagerTensor which is'nt supported."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "serving_function_image:post"
      },
      "source": [
        "### Serving function for image data\n",
        "\n",
        "#### Preprocessing\n",
        "\n",
        "To pass images to the prediction service, you encode the compressed (e.g., JPEG) image bytes into base 64 -- which makes the content safe from modification while transmitting binary data over the network. Since this deployed model expects input data as raw (uncompressed) bytes, you need to ensure that the base 64 encoded data gets converted back to raw bytes, and then preprocessed to match the model input requirements, before it's passed as input to the deployed model.\n",
        "\n",
        "To resolve this, you define a serving function (`serving_fn`) and attach it to the model as a preprocessing step. Add a `@tf.function` decorator so the serving function is fused to the underlying model (instead of upstream on a CPU).\n",
        "\n",
        "When you send a prediction or explanation request, the content of the request is base 64 decoded into a Tensorflow string (`tf.string`), which is passed to the serving function (`serving_fn`). The serving function preprocesses the `tf.string` into raw (uncompressed) numpy bytes (`preprocess_fn`) to match the input requirements of the model:\n",
        "\n",
        "- `io.decode_jpeg`- Decompresses the JPG image which is returned as a Tensorflow tensor with three channels (RGB).\n",
        "- `image.convert_image_dtype` - Changes integer pixel values to float 32, and rescales pixel data between 0 and 1.\n",
        "- `image.resize` - Resizes the image to match the input shape for the model.\n",
        "\n",
        "At this point, the data can be passed to the model (`m_call`), via a concrete function. The serving function is a static graph, while the model is a dynamic graph. The concrete function performs the tasks of marshalling the input data from the serving function to the model, and marshalling the prediction result from the model back to the serving function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "serving_function_image"
      },
      "outputs": [],
      "source": [
        "CONCRETE_INPUT = \"numpy_inputs\"\n",
        "\n",
        "\n",
        "def _preprocess(bytes_input):\n",
        "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
        "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
        "    resized = tf.image.resize(decoded, size=(224, 224))\n",
        "    return resized\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def preprocess_fn(bytes_inputs):\n",
        "    decoded_images = tf.map_fn(\n",
        "        _preprocess, bytes_inputs, dtype=tf.float32, back_prop=False\n",
        "    )\n",
        "    return {\n",
        "        CONCRETE_INPUT: decoded_images\n",
        "    }  # User needs to make sure the key matches model's input\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def serving_fn(bytes_inputs):\n",
        "    images = preprocess_fn(bytes_inputs)\n",
        "    prob = m_call(**images)\n",
        "    return prob\n",
        "\n",
        "\n",
        "m_call = tf.function(tfhub_model_icn.call).get_concrete_function(\n",
        "    [tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32, name=CONCRETE_INPUT)]\n",
        ")\n",
        "\n",
        "tf.saved_model.save(\n",
        "    tfhub_model_icn, MODEL_ICN_DIR, signatures={\"serving_default\": serving_fn}\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "serving_function_signature:image"
      },
      "source": [
        "### Get the serving function signature\n",
        "\n",
        "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
        "\n",
        "For your purpose, you need the signature of the serving function. Why? Well, when you send your data for prediction as a HTTP request packet, the image data is base64 encoded, and your TF.Keras model takes numpy input. Your serving function does the conversion from base64 to a numpy array.\n",
        "\n",
        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you use later when you make a prediction request."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "serving_function_signature:image"
      },
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load(MODEL_ICN_DIR)\n",
        "\n",
        "serving_input_icn = list(\n",
        "    loaded.signatures[\"serving_default\"].structured_input_signature[1].keys()\n",
        ")[0]\n",
        "print(\"Serving function input:\", serving_input_icn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e8ce91147c93"
      },
      "source": [
        "### Upload the TensorFlow Hub model to a `Vertex AI Model` resource\n",
        "\n",
        "Finally, you upload the model artifacts from the TFHub model and serving function into a `Vertex AI Model` resource.\n",
        "\n",
        "*Note:* When you upload the model artifacts to a `Vertex AI Model` resource, you specify the corresponding deployment container image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ad61e1429512"
      },
      "outputs": [],
      "source": [
        "model_icn = aiplatform.Model.upload(\n",
        "    display_name=\"icn\",\n",
        "    artifact_uri=MODEL_ICN_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        ")\n",
        "\n",
        "print(model_icn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1e893c95868b"
      },
      "source": [
        "### Download the pretrained sentence encoder model\n",
        "\n",
        "Next, you download the pretrained text sentence encoder model from TensorFlow Hub. The model gets downloaded as a TF.Keras layer. To finalize the model, in this example, you create a `Sequential()` model with the downloaded TFHub model as a layer, and specify the input shape to the model. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a3760344b764"
      },
      "outputs": [],
      "source": [
        "tfhub_model_use = tf.keras.Sequential(\n",
        "    [hub.KerasLayer(\"https://tfhub.dev/google/universal-sentence-encoder/4\")]\n",
        ")\n",
        "\n",
        "# force the model to build\n",
        "tfhub_model_use.predict([\"foo\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63de49055083"
      },
      "source": [
        "### Save the model artifacts\n",
        "\n",
        "At this point, the model is in memory. Next, you save the model artifacts to a Cloud Storage location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "64618c713db9"
      },
      "outputs": [],
      "source": [
        "MODEL_USE_DIR = BUCKET_URI + \"/model_use\"\n",
        "tfhub_model_use.save(MODEL_USE_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "serving_function_signature:image"
      },
      "source": [
        "## Get the serving function signature\n",
        "\n",
        "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
        "\n",
        "For your purpose, you need the signature of the serving function. \n",
        "\n",
        "When making a prediction request, you need to route the request to the serving function instead of the model, so you need to know the input layer name of the serving function -- which you use later when you make a prediction request."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "serving_function_signature:image"
      },
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load(MODEL_USE_DIR)\n",
        "\n",
        "serving_input_use = list(\n",
        "    loaded.signatures[\"serving_default\"].structured_input_signature[1].keys()\n",
        ")[0]\n",
        "print(\"Serving function input:\", serving_input_use)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e8ce91147c93"
      },
      "source": [
        "### Upload the TensorFlow Hub model to a `Vertex AI Model` resource\n",
        "\n",
        "Finally, you upload the model artifacts from the TFHub model and serving function into a `Vertex AI Model` resource.\n",
        "\n",
        "*Note:* When you upload the model artifacts to a `Vertex AI Model` resource, you specify the corresponding deployment container image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ad61e1429512"
      },
      "outputs": [],
      "source": [
        "model_use = aiplatform.Model.upload(\n",
        "    display_name=\"use\",\n",
        "    artifact_uri=MODEL_USE_DIR,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE,\n",
        ")\n",
        "\n",
        "print(model_use)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gk9tr92tSh-f"
      },
      "source": [
        "## Creating a deployment resource pool\n",
        "\n",
        "Currently, creating deploynent resource pools is only supported via the REST-based API (e.g., CURL).\n",
        "\n",
        "Use `CreateDeploymentResourcePool` API to create a resource pool, with the following configuration:\n",
        "\n",
        "- `dedicated_resources`: Compute (HW) resources to allocate for the shared vm.\n",
        "- `min_replica_count`: Auto-scaling, the minimum number of compute nodes.\n",
        "- `max_replica_count`: Auto-scaling, the maximum number of compute nodes.\n",
        "\n",
        "Learn more about [Deployment Resource Pools](https://cloud.google.com/vertex-ai/docs/predictions/model-co-hosting)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YiBmoiWYcMQt"
      },
      "outputs": [],
      "source": [
        "DEPLOYMENT_RESOURCE_POOL_ID = \"shared-vm\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0CHPJ4h-Slgs"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import pprint\n",
        "pp = pprint.PrettyPrinter(indent=4)\n",
        "\n",
        "MIN_NODES = 1\n",
        "MAX_NODES = 2\n",
        "\n",
        "CREATE_RP_PAYLOAD = {\n",
        "  \"deployment_resource_pool\":{\n",
        "    \"dedicated_resources\":{\n",
        "      \"machine_spec\":{\n",
        "        \"machine_type\": DEPLOY_COMPUTE\n",
        "      },\n",
        "      \"min_replica_count\": MIN_NODES, \n",
        "      \"max_replica_count\": MAX_NODES\n",
        "    }\n",
        "  },\n",
        "  \"deployment_resource_pool_id\":DEPLOYMENT_RESOURCE_POOL_ID\n",
        "}\n",
        "CREATE_RP_REQUEST=json.dumps(CREATE_RP_PAYLOAD)\n",
        "pp.pprint(\"CREATE_RP_REQUEST: \" + CREATE_RP_REQUEST)\n",
        "\n",
        "! curl \\\n",
        "-X POST \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools \\\n",
        "-d '{CREATE_RP_REQUEST}'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3WP4WV_bZzDh"
      },
      "source": [
        "## Get a deployment resource pool\n",
        "\n",
        "Use `GetDeploymentResourcePool` API to check out the deploynent resource pool that you created. \n",
        "\n",
        "Learn more about [Get Deployment Resource Pool](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/projects.locations.deploymentResourcePools/get)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6wTLyhPraFah"
      },
      "outputs": [],
      "source": [
        "! curl -X GET \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools/{DEPLOYMENT_RESOURCE_POOL_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gBQyGl-7aaZC"
      },
      "source": [
        "## List all deployment resource pools\n",
        "\n",
        "Use `ListDeploymentResourcePools` API to list all the deployment resource pools. \n",
        "\n",
        "Learn more about [Listing Deployment Resource Pools](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/projects.locations.deploymentResourcePools/list)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pxls4sNnaltU"
      },
      "outputs": [],
      "source": [
        "! curl -X GET \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "628de0914ba1"
      },
      "source": [
        "## Creating two `Endpoint` resource\n",
        "\n",
        "Next, you create two `Endpoint` resources using the `Endpoint.create()` method. At a minimum, you specify the display name for the endpoint. Optionally, you can specify the project and location (location); otherwise the settings are inherited by the values you set when you initialized the Vertex AI SDK for Python with the `init()` method.\n",
        "\n",
        "In this example, the following parameters are specified:\n",
        "\n",
        "- `display_name`: A human readable name for the `Endpoint` resource.\n",
        "\n",
        "This method returns an `Endpoint` object.\n",
        "\n",
        "Learn more about [Vertex AI Endpoints](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0ea443f9593b"
      },
      "outputs": [],
      "source": [
        "endpoint_icn = aiplatform.Endpoint.create(display_name=\"icn\")\n",
        "\n",
        "print(endpoint_icn)\n",
        "\n",
        "endpoint_use = aiplatform.Endpoint.create(display_name=\"use\")\n",
        "\n",
        "print(endpoint_use)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "azXkQSkjb3tv"
      },
      "source": [
        "## Deploy Model in a Deployment Resource Pool\n",
        "\n",
        "After you have created a Model and an Endpoint, you're ready to deploy using the DeployModel API. See an example of the CURL command below. Notice how you specified the `shared_resources` of DeployedModel with the deployment resource name of the resource pool that was created. \n",
        "\n",
        "Model deployments for the same deployment resource pool can be started concurrently.\n",
        "\n",
        "### Deploy the image classification model\n",
        "\n",
        "Next, you deploy the image classification model to an `Endpoint` using your deployment resource pool."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bEcQp8eUdHok"
      },
      "outputs": [],
      "source": [
        "SHARED_RESOURCE = \"projects/{project_id}/locations/{location}/deploymentResourcePools/{deployment_resource_pool_id}\".format(\n",
        "    project_id=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    deployment_resource_pool_id=DEPLOYMENT_RESOURCE_POOL_ID,\n",
        ")\n",
        "\n",
        "DEPLOY_MODEL_PAYLOAD = {\n",
        "    \"deployedModel\": {\n",
        "        \"model\": model_icn.resource_name,\n",
        "        \"shared_resources\": SHARED_RESOURCE,\n",
        "        \"enable_container_logging\": True,\n",
        "    },\n",
        "    \"trafficSplit\": {\"0\": 100},\n",
        "}\n",
        "DEPLOY_MODEL_REQUEST = json.dumps(DEPLOY_MODEL_PAYLOAD)\n",
        "pp.pprint(\"DEPLOY_MODEL_REQUEST: \" + DEPLOY_MODEL_REQUEST)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "92c2036684b3"
      },
      "outputs": [],
      "source": [
        "ENDPOINT_ID = endpoint_icn.name\n",
        "\n",
        "output = ! curl -X POST \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel \\\n",
        "-d '{DEPLOY_MODEL_REQUEST}'\n",
        "\n",
        "for line in output:\n",
        "    if '\"name\"' in line:\n",
        "        operation_id = line.split(\":\")[-1].strip()[:-1]\n",
        "        print(operation_id)\n",
        "        break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a5053c9db02a"
      },
      "source": [
        "### Wait for deployment to complete\n",
        "\n",
        "Next, you query the status of the operation, waiting for the operation state `done` to be set to `true`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "563db25caa5c"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "done = False\n",
        "while done != '\"done\": true':\n",
        "    status = ! curl -X GET \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{operation_id}\n",
        "    for line in status:\n",
        "        if '\"done\"' in line.strip():\n",
        "            done = line.strip()[0:-1]\n",
        "    print(\"DONE status:\", done)\n",
        "    time.sleep(30)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68ceccf7a8a7"
      },
      "source": [
        "### Deploy the text sentence encoder model\n",
        "\n",
        "Next, you deploy the text sentence encoder model to an `Endpoint` using your deployment resource pool."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bEcQp8eUdHok"
      },
      "outputs": [],
      "source": [
        "DEPLOY_MODEL_PAYLOAD = {\n",
        "    \"deployedModel\": {\n",
        "        \"model\": model_use.resource_name,\n",
        "        \"shared_resources\": SHARED_RESOURCE,\n",
        "        \"enable_container_logging\": True,\n",
        "    },\n",
        "    \"trafficSplit\": {\"0\": 100},\n",
        "}\n",
        "DEPLOY_MODEL_REQUEST = json.dumps(DEPLOY_MODEL_PAYLOAD)\n",
        "pp.pprint(\"DEPLOY_MODEL_REQUEST: \" + DEPLOY_MODEL_REQUEST)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "288acc79c173"
      },
      "outputs": [],
      "source": [
        "ENDPOINT_ID = endpoint_use.name\n",
        "\n",
        "output = ! curl -X POST \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:deployModel \\\n",
        "-d '{DEPLOY_MODEL_REQUEST}'\n",
        "\n",
        "for line in output:\n",
        "    if '\"name\"' in line:\n",
        "        operation_id = line.split(\":\")[-1].strip()[:-1]\n",
        "        print(operation_id)\n",
        "        break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87ba1aa4e0b0"
      },
      "source": [
        "### Wait for deployment to complete\n",
        "\n",
        "Next, you query the status of the operation, waiting for the operation state `done` to be set to `true`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5b075d6ec099"
      },
      "outputs": [],
      "source": [
        "done = False\n",
        "while done != '\"done\": true':\n",
        "    status = ! curl -X GET \\\n",
        " -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        " -H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{operation_id}\n",
        "    for line in status:\n",
        "        if '\"done\"' in line.strip():\n",
        "            done = line.strip()[0:-1]\n",
        "    print(\"DONE status:\", done)\n",
        "    time.sleep(30)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "033a5bfb969c"
      },
      "source": [
        "### Create test example the image classification model\n",
        "\n",
        "Next, you test your deployed image classification model. First, you encode your test data for the serving function, which is in the format:\n",
        "\n",
        "`{ serving_input: { 'b64': base64_encoded_bytes } }`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "54947579cdf7"
      },
      "outputs": [],
      "source": [
        "! gsutil cp gs://cloud-ml-data/img/flower_photos/daisy/100080576_f52e8ee070_n.jpg test.jpg\n",
        "\n",
        "import base64\n",
        "\n",
        "with open(\"test.jpg\", \"rb\") as f:\n",
        "    data = f.read()\n",
        "b64str = base64.b64encode(data).decode(\"utf-8\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fdf3dd237d32"
      },
      "source": [
        "### Make the prediction request for the image classification model\n",
        "\n",
        "Finally, you make a prediction request. Since the model was trained on ImageNet, the prediction returns the probabilities for the corresponding 1000 classes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "predict_request:mbsdk,custom,icn"
      },
      "outputs": [],
      "source": [
        "# The format of each instance should conform to the deployed model's prediction input schema.\n",
        "instances = [{serving_input_icn: {\"b64\": b64str}}]\n",
        "\n",
        "prediction = endpoint_icn.predict(instances=instances)\n",
        "\n",
        "print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "033a5bfb969c"
      },
      "source": [
        "### Create test example the text sentence encoder model\n",
        "\n",
        "Next, you test your deployed text sentence encoder model. First, you encode your test data for the serving function, which is in the format:\n",
        "\n",
        "`\"word1 word2 ... wordN\"`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3dd85e9ce024"
      },
      "outputs": [],
      "source": [
        "instance = \"the brown fox jumped over the laxy dog\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e3185c34f59c"
      },
      "source": [
        "### Make the prediction request for the text sentence encoder model\n",
        "\n",
        "Finally, you make a prediction request. The prediction returns an embedding which is a 500 element vector."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "96877270ec07"
      },
      "outputs": [],
      "source": [
        "endpoint_use.predict([instance])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cleanup"
      },
      "outputs": [],
      "source": [
        "# Set this to true only if you'd like to delete your bucket\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "\n",
        "# Undeploy the models\n",
        "# When you're done doing predictions, you undeploy the model from the `Endpoint` resouce. \n",
        "# This deprovisions all compute resources and ends billing for the deployed model.\n",
        "endpoint_icn.undeploy_all()\n",
        "endpoint_use.undeploy_all()\n",
        "\n",
        "# Delete the `Endpoint` resources\n",
        "# The method 'delete()' deletes the model.\n",
        "model_icn.delete()\n",
        "model_use.delete()\n",
        "\n",
        "# Delete the `Endpoint` resources\n",
        "# The method 'delete()' deletes the endpoint.\n",
        "endpoint_icn.delete()\n",
        "endpoint_use.delete()\n",
        "\n",
        "# Delete the `DeploymentResourcePool`\n",
        "# The method 'delete()' deletes your deployment resource pool.\n",
        "! curl -X DELETE \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/deploymentResourcePools/{DEPLOYMENT_RESOURCE_POOL_ID}\n",
        "\n",
        "!rm -f test.jpg"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_vertex_endpoint_and_shared_vm.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}