examples/getting-started-session-based/03-serving-session-based-model-torch-backend.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "97250792",
   "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": "0a2228da",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com//notebooks/dlsw-notebooks/remtting-started-session-based-03-serving-session-based-model-torch-backend/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Serving a Session-based Recommendation model with Torch Backend"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4127de28-a7ce-4ff7-8dcc-1575a70ca7c8",
   "metadata": {},
   "source": [
    "This notebook is created using the latest stable [merlin-pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-pytorch/tags) container.\n",
    "\n",
    "At this point, when you reach out to this notebook, we expect that you have already executed the `01-ETL-with-NVTabular.ipynb` and `02-session-based-XLNet-with-PyT.ipynb` notebooks, and saved the NVT workflow and the trained session-based model.\n",
    "\n",
    "In this notebook, you are going to learn how you can serve a trained Transformer-based PyTorch model on NVIDIA [Triton Inference Server](https://github.com/triton-inference-server/server)  (TIS) with Torch backend using [Merlin systems](https://github.com/NVIDIA-Merlin/systems) library. One common way to do inference with a trained model is to use TorchScript, an intermediate representation of a PyTorch model that can be run in Python as well as in a high performance environment like C++. [TorchScript](https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html) is actually the recommended model format for scaled inference and deployment. TIS [PyTorch (LibTorch) backend](https://github.com/triton-inference-server/pytorch_backend) is designed to run TorchScript models using the PyTorch C++ API.\n",
    "\n",
    "[Triton Inference Server](https://github.com/triton-inference-server/server) (TIS) simplifies the deployment of AI models at scale in production. TIS provides a cloud and edge inferencing solution optimized for both CPUs and GPUs. It supports a number of different machine learning frameworks such as TensorFlow and PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "599efc90",
   "metadata": {},
   "source": [
    "### Import required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3ba89970",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/dtypes/mappings/tf.py:52: UserWarning: Tensorflow dtype mappings did not load successfully due to an error: No module named 'tensorflow'\n",
      "  warn(f\"Tensorflow dtype mappings did not load successfully due to an error: {exc.msg}\")\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "import cudf\n",
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch \n",
    "\n",
    "from transformers4rec import torch as tr\n",
    "from merlin.io import Dataset\n",
    "\n",
    "from merlin.core.dispatch import make_df \n",
    "from merlin.systems.dag import Ensemble  \n",
    "from merlin.systems.dag.ops.pytorch import PredictPyTorch \n",
    "from merlin.systems.dag.ops.workflow import TransformWorkflow "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac2aebe1-66a9-4c71-9bbc-1874625bc4e8",
   "metadata": {
    "tags": []
   },
   "source": [
    "We define the paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5bc4a16b-32be-42b7-9b6f-e4f4cde3f345",
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DATA_DIR = os.environ.get(\"INPUT_DATA_DIR\", \"/workspace/data\")\n",
    "OUTPUT_DIR = os.environ.get(\"OUTPUT_DIR\", f\"{INPUT_DATA_DIR}/sessions_by_day\")\n",
    "model_path= os.environ.get(\"model_path\", f\"{INPUT_DATA_DIR}/saved_model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a510b6ef",
   "metadata": {},
   "source": [
    "### Set the schema object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30a0518a-eb01-4ac4-9c6d-36b328985765",
   "metadata": {},
   "source": [
    "We create the schema object by reading the processed train parquet file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9d1299fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from merlin.schema import Schema\n",
    "from merlin.io import Dataset\n",
    "\n",
    "train = Dataset(os.path.join(INPUT_DATA_DIR, \"processed_nvt/part_0.parquet\"))\n",
    "schema = train.schema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81764b19-4495-45c0-9cb5-6d937962d2bc",
   "metadata": {},
   "source": [
    "We need to load the saved model to be able to serve it on TIS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3bbc34a4-2277-4961-8061-59aadaa5116c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cloudpickle\n",
    "loaded_model = cloudpickle.load(\n",
    "                open(os.path.join(model_path, \"t4rec_model_class.pkl\"), \"rb\")\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c68ab3c7-a576-4fa0-a6f3-fe2bd11effc9",
   "metadata": {},
   "source": [
    "Switch the model to eval mode. We call `model.eval()` before tracing to set dropout and batch normalization layers to evaluation mode before running inference. Failing to do this might yield inconsistent inference results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e516a78d-2e1a-4124-ba46-f60b245d3329",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (heads): ModuleList(\n",
       "    (0): Head(\n",
       "      (body): SequentialBlock(\n",
       "        (0): TabularSequenceFeatures(\n",
       "          (to_merge): ModuleDict(\n",
       "            (continuous_module): SequentialBlock(\n",
       "              (0): ContinuousFeatures(\n",
       "                (filter_features): FilterFeatures()\n",
       "                (_aggregation): ConcatFeatures()\n",
       "              )\n",
       "              (1): SequentialBlock(\n",
       "                (0): DenseBlock(\n",
       "                  (0): Linear(in_features=2, out_features=64, bias=True)\n",
       "                  (1): ReLU(inplace=True)\n",
       "                )\n",
       "              )\n",
       "              (2): AsTabular()\n",
       "            )\n",
       "            (categorical_module): SequenceEmbeddingFeatures(\n",
       "              (filter_features): FilterFeatures()\n",
       "              (embedding_tables): ModuleDict(\n",
       "                (item_id-list): Embedding(495, 64, padding_idx=0)\n",
       "                (category-list): Embedding(172, 64, padding_idx=0)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (_aggregation): ConcatFeatures()\n",
       "          (projection_module): SequentialBlock(\n",
       "            (0): DenseBlock(\n",
       "              (0): Linear(in_features=192, out_features=100, bias=True)\n",
       "              (1): ReLU(inplace=True)\n",
       "            )\n",
       "          )\n",
       "          (_masking): MaskedLanguageModeling()\n",
       "        )\n",
       "        (1): SequentialBlock(\n",
       "          (0): DenseBlock(\n",
       "            (0): Linear(in_features=100, out_features=64, bias=True)\n",
       "            (1): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (2): TansformerBlock(\n",
       "          (transformer): XLNetModel(\n",
       "            (word_embedding): Embedding(1, 64)\n",
       "            (layer): ModuleList(\n",
       "              (0): XLNetLayer(\n",
       "                (rel_attn): XLNetRelativeAttention(\n",
       "                  (layer_norm): LayerNorm((64,), eps=0.03, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.3, inplace=False)\n",
       "                )\n",
       "                (ff): XLNetFeedForward(\n",
       "                  (layer_norm): LayerNorm((64,), eps=0.03, elementwise_affine=True)\n",
       "                  (layer_1): Linear(in_features=64, out_features=256, bias=True)\n",
       "                  (layer_2): Linear(in_features=256, out_features=64, bias=True)\n",
       "                  (dropout): Dropout(p=0.3, inplace=False)\n",
       "                )\n",
       "                (dropout): Dropout(p=0.3, inplace=False)\n",
       "              )\n",
       "              (1): XLNetLayer(\n",
       "                (rel_attn): XLNetRelativeAttention(\n",
       "                  (layer_norm): LayerNorm((64,), eps=0.03, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.3, inplace=False)\n",
       "                )\n",
       "                (ff): XLNetFeedForward(\n",
       "                  (layer_norm): LayerNorm((64,), eps=0.03, elementwise_affine=True)\n",
       "                  (layer_1): Linear(in_features=64, out_features=256, bias=True)\n",
       "                  (layer_2): Linear(in_features=256, out_features=64, bias=True)\n",
       "                  (dropout): Dropout(p=0.3, inplace=False)\n",
       "                )\n",
       "                (dropout): Dropout(p=0.3, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (dropout): Dropout(p=0.3, inplace=False)\n",
       "          )\n",
       "          (masking): MaskedLanguageModeling()\n",
       "        )\n",
       "      )\n",
       "      (prediction_task_dict): ModuleDict(\n",
       "        (next-item): NextItemPredictionTask(\n",
       "          (sequence_summary): SequenceSummary(\n",
       "            (summary): Identity()\n",
       "            (activation): Identity()\n",
       "            (first_dropout): Identity()\n",
       "            (last_dropout): Identity()\n",
       "          )\n",
       "          (metrics): ModuleList(\n",
       "            (0): NDCGAt()\n",
       "            (1): RecallAt()\n",
       "          )\n",
       "          (loss): CrossEntropyLoss()\n",
       "          (embeddings): SequenceEmbeddingFeatures(\n",
       "            (filter_features): FilterFeatures()\n",
       "            (embedding_tables): ModuleDict(\n",
       "              (item_id-list): Embedding(495, 64, padding_idx=0)\n",
       "              (category-list): Embedding(172, 64, padding_idx=0)\n",
       "            )\n",
       "          )\n",
       "          (item_embedding_table): Embedding(495, 64, padding_idx=0)\n",
       "          (masking): MaskedLanguageModeling()\n",
       "          (pre): Block(\n",
       "            (module): NextItemPredictionTask(\n",
       "              (item_embedding_table): Embedding(495, 64, padding_idx=0)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = loaded_model.cuda()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "409c17c2-c81d-4f41-9577-bd380ae10921",
   "metadata": {},
   "source": [
    "### Trace the model\n",
    "\n",
    "We serve the model with the PyTorch backend that is used to execute TorchScript models. All models created in PyTorch using the python API must be traced/scripted to produce a TorchScript model. For tracing the model, we use [torch.jit.trace](https://pytorch.org/docs/stable/generated/torch.jit.trace.html) api that takes the model as a Python function or torch.nn.Module, and an example input  that will be passed to the function while tracing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6273c8e5-db62-4cc6-a4f3-945155a463d6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_paths = os.path.join(OUTPUT_DIR, f\"{1}/train.parquet\")\n",
    "dataset = Dataset(train_paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ddb4b5f-9b8c-487b-9fa1-cd4b6d3c090f",
   "metadata": {},
   "source": [
    "Create a dict of tensors to feed it as example inputs in the `torch.jit.trace()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d8df9e4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from merlin.table import TensorTable, TorchColumn\n",
    "from merlin.table.conversions import convert_col\n",
    "\n",
    "df = cudf.read_parquet(train_paths, columns=model.input_schema.column_names)\n",
    "table = TensorTable.from_df(df.loc[:100])\n",
    "for column in table.columns:\n",
    "    table[column] = convert_col(table[column], TorchColumn)\n",
    "model_input_dict = table.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6b4731f0-e0df-48da-b405-3fead25cadd0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([306,   5,  40,  17,  43,  20,  69,   8,  57, 137,  35,  37,  85,  65,\n",
       "          5,  28,   9, 153,  74,  53,  15, 173,  59,  32,  11,  21,  23,  23,\n",
       "          9,  15,  12,  69,  37,  16,   6,  22,  39,  20,  22,  95,  40,   7,\n",
       "         25,  32,  17,   8,  26,  32,  33,  18,  12,  10,  41,  14,  28,  56,\n",
       "         30,  21,  16,  42,  13,  83,  65,  46, 105,  38,  11,   3,   3,  14,\n",
       "          9,  36, 116,  15,  15,  23,   8,  16,  68, 151,  60,  18,  48,  19,\n",
       "         16,   4,  37, 246, 169,  21,  16, 116,  27,   4,  19,  76,   6,  31,\n",
       "        153,  38,  35,  11,  38,   3,  73,  38,  74,   6,   7,  12,  18,  10,\n",
       "         54,  11,  29,   5,  24,  11,  20,   3,  17,  42,  26,  24,  30,  26,\n",
       "         62,  89,  12,  38,  18,   3,  10,  18,  15, 131,  19,   6,  51,  60,\n",
       "         10,   3,  14,  22,  21,  39,  44, 221,  88,  14,  16,  80,   5,  16,\n",
       "         21,  81,  27,   8,  20,  49,  32,  83,  49,  19,   3,  17,   8,  10,\n",
       "         29,  62,  94,  38,  15,  11,  12,  16,  10,  31,   7,  53,   3,  42,\n",
       "         38,  25,   5,  62,  20,  73,  48,   6,  12,  19,  15,  38,  30,   9,\n",
       "         82,  31,  49,  64,  22,  38,  10,  56,  11,  13,   3,  14,  39,  18,\n",
       "         47,  65,  18,  15,   9,  74,   3,  50,  37,  22,  66,  47,  23,  17,\n",
       "          8,  21,  35,   7,  12,  16,  21,  26,  31,  13,  20,   9, 193,  49,\n",
       "          9,  62,  51,  45,  90,  14,  47,   9,  73,  16,   3,  62,  24,  82,\n",
       "          7,  14,  37,  29,  26,  42,   6,  90,   3,  10,  33,   7,   7,  10,\n",
       "         31,  10,  12,  21,  55,  25,  21,   3,  20,  24,  25,   4,   3,  52,\n",
       "          5,   5,  10,  12,  37, 162,  31,   5, 119,   5,  24,  65,   4,  10,\n",
       "         46,  86,   5,  58,  15,  48,  66,  14,  23,  12,  13,   6,  48,   8,\n",
       "         22,  95,   5,  42,  86, 108,  26,   7,  80,  54,  63,  12, 147, 177,\n",
       "         17,  18,  24,  15,  40,   5,  40,   7,   6,  63,   4,  18, 123,  33,\n",
       "         36,  25,  40,  18,  16,  10,  18,  26,  21,  59,  44,  12,  28,  30,\n",
       "        134,   7,  21,   8,   7,  32,  41,  60,  52,  25,  36,   6,  45,  39,\n",
       "         16,  20,  95,   8,  56,  53,  48,  17,  14,   3,  46,  35,  17,  12,\n",
       "         30,   8,   5,  54,  75,  96,   4,  43,   8,  61,   4,   8,  34,  30,\n",
       "         34,  49,  29,  92,   6,  28,  26,  22,  46,  20,  11,  14,  13,  75,\n",
       "         22,  21,  17, 166,   4,  87,   5,  11,  37,  26,  23],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_input_dict['item_id-list__values']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "43e789a3-2423-44eb-ae5d-0c154557424f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "traced_model = torch.jit.trace(model, model_input_dict, strict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a404bdb6-6e19-4899-8172-214adef384a8",
   "metadata": {},
   "source": [
    "Generate model input and output schemas to feed in the `PredictPyTorch` operator below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c6a814aa-4954-404b-bd2d-161ed8066f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_schema = model.input_schema\n",
    "output_schema = model.output_schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "757cd0c5-f581-488b-a8de-b8d1188820d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>tags</th>\n",
       "      <th>dtype</th>\n",
       "      <th>is_list</th>\n",
       "      <th>is_ragged</th>\n",
       "      <th>properties.value_count.min</th>\n",
       "      <th>properties.value_count.max</th>\n",
       "      <th>properties.num_buckets</th>\n",
       "      <th>properties.freq_threshold</th>\n",
       "      <th>properties.max_size</th>\n",
       "      <th>properties.cat_path</th>\n",
       "      <th>properties.embedding_sizes.cardinality</th>\n",
       "      <th>properties.embedding_sizes.dimension</th>\n",
       "      <th>properties.domain.min</th>\n",
       "      <th>properties.domain.max</th>\n",
       "      <th>properties.domain.name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>weekday_sin-list</td>\n",
       "      <td>(Tags.LIST, Tags.CONTINUOUS)</td>\n",
       "      <td>DType(name='float32', element_type=&lt;ElementTyp...</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>2</td>\n",
       "      <td>16</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>age_days-list</td>\n",
       "      <td>(Tags.LIST, Tags.CONTINUOUS)</td>\n",
       "      <td>DType(name='float32', element_type=&lt;ElementTyp...</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>2</td>\n",
       "      <td>16</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>item_id-list</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.LIST, Tags.ITEM, Tags.ID)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>2</td>\n",
       "      <td>16</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.item_id.parquet</td>\n",
       "      <td>495.0</td>\n",
       "      <td>52.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>494.0</td>\n",
       "      <td>item_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>category-list</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.LIST)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>2</td>\n",
       "      <td>16</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.category.parquet</td>\n",
       "      <td>172.0</td>\n",
       "      <td>29.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>171.0</td>\n",
       "      <td>category</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'weekday_sin-list', 'tags': {<Tags.LIST: 'list'>, <Tags.CONTINUOUS: 'continuous'>}, 'properties': {'value_count': {'min': 2, 'max': 16}}, 'dtype': DType(name='float32', element_type=<ElementType.Float: 'float'>, element_size=32, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None), Dimension(min=2, max=16)))), 'is_list': True, 'is_ragged': True}, {'name': 'age_days-list', 'tags': {<Tags.LIST: 'list'>, <Tags.CONTINUOUS: 'continuous'>}, 'properties': {'value_count': {'min': 2, 'max': 16}}, 'dtype': DType(name='float32', element_type=<ElementType.Float: 'float'>, element_size=32, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None), Dimension(min=2, max=16)))), 'is_list': True, 'is_ragged': True}, {'name': 'item_id-list', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.LIST: 'list'>, <Tags.ITEM: 'item'>, <Tags.ID: 'id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'cat_path': './/categories/unique.item_id.parquet', 'embedding_sizes': {'cardinality': 495.0, 'dimension': 52.0}, 'domain': {'min': 0, 'max': 494, 'name': 'item_id'}, 'value_count': {'min': 2, 'max': 16}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None), Dimension(min=2, max=16)))), 'is_list': True, 'is_ragged': True}, {'name': 'category-list', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.LIST: 'list'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'cat_path': './/categories/unique.category.parquet', 'embedding_sizes': {'cardinality': 172.0, 'dimension': 29.0}, 'domain': {'min': 0, 'max': 171, 'name': 'category'}, 'value_count': {'min': 2, 'max': 16}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None), Dimension(min=2, max=16)))), 'is_list': True, 'is_ragged': True}]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_schema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f8a8bfc-c6e1-44dd-ab65-cbdab2135e8e",
   "metadata": {},
   "source": [
    "Let's create a folder that we can store the exported models and the config files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6b2deb2b-e223-4b5d-b655-810e1aefa7e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "ens_model_path = os.environ.get(\"ens_model_path\", f\"{INPUT_DATA_DIR}/models\")\n",
    "# Make sure we have a clean stats space for Dask\n",
    "if os.path.isdir(ens_model_path):\n",
    "    shutil.rmtree(ens_model_path)\n",
    "os.mkdir(ens_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "845c935c-8bfc-4c92-bba2-2a915edbfc68",
   "metadata": {},
   "source": [
    "We want to serve NVT model and our trained session-based model together as an ensemble to the Triton Inference Server. That way we can send raw requests to Triton and return back item scores per session. For that we need to load our save workflow first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0360f457-a198-42c0-b860-3356f7956ab9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['item_id', 'category', 'day', 'age_days', 'weekday_sin', 'session_id']\n"
     ]
    }
   ],
   "source": [
    "from nvtabular.workflow import Workflow\n",
    "workflow = Workflow.load(os.path.join(INPUT_DATA_DIR, \"workflow_etl\"))\n",
    "print(workflow.input_schema.column_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3449615-2120-402d-b5c3-1544ee3224dd",
   "metadata": {},
   "source": [
    "For transforming the raw input features during inference, we use [TransformWorkflow](https://github.com/NVIDIA-Merlin/systems/blob/stable/merlin/systems/dag/ops/workflow.py) operator that ensures the workflow is correctly saved and packaged with the required config so the server will know how to load it. We use [PredictPyTorch](https://github.com/NVIDIA-Merlin/systems/blob/stable/merlin/systems/dag/ops/pytorch.py) operator that takes a pytorch model and packages it correctly for tritonserver to run on the PyTorch backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4f96597c-1c05-4fb0-ad3e-c55c21599158",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/merlin/systems/dag/node.py:100: UserWarning: Operator 'TransformWorkflow' is producing the output column 'session_id', which is not being used by any downstream operator in the ensemble graph.\n",
      "  warnings.warn(\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/systems/dag/node.py:100: UserWarning: Operator 'TransformWorkflow' is producing the output column 'day-first', which is not being used by any downstream operator in the ensemble graph.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "torch_op = workflow.input_schema.column_names >> TransformWorkflow(workflow) >> PredictPyTorch(\n",
    "    traced_model, input_schema, output_schema\n",
    ")\n",
    "\n",
    "ensemble = Ensemble(torch_op, workflow.input_schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f36b0f90-5392-4d7a-9a48-57737cf63cd1",
   "metadata": {},
   "source": [
    "The last step is to create the ensemble artifacts that Triton Inference Server can consume. To make these artifacts, we import the Ensemble class. The class is responsible for interpreting the graph and exporting the correct files for the server.\n",
    "\n",
    "When we create an `Ensemble` object we supply the graph and a schema representing the starting input of the graph. The inputs to the ensemble graph are the inputs to the first operator of out graph. After we created the Ensemble we export the graph, supplying an export path for the `ensemble.export` function. This returns an ensemble config which represents the entire inference pipeline and a list of node-specific configs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5b0d14bb-7765-45e8-8fd0-9d508dc3ec14",
   "metadata": {},
   "outputs": [],
   "source": [
    "ens_config, node_configs = ensemble.export(ens_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a3ba86eb-ca25-4a0c-9daf-61c9911b29ab",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>tags</th>\n",
       "      <th>dtype</th>\n",
       "      <th>is_list</th>\n",
       "      <th>is_ragged</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>item_id</td>\n",
       "      <td>()</td>\n",
       "      <td>DType(name='int32', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>category</td>\n",
       "      <td>()</td>\n",
       "      <td>DType(name='int32', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>day</td>\n",
       "      <td>()</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>age_days</td>\n",
       "      <td>()</td>\n",
       "      <td>DType(name='float32', element_type=&lt;ElementTyp...</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>weekday_sin</td>\n",
       "      <td>()</td>\n",
       "      <td>DType(name='float32', element_type=&lt;ElementTyp...</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>session_id</td>\n",
       "      <td>()</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'item_id', 'tags': set(), 'properties': {}, 'dtype': DType(name='int32', element_type=<ElementType.Int: 'int'>, element_size=32, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'category', 'tags': set(), 'properties': {}, 'dtype': DType(name='int32', element_type=<ElementType.Int: 'int'>, element_size=32, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'day', 'tags': set(), 'properties': {}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'age_days', 'tags': set(), 'properties': {}, 'dtype': DType(name='float32', element_type=<ElementType.Float: 'float'>, element_size=32, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'weekday_sin', 'tags': set(), 'properties': {}, 'dtype': DType(name='float32', element_type=<ElementType.Float: 'float'>, element_size=32, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'session_id', 'tags': set(), 'properties': {}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ensemble.input_schema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a36169a5-f218-44b5-b034-7d299ce718ed",
   "metadata": {},
   "source": [
    "## Starting Triton Server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb507766-ac9b-4c8c-8339-c1c951648428",
   "metadata": {},
   "source": [
    "It is time to deploy all the models as an ensemble model to Triton Inference Serve TIS. After we export the ensemble, we are ready to start the TIS. You can start triton server by using the following command on your terminal:\n",
    "\n",
    "`tritonserver --model-repository=<ensemble_export_path>`\n",
    "\n",
    "For the `--model-repository` argument, specify the same path as the export_path that you specified previously in the `ensemble.export` method. This command will launch the server and load all the models to the server. Once all the models are loaded successfully, you should see READY status printed out in the terminal for each loaded model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "46a86c8d-9ec1-4422-8f8c-4d49e83f6783",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "client created.\n"
     ]
    }
   ],
   "source": [
    "import tritonclient.http as client\n",
    "\n",
    "# Create a triton client\n",
    "try:\n",
    "    triton_client = client.InferenceServerClient(url=\"localhost:8000\", verbose=True)\n",
    "    print(\"client created.\")\n",
    "except Exception as e:\n",
    "    print(\"channel creation failed: \" + str(e))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bab41dd0-0155-4d97-8bdf-2451177d46f1",
   "metadata": {},
   "source": [
    "After we create the client and verified it is connected to the server instance, we can communicate with the server and ensure all the models are loaded correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dda3f852-a019-4bf1-831b-f63b750a1192",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GET /v2/health/live, headers None\n",
      "<HTTPSocketPoolResponse status=200 headers={'content-length': '0', 'content-type': 'text/plain'}>\n",
      "POST /v2/repository/index, headers None\n",
      "\n",
      "<HTTPSocketPoolResponse status=200 headers={'content-type': 'application/json', 'content-length': '188'}>\n",
      "bytearray(b'[{\"name\":\"0_transformworkflowtriton\",\"version\":\"1\",\"state\":\"READY\"},{\"name\":\"1_predictpytorchtriton\",\"version\":\"1\",\"state\":\"READY\"},{\"name\":\"executor_model\",\"version\":\"1\",\"state\":\"READY\"}]')\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'name': '0_transformworkflowtriton', 'version': '1', 'state': 'READY'},\n",
       " {'name': '1_predictpytorchtriton', 'version': '1', 'state': 'READY'},\n",
       " {'name': 'executor_model', 'version': '1', 'state': 'READY'}]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ensure triton is in a good state\n",
    "triton_client.is_server_live()\n",
    "triton_client.get_model_repository_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "415146b1-9e9c-4c72-a9b9-697041ba27ef",
   "metadata": {},
   "source": [
    "### Send request to Triton and get the response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f218f19-63d8-498f-aec8-b5dfa56ca3f3",
   "metadata": {},
   "source": [
    "The last step of a machine learning (ML)/deep learning (DL) pipeline is to deploy the model to production, and get responses for a given query or a set of queries.\n",
    "In this section, we generate a dataframe that we can serve as a request to TIS. We do serve the raw dataframe and in the production setting, we want to transform the input data as done during training (ETL). We need to apply the same mean/std for continuous features and use the same categorical mapping to convert the categories to continuous integer before we use the deployed DL model for a prediction."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e3abf91-67e8-4f04-842c-f128a3379569",
   "metadata": {},
   "source": [
    "Let's generate a dataframe with raw input values. We can send this dataframe to Triton as a request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0acd5649-31fe-4f3f-87a2-2607477638b5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   session_id  item_id  category  age_days  weekday_sin  day\n",
      "0       79856        3         2  0.327276     0.080060    2\n",
      "1       74117        6         4  0.012172     0.147716    1\n"
     ]
    }
   ],
   "source": [
    "NUM_ROWS =1000\n",
    "long_tailed_item_distribution = np.clip(np.random.lognormal(3., 1., int(NUM_ROWS)).astype(np.int32), 1, 50000)\n",
    "# generate random item interaction features \n",
    "df = pd.DataFrame(np.random.randint(70000, 90000, int(NUM_ROWS)), columns=['session_id'])\n",
    "df['item_id'] = long_tailed_item_distribution\n",
    "\n",
    "# generate category mapping for each item-id\n",
    "df['category'] = pd.cut(df['item_id'], bins=334, labels=np.arange(1, 335)).astype(np.int32)\n",
    "df['age_days'] = np.random.uniform(0, 1, int(NUM_ROWS)).astype(np.float32)\n",
    "df['weekday_sin']= np.random.uniform(0, 1, int(NUM_ROWS)).astype(np.float32)\n",
    "\n",
    "# generate day mapping for each session \n",
    "map_day = dict(zip(df.session_id.unique(), np.random.randint(1, 10, size=(df.session_id.nunique()))))\n",
    "df['day'] =  df.session_id.map(map_day)\n",
    "\n",
    "print(df.head(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52832509-6444-4fd2-8834-a02eec429e14",
   "metadata": {},
   "source": [
    "Once our models are successfully loaded to the TIS, we can now easily send a request to TIS and get a response for our query with send_triton_request utility function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "42091c25-7676-414e-bb8c-8432aeb58297",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'next-item': array([[-3.9399953, -2.632081 , -4.2211075, ..., -3.6699016, -3.673493 ,\n",
       "         -3.1244578],\n",
       "        [-3.940445 , -2.6335964, -4.2203593, ..., -3.671566 , -3.6745713,\n",
       "         -3.1240335],\n",
       "        [-3.9393594, -2.6300201, -4.222065 , ..., -3.6674871, -3.672068 ,\n",
       "         -3.1251097],\n",
       "        ...,\n",
       "        [-3.9396427, -2.6304667, -4.2218847, ..., -3.6677885, -3.6724825,\n",
       "         -3.1250875],\n",
       "        [-3.939829 , -2.6316376, -4.221267 , ..., -3.6693997, -3.6732295,\n",
       "         -3.1245873],\n",
       "        [-3.9399223, -2.631995 , -4.2210817, ..., -3.669589 , -3.6734715,\n",
       "         -3.1244512]], dtype=float32)}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from merlin.systems.triton.utils import send_triton_request\n",
    "response = send_triton_request(workflow.input_schema, df, output_schema.column_names, endpoint=\"localhost:8001\")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0fa425a4-9c00-45ed-a4b1-fd75ca4bf819",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(28, 495)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response['next-item'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0581195-d496-4536-a93b-3071ed9088ea",
   "metadata": {},
   "source": [
    "We return a response for each request in the df. Each row in the `response['next-item']` array corresponds to the logit values per item in the catalog, and one logit score corresponding to the null, OOV and padded items. The first score of each array in each row corresponds to the score for the padded item, OOV or null item. Note that we dont have OOV or null items in our syntheticall generated datasets.\n",
    "\n",
    "This is the end of this suit of examples. You successfully performed feature engineering with NVTabular trained transformer architecture based session-based recommendation models with Transformers4Rec deployed the saved workflow and the trained model to Triton Inference Server with Torch backend, sent request and got responses from the server. If you would like to learn how to serve a TF4Rec model with Python backend please visit this [example](https://github.com/NVIDIA-Merlin/Transformers4Rec/blob/stable/examples/end-to-end-session-based/02-End-to-end-session-based-with-Yoochoose-PyT.ipynb)."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "7b543a88d374ac88bf8df97911b380f671b13649694a5b49eb21e60fd27eb479"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}