examples/scaling-criteo/04-Triton-Inference-with-Merlin-Models-TensorFlow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c9026a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2022 NVIDIA Corporation. All Rights Reserved.\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",
    "#     http://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.\n",
    "# ==============================================================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f6f1d42",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_merlin_scaling-criteo-04-triton-inference-with-merlin-models-tensorflow/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Scaling Criteo: Triton Inference with Merlin Models TensorFlow\n",
    "This notebook is created using the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) container.\n",
    "\n",
    "## Overview\n",
    "In the previous notebook, we processed the [criteo dataset with NVTabular](02-ETL-with-NVTabular.ipynb) and [trained a DLRM model with Merlin Model Tensorflow](04-Triton-Inference-with-Merlin-Models-TensorFlow.ipynb). Finally, we want to deploy our pipeline to [Triton Inference Server (TIS)](https://github.com/triton-inference-server/server), which can serve our model in a production environment.\n",
    "\n",
    "We can send raw data to the API endpoint. TIS will execute the same NVTabular workflow for feature engineering and predict the processed data with Merlin Models TensorFlow. We deploy the pipeline as an ensemble and receive the predict scores. This notebook is based on the Example, [Serving Ranking Models With Merlin Systems](https://github.com/NVIDIA-Merlin/systems/blob/stable/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb), in Merlin systems. If you are interested in more details, we recommend to go through the example, first.\n",
    "\n",
    "## Learning objectives\n",
    "- Deploy an ensemble pipeline of NVTabular and Merlin Models TensorFlow to Triton Inference Server\n",
    "- Get prediction from Triton Inference Server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c13f81",
   "metadata": {},
   "source": [
    "## Saved NVTabular workflow and Merlin Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5583915",
   "metadata": {},
   "source": [
    "We load the required libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e8c67a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "os.environ[\"TF_GPU_ALLOCATOR\"] = \"cuda_malloc_async\"\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "import tritonclient.grpc as grpcclient\n",
    "from nvtabular.workflow import Workflow\n",
    "# The following import is required even though it is not referenced in the program.\n",
    "# It loads artifacts that affect the schema and how the model is saved on disk.\n",
    "import merlin.models.tf as mm  # noqa: F401\n",
    "from merlin.schema.tags import Tags\n",
    "from merlin.systems.dag.ops.workflow import TransformWorkflow\n",
    "from merlin.systems.dag.ops.tensorflow import PredictTensorflow\n",
    "from merlin.systems.dag.ensemble import Ensemble\n",
    "from merlin.systems.triton import convert_df_to_triton_input\n",
    "from merlin.core.dispatch import get_lib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25a6ac04",
   "metadata": {},
   "source": [
    "We define the path for the saved workflow and model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9356fcd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_DIR = os.environ.get(\"BASE_DIR\", '/raid/data/criteo/test_dask/output/')\n",
    "original_data_path = os.environ.get(\"INPUT_FOLDER\", \"/raid/data/criteo/converted/criteo\")\n",
    "input_path = BASE_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4faf1583",
   "metadata": {},
   "source": [
    "We load the NVTabular workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "577fa9ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = Workflow.load(os.path.join(input_path, \"workflow\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e34b05",
   "metadata": {},
   "source": [
    "We need to remove the target columns from the workflow. The target columns are required to train our model. However, we do not know the targets during inference in the production environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fff2ccf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_columns = workflow.output_schema.select_by_tag(Tags.TARGET).column_names\n",
    "workflow.remove_inputs(label_columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2277937",
   "metadata": {},
   "source": [
    "We load the saved Merlin Models TensorFlow model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3de3bb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = tf.keras.models.load_model(os.path.join(input_path, \"dlrm\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fa11ab9",
   "metadata": {},
   "source": [
    "## Deploying Ensemble to Triton Inference Server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911541f6",
   "metadata": {},
   "source": [
    "We create our prediction pipeline:\n",
    "- the NVTabular workflow is executed via TransformWorkflow()\n",
    "- the TensorFlow model predict the output of the NVTabular workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe8a4ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_operators = (\n",
    "    workflow.input_schema.column_names >>\n",
    "    TransformWorkflow(workflow) >>\n",
    "    PredictTensorflow(model)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc0400cc",
   "metadata": {},
   "source": [
    "We create the Ensemble graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afc21189",
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble = Ensemble(serving_operators, workflow.input_schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1c899fd",
   "metadata": {},
   "source": [
    "We generate the Triton Inference Server artifacts and export them in the `export_path` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e27246ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "export_path = os.path.join(input_path, \"ensemble\")\n",
    "ens_conf, node_confs = ensemble.export(export_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ed386d",
   "metadata": {},
   "source": [
    "After we export the ensemble, we are ready to start the Triton Inference Server. The server is installed in the merlin-tensorflow-container. If you are not using one of our containers, then ensure it is installed in your environment. For more information, see the Triton Inference Server [documentation](https://github.com/triton-inference-server/server/blob/r22.03/README.md#documentation). \n",
    "\n",
    "You can start the server by running the following command:\n",
    "\n",
    "```shell\n",
    "tritonserver --model-repository=/raid/data/criteo/test_dask/output/ensemble --backend-config=tensorflow,version=2\n",
    "```\n",
    "\n",
    "For the `--model-repository` argument, specify the same value as the `export_path` that you specified previously in the `ensemble.export` method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "006b2581",
   "metadata": {},
   "source": [
    "## Get prediction from Triton Inference Server\n",
    "\n",
    "After we started Triton Inference Server and it loaded all models, we will send raw data as a request and receive the predictions.\n",
    "\n",
    "We read 3 example rows from the last parquet file from the raw data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93d1b8cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_lib = get_lib()\n",
    "input_cols = workflow.input_schema.column_names\n",
    "# read in data for request\n",
    "data = df_lib.read_parquet(\n",
    "    os.path.join(sorted(glob.glob(original_data_path + \"/*.parquet\"))[-1]),\n",
    "    columns=input_cols\n",
    ")\n",
    "batch = data[:3]\n",
    "batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60486140",
   "metadata": {},
   "source": [
    "We generate a Triton Inference Server request object. \n",
    "\n",
    "Currently, NA/None values are not supported for `int32` columns. As a workaround, we will NA values with 0. This will be updated in the future. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b833997d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create inputs and outputs\n",
    "inputs = convert_df_to_triton_input(workflow.input_schema, batch.fillna(0), grpcclient.InferInput)\n",
    "output_cols = ensemble.output_schema.column_names\n",
    "outputs = [\n",
    "    grpcclient.InferRequestedOutput(col)\n",
    "    for col in output_cols\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51a3c4ef",
   "metadata": {},
   "source": [
    "We send the request to Triton Inference Server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d29bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# send request to tritonserver\n",
    "with grpcclient.InferenceServerClient(\"localhost:8001\") as client:\n",
    "    response = client.infer(\"executor_model\", inputs, request_id=\"1\", outputs=outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bb2a8bd",
   "metadata": {},
   "source": [
    "We print out the predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "624d4ae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for col in ensemble.output_schema.column_names:\n",
    "    print(col, response.as_numpy(col), response.as_numpy(col).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa249750",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this example, we deployed a recommender system pipeline as an ensemble. First, NVTabular created features and afterwards, Merlin Models TensorFlow predicted the processed data. The DLRM architecture was used as a model. This process ensures that the training and production environments use the same feature engineering.\n",
    "\n",
    "## Next steps\n",
    "\n",
    "If you are interested in more details of the pipeline, we recommend to try out the [Merlin System example](https://github.com/NVIDIA-Merlin/systems/blob/stable/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb).\n",
    "\n",
    "In our Merlin repository, we provide [another end-to-end example](../Building-and-deploying-multi-stage-RecSys/) using a candidate retrieval and ranking model. In addition, we use approximate nearest neighbor and a feature store."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}