hps_trt/notebooks/demo_for_tf_trained_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7cb13422",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_hugectr_hps-hps-tensorflow-triton-deployment/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# HPS TensorRT Plugin Demo for TensorFlow Trained Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc7b66b1",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This notebook demonstrates how to build and deploy the HPS-integrated TensorRT engine for the model trained with TensorFlow.\n",
    "\n",
    "For more details about HPS, please refer to [HugeCTR Hierarchical Parameter Server (HPS)](https://nvidia-merlin.github.io/HugeCTR/master/hierarchical_parameter_server/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95b12df6",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "### Use NGC\n",
    "\n",
    "The HPS TensorRT plugin is preinstalled in the 24.06 and later [Merlin HugeCTR Container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-hugectr): `nvcr.io/nvidia/merlin/merlin-hugectr:24.06`.\n",
    "\n",
    "You can check the existence of the required libraries by running the following Python code after launching this container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6cd5edcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ctypes\n",
    "plugin_lib_name = \"/usr/local/hps_trt/lib/libhps_plugin.so\"\n",
    "plugin_handle = ctypes.CDLL(plugin_lib_name, mode=ctypes.RTLD_GLOBAL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605963dc",
   "metadata": {},
   "source": [
    "## Configurations\n",
    "\n",
    "First of all we specify the required configurations, e.g., the arguments needed for generating the dataset, the model parameters and the paths to save the model. We will use DLRM model which has one embedding table, bottom MLP layers, interaction layer and top MLP layers. Please note that the input to the embedding layer will be a dense key tensor of int32."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eed1333f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-21 03:16:46.032517: I tensorflow/core/platform/cpu_feature_guard.cc:183] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import struct\n",
    "\n",
    "args = dict()\n",
    "\n",
    "args[\"gpu_num\"] = 1                               # the number of available GPUs\n",
    "args[\"iter_num\"] = 50                             # the number of training iteration\n",
    "args[\"slot_num\"] = 26                             # the number of feature fields in this embedding layer\n",
    "args[\"embed_vec_size\"] = 128                      # the dimension of embedding vectors\n",
    "args[\"dense_dim\"] = 13                            # the dimension of dense features\n",
    "args[\"global_batch_size\"] = 1024                  # the globally batchsize for all GPUs\n",
    "args[\"max_vocabulary_size\"] = 260000\n",
    "args[\"vocabulary_range_per_slot\"] = [[i*10000, (i+1)*10000] for i in range(26)]\n",
    "args[\"combiner\"] = \"mean\"\n",
    "\n",
    "args[\"ps_config_file\"] = \"dlrm_tf.json\"\n",
    "args[\"embedding_table_path\"] = \"dlrm_tf_sparse.model\"\n",
    "args[\"saved_path\"] = \"dlrm_tf_saved_model\"\n",
    "args[\"np_key_type\"] = np.int32\n",
    "args[\"np_vector_type\"] = np.float32\n",
    "args[\"tf_key_type\"] = tf.int32\n",
    "args[\"tf_vector_type\"] = tf.float32\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \",\".join(map(str, range(args[\"gpu_num\"])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a127fc7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_random_samples(num_samples, vocabulary_range_per_slot, dense_dim, key_dtype = args[\"np_key_type\"]):\n",
    "    keys = list()\n",
    "    for vocab_range in vocabulary_range_per_slot:\n",
    "        keys_per_slot = np.random.randint(low=vocab_range[0], high=vocab_range[1], size=(num_samples, 1), dtype=key_dtype)\n",
    "        keys.append(keys_per_slot)\n",
    "    keys = np.concatenate(np.array(keys), axis = 1)\n",
    "    numerical_features = np.random.random((num_samples, dense_dim)).astype(np.float32)\n",
    "    labels = np.random.randint(low=0, high=2, size=(num_samples, 1))\n",
    "    return keys, numerical_features, labels\n",
    "\n",
    "def tf_dataset(keys, numerical_features, labels, batchsize):\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((keys, numerical_features, labels))\n",
    "    dataset = dataset.batch(batchsize, drop_remainder=True)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d597a062",
   "metadata": {},
   "source": [
    "## Train with native TF layers\n",
    "\n",
    "We define the model graph for training with native TF layers, i.e., `tf.nn.embedding_lookup`, `tf.keras.layers.Dense` and so on. We can then train the model and extract the trained weights of the embedding table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a8e17a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(tf.keras.layers.Layer):\n",
    "    def __init__(self,\n",
    "                arch,\n",
    "                activation='relu',\n",
    "                out_activation=None,\n",
    "                **kwargs):\n",
    "        super(MLP, self).__init__(**kwargs)\n",
    "        self.layers = []\n",
    "        index = 0\n",
    "        for units in arch[:-1]:\n",
    "            self.layers.append(tf.keras.layers.Dense(units, activation=activation, name=\"{}_{}\".format(kwargs['name'], index)))\n",
    "            index+=1\n",
    "        self.layers.append(tf.keras.layers.Dense(arch[-1], activation=out_activation, name=\"{}_{}\".format(kwargs['name'], index)))\n",
    "\n",
    "            \n",
    "    def call(self, inputs, training=True):\n",
    "        x = self.layers[0](inputs)\n",
    "        for layer in self.layers[1:]:\n",
    "            x = layer(x)\n",
    "        return x\n",
    "\n",
    "class SecondOrderFeatureInteraction(tf.keras.layers.Layer):\n",
    "    def __init__(self):\n",
    "        super(SecondOrderFeatureInteraction, self).__init__()\n",
    "\n",
    "    def call(self, inputs, num_feas):\n",
    "        dot_products = tf.reshape(tf.matmul(inputs, inputs, transpose_b=True), (-1, num_feas * num_feas))\n",
    "        indices = tf.constant([i * num_feas + j for j in range(1, num_feas) for i in range(j)])\n",
    "        flat_interactions = tf.gather(dot_products, indices, axis=1)\n",
    "        return flat_interactions\n",
    "\n",
    "class DLRM(tf.keras.models.Model):\n",
    "    def __init__(self,\n",
    "                 init_tensors,\n",
    "                 embed_vec_size,\n",
    "                 slot_num,\n",
    "                 dense_dim,\n",
    "                 arch_bot,\n",
    "                 arch_top,\n",
    "                 **kwargs):\n",
    "        super(DLRM, self).__init__(**kwargs)\n",
    "        \n",
    "\n",
    "        self.init_tensors = init_tensors\n",
    "        self.params = tf.Variable(initial_value=tf.concat(self.init_tensors, axis=0))\n",
    "        \n",
    "        self.embed_vec_size = embed_vec_size\n",
    "        self.slot_num = slot_num\n",
    "        self.dense_dim = dense_dim\n",
    "    \n",
    "        self.bot_nn = MLP(arch_bot, name = \"bottom\", out_activation='relu')\n",
    "        self.top_nn = MLP(arch_top, name = \"top\", out_activation='sigmoid')\n",
    "        self.interaction_op = SecondOrderFeatureInteraction()\n",
    "\n",
    "        self.interaction_out_dim = self.slot_num * (self.slot_num+1) // 2\n",
    "        self.reshape_layer1 = tf.keras.layers.Reshape((1, arch_bot[-1]), name = \"reshape1\")\n",
    "        self.concat1 = tf.keras.layers.Concatenate(axis=1, name = \"concat1\")\n",
    "        self.concat2 = tf.keras.layers.Concatenate(axis=1, name = \"concat2\")\n",
    "            \n",
    "    def call(self, inputs, training=True):\n",
    "        categorical_features = inputs[\"keys\"]\n",
    "        numerical_features = inputs[\"numerical_features\"]\n",
    "        \n",
    "        embedding_vector = tf.nn.embedding_lookup(params=self.params, ids=categorical_features)\n",
    "        dense_x = self.bot_nn(numerical_features)\n",
    "        concat_features = self.concat1([embedding_vector, self.reshape_layer1(dense_x)])\n",
    "        \n",
    "        Z = self.interaction_op(concat_features, self.slot_num+1)\n",
    "        z = self.concat2([dense_x, Z])\n",
    "        logit = self.top_nn(z)\n",
    "        return logit\n",
    "\n",
    "    def summary(self):\n",
    "        inputs = {\"keys\": tf.keras.Input(shape=(self.slot_num, ), dtype=args[\"tf_key_type\"], name=\"keys\"), \n",
    "                  \"numerical_features\": tf.keras.Input(shape=(self.dense_dim, ), dtype=tf.float32, name=\"numrical_features\")}\n",
    "        model = tf.keras.models.Model(inputs=inputs, outputs=self.call(inputs))\n",
    "        return model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fe1b2fbe",
   "metadata": {},
   "outputs": [],
   "source": [
    " def train(args):\n",
    "    init_tensors = np.ones(shape=[args[\"max_vocabulary_size\"], args[\"embed_vec_size\"]], dtype=args[\"np_vector_type\"])\n",
    "    \n",
    "    model = DLRM(init_tensors, args[\"embed_vec_size\"], args[\"slot_num\"], args[\"dense_dim\"],\n",
    "                arch_bot = [512, 256, args[\"embed_vec_size\"]],\n",
    "                arch_top = [1024, 1024, 512, 256, 1],\n",
    "                name = \"dlrm\")\n",
    "    model.summary()\n",
    "    optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)\n",
    "    loss_fn = tf.keras.losses.BinaryCrossentropy()\n",
    "    \n",
    "    def _train_step(inputs, labels):\n",
    "        with tf.GradientTape() as tape:\n",
    "            logit = model(inputs)\n",
    "            loss = loss_fn(labels, logit)\n",
    "        grads = tape.gradient(loss, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
    "        return loss, logit\n",
    "\n",
    "    keys, numerical_features, labels = generate_random_samples(args[\"global_batch_size\"]  * args[\"iter_num\"], args[\"vocabulary_range_per_slot\"], args[\"dense_dim\"], args[\"np_key_type\"])\n",
    "    dataset = tf_dataset(keys, numerical_features, labels, args[\"global_batch_size\"])\n",
    "    for i, (keys, numerical_features, labels) in enumerate(dataset):\n",
    "        inputs = {\"keys\": keys, \"numerical_features\": numerical_features}\n",
    "        loss, logit = _train_step(inputs, labels)\n",
    "        print(\"-\"*20, \"Step {}, loss: {}\".format(i, loss),  \"-\"*20)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bc8ecf24",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-21 03:16:55.963734: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1638] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 30974 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB, pci bus id: 0000:06:00.0, compute capability: 7.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:The following Variables were used in a Lambda layer's call (tf.compat.v1.nn.embedding_lookup), but are not present in its tracked objects:   <tf.Variable 'Variable:0' shape=(260000, 128) dtype=float32>. This is a strong indication that the Lambda layer should be rewritten as a subclassed Layer.\n",
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " numrical_features (InputLayer)  [(None, 13)]        0           []                               \n",
      "                                                                                                  \n",
      " bottom (MLP)                   (None, 128)          171392      ['numrical_features[0][0]']      \n",
      "                                                                                                  \n",
      " keys (InputLayer)              [(None, 26)]         0           []                               \n",
      "                                                                                                  \n",
      " tf.compat.v1.nn.embedding_look  (None, 26, 128)     0           ['keys[0][0]']                   \n",
      " up (TFOpLambda)                                                                                  \n",
      "                                                                                                  \n",
      " reshape1 (Reshape)             (None, 1, 128)       0           ['bottom[0][0]']                 \n",
      "                                                                                                  \n",
      " concat1 (Concatenate)          (None, 27, 128)      0           ['tf.compat.v1.nn.embedding_looku\n",
      "                                                                 p[0][0]',                        \n",
      "                                                                  'reshape1[0][0]']               \n",
      "                                                                                                  \n",
      " second_order_feature_interacti  (None, 351)         0           ['concat1[0][0]']                \n",
      " on (SecondOrderFeatureInteract                                                                   \n",
      " ion)                                                                                             \n",
      "                                                                                                  \n",
      " concat2 (Concatenate)          (None, 479)          0           ['bottom[0][0]',                 \n",
      "                                                                  'second_order_feature_interactio\n",
      "                                                                 n[0][0]']                        \n",
      "                                                                                                  \n",
      " top (MLP)                      (None, 1)            2197505     ['concat2[0][0]']                \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2,368,897\n",
      "Trainable params: 2,368,897\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-21 03:16:57.578464: I tensorflow/core/common_runtime/executor.cc:1209] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_2' with dtype int64 and shape [51200,1]\n",
      "\t [[{{node Placeholder/_2}}]]\n",
      "2023-08-21 03:16:58.892396: I tensorflow/compiler/xla/service/service.cc:169] XLA service 0x55e0fdfeb330 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:\n",
      "2023-08-21 03:16:58.892450: I tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (0): Tesla V100-SXM2-32GB, Compute Capability 7.0\n",
      "2023-08-21 03:16:58.897903: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.\n",
      "2023-08-21 03:16:59.379151: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:424] Loaded cuDNN version 8902\n",
      "2023-08-21 03:16:59.502058: I ./tensorflow/compiler/jit/device_compiler.h:180] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:5 out of the last 5 calls to <function _BaseOptimizer._update_step_xla at 0x7fa9660adab0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "WARNING:tensorflow:6 out of the last 6 calls to <function _BaseOptimizer._update_step_xla at 0x7fa9660adab0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "-------------------- Step 0, loss: 39.68028259277344 --------------------\n",
      "-------------------- Step 1, loss: 2571352064.0 --------------------\n",
      "-------------------- Step 2, loss: 639234.5 --------------------\n",
      "-------------------- Step 3, loss: 4132346.75 --------------------\n",
      "-------------------- Step 4, loss: 20792958.0 --------------------\n",
      "-------------------- Step 5, loss: 5957.8994140625 --------------------\n",
      "-------------------- Step 6, loss: 231005.96875 --------------------\n",
      "-------------------- Step 7, loss: 185315.3125 --------------------\n",
      "-------------------- Step 8, loss: 151740.75 --------------------\n",
      "-------------------- Step 9, loss: 43695.6640625 --------------------\n",
      "-------------------- Step 10, loss: 45556.24609375 --------------------\n",
      "-------------------- Step 11, loss: 131654.78125 --------------------\n",
      "-------------------- Step 12, loss: 1.8805829286575317 --------------------\n",
      "-------------------- Step 13, loss: 49121.47265625 --------------------\n",
      "-------------------- Step 14, loss: 60609.62109375 --------------------\n",
      "-------------------- Step 15, loss: 676294.375 --------------------\n",
      "-------------------- Step 16, loss: 31208.66015625 --------------------\n",
      "-------------------- Step 17, loss: 156789.65625 --------------------\n",
      "-------------------- Step 18, loss: 103213.1015625 --------------------\n",
      "-------------------- Step 19, loss: 22.394046783447266 --------------------\n",
      "-------------------- Step 20, loss: 10789.5703125 --------------------\n",
      "-------------------- Step 21, loss: 2716.05859375 --------------------\n",
      "-------------------- Step 22, loss: 139559.96875 --------------------\n",
      "-------------------- Step 23, loss: 130419.9453125 --------------------\n",
      "-------------------- Step 24, loss: 13583.6923828125 --------------------\n",
      "-------------------- Step 25, loss: 7378.22802734375 --------------------\n",
      "-------------------- Step 26, loss: 81185.40625 --------------------\n",
      "-------------------- Step 27, loss: 18370.255859375 --------------------\n",
      "-------------------- Step 28, loss: 3314.90478515625 --------------------\n",
      "-------------------- Step 29, loss: 15871.3154296875 --------------------\n",
      "-------------------- Step 30, loss: 545.2841796875 --------------------\n",
      "-------------------- Step 31, loss: 1281.3038330078125 --------------------\n",
      "-------------------- Step 32, loss: 52890.65625 --------------------\n",
      "-------------------- Step 33, loss: 2550.232177734375 --------------------\n",
      "-------------------- Step 34, loss: 4526.03759765625 --------------------\n",
      "-------------------- Step 35, loss: 25.5832462310791 --------------------\n",
      "-------------------- Step 36, loss: 22.22301483154297 --------------------\n",
      "-------------------- Step 37, loss: 17.7525691986084 --------------------\n",
      "-------------------- Step 38, loss: 9.034607887268066 --------------------\n",
      "-------------------- Step 39, loss: 1.6510401964187622 --------------------\n",
      "-------------------- Step 40, loss: 6.275766372680664 --------------------\n",
      "-------------------- Step 41, loss: 3.707094430923462 --------------------\n",
      "-------------------- Step 42, loss: 0.7623991966247559 --------------------\n",
      "-------------------- Step 43, loss: 1.5783321857452393 --------------------\n",
      "-------------------- Step 44, loss: 0.8166252374649048 --------------------\n",
      "-------------------- Step 45, loss: 0.885994553565979 --------------------\n",
      "-------------------- Step 46, loss: 0.912842869758606 --------------------\n",
      "-------------------- Step 47, loss: 0.7323049902915955 --------------------\n",
      "-------------------- Step 48, loss: 0.7469371557235718 --------------------\n",
      "-------------------- Step 49, loss: 0.8475004434585571 --------------------\n",
      "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n",
      "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-21 03:17:12.248789: I tensorflow/core/common_runtime/executor.cc:1209] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'inputs' with dtype float and shape [?,128]\n",
      "\t [[{{node inputs}}]]\n",
      "2023-08-21 03:17:12.721088: I tensorflow/core/common_runtime/executor.cc:1209] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'inputs' with dtype float and shape [?,128]\n",
      "\t [[{{node inputs}}]]\n",
      "WARNING:absl:Found untraced functions such as bottom_0_layer_call_fn, bottom_0_layer_call_and_return_conditional_losses, bottom_1_layer_call_fn, bottom_1_layer_call_and_return_conditional_losses, bottom_2_layer_call_fn while saving (showing 5 of 16). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: dlrm_tf_saved_model/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: dlrm_tf_saved_model/assets\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Model's `__init__()` arguments contain non-serializable objects. Please implement a `get_config()` method in the subclassed Model for proper saving and loading. Defaulting to empty config.\n"
     ]
    }
   ],
   "source": [
    "trained_model = train(args)\n",
    "weights_list = trained_model.get_weights()\n",
    "embedding_weights = weights_list[-1]\n",
    "trained_model.save(args[\"saved_path\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3fbf1e53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Release the occupied GPU memory by TensorFlow and Keras\n",
    "from numba import cuda\n",
    "cuda.select_device(0)\n",
    "cuda.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bbd4051",
   "metadata": {},
   "source": [
    "## Build the HPS-integrated TensorRT engine\n",
    "In order to use HPS in the inference stage, we need to convert the embedding weights to the formats required by HPS first and create JSON configuration file for HPS.\n",
    "\n",
    "Then we convert the TF saved model to ONNX, and employ the ONNX GraphSurgoen tool to replace the native TF embedding lookup layer with the placeholder of HPS TensorRT plugin layer.\n",
    "\n",
    "After that, we can build the TensorRT engine, which is comprised of the HPS TensorRT plugin layer and the dense network."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5510a6b",
   "metadata": {},
   "source": [
    "### Step1: Prepare sparse model and JSON configuration file for HPS\n",
    "\n",
    "Please note that the storage format in the `dlrm_tf_sparse.model/key` file is int64, while the HPS TensorRT plugin currently only support int32 when loading the keys into memory. There is no overflow since the key value range is 0~260000."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "99d772bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_sparse_model(embeddings_weights, embedding_table_path, embedding_vec_size):\n",
    "    os.system(\"mkdir -p {}\".format(embedding_table_path))\n",
    "    with open(\"{}/key\".format(embedding_table_path), 'wb') as key_file, \\\n",
    "        open(\"{}/emb_vector\".format(embedding_table_path), 'wb') as vec_file:\n",
    "      for key in range(embeddings_weights.shape[0]):\n",
    "        vec = embeddings_weights[key]\n",
    "        key_struct = struct.pack('q', key)\n",
    "        vec_struct = struct.pack(str(embedding_vec_size) + \"f\", *vec)\n",
    "        key_file.write(key_struct)\n",
    "        vec_file.write(vec_struct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6baeb83a",
   "metadata": {},
   "outputs": [],
   "source": [
    "convert_to_sparse_model(embedding_weights, args[\"embedding_table_path\"], args[\"embed_vec_size\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9248e956",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing dlrm_tf.json\n"
     ]
    }
   ],
   "source": [
    "%%writefile dlrm_tf.json\n",
    "{\n",
    "    \"supportlonglong\": false,\n",
    "    \"models\": [{\n",
    "        \"model\": \"dlrm\",\n",
    "        \"sparse_files\": [\"dlrm_tf_sparse.model\"],\n",
    "        \"num_of_worker_buffer_in_pool\": 3,\n",
    "        \"embedding_table_names\":[\"sparse_embedding0\"],\n",
    "        \"embedding_vecsize_per_table\": [128],\n",
    "        \"maxnum_catfeature_query_per_table_per_sample\": [26],\n",
    "        \"default_value_for_each_table\": [1.0],\n",
    "        \"deployed_device_list\": [0],\n",
    "        \"max_batch_size\": 1024,\n",
    "        \"cache_refresh_percentage_per_iteration\": 0.2,\n",
    "        \"hit_rate_threshold\": 1.0,\n",
    "        \"gpucacheper\": 1.0,\n",
    "        \"gpucache\": true\n",
    "        }\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad232f01",
   "metadata": {},
   "source": [
    "### Step2: Convert to ONNX and do ONNX graph surgery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f9631307",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3.10/runpy.py:126: RuntimeWarning: 'tf2onnx.convert' found in sys.modules after import of package 'tf2onnx', but prior to execution of 'tf2onnx.convert'; this may result in unpredictable behaviour\n",
      "  warn(RuntimeWarning(msg))\n",
      "2023-08-21 03:17:49,926 - WARNING - ***IMPORTANT*** Installed protobuf is not cpp accelerated. Conversion will be extremely slow. See https://github.com/onnx/tensorflow-onnx/issues/1557\n",
      "2023-08-21 03:17:50,868 - WARNING - '--tag' not specified for saved_model. Using --tag serve\n",
      "2023-08-21 03:17:56,302 - INFO - Signatures found in model: [serving_default].\n",
      "2023-08-21 03:17:56,302 - WARNING - '--signature_def' not specified, using first signature: serving_default\n",
      "2023-08-21 03:17:56,302 - INFO - Output names: ['output_1']\n",
      "2023-08-21 03:18:02,064 - INFO - Using tensorflow=2.12.0, onnx=1.14.0, tf2onnx=1.14.0/8f8d49\n",
      "2023-08-21 03:18:02,064 - INFO - Using opset <onnx, 15>\n",
      "2023-08-21 03:18:03,255 - INFO - Computed 0 values for constant folding\n",
      "2023-08-21 03:18:04,203 - INFO - Optimizing ONNX model\n",
      "2023-08-21 03:18:04,624 - INFO - After optimization: Cast -3 (3->0), Concat -1 (3->2), Const -15 (35->20), Identity -2 (2->0), Shape -1 (1->0), Slice -1 (1->0), Squeeze -1 (1->0), Unsqueeze -3 (3->0)\n",
      "2023-08-21 03:18:07,745 - INFO - \n",
      "2023-08-21 03:18:07,745 - INFO - Successfully converted TensorFlow model dlrm_tf_saved_model to ONNX\n",
      "2023-08-21 03:18:07,745 - INFO - Model inputs: ['keys', 'numerical_features']\n",
      "2023-08-21 03:18:07,745 - INFO - Model outputs: ['output_1']\n",
      "2023-08-21 03:18:07,745 - INFO - ONNX model is saved at dlrm_tf.onnx\n"
     ]
    }
   ],
   "source": [
    "# convert TF SavedModel to ONNX\n",
    "!python -m tf2onnx.convert --saved-model dlrm_tf_saved_model --output dlrm_tf.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7c304f7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ONNX graph surgery to insert HPS the TensorRT plugin placeholder\n",
    "import onnx_graphsurgeon as gs\n",
    "from onnx import  shape_inference\n",
    "import numpy as np\n",
    "import onnx\n",
    "\n",
    "graph = gs.import_onnx(onnx.load(\"dlrm_tf.onnx\"))\n",
    "saved = []\n",
    "\n",
    "for node in graph.nodes:\n",
    "    if node.name == \"StatefulPartitionedCall/dlrm/embedding_lookup\":\n",
    "        categorical_features = gs.Variable(name=\"categorical_features\", dtype=np.int32, shape=(\"unknown\", 26))\n",
    "        hps_node = gs.Node(op=\"HPS_TRT\", attrs={\"ps_config_file\": \"dlrm_tf.json\\0\", \"model_name\": \"dlrm\\0\", \"table_id\": 0, \"emb_vec_size\": 128}, \n",
    "                           inputs=[categorical_features], outputs=[node.outputs[0]])\n",
    "        graph.nodes.append(hps_node)\n",
    "        saved.append(categorical_features)\n",
    "        node.outputs.clear()\n",
    "for i in graph.inputs:\n",
    "    if i.name == \"numerical_features\":\n",
    "        saved.append(i)\n",
    "graph.inputs = saved\n",
    "\n",
    "graph.cleanup().toposort()\n",
    "onnx.save(gs.export_onnx(graph), \"dlrm_tf_with_hps.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52d5f090",
   "metadata": {},
   "source": [
    "### Step3: Build the TensorRT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8da1a8fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[08/21/2023-03:18:16] [TRT] [I] [MemUsageChange] Init CUDA: CPU +2013, GPU +0, now: CPU 4018, GPU 721 (MiB)\n",
      "[08/21/2023-03:18:22] [TRT] [I] [MemUsageChange] Init builder kernel library: CPU +421, GPU +72, now: CPU 4516, GPU 793 (MiB)\n",
      "[08/21/2023-03:18:22] [TRT] [W] CUDA lazy loading is not enabled. Enabling it can significantly reduce device memory usage and speed up TensorRT initialization. See \"Lazy Loading\" section of CUDA documentation https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#lazy-loading\n",
      "[08/21/2023-03:18:22] [TRT] [W] onnx2trt_utils.cpp:374: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.\n",
      "[08/21/2023-03:18:22] [TRT] [I] No importer registered for op: HPS_TRT. Attempting to import as plugin.\n",
      "[08/21/2023-03:18:22] [TRT] [I] Searching for plugin: HPS_TRT, plugin_version: 1, plugin_namespace: \n",
      "=====================================================HPS Parse====================================================\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: fuse_embedding_table is not specified using default: 0\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: dense_file is not specified using default: \n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: num_of_refresher_buffer_in_pool is not specified using default: 1\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: maxnum_des_feature_per_sample is not specified using default: 26\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: refresh_delay is not specified using default: 0\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: refresh_interval is not specified using default: 0\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: use_static_table is not specified using default: 0\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: use_context_stream is not specified using default: 1\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: use_hctr_cache_implementation is not specified using default: 1\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: thread_pool_size is not specified using default: 16\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: init_ec is not specified using default: 1\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: enable_pagelock is not specified using default: 0\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: fp8_quant is not specified using default: 0\n",
      "[HCTR][03:18:22.774][INFO][RK0][main]: HPS plugin uses context stream for model dlrm: True\n",
      "====================================================HPS Create====================================================\n",
      "[HCTR][03:18:22.775][INFO][RK0][main]: Creating HashMap CPU database backend...\n",
      "[HCTR][03:18:22.775][DEBUG][RK0][main]: Created blank database backend in local memory!\n",
      "[HCTR][03:18:22.775][INFO][RK0][main]: Volatile DB: initial cache rate = 1\n",
      "[HCTR][03:18:22.775][INFO][RK0][main]: Volatile DB: cache missed embeddings = 0\n",
      "[HCTR][03:18:22.775][DEBUG][RK0][main]: Created raw model loader in local memory!\n",
      "[HCTR][03:18:24.860][INFO][RK0][main]: Table: hps_et.dlrm.sparse_embedding0; cached 260000 / 260000 embeddings in volatile database (HashMapBackend); load: 260000 / 18446744073709551615 (0.00%).\n",
      "[HCTR][03:18:24.863][DEBUG][RK0][main]: Real-time subscribers created!\n",
      "[HCTR][03:18:24.864][INFO][RK0][main]: Creating embedding cache in device 0.\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Model name: dlrm\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Max batch size: 1024\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Fuse embedding tables: False\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Number of embedding tables: 1\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 1.000000\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Embedding cache type: dynamic\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Use I64 input key: False\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: The size of worker memory pool: 3\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][03:18:24.869][INFO][RK0][main]: The refresh percentage : 0.200000\n",
      "[HCTR][03:18:24.902][INFO][RK0][main]: Initialize the embedding cache by by inserting the same size model file with embedding cache from beginning\n",
      "[HCTR][03:18:24.902][DEBUG][RK0][main]: Created raw model loader in local memory!\n",
      "[HCTR][03:18:24.902][INFO][RK0][main]: EC initialization on device 0 for hps_et.dlrm.sparse_embedding0\n",
      "[HCTR][03:18:24.947][INFO][RK0][main]: Initialize the embedding table 0 for iteration 0 with number of 51968 keys.\n",
      "[HCTR][03:18:24.992][INFO][RK0][main]: Initialize the embedding table 0 for iteration 1 with number of 51968 keys.\n",
      "[HCTR][03:18:25.018][INFO][RK0][main]: Initialize the embedding table 0 for iteration 2 with number of 51968 keys.\n",
      "[HCTR][03:18:25.047][INFO][RK0][main]: Initialize the embedding table 0 for iteration 3 with number of 51968 keys.\n",
      "[HCTR][03:18:25.069][INFO][RK0][main]: Initialize the embedding table 0 for iteration 4 with number of 51968 keys.\n",
      "[HCTR][03:18:25.077][INFO][RK0][main]: LookupSession i64_input_key: False\n",
      "[HCTR][03:18:25.077][INFO][RK0][main]: Creating lookup session for dlrm on device: 0\n",
      "[08/21/2023-03:18:25] [TRT] [I] Successfully created plugin: HPS_TRT\n",
      "[08/21/2023-03:18:25] [TRT] [I] BuilderFlag::kTF32 is set but hardware does not support TF32. Disabling TF32.\n",
      "[08/21/2023-03:18:25] [TRT] [I] Graph optimization time: 0.034216 seconds.\n",
      "[08/21/2023-03:18:25] [TRT] [I] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +1, GPU +8, now: CPU 8710, GPU 1051 (MiB)\n",
      "[08/21/2023-03:18:25] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +1, GPU +10, now: CPU 8711, GPU 1061 (MiB)\n",
      "[08/21/2023-03:18:25] [TRT] [I] BuilderFlag::kTF32 is set but hardware does not support TF32. Disabling TF32.\n",
      "[08/21/2023-03:18:25] [TRT] [I] Local timing cache in use. Profiling results in this builder pass will not be stored.\n",
      "[08/21/2023-03:18:27] [TRT] [I] Detected 2 inputs and 1 output network tensors.\n",
      "[08/21/2023-03:18:27] [TRT] [I] Total Host Persistent Memory: 144\n",
      "[08/21/2023-03:18:27] [TRT] [I] Total Device Persistent Memory: 0\n",
      "[08/21/2023-03:18:27] [TRT] [I] Total Scratch Memory: 18350080\n",
      "[08/21/2023-03:18:27] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 0 MiB, GPU 41 MiB\n",
      "[08/21/2023-03:18:27] [TRT] [I] [BlockAssignment] Started assigning block shifts. This will take 2 steps to complete.\n",
      "[08/21/2023-03:18:27] [TRT] [I] [BlockAssignment] Algorithm ShiftNTopDown took 0.007954ms to assign 2 blocks to 2 nodes requiring 31981568 bytes.\n",
      "[08/21/2023-03:18:27] [TRT] [I] Total Activation Memory: 31981568\n",
      "[08/21/2023-03:18:27] [TRT] [I] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +8, now: CPU 8764, GPU 1091 (MiB)\n",
      "[08/21/2023-03:18:27] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +0, GPU +10, now: CPU 8764, GPU 1101 (MiB)\n",
      "[08/21/2023-03:18:27] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +16, now: CPU 0, GPU 16 (MiB)\n",
      "Successfully build the TensorRT engine\n"
     ]
    }
   ],
   "source": [
    "# build the TensorRT engine based on dlrm_tf_with_hps.onnx\n",
    "import tensorrt as trt\n",
    "import ctypes\n",
    "\n",
    "plugin_lib_name = \"/usr/local/hps_trt/lib/libhps_plugin.so\"\n",
    "handle = ctypes.CDLL(plugin_lib_name, mode=ctypes.RTLD_GLOBAL)\n",
    "\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n",
    "EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)\n",
    "\n",
    "def build_engine_from_onnx(onnx_model_path):\n",
    "    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser, builder.create_builder_config() as builder_config:        \n",
    "        model = open(onnx_model_path, 'rb')\n",
    "        parser.parse(model.read())\n",
    "\n",
    "        profile = builder.create_optimization_profile()        \n",
    "        profile.set_shape(\"categorical_features\", (1, 26), (1024, 26), (1024, 26))    \n",
    "        profile.set_shape(\"numerical_features\", (1, 13), (1024, 13), (1024, 13))\n",
    "        builder_config.add_optimization_profile(profile)\n",
    "        engine = builder.build_serialized_network(network, builder_config)\n",
    "        return engine\n",
    "\n",
    "serialized_engine = build_engine_from_onnx(\"dlrm_tf_with_hps.onnx\")\n",
    "with open(\"dlrm_tf_with_hps.trt\", \"wb\") as fout:\n",
    "    fout.write(serialized_engine)\n",
    "print(\"Successfully build the TensorRT engine\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd1437e3",
   "metadata": {},
   "source": [
    "## Deploy HPS-integrated TensorRT engine with Triton on multiple GPUs\n",
    "\n",
    "In order to deploy the TensorRT engine with the Triton TensorRT backend, we need to create the model repository and define the `config.pbtxt` first. Since we are deploy model instances on multiple GPUs, we need to modify the `\"deployed_device_list\"` entry in `dlrm_tf.json` accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9947c0a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p model_repo/dlrm_tf_with_hps/1\n",
    "!mv dlrm_tf_with_hps.trt model_repo/dlrm_tf_with_hps/1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b73f9b05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting model_repo/dlrm_tf_with_hps/config.pbtxt\n"
     ]
    }
   ],
   "source": [
    "%%writefile model_repo/dlrm_tf_with_hps/config.pbtxt\n",
    "\n",
    "platform: \"tensorrt_plan\"\n",
    "default_model_filename: \"dlrm_tf_with_hps.trt\"\n",
    "backend: \"tensorrt\"\n",
    "max_batch_size: 0\n",
    "input [\n",
    "  {\n",
    "    name: \"categorical_features\"\n",
    "    data_type: TYPE_INT32\n",
    "    dims: [-1,26]\n",
    "  },\n",
    "  {\n",
    "    name: \"numerical_features\"\n",
    "    data_type: TYPE_FP32\n",
    "    dims: [-1,13]\n",
    "  }\n",
    "]\n",
    "output [\n",
    "  {\n",
    "      name: \"output_1\"\n",
    "      data_type: TYPE_FP32\n",
    "      dims: [-1,1]\n",
    "  }\n",
    "]\n",
    "instance_group [\n",
    "  {\n",
    "    count: 1\n",
    "    kind: KIND_GPU\n",
    "    gpus:[0,1,2,3]\n",
    "  }\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "55ac01fe-d769-46ee-a2e3-565a913f195f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting dlrm_tf.json\n"
     ]
    }
   ],
   "source": [
    "%%writefile dlrm_tf.json\n",
    "{\n",
    "    \"supportlonglong\": false,\n",
    "    \"models\": [{\n",
    "        \"model\": \"dlrm\",\n",
    "        \"sparse_files\": [\"dlrm_tf_sparse.model\"],\n",
    "        \"num_of_worker_buffer_in_pool\": 3,\n",
    "        \"embedding_table_names\":[\"sparse_embedding0\"],\n",
    "        \"embedding_vecsize_per_table\": [128],\n",
    "        \"maxnum_catfeature_query_per_table_per_sample\": [26],\n",
    "        \"default_value_for_each_table\": [1.0],\n",
    "        \"deployed_device_list\": [0,1,2,3],\n",
    "        \"max_batch_size\": 1024,\n",
    "        \"cache_refresh_percentage_per_iteration\": 0.2,\n",
    "        \"hit_rate_threshold\": 1.0,\n",
    "        \"gpucacheper\": 1.0,\n",
    "        \"gpucache\": true\n",
    "        }\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "de7f8b12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34mmodel_repo/dlrm_tf_with_hps\u001b[0m\n",
      "├── \u001b[01;34m1\u001b[0m\n",
      "│   └── \u001b[00mdlrm_tf_with_hps.trt\u001b[0m\n",
      "└── \u001b[00mconfig.pbtxt\u001b[0m\n",
      "\n",
      "1 directory, 2 files\n"
     ]
    }
   ],
   "source": [
    "!tree model_repo/dlrm_tf_with_hps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57235f3f",
   "metadata": {},
   "source": [
    "We can then launch the Triton inference server using the TensorRT backend. Please note that `LD_PRELOAD` is utilized to load the custom TensorRT plugin (i.e., HPS TensorRT plugin) into Triton.\n",
    "\n",
    "Note: `Since Background processes not supported by Jupyter, please launch the Triton Server according to the following command independently in the background`.\n",
    "\n",
    "> **LD_PRELOAD=/usr/local/hps_trt/lib/libhps_plugin.so tritonserver --model-repository=/hugectr/hps_trt/notebooks/model_repo/ --load-model=dlrm_tf_with_hps --model-control-mode=explicit**\n",
    "\n",
    "If you successfully started tritonserver, you should see a log similar to following:\n",
    "\n",
    "\n",
    "```bash\n",
    "TRITONBACKEND_ModelInstanceInitialize: dlrm_tf_with_hps_0 (GPU device 0)\n",
    "TRITONBACKEND_ModelInstanceInitialize: dlrm_tf_with_hps_0 (GPU device 1)\n",
    "TRITONBACKEND_ModelInstanceInitialize: dlrm_tf_with_hps_0 (GPU device 2)\n",
    "TRITONBACKEND_ModelInstanceInitialize: dlrm_tf_with_hps_0 (GPU device 3)\n",
    "\n",
    "\n",
    "+----------+--------------------------------+--------------------------------+\n",
    "| Backend  | Path                           | Config                         |\n",
    "+----------+--------------------------------+--------------------------------+\n",
    "| tensorrt | /opt/tritonserver/backends/ten | {\"cmdline\":{\"auto-complete-con |\n",
    "|          | sorrt/libtriton_tensorrt.so    | fig\":\"true\",\"min-compute-capab |\n",
    "|          |                                | ility\":\"6.000000\",\"backend-dir |\n",
    "|          |                                | ectory\":\"/opt/tritonserver/bac |\n",
    "|          |                                | kends\",\"default-max-batch-size |\n",
    "|          |                                | \":\"4\"}}                        |\n",
    "|          |                                |                                |\n",
    "+----------+--------------------------------+--------------------------------+\n",
    "\n",
    "+------------------+---------+--------+\n",
    "| Model            | Version | Status |\n",
    "+------------------+---------+--------+\n",
    "| dlrm_tf_with_hps | 1       | READY  |\n",
    "+------------------+---------+--------+\n",
    "```\n",
    "\n",
    "We can then send the requests to the Triton inference server using the HTTP client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "398dc8ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction result is \n",
      "[[0.34091672]\n",
      " [0.34091672]\n",
      " [0.34091672]\n",
      " ...\n",
      " [0.34091672]\n",
      " [0.34091672]\n",
      " [0.34091672]]\n",
      "Response details:\n",
      "{'model_name': 'dlrm_tf_with_hps', 'model_version': '1', 'outputs': [{'name': 'output_1', 'datatype': 'FP32', 'shape': [1024, 1], 'parameters': {'binary_data_size': 4096}}]}\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import tritonclient.http as httpclient\n",
    "from tritonclient.utils import *\n",
    "\n",
    "BATCH_SIZE = 1024\n",
    "\n",
    "categorical_feature = np.random.randint(0,260000,size=(BATCH_SIZE,26)).astype(np.int32)\n",
    "numerical_feature = np.random.random((BATCH_SIZE, 13)).astype(np.float32)\n",
    "\n",
    "inputs = [\n",
    "    httpclient.InferInput(\"categorical_features\", \n",
    "                          categorical_feature.shape,\n",
    "                          np_to_triton_dtype(np.int32)),\n",
    "    httpclient.InferInput(\"numerical_features\", \n",
    "                          numerical_feature.shape,\n",
    "                          np_to_triton_dtype(np.float32)),                          \n",
    "]\n",
    "inputs[0].set_data_from_numpy(categorical_feature)\n",
    "inputs[1].set_data_from_numpy(numerical_feature)\n",
    "\n",
    "\n",
    "outputs = [\n",
    "    httpclient.InferRequestedOutput(\"output_1\")\n",
    "]\n",
    "\n",
    "model_name = \"dlrm_tf_with_hps\"\n",
    "\n",
    "with httpclient.InferenceServerClient(\"localhost:8000\") as client:\n",
    "    response = client.infer(model_name,\n",
    "                            inputs,\n",
    "                            outputs=outputs)\n",
    "    result = response.get_response()\n",
    "    \n",
    "    print(\"Prediction result is \\n{}\".format(response.as_numpy(\"output_1\")))\n",
    "    print(\"Response details:\\n{}\".format(result))\n"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}