03-model-inference-hugectr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f7086876-72ed-4ac8-aeba-f05dbc641798",
   "metadata": {},
   "source": [
    "# Serving models using NVIDIA Triton Inference Server and Vertex AI Prediction\n",
    "\n",
    "This notebook demonstrates how to serve NVIDIA Merlin HugeCTR deep learning models using [NVIDIA Triton Inference Server](https://github.com/triton-inference-server/server) and [Vertex AI Prediction](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions).\n",
    "The notebook compiles prescriptive guidance for the following tasks:\n",
    "\n",
    "1. Creating Triton ensemble models that combine NVTabular preprocessing workflows and Merlin HugeCTR models\n",
    "2. Building a Vertex Prediction custom serving container image for serving the ensemble models with NVIDIA Triton Inference server. \n",
    "3. Uploading the ensemble model and its metadata to Vertex AI Model Resources.\n",
    "4. Deploying the model with the Triton Inference Server container to Vertex Prediction Endpoints.\n",
    "5. Getting online predictions from the deployed ensemble.\n",
    "\n",
    "To fully benefit from the content covered in this notebook, you should have a solid understanding of key Vertex AI Prediction concepts like models, endpoints, and model deployments. We strongly recommend reviewing [Vertex AI Prediction documentation](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) before proceeding.\n",
    "\n",
    "## NVIDIA Triton Inference Server Overview\n",
    "\n",
    "[Triton Inference Server](https://github.com/triton-inference-server/server) provides an inferencing solution optimized for both CPUs and GPUs. Triton can run multiple models from the same or different frameworks concurrently on a single GPU or CPU. In a multi-GPU server, it automatically creates an instance of each model on each GPU to increase utilization without extra coding. It supports real-time inferencing, batch inferencing to maximize GPU/CPU utilization, and streaming inference with built-in support for audio streaming input. It also supports model ensembles for use cases that require multiple models to perform end-to-end inference.\n",
    "\n",
    "The following figure shows the Triton Inference Server high-level architecture.\n",
    "\n",
    "<img src=\"images/triton-architecture.png\" alt=\"Triton Architecture\" style=\"width:50%\"/>\n",
    "\n",
    "\n",
    "- The model repository is a file-system based repository of the models that Triton will make available for inferencing. \n",
    "- Inference requests arrive at the server via either HTTP/REST or gRPC and are then routed to the appropriate per-model scheduler. \n",
    "- Triton implements multiple scheduling and batching algorithms that can be configured on a model-by-model basis.\n",
    "- The backend performs inferencing using the inputs provided in the batched requests to produce the requested outputs.\n",
    "\n",
    "\n",
    "Triton server provides readiness and liveness health endpoints, as well as utilization, throughput, and latency metrics, which enable the integration of Triton into deployment environments, such as Vertex AI Prediction.\n",
    "\n",
    "Refer to [Triton Inference Server Architecture](https://github.com/triton-inference-server/server/blob/main/docs/architecture.md) for more detailed information.\n",
    "\n",
    "## Triton Inference Server on Vertex AI Prediction\n",
    "\n",
    "\n",
    "\n",
    "In this section, we describe the deployment of Triton Inference Server on Vertex AI Prediction. Although, the focus of this notebook is on demonstrating how to serve an ensemble of an NVTabular preprocessing workflow and a HugeCTR model, the outlined design patterns are applicable to a wider set of serving scenarios.  The following figure shows a deployment architecture.\n",
    "\n",
    "<img src=\"./images/triton-in-vertex.png\" alt=\"Triton on Vertex AI Prediction\" style=\"width:70%\"/>\n",
    "\n",
    "\n",
    "Triton Inference Server runs inside a container based on a custom serving image. The custom container image is built on top of [NVIDIA Merlin Inference image](https://ngc.nvidia.com/catalog/containers/nvidia:merlin:merlin-inference) and adds packages and configurations to align with Vertex AI [requirements for custom serving container images](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements). \n",
    "\n",
    "An ensemble to be served by Triton is registered with Vertex AI Prediction as a `Model`. The `Model`'s metadata reference a location of the ensemble artifacts in Google Cloud Storage and the custom serving container and its configurations. \n",
    "\n",
    "After the model is deployed to a Vertex AI Prediction endpoint, the entrypoint script of the custom container copies the ensemble's artifacts from the GCS location to a local file system in the container. It then starts Triton, referencing a local copy of the ensemble as Triton's model repository. \n",
    "\n",
    "Triton loads the models comprising the ensemble and exposes inference, health, and model management REST endpoints using [standard inference protocols](https://github.com/kserve/kserve/tree/master/docs/predict-api/v2). The Triton's inference endpoint - `/v2/models/{ENSEMBLE_NAME}/infer` is mapped to Vertex AI Prediction predict route and exposed to external clients through Vertex Prediction endpoint. The Triton's health endpoint - `/v2/health/ready` - is mapped to Vertex AI Prediction health route and used by Vertex AI Prediction for [health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#health)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c28b3137-249c-49a3-a8be-f459b3774ea2",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "In this section of the notebook you configure your environment settings, including a GCP project, a Vertex AI compute region, and a Vertex AI staging GCS bucket. You also set the locations of the fitted NVTaubular workflow and the trained HugeCTR model created by the [01-dataset-preprocessing.ipynb](https://github.com/GoogleCloudPlatform/nvidia-merlin-on-vertex-ai/blob/main/01-dataset-preprocessing.ipynb) and [02-model-training-hugectr.ipynb](https://github.com/GoogleCloudPlatform/nvidia-merlin-on-vertex-ai/blob/main/02-model-training-hugectr.ipynb) notebooks, as well as the Vertex AI model name, description, and endpoint name.\n",
    "\n",
    "Make sure to update the below cells with the values reflecting your environment.\n",
    "\n",
    "First import all the necessary python packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9791d621-22e8-41bc-8681-a48aec45ef6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import shutil\n",
    "import time\n",
    "\n",
    "from pathlib import Path\n",
    "from src.serving import export\n",
    "\n",
    "from google.cloud import aiplatform as vertex_ai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "023b5454",
   "metadata": {},
   "source": [
    "Change the following variables according to your definitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd0098a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Project definitions\n",
    "PROJECT_ID = '<YOUR PROJECT ID>' # Change to your project id.\n",
    "REGION = '<LOCATION OF RESOURCES>'  # Change to your region.\n",
    "\n",
    "# Bucket definitions\n",
    "STAGING_BUCKET = '<YOUR BUCKET NAME>' # Change to your bucket.\n",
    "WORKFLOW_MODEL_PATH = \"gs://...\" # Change to GCS path of the nvt workflow\n",
    "HUGECTR_MODEL_PATH = \"gs://...\" # Change to GCS path of the hugectr trained model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d4783df",
   "metadata": {},
   "source": [
    "Change the following variables ONLY if necessary.  \n",
    "You can leave the default variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e20da19-0dda-4a58-b606-341e50c9bb5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_ARTIFACTS_REPOSITORY = f'gs://{STAGING_BUCKET}/recsys-models'\n",
    "\n",
    "MODEL_NAME = 'deepfm'\n",
    "MODEL_VERSION = 'v01'\n",
    "MODEL_DISPLAY_NAME = f'criteo-hugectr-{MODEL_NAME}-{MODEL_VERSION}'\n",
    "MODEL_DESCRIPTION = 'HugeCTR DeepFM model'\n",
    "ENDPOINT_DISPLAY_NAME = f'hugectr-{MODEL_NAME}-{MODEL_VERSION}'\n",
    "\n",
    "LOCAL_WORKSPACE = '/home/jupyter/serving_notebook'\n",
    "IMAGE_NAME = 'triton-serving'\n",
    "IMAGE_URI = f'gcr.io/{PROJECT_ID}/{IMAGE_NAME}'\n",
    "DOCKERNAME = 'triton'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b8e3f45-2bab-48fb-9b81-4e0e65635ece",
   "metadata": {},
   "source": [
    "### Initialize Vertex AI SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2e0cca3-11e6-4454-af41-c5c632dbbeec",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_ai.init(\n",
    "    project=PROJECT_ID,\n",
    "    location=REGION,\n",
    "    staging_bucket=STAGING_BUCKET\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51f9c710-537a-49c9-86d5-fae7e9f2b4be",
   "metadata": {},
   "source": [
    "## 1. Exporting Triton ensemble model\n",
    "\n",
    "A Triton ensemble model represents a pipeline of one or more models and the connection of input and output tensors between these models. Ensemble models are intended to encapsulate inference pipelines that involves multiple steps, each performed by a different model. For example, a common  \"data preprocessing -> inference -> data postprocessing\" pattern. Using ensemble models for this purpose can avoid the overhead of transferring intermediate tensors between client and serving endpoints and minimize the number of requests that must be sent to Triton. \n",
    "\n",
    "In our case, an inference pipeline comprises two steps: input preprocessing using a fitted NVTabular workflow and generating predictions using a HugeCTR ranking model.\n",
    "\n",
    "An ensemble model is not an actual serialized model. There are no addtional model artifacts created when an ensemble is defined. It is a configuration that specifies which actual models comprise the ensemble, the execution flow when processing an inference request and the flow of data between inputs and outputs of the component models. This configuration is defined using the same [protocol buffer](https://developers.google.com/protocol-buffers) based configuration format as used for serving other model types in Triton. Refer to [Triton Inference Server Model Configuration guide](https://github.com/triton-inference-server/server/blob/main/docs/model_configuration.md) for detailed information about configuring models and model ensembles.\n",
    "\n",
    "You can create an ensemble model manually by arranging the component models into the prescribed folder structure and editing the required configuration files. For ensemble models that utilize the \"NVTabular workflow -> Inference Model\" processing pattern you can utilize a set of utility functions provided by the `nvtabular.inference.triton` module. Specifically to create a \"NVTabular workflow -> HugeCTR model\" ensemble, as utilized in this notebook, you can use the `nvtabular.inference.triton.export_hugectr_ensemble` function.\n",
    "\n",
    "We have encapsulated the ensemble export logic in the `src.serving.export_ensemble` function. In addition to calling `nvtabular.inference.triton.export_hugectr_ensemble`, the function also creates a JSON configuration file required by Triton when serving HugeCTR models. This file - `ps.json` - specifies the locations of different components comprising a saved HugeCTR model and is used by Triton HugeCTR backend to correctly load the saved model and prepare it for serving. \n",
    "\n",
    "Recall that the entrypoint script in the custom serving container copies the ensemble's models artifacts from a source GCS location as prepared by Vertex AI Prediction into the serving container's local file systems. The `ps.json` file needs to use the paths that correctly point to saved model artifacts in the container's file system. Also some of the paths embedded in the configs generated by `nvtabular.inference.triton.export_hugectr_ensemble` use absolute paths and need to be properly set. The `src.serving.export_ensemble` function handles all of that. You can specify the target root folder in the containers local file system using the `model_repository_path` parameter and all the paths will be adjusted accordingly.\n",
    "\n",
    "\n",
    "\n",
    "### Copy a HugeCTR saved model and a fitted NVTabular workflow to a local staging folder\n",
    "\n",
    "The `nvtabular.inference.triton.export_hugectr_ensemble` does not support GCS. As such you need to copy NVTabular workflow and HugeCTR model artifacts to a local file system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc4fc36f-621e-4471-97ac-0580d2bca48c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.isdir(LOCAL_WORKSPACE):\n",
    "    shutil.rmtree(LOCAL_WORKSPACE)\n",
    "os.makedirs(LOCAL_WORKSPACE)\n",
    "\n",
    "!gsutil -m cp -r {WORKFLOW_MODEL_PATH} {LOCAL_WORKSPACE}\n",
    "!gsutil -m cp -r {HUGECTR_MODEL_PATH} {LOCAL_WORKSPACE}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45757f1d-2e9b-4035-b982-d4c289e67dec",
   "metadata": {},
   "source": [
    "### Export the ensemble model\n",
    "\n",
    "The `src.export.export_ensemble` utility function takes a number of arguments that are required to set up a proper flow of tensors between inputs and outputs of the NVTabular workflow and the HugeCTR model.\n",
    "\n",
    "- `model_name` - The model name that will be used as a prefix for the generated ensemble artifacts.\n",
    "- `workflow_path` - The local path to the NVTabular workflow\n",
    "- `saved_model_path` - The local path to the saved HugeCTR model\n",
    "- `output_path` - The local path to the location where an ensemble will be exported\n",
    "- `model_repository_path` - The path to use as a root  in `ps.json` and other config files\n",
    "- `max_batch` - The maximum size of a serving batch that will be supported by the ensemble \n",
    "\n",
    "\n",
    "The following settings should match the settings of the NVTabular workflow\n",
    "\n",
    "- `categorical_columns` - The list of names of categorical input features to the NVTabular workflow\n",
    "- `continuous_columns` - The list of names of continuous input features to the NVTabular workflow\n",
    "\n",
    "\n",
    "The following settings should match the respective settings in the HugeCTR model\n",
    "\n",
    "- `num_outputs` - The number of outputs from the HugeCTR model\n",
    "- `embedding_vector_size` - The size of an embedding vector used by the HugeCTR model\n",
    "- `num_slots` - The number of slots used for sparse features of the HugeCTR model\n",
    "- `max_nnz` - This value controls how sparse features are coded in the embedding arrays \n",
    "\n",
    "\n",
    "As noted before, in this notebook we assume that you generated the NVTabular workflow and the HugeCTR model using the the [01-dataset-preprocessing.ipynb](https://github.com/GoogleCloudPlatform/nvidia-merlin-on-vertex-ai/blob/main/01-dataset-preprocessing.ipynb) and [02-model-training-hugectr.ipynb](https://github.com/GoogleCloudPlatform/nvidia-merlin-on-vertex-ai/blob/main/02-model-training-hugectr.ipynb) notebooks. The workflow captures the preprocessing logic for the Criteo dataset and the HugeCTR model is an implementation of [the DeepFM CTR model](https://arxiv.org/abs/1703.04247). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b84b102-e68f-4cdb-9341-362920556feb",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_SLOTS = 26\n",
    "MAX_NNZ = 2\n",
    "EMBEDDING_VECTOR_SIZE = 11\n",
    "MAX_BATCH_SIZE = 64\n",
    "NUM_OUTPUTS = 1\n",
    "\n",
    "continuous_columns = [\"I\" + str(x) for x in range(1, 14)]\n",
    "categorical_columns = [\"C\" + str(x) for x in range(1, 27)]\n",
    "label_columns = [\"label\"]\n",
    "\n",
    "local_workflow_path = str(Path(LOCAL_WORKSPACE) / Path(WORKFLOW_MODEL_PATH).parts[-1])\n",
    "local_saved_model_path = str(Path(LOCAL_WORKSPACE) / Path(HUGECTR_MODEL_PATH).parts[-1])\n",
    "local_ensemble_path = str(Path(LOCAL_WORKSPACE) / f'triton-ensemble-{time.strftime(\"%Y%m%d%H%M%S\")}')\n",
    "model_repository_path = '/model'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76abc820-3603-4b8b-9dec-a22c07d31d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "export.export_ensemble(\n",
    "    model_name=MODEL_NAME,\n",
    "    workflow_path=local_workflow_path,\n",
    "    saved_model_path=local_saved_model_path,\n",
    "    output_path=local_ensemble_path,\n",
    "    categorical_columns=categorical_columns,\n",
    "    continuous_columns=continuous_columns,\n",
    "    label_columns=label_columns,\n",
    "    num_slots=NUM_SLOTS,\n",
    "    max_nnz=MAX_NNZ,\n",
    "    num_outputs=NUM_OUTPUTS,\n",
    "    embedding_vector_size=EMBEDDING_VECTOR_SIZE,\n",
    "    max_batch_size=MAX_BATCH_SIZE,\n",
    "    model_repository_path=model_repository_path\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46c48c08-1267-4117-ae54-1331d7afe4ed",
   "metadata": {},
   "source": [
    "The previous cell created the following local folder structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d654b85e-db4a-47c2-97d2-960bc802fef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls -la {local_ensemble_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdd6401c-dca4-4e45-ad64-e7dd6c8e9e4f",
   "metadata": {},
   "source": [
    "The `deepfm` folder contains artifacts and configurations for the HugeCTR model. The `deepfm_ens` folder contains a configuration for the ensemble model. And the `deepfm_nvt` contains artifacts and configurations for the NVTabular preprocessing workflow. The `ps.json` file contains information required by the Triton's HugeCTR backend.\n",
    "\n",
    "Notice that the file paths in `ps.json`  use the value from `model_repository_path`. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1671792c-6566-4982-b2c3-7aab7a92e6d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "! cat {local_ensemble_path}/ps.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481fdc6e-1047-44a0-b6fa-2bd7c5b6e232",
   "metadata": {},
   "source": [
    "### Upload the ensemble to GCS\n",
    "\n",
    "In the later steps you will register the exported ensemble model as a Vertex AI Prediction model resource. Before doing that we need to move the ensemble to GCS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad435708-b92b-40d6-83f9-3b685bfeaba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcs_ensemble_path = '{}/{}'.format(MODEL_ARTIFACTS_REPOSITORY, Path(local_ensemble_path).parts[-1])\n",
    "\n",
    "!gsutil -m cp -r {local_ensemble_path}/* {gcs_ensemble_path}/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f662a18-0b51-4b30-8e3f-81daf64ca5a9",
   "metadata": {},
   "source": [
    "## 2. Building a custom serving container \n",
    "\n",
    "The custom serving container is derived from the NVIDIA NGC Merlin inference container. It adds Google Cloud SDK and an entrypoint script that executes the tasks described in detail in the overview."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19efaa3e-bace-47d6-8023-dc5a688ae7a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "! cat src/Dockerfile.triton"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ac0c951-23b5-4f92-b809-b6fa323beec2",
   "metadata": {},
   "source": [
    "As described in detail in the overview, the entry point script copies the ensemble artifacts to the serving container's local file system and starts Triton."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24cd7086-111a-4336-a3c6-108a63a994b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "! cat src/serving/entrypoint.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b113e891-7324-4bec-9381-f775dd6df4d5",
   "metadata": {},
   "source": [
    "You use [Cloud Build](https://cloud.google.com/build) to build the serving container and push it to your projects [Container Registry](https://cloud.google.com/container-registry#:~:text=Container%20Registry%20is%20a%20single,pipelines%20to%20get%20fast%20feedback.)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed3fff22-3541-466a-9d2c-8e007cfb4940",
   "metadata": {},
   "outputs": [],
   "source": [
    "FILE_LOCATION = './src'\n",
    "! gcloud builds submit --config src/cloudbuild.yaml --substitutions _DOCKERNAME=$DOCKERNAME,_IMAGE_URI=$IMAGE_URI,_FILE_LOCATION=$FILE_LOCATION --timeout=2h --machine-type=e2-highcpu-8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "930f7ae3-76c4-4bfe-8a80-44d3567b0b6d",
   "metadata": {},
   "source": [
    "## 3. Uploading the model and its metadata to Vertex Models.\n",
    "\n",
    "In the following cell you will register (upload) the ensemble model as a Vertex AI Prediction `Model` resource. \n",
    "\n",
    "Refer to [Use a custom container for prediction guide](https://cloud.google.com/vertex-ai/docs/predictions/use-custom-container) for detailed information about creating Vertex AI Prediction `Model` resources.\n",
    "\n",
    "Notice that the value of  `model_repository_path` that was used when exporting the ensemble is passed as a command line parameter to the serving container. The entrypoint script in the container will copy the ensemble artifacts to this location when the container starts. This ensures that the locations of the artifacts in the container's local file system and the paths in the `ps.json` and other configuration files used by Triton match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9e3c2ac-90fb-4175-b185-d2981cbdebf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_container_args = [model_repository_path]\n",
    "\n",
    "model = vertex_ai.Model.upload(\n",
    "    display_name=MODEL_DISPLAY_NAME,\n",
    "    description=MODEL_DESCRIPTION,\n",
    "    serving_container_image_uri=IMAGE_URI,\n",
    "    artifact_uri=gcs_ensemble_path,\n",
    "    serving_container_args=serving_container_args,\n",
    "    sync=True\n",
    ")\n",
    "\n",
    "model.resource_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c572a7c-c5e5-4383-bbd5-ce11b6659945",
   "metadata": {},
   "source": [
    "## 4. Deploying the model to Vertex AI Prediction.\n",
    "\n",
    "Deploying a Vertex AI Prediction `Model` is a two step process. First you create an endpoint that will expose an external interface to clients consuming the model. After the endpoint is ready you can deploy multiple versions of a model to the endpoint.\n",
    "\n",
    "Refer to [Deploy a model using the Vertex AI API guide](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) for more information about the APIs used in the following cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "763abf60-c325-4222-b550-0aca98492d15",
   "metadata": {},
   "source": [
    "### Create the Vertex Endpoint\n",
    "\n",
    "Before deploying the ensemble model you need to create a Vertex AI Prediction endpoint. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2ef58a0-1f13-40a7-b545-f6a75ee0fcb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = vertex_ai.Endpoint.create(\n",
    "    display_name=ENDPOINT_DISPLAY_NAME\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "866fd6fc-2777-457e-b446-8b1532984746",
   "metadata": {},
   "source": [
    "### Deploy the model to Vertex Prediction endpoint\n",
    "\n",
    "After the endpoint is ready, you can deploy your ensemble model to the endpoint. You will run the ensemble on a GPU node equipped with the NVIDIA Tesla T4 GPUs. \n",
    "\n",
    "Refer to [Deploy a model using the Vertex AI API guide](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf2bb85f-330c-4ce6-9a8f-d5b96772ab2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "traffic_percentage = 100\n",
    "machine_type = \"n1-standard-8\"\n",
    "accelerator_type=\"NVIDIA_TESLA_T4\"\n",
    "accelerator_count = 1\n",
    "min_replica_count = 1\n",
    "max_replica_count = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fd27dda-7c1b-4ce5-9477-e62a9f1199bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.deploy(\n",
    "    endpoint=endpoint,\n",
    "    deployed_model_display_name=MODEL_DISPLAY_NAME,\n",
    "    machine_type=machine_type,\n",
    "    min_replica_count=min_replica_count,\n",
    "    max_replica_count=max_replica_count,\n",
    "    traffic_percentage=traffic_percentage,\n",
    "    accelerator_type=accelerator_type,\n",
    "    accelerator_count=accelerator_count,\n",
    "    sync=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca6c71d4-56c4-4f24-aee6-d1b96f6e4e8a",
   "metadata": {},
   "source": [
    "## 5. Invoking the model\n",
    "\n",
    "To invoke the ensemble through Vertex AI Prediction endpoint you need to format your request using a [standard Inference Request JSON Object](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#inference) or a [Inference Request JSON Object with a binary extension](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_binary_data.md) and submit a request to Vertex AI Prediction [REST rawPredict endpoint](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/projects.locations.endpoints/rawPredict). You need to use the `rawPredict` rather than `predict` endpoint because inference request formats used by Triton are not compatible with the Vertex AI Prediction [standard input format](https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-custom-models#formatting-prediction-input).\n",
    "\n",
    "The below cell shows a sample request body formatted as a [standard Inference Request JSON Object](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#inference). The request encapsulates a batch of three records from the Criteo dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924669cf-5442-4454-b881-e2a792a96f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = {\n",
    "    'id': '1',\n",
    "    'inputs': [\n",
    "        {'name': 'I1','shape': [3, 1], 'datatype': 'INT32', 'data': [5, 32, 0]},\n",
    "        {'name': 'I2', 'shape': [3, 1], 'datatype': 'INT32', 'data': [110, 3, 233]},\n",
    "        {'name': 'I3', 'shape': [3, 1], 'datatype': 'INT32', 'data': [0, 5, 1]},\n",
    "        {'name': 'I4', 'shape': [3, 1], 'datatype': 'INT32', 'data': [16, 0, 146]},\n",
    "        {'name': 'I5', 'shape': [3, 1], 'datatype': 'INT32', 'data': [0, 1, 1]},\n",
    "        {'name': 'I6', 'shape': [3, 1], 'datatype': 'INT32', 'data': [1, 0, 0]},\n",
    "        {'name': 'I7', 'shape': [3, 1], 'datatype': 'INT32', 'data': [0, 0, 0]},\n",
    "        {'name': 'I8', 'shape': [3, 1], 'datatype': 'INT32', 'data': [14, 61, 99]},\n",
    "        {'name': 'I9', 'shape': [3, 1], 'datatype': 'INT32', 'data': [7, 5, 7]},\n",
    "        {'name': 'I10', 'shape': [3, 1], 'datatype': 'INT32', 'data': [1, 0, 0]},\n",
    "        {'name': 'I11', 'shape': [3, 1], 'datatype': 'INT32', 'data': [0, 1, 1]},\n",
    "        {'name': 'I12', 'shape': [3, 1], 'datatype': 'INT32', 'data': [306, 3157, 3101]},\n",
    "        {'name': 'I13', 'shape': [3, 1], 'datatype': 'INT32', 'data': [0, 5, 1]},\n",
    "        {'name': 'C1', 'shape': [3, 1], 'datatype': 'INT32', 'data': [1651969401, -436994675, 1651969401]},\n",
    "        {'name': 'C2', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-501260968, -1599406170, -1382530557]},\n",
    "        {'name': 'C3', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1343601617, 1873417685, 1656669709]},\n",
    "        {'name': 'C4', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1805877297, -628476895, 946620910]},\n",
    "        {'name': 'C5', 'shape': [3, 1], 'datatype': 'INT32', 'data': [951068488, 1020698403, -413858227]},\n",
    "        {'name': 'C6', 'shape': [3, 1], 'datatype': 'INT32', 'data': [1875733963, 1875733963, 1875733963]},\n",
    "        {'name': 'C7', 'shape': [3, 1], 'datatype': 'INT32', 'data': [897624609, -1424560767, -1242174622]},\n",
    "        {'name': 'C8', 'shape': [3, 1], 'datatype': 'INT32', 'data': [679512323, 1128426537, -772617077]},\n",
    "        {'name': 'C9', 'shape': [3, 1], 'datatype': 'INT32', 'data': [1189011366, 502653268, 776897055]},\n",
    "        {'name': 'C10', 'shape': [3, 1], 'datatype': 'INT32', 'data': [771915201, 2112471209, 771915201]},\n",
    "        {'name': 'C11', 'shape': [3, 1], 'datatype': 'INT32', 'data': [209470001, 1716706404, 209470001]},\n",
    "        {'name': 'C12', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1785193185, -1712632281, 309420420]},\n",
    "        {'name': 'C13', 'shape': [3, 1], 'datatype': 'INT32', 'data': [12976055, 12976055, 12976055]},\n",
    "        {'name': 'C14', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1102125769, -1102125769, -1102125769]},\n",
    "        {'name': 'C15', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1978960692, -205783399, -150008565]},\n",
    "        {'name': 'C16', 'shape': [3, 1], 'datatype': 'INT32', 'data': [1289502458, 1289502458, 1289502458]},\n",
    "        {'name': 'C17', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-771205462, -771205462, -771205462]},\n",
    "        {'name': 'C18', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1206449222, -1578429167, 1653545869]},\n",
    "        {'name': 'C19', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1793932789, -1793932789, -1793932789]},\n",
    "        {'name': 'C20', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-1014091992, -20981661, -1014091992]},\n",
    "        {'name': 'C21', 'shape': [3, 1], 'datatype': 'INT32', 'data': [351689309, -1556988767, 351689309]},\n",
    "        {'name': 'C22', 'shape': [3, 1], 'datatype': 'INT32', 'data': [632402057, -924717482, 632402057]},\n",
    "        {'name': 'C23', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-675152885, 391309800, -675152885]},\n",
    "        {'name': 'C24', 'shape': [3, 1], 'datatype': 'INT32', 'data': [2091868316, 1966410890, 883538181]},\n",
    "        {'name': 'C25', 'shape': [3, 1], 'datatype': 'INT32', 'data': [809724924, -1726799382, -10139646]},\n",
    "        {'name': 'C26', 'shape': [3, 1], 'datatype': 'INT32', 'data': [-317696227, -1218975401, -317696227]}]\n",
    "}\n",
    "\n",
    "with open('criteo_payload.json', 'w') as f:\n",
    "    json.dump(payload, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27197c74-94c7-454f-a469-151c680de97b",
   "metadata": {},
   "source": [
    "You can invoke the Vertex AI Prediction `rawPredict` endpoint using any HTTP tool or library, including `curl`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f98d37f-87bf-47d8-951e-a88dee48421e",
   "metadata": {},
   "outputs": [],
   "source": [
    "uri = f'https://{REGION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:rawPredict'\n",
    "\n",
    "! curl -X POST \\\n",
    "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
    "-H \"Content-Type: application/json\"  \\\n",
    "{uri} \\\n",
    "-d @criteo_payload.json"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "",
   "name": "managed-notebooks.m87",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/managed-notebooks:m87"
  },
  "interpreter": {
   "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}