examples/sagemaker-tensorflow/sagemaker-merlin-tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "71d18943-381f-4096-b8c2-349cd12f57bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2022, NVIDIA CORPORATION.\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",
    "# 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": "1e73569d-99f1-4449-bfd1-8331fcde362d",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_merlin_getting-started-movielens-01-download-convert/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Training and Serving Merlin on AWS SageMaker\n",
    "\n",
    "This notebook is created using the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) container.\n",
    "Note that AWS libraries in this notebook require AWS credentials, and if you are running this notebook in a container, you might need to restart the container with the AWS credentials mounted, e.g., `-v $HOME/.aws:$HOME/.aws`.\n",
    "\n",
    "\n",
    "With AWS Sagemaker, you can package your own models that can then be trained and deployed in the SageMaker environment. This notebook shows you how to use Merlin for training and inference in the SageMaker environment.\n",
    "\n",
    "It assumes that readers are familiar wtth some basic concepts in NVIDIA Merlin,\n",
    "such as:\n",
    "\n",
    "- Using NVTabular to GPU-accelerate preprocessing and feature engineering,\n",
    "- Training a ranking model using Merlin Models, and\n",
    "- Inference with the Triton Inference Server and Merlin Models for Tensorflow.\n",
    "\n",
    "To learn more about these concepts in NVIDIA Merlin, see for example\n",
    "[Deploying a Multi-Stage Recommender System](../Building-and-deploying-multi-stage-RecSys/README.md)\n",
    "in this repository or example notebooks in\n",
    "[Merlin Models](https://github.com/NVIDIA-Merlin/models/tree/stable/examples).\n",
    "\n",
    "To run this notebook, you need to have [Amazon SageMaker Python SDK](https://sagemaker.readthedocs.io/en/stable/) installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1fa66a78-2243-433c-add5-5a49a22b4718",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting sagemaker\n",
      "  Downloading sagemaker-2.116.0.tar.gz (592 kB)\n",
      "\u001b[K     |████████████████████████████████| 592 kB 4.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: attrs<23,>=20.3.0 in /usr/local/lib/python3.8/dist-packages (from sagemaker) (22.1.0)\n",
      "Requirement already satisfied: boto3<2.0,>=1.20.21 in /usr/local/lib/python3.8/dist-packages (from sagemaker) (1.25.2)\n",
      "Requirement already satisfied: google-pasta in /usr/local/lib/python3.8/dist-packages (from sagemaker) (0.2.0)\n",
      "Collecting importlib-metadata<5.0,>=1.4.0\n",
      "  Downloading importlib_metadata-4.13.0-py3-none-any.whl (23 kB)\n",
      "Requirement already satisfied: numpy<2.0,>=1.9.0 in /usr/local/lib/python3.8/dist-packages (from sagemaker) (1.22.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from sagemaker) (21.3)\n",
      "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (from sagemaker) (1.3.5)\n",
      "Collecting pathos\n",
      "  Downloading pathos-0.3.0-py3-none-any.whl (79 kB)\n",
      "\u001b[K     |████████████████████████████████| 79 kB 10.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting protobuf3-to-dict<1.0,>=0.1.5\n",
      "  Downloading protobuf3-to-dict-0.1.5.tar.gz (3.5 kB)\n",
      "Requirement already satisfied: protobuf<4.0,>=3.1 in /usr/local/lib/python3.8/dist-packages (from sagemaker) (3.19.6)\n",
      "Collecting schema\n",
      "  Downloading schema-0.7.5-py2.py3-none-any.whl (17 kB)\n",
      "Collecting smdebug_rulesconfig==1.0.1\n",
      "  Downloading smdebug_rulesconfig-1.0.1-py2.py3-none-any.whl (20 kB)\n",
      "Requirement already satisfied: s3transfer<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from boto3<2.0,>=1.20.21->sagemaker) (0.6.0)\n",
      "Requirement already satisfied: botocore<1.29.0,>=1.28.2 in /usr/local/lib/python3.8/dist-packages (from boto3<2.0,>=1.20.21->sagemaker) (1.28.2)\n",
      "Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /usr/local/lib/python3.8/dist-packages (from boto3<2.0,>=1.20.21->sagemaker) (1.0.1)\n",
      "Requirement already satisfied: six in /usr/lib/python3/dist-packages (from google-pasta->sagemaker) (1.14.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata<5.0,>=1.4.0->sagemaker) (3.10.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging>=20.0->sagemaker) (3.0.9)\n",
      "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->sagemaker) (2022.5)\n",
      "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.8/dist-packages (from pandas->sagemaker) (2.8.2)\n",
      "Collecting pox>=0.3.2\n",
      "  Downloading pox-0.3.2-py3-none-any.whl (29 kB)\n",
      "Collecting dill>=0.3.6\n",
      "  Downloading dill-0.3.6-py3-none-any.whl (110 kB)\n",
      "\u001b[K     |████████████████████████████████| 110 kB 17.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting multiprocess>=0.70.14\n",
      "  Downloading multiprocess-0.70.14-py38-none-any.whl (132 kB)\n",
      "\u001b[K     |████████████████████████████████| 132 kB 17.8 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting ppft>=1.7.6.6\n",
      "  Downloading ppft-1.7.6.6-py3-none-any.whl (52 kB)\n",
      "\u001b[K     |████████████████████████████████| 52 kB 2.9 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting contextlib2>=0.5.5\n",
      "  Downloading contextlib2-21.6.0-py2.py3-none-any.whl (13 kB)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.25.4 in /usr/local/lib/python3.8/dist-packages (from botocore<1.29.0,>=1.28.2->boto3<2.0,>=1.20.21->sagemaker) (1.26.12)\n",
      "Building wheels for collected packages: sagemaker, protobuf3-to-dict\n",
      "  Building wheel for sagemaker (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for sagemaker: filename=sagemaker-2.116.0-py2.py3-none-any.whl size=809052 sha256=f446dd6eed6d268b7f3f2709f8f11c1ba153e382fbea9b2caedd517c1fb71215\n",
      "  Stored in directory: /root/.cache/pip/wheels/3e/cb/b1/5b13ff7b150aa151e4a11030a6c41b1e457c31a52ea1ef11b0\n",
      "  Building wheel for protobuf3-to-dict (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for protobuf3-to-dict: filename=protobuf3_to_dict-0.1.5-py3-none-any.whl size=4029 sha256=8f99baaa875ba544d54f624f95dfbf4fd52ca96d52ce8af6d05c1ff2bb8435b2\n",
      "  Stored in directory: /root/.cache/pip/wheels/fc/10/27/2d1e23d8b9a9013a83fbb418a0b17b1e6f81c8db8f53b53934\n",
      "Successfully built sagemaker protobuf3-to-dict\n",
      "Installing collected packages: importlib-metadata, pox, dill, multiprocess, ppft, pathos, protobuf3-to-dict, contextlib2, schema, smdebug-rulesconfig, sagemaker\n",
      "  Attempting uninstall: importlib-metadata\n",
      "    Found existing installation: importlib-metadata 5.0.0\n",
      "    Uninstalling importlib-metadata-5.0.0:\n",
      "      Successfully uninstalled importlib-metadata-5.0.0\n",
      "Successfully installed contextlib2-21.6.0 dill-0.3.6 importlib-metadata-4.13.0 multiprocess-0.70.14 pathos-0.3.0 pox-0.3.2 ppft-1.7.6.6 protobuf3-to-dict-0.1.5 sagemaker-2.116.0 schema-0.7.5 smdebug-rulesconfig-1.0.1\n"
     ]
    }
   ],
   "source": [
    "! python -m pip install sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe6eadc-10ab-41c4-ba55-67ea59ac6fb5",
   "metadata": {},
   "source": [
    "## Part 1: Generating Dataset and Docker image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6256e2bd-899a-48fa-9111-f594f6b3b2a1",
   "metadata": {},
   "source": [
    "### Generating Dataset\n",
    "\n",
    "In this notebook, we use the synthetic train and test datasets generated by mimicking the real [Ali-CCP](https://tianchi.aliyun.com/dataset/dataDetail?dataId=408#1): Alibaba Click and Conversion Prediction dataset to build our recommender system ranking models. The Ali-CCP is a dataset gathered from real-world traffic logs of the recommender system in Taobao, the largest online retail platform in the world.\n",
    "\n",
    "If you would like to use real Ali-CCP dataset instead, you can download the training and test datasets on [tianchi.aliyun.com](https://tianchi.aliyun.com/dataset/dataDetail?dataId=408#1). You can then use [get_aliccp()](https://github.com/NVIDIA-Merlin/models/blob/stable/merlin/datasets/ecommerce/aliccp/dataset.py#L43) function to curate the raw csv files and save them as parquet files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e96e35df-27b4-4979-a3f5-676593c87106",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.USER_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.USER: 'user'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.ITEM_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.ITEM: 'item'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "from merlin.datasets.synthetic import generate_data\n",
    "\n",
    "DATA_FOLDER = os.environ.get(\"DATA_FOLDER\", \"/workspace/data/\")\n",
    "NUM_ROWS = os.environ.get(\"NUM_ROWS\", 1_000_000)\n",
    "SYNTHETIC_DATA = eval(os.environ.get(\"SYNTHETIC_DATA\", \"True\"))\n",
    "BATCH_SIZE = int(os.environ.get(\"BATCH_SIZE\", 512))\n",
    "\n",
    "if SYNTHETIC_DATA:\n",
    "    train, valid = generate_data(\"aliccp-raw\", int(NUM_ROWS), set_sizes=(0.7, 0.3))\n",
    "    # save the datasets as parquet files\n",
    "    train.to_ddf().to_parquet(os.path.join(DATA_FOLDER, \"train\"))\n",
    "    valid.to_ddf().to_parquet(os.path.join(DATA_FOLDER, \"valid\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3eb0b2e-b865-4b48-8182-ce6f8c4f0789",
   "metadata": {},
   "source": [
    "### Training Script\n",
    "\n",
    "The training script [train.py](./train.py) in this example starts with the synthetic dataset we have created in the previous cell and produces a ranking model by performing the following tasks:\n",
    "- Perform feature engineering and preprocessing with [NVTabular](https://github.com/NVIDIA-Merlin/NVTabular). NVTabular implements common feature engineering and preprocessing operators in easy-to-use, high-level APIs.\n",
    "- Use [Merlin Models](https://github.com/NVIDIA-Merlin/models/) to train [Facebook's DLRM model](https://arxiv.org/pdf/1906.00091.pdf) in Tensorflow.\n",
    "- Prepares [ensemble models](https://github.com/triton-inference-server/server/blob/main/docs/user_guide/architecture.md#ensemble-models) for serving on [Triton Inference Server](https://github.com/triton-inference-server/server).\n",
    "The training script outputs to `model_dir` the final NVTabular workflow and the trained DLRM model as an ensemble model. You want to make sure that your script generates any artifacts within `model_dir`, since SageMaker packages any files in this directory into a compressed tar archive and made available at the S3 location. The ensemble model that is uploaded to S3 will be used later to handle predictions in Triton inference server later in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "60996955-1d6c-43e7-b29a-ad451b9aee91",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile train.py\n",
    "#\n",
    "# Copyright (c) 2022, NVIDIA CORPORATION.\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",
    "import argparse\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import tempfile\n",
    "\n",
    "# We can control how much memory to give tensorflow with this environment variable\n",
    "# IMPORTANT: make sure you do this before you initialize TF's runtime, otherwise\n",
    "# TF will have claimed all free GPU memory\n",
    "os.environ[\"TF_MEMORY_ALLOCATION\"] = \"0.7\"  # fraction of free memory\n",
    "\n",
    "import merlin.io\n",
    "import merlin.models.tf as mm\n",
    "import nvtabular as nvt\n",
    "import tensorflow as tf\n",
    "from merlin.schema.tags import Tags\n",
    "from merlin.systems.dag.ops.workflow import TransformWorkflow\n",
    "from merlin.systems.dag.ops.tensorflow import PredictTensorflow\n",
    "from merlin.systems.dag.ensemble import Ensemble\n",
    "import numpy as np\n",
    "from nvtabular.ops import *\n",
    "\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "logger.setLevel(logging.DEBUG)\n",
    "logger.addHandler(logging.StreamHandler(sys.stdout))\n",
    "\n",
    "\n",
    "def parse_args():\n",
    "    \"\"\"\n",
    "    Parse arguments passed from the SageMaker API to the container.\n",
    "    \"\"\"\n",
    "\n",
    "    parser = argparse.ArgumentParser()\n",
    "\n",
    "    # Hyperparameters sent by the client are passed as command-line arguments to the script\n",
    "    parser.add_argument(\"--epochs\", type=int, default=1)\n",
    "    parser.add_argument(\"--batch_size\", type=int, default=1024)\n",
    "\n",
    "    # Data directories\n",
    "    parser.add_argument(\n",
    "        \"--train_dir\", type=str, default=os.environ.get(\"SM_CHANNEL_TRAIN\")\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--valid_dir\", type=str, default=os.environ.get(\"SM_CHANNEL_VALID\")\n",
    "    )\n",
    "\n",
    "    # Model directory: we will use the default set by SageMaker, /opt/ml/model\n",
    "    parser.add_argument(\"--model_dir\", type=str, default=os.environ.get(\"SM_MODEL_DIR\"))\n",
    "\n",
    "    return parser.parse_known_args()\n",
    "\n",
    "\n",
    "def create_nvtabular_workflow(train_path, valid_path):\n",
    "    user_id = [\"user_id\"] >> Categorify() >> TagAsUserID()\n",
    "    item_id = [\"item_id\"] >> Categorify() >> TagAsItemID()\n",
    "    targets = [\"click\"] >> AddMetadata(tags=[Tags.BINARY_CLASSIFICATION, \"target\"])\n",
    "\n",
    "    item_features = (\n",
    "        [\"item_category\", \"item_shop\", \"item_brand\"]\n",
    "        >> Categorify()\n",
    "        >> TagAsItemFeatures()\n",
    "    )\n",
    "\n",
    "    user_features = (\n",
    "        [\n",
    "            \"user_shops\",\n",
    "            \"user_profile\",\n",
    "            \"user_group\",\n",
    "            \"user_gender\",\n",
    "            \"user_age\",\n",
    "            \"user_consumption_2\",\n",
    "            \"user_is_occupied\",\n",
    "            \"user_geography\",\n",
    "            \"user_intentions\",\n",
    "            \"user_brands\",\n",
    "            \"user_categories\",\n",
    "        ]\n",
    "        >> Categorify()\n",
    "        >> TagAsUserFeatures()\n",
    "    )\n",
    "\n",
    "    outputs = user_id + item_id + item_features + user_features + targets\n",
    "\n",
    "    workflow = nvt.Workflow(outputs)\n",
    "\n",
    "    return workflow\n",
    "\n",
    "\n",
    "def create_ensemble(workflow, model):\n",
    "    serving_operators = (\n",
    "        workflow.input_schema.column_names\n",
    "        >> TransformWorkflow(workflow)\n",
    "        >> PredictTensorflow(model)\n",
    "    )\n",
    "    ensemble = Ensemble(serving_operators, workflow.input_schema)\n",
    "    return ensemble\n",
    "\n",
    "\n",
    "def train():\n",
    "    \"\"\"\n",
    "    Train the Merlin model.\n",
    "    \"\"\"\n",
    "    train_path = os.path.join(args.train_dir, \"*.parquet\")\n",
    "    valid_path = os.path.join(args.valid_dir, \"*.parquet\")\n",
    "\n",
    "    workflow = create_nvtabular_workflow(\n",
    "        train_path=train_path,\n",
    "        valid_path=valid_path,\n",
    "    )\n",
    "\n",
    "    train_dataset = nvt.Dataset(train_path)\n",
    "    valid_dataset = nvt.Dataset(valid_path)\n",
    "\n",
    "    output_path = tempfile.mkdtemp()\n",
    "    workflow_path = os.path.join(output_path, \"workflow\")\n",
    "\n",
    "    workflow.fit(train_dataset)\n",
    "    workflow.transform(train_dataset).to_parquet(\n",
    "        output_path=os.path.join(output_path, \"train\")\n",
    "    )\n",
    "    workflow.transform(valid_dataset).to_parquet(\n",
    "        output_path=os.path.join(output_path, \"valid\")\n",
    "    )\n",
    "\n",
    "    workflow.save(workflow_path)\n",
    "    logger.info(f\"Workflow saved to {workflow_path}.\")\n",
    "\n",
    "    train_data = merlin.io.Dataset(os.path.join(output_path, \"train\", \"*.parquet\"))\n",
    "    valid_data = merlin.io.Dataset(os.path.join(output_path, \"valid\", \"*.parquet\"))\n",
    "\n",
    "    schema = train_data.schema\n",
    "    target_column = schema.select_by_tag(Tags.TARGET).column_names[0]\n",
    "\n",
    "    model = mm.DLRMModel(\n",
    "        schema,\n",
    "        embedding_dim=64,\n",
    "        bottom_block=mm.MLPBlock([128, 64]),\n",
    "        top_block=mm.MLPBlock([128, 64, 32]),\n",
    "        prediction_tasks=mm.BinaryClassificationTask(target_column),\n",
    "    )\n",
    "\n",
    "    model.compile(\"adam\", run_eagerly=False, metrics=[tf.keras.metrics.AUC()])\n",
    "\n",
    "    batch_size = args.batch_size\n",
    "    epochs = args.epochs\n",
    "    logger.info(f\"batch_size = {batch_size}, epochs = {epochs}\")\n",
    "\n",
    "    model.fit(\n",
    "        train_data,\n",
    "        validation_data=valid_data,\n",
    "        batch_size=args.batch_size,\n",
    "        epochs=epochs,\n",
    "        verbose=2,\n",
    "    )\n",
    "\n",
    "    model_path = os.path.join(output_path, \"dlrm\")\n",
    "    model.save(model_path)\n",
    "    logger.info(f\"Model saved to {model_path}.\")\n",
    "\n",
    "    # We remove the label columns from its inputs.\n",
    "    # This removes all columns with the TARGET tag from the workflow.\n",
    "    # We do this because we need to set the workflow to only require the\n",
    "    # features needed to predict, not train, when creating an inference\n",
    "    # pipeline.\n",
    "    label_columns = workflow.output_schema.select_by_tag(Tags.TARGET).column_names\n",
    "    workflow.remove_inputs(label_columns)\n",
    "\n",
    "    ensemble = create_ensemble(workflow, model)\n",
    "    ensemble_path = args.model_dir\n",
    "    ensemble.export(ensemble_path)\n",
    "    logger.info(f\"Ensemble graph saved to {ensemble_path}.\")\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    args, _ = parse_args()\n",
    "    train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "361e6bb6-1d68-4797-a2b5-d87d164995cc",
   "metadata": {},
   "source": [
    "### Create the Dockerfile\n",
    "\n",
    "The `Dockerfile` describes the image that will be used on SageMaker for training and inference.\n",
    "We start from the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) docker image and install the [sagemaker-training-toolkit](https://github.com/aws/sagemaker-training-toolkit) library, which makes the image compatible with Sagemaker for training models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "26ba6bc5-f6ea-4d35-88f9-23683826e45b",
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting container/Dockerfile\n"
     ]
    }
   ],
   "source": [
    "%%writefile container/Dockerfile\n",
    "\n",
    "FROM nvcr.io/nvidia/merlin/merlin-tensorflow:22.10\n",
    "\n",
    "RUN pip3 install sagemaker-training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb9afdca-19f6-45e6-8a73-3878c87a3577",
   "metadata": {},
   "source": [
    "### Building and registering the container\n",
    "\n",
    "The following shell code shows how to build the container image using `docker build` and push the container image to ECR using `docker push`. This code is available as the shell script `build_and_push_image.sh`. If you are running this notebook inside the [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) docker container, you probably need to execute the script outside the container (e.g., in your terminal where you can run the `docker` command).\n",
    "\n",
    "You need to have the AWS CLI installed to run this code. To install the AWS CLI, see [Installing or updating the latest version of the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html#getting-started-install-instructions).\n",
    "\n",
    "This code looks for an ECR repository in the account you're using and the current default region (if you're using a SageMaker notebook instance, this is the region where the notebook instance was created). If the repository doesn't exist, the script will create it.\n",
    "\n",
    "Note that running the following script requires permissions to create new repositories on Amazon ECR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2e40ffca-b651-413f-ac64-43f44a01e7ef",
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./build_and_push_image.sh\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./build_and_push_image.sh\n",
    "\n",
    "#!/bin/bash\n",
    "\n",
    "set -euo pipefail\n",
    "\n",
    "# The name of our algorithm\n",
    "ALGORITHM_NAME=sagemaker-merlin-tensorflow\n",
    "REGION=us-east-1\n",
    "\n",
    "cd container\n",
    "\n",
    "ACCOUNT=$(aws sts get-caller-identity --query Account --output text --region ${REGION})\n",
    "\n",
    "# Get the region defined in the current configuration (default to us-west-2 if none defined)\n",
    "\n",
    "REPOSITORY=\"${ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com\"\n",
    "IMAGE_URI=\"${REPOSITORY}/${ALGORITHM_NAME}:latest\"\n",
    "\n",
    "# Get the login command from ECR and execute it directly\n",
    "aws ecr get-login-password --region ${REGION} | docker login --username AWS --password-stdin ${REPOSITORY}\n",
    "\n",
    "# If the repository doesn't exist in ECR, create it.\n",
    "\n",
    "aws ecr describe-repositories --repository-names \"${ALGORITHM_NAME}\" --region ${REGION} > /dev/null 2>&1\n",
    "\n",
    "if [ $? -ne 0 ]\n",
    "then\n",
    "    aws ecr create-repository --repository-name \"${ALGORITHM_NAME}\" --region ${REGION} > /dev/null\n",
    "fi\n",
    "\n",
    "# Build the docker image locally with the image name and then push it to ECR\n",
    "# with the full name.\n",
    "\n",
    "docker build  -t ${ALGORITHM_NAME} .\n",
    "docker tag ${ALGORITHM_NAME} ${IMAGE_URI}\n",
    "\n",
    "docker push ${IMAGE_URI}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9a92a0fd-4bfb-4da7-b598-073b178f3dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you are able to run `docker` from the notebook environment, you can uncomment and run the below script.\n",
    "# ! ./build_and_push_image.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d525205-b00a-4177-ab14-8d3a218d4ac9",
   "metadata": {},
   "source": [
    "## Part 2: Training your Merlin model on Sagemaker\n",
    "\n",
    "To deploy the training script onto Sagemaker, we use the Sagemaker Python SDK.\n",
    "Here, we create a Sagemaker session that we will use to perform our Sagemaker operations, specify the bucket to use, and the role for working with Sagemaker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2b62f39e-af41-4aec-857f-0541038d9c5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Couldn't call 'get_role' to get Role ARN from role name AWSOS-AD-Engineer to get Role path.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arn:aws:iam::843263297212:role/AWSOS-AD-Engineer\n"
     ]
    }
   ],
   "source": [
    "import sagemaker\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "\n",
    "# S3 prefix\n",
    "prefix = \"DEMO-merlin-tensorflow-aliccp\"\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "\n",
    "print(role)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6e5d758-e4f3-4f44-ba4f-9823f74a04e0",
   "metadata": {},
   "source": [
    "We can use the Sagemaker Python SDK to upload the Ali-CCP synthetic data to our S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cd9109d1-73cf-4177-b896-09cff9289bf8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-843263297212/DEMO-merlin-tensorflow-aliccp\n"
     ]
    }
   ],
   "source": [
    "data_location = sess.upload_data(DATA_FOLDER, key_prefix=prefix)\n",
    "\n",
    "print(data_location)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "922dbb12-76fb-41c0-ad90-32b7d1328e62",
   "metadata": {},
   "source": [
    "### Training on Sagemaker using the Python SDK\n",
    "\n",
    "Sagemaker provides the Python SDK for training a model on Sagemaker.\n",
    "\n",
    "Here, we start by using the ECR image URL of the image we pushed in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "14b9fc54-bafa-4487-8f34-2f9448c0b3cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "843263297212.dkr.ecr.us-east-1.amazonaws.com/sagemaker-merlin-tensorflow:latest\n"
     ]
    }
   ],
   "source": [
    "import boto3\n",
    "\n",
    "sts_client = boto3.client(\"sts\")\n",
    "account = sts_client.get_caller_identity()[\"Account\"]\n",
    "\n",
    "my_session = boto3.session.Session()\n",
    "region = my_session.region_name\n",
    "\n",
    "algorithm_name = \"sagemaker-merlin-tensorflow\"\n",
    "\n",
    "ecr_image = \"{}.dkr.ecr.{}.amazonaws.com/{}:latest\".format(\n",
    "    account, region, algorithm_name\n",
    ")\n",
    "\n",
    "print(ecr_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c51cafe-8685-4789-b348-e70b2f811035",
   "metadata": {},
   "source": [
    "We can call `Estimator.fit()` to start training on Sagemaker. Here, we use a `g4dn` GPU instance that are equipped with NVIDIA T4 GPUs.\n",
    "Our training script `train.py` is passed to the Estimator through the `entry_point` parameter.\n",
    "Behind the scenes, the Sagemaker Python SDK will upload the training script specified in the`entry_point` field (`train.py` in our case)\n",
    "to the S3 bucket and set the `SAGEMAKER_PROGRAM` environment variable in the training instance to the S3 location so that the training instance\n",
    "can download the training script on S3 to the training instance.\n",
    "We also adjust our hyperparameters in the `hyperparameters` field.\n",
    "We have uploaded our training dataset to our S3 bucket in the previous code cell, and the S3 URLs to our training and validation sets are passed into the `fit()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f238c8c6-656e-4b3d-96c6-32654357fc13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2022-11-09 10:18:31 Starting - Starting the training job...\n",
      "2022-11-09 10:18:54 Starting - Preparing the instances for trainingProfilerReport-1667989110: InProgress\n",
      "......\n",
      "2022-11-09 10:19:54 Downloading - Downloading input data...\n",
      "2022-11-09 10:20:34 Training - Downloading the training image..................................\u001b[34m==================================\u001b[0m\n",
      "\u001b[34m== Triton Inference Server Base ==\u001b[0m\n",
      "\u001b[34m==================================\u001b[0m\n",
      "\u001b[34mNVIDIA Release 22.08 (build 42766143)\u001b[0m\n",
      "\u001b[34mCopyright (c) 2018-2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.\u001b[0m\n",
      "\u001b[34mVarious files include modifications (c) NVIDIA CORPORATION & AFFILIATES.  All rights reserved.\u001b[0m\n",
      "\u001b[34mThis container image and its contents are governed by the NVIDIA Deep Learning Container License.\u001b[0m\n",
      "\u001b[34mBy pulling and using the container, you accept the terms and conditions of this license:\u001b[0m\n",
      "\u001b[34mhttps://developer.nvidia.com/ngc/nvidia-deep-learning-container-license\u001b[0m\n",
      "\u001b[34mNOTE: CUDA Forward Compatibility mode ENABLED.\n",
      "  Using CUDA 11.7 driver version 515.65.01 with kernel driver version 510.47.03.\n",
      "  See https://docs.nvidia.com/deploy/cuda-compatibility/ for details.\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:03,405 sagemaker-training-toolkit INFO     No Neurons detected (normal if no neurons installed)\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:03,438 sagemaker-training-toolkit INFO     No Neurons detected (normal if no neurons installed)\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:03,473 sagemaker-training-toolkit INFO     No Neurons detected (normal if no neurons installed)\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:03,485 sagemaker-training-toolkit INFO     Invoking user script\u001b[0m\n",
      "\u001b[34mTraining Env:\u001b[0m\n",
      "\u001b[34m{\n",
      "    \"additional_framework_parameters\": {},\n",
      "    \"channel_input_dirs\": {\n",
      "        \"train\": \"/opt/ml/input/data/train\",\n",
      "        \"valid\": \"/opt/ml/input/data/valid\"\n",
      "    },\n",
      "    \"current_host\": \"algo-1\",\n",
      "    \"current_instance_group\": \"homogeneousCluster\",\n",
      "    \"current_instance_group_hosts\": [\n",
      "        \"algo-1\"\n",
      "    ],\n",
      "    \"current_instance_type\": \"ml.g4dn.xlarge\",\n",
      "    \"distribution_hosts\": [],\n",
      "    \"distribution_instance_groups\": [],\n",
      "    \"framework_module\": null,\n",
      "    \"hosts\": [\n",
      "        \"algo-1\"\n",
      "    ],\n",
      "    \"hyperparameters\": {\n",
      "        \"batch_size\": 1024,\n",
      "        \"epoch\": 10\n",
      "    },\n",
      "    \"input_config_dir\": \"/opt/ml/input/config\",\n",
      "    \"input_data_config\": {\n",
      "        \"train\": {\n",
      "            \"TrainingInputMode\": \"File\",\n",
      "            \"S3DistributionType\": \"FullyReplicated\",\n",
      "            \"RecordWrapperType\": \"None\"\n",
      "        },\n",
      "        \"valid\": {\n",
      "            \"TrainingInputMode\": \"File\",\n",
      "            \"S3DistributionType\": \"FullyReplicated\",\n",
      "            \"RecordWrapperType\": \"None\"\n",
      "        }\n",
      "    },\n",
      "    \"input_dir\": \"/opt/ml/input\",\n",
      "    \"instance_groups\": [\n",
      "        \"homogeneousCluster\"\n",
      "    ],\n",
      "    \"instance_groups_dict\": {\n",
      "        \"homogeneousCluster\": {\n",
      "            \"instance_group_name\": \"homogeneousCluster\",\n",
      "            \"instance_type\": \"ml.g4dn.xlarge\",\n",
      "            \"hosts\": [\n",
      "                \"algo-1\"\n",
      "            ]\n",
      "        }\n",
      "    },\n",
      "    \"is_hetero\": false,\n",
      "    \"is_master\": true,\n",
      "    \"is_modelparallel_enabled\": null,\n",
      "    \"is_smddpmprun_installed\": false,\n",
      "    \"job_name\": \"sagemaker-merlin-tensorflow-2022-11-09-10-18-29-376\",\n",
      "    \"log_level\": 20,\n",
      "    \"master_hostname\": \"algo-1\",\n",
      "    \"model_dir\": \"/opt/ml/model\",\n",
      "    \"module_dir\": \"s3://sagemaker-us-east-1-843263297212/sagemaker-merlin-tensorflow-2022-11-09-10-18-29-376/source/sourcedir.tar.gz\",\n",
      "    \"module_name\": \"train\",\n",
      "    \"network_interface_name\": \"eth0\",\n",
      "    \"num_cpus\": 4,\n",
      "    \"num_gpus\": 1,\n",
      "    \"num_neurons\": 0,\n",
      "    \"output_data_dir\": \"/opt/ml/output/data\",\n",
      "    \"output_dir\": \"/opt/ml/output\",\n",
      "    \"output_intermediate_dir\": \"/opt/ml/output/intermediate\",\n",
      "    \"resource_config\": {\n",
      "        \"current_host\": \"algo-1\",\n",
      "        \"current_instance_type\": \"ml.g4dn.xlarge\",\n",
      "        \"current_group_name\": \"homogeneousCluster\",\n",
      "        \"hosts\": [\n",
      "            \"algo-1\"\n",
      "        ],\n",
      "        \"instance_groups\": [\n",
      "            {\n",
      "                \"instance_group_name\": \"homogeneousCluster\",\n",
      "                \"instance_type\": \"ml.g4dn.xlarge\",\n",
      "                \"hosts\": [\n",
      "                    \"algo-1\"\n",
      "                ]\n",
      "            }\n",
      "        ],\n",
      "        \"network_interface_name\": \"eth0\"\n",
      "    },\n",
      "    \"user_entry_point\": \"train.py\"\u001b[0m\n",
      "\u001b[34m}\u001b[0m\n",
      "\u001b[34mEnvironment variables:\u001b[0m\n",
      "\u001b[34mSM_HOSTS=[\"algo-1\"]\u001b[0m\n",
      "\u001b[34mSM_NETWORK_INTERFACE_NAME=eth0\u001b[0m\n",
      "\u001b[34mSM_HPS={\"batch_size\":1024,\"epoch\":10}\u001b[0m\n",
      "\u001b[34mSM_USER_ENTRY_POINT=train.py\u001b[0m\n",
      "\u001b[34mSM_FRAMEWORK_PARAMS={}\u001b[0m\n",
      "\u001b[34mSM_RESOURCE_CONFIG={\"current_group_name\":\"homogeneousCluster\",\"current_host\":\"algo-1\",\"current_instance_type\":\"ml.g4dn.xlarge\",\"hosts\":[\"algo-1\"],\"instance_groups\":[{\"hosts\":[\"algo-1\"],\"instance_group_name\":\"homogeneousCluster\",\"instance_type\":\"ml.g4dn.xlarge\"}],\"network_interface_name\":\"eth0\"}\u001b[0m\n",
      "\u001b[34mSM_INPUT_DATA_CONFIG={\"train\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"},\"valid\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"}}\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_DATA_DIR=/opt/ml/output/data\u001b[0m\n",
      "\u001b[34mSM_CHANNELS=[\"train\",\"valid\"]\u001b[0m\n",
      "\u001b[34mSM_CURRENT_HOST=algo-1\u001b[0m\n",
      "\u001b[34mSM_CURRENT_INSTANCE_TYPE=ml.g4dn.xlarge\u001b[0m\n",
      "\u001b[34mSM_CURRENT_INSTANCE_GROUP=homogeneousCluster\u001b[0m\n",
      "\u001b[34mSM_CURRENT_INSTANCE_GROUP_HOSTS=[\"algo-1\"]\u001b[0m\n",
      "\u001b[34mSM_INSTANCE_GROUPS=[\"homogeneousCluster\"]\u001b[0m\n",
      "\u001b[34mSM_INSTANCE_GROUPS_DICT={\"homogeneousCluster\":{\"hosts\":[\"algo-1\"],\"instance_group_name\":\"homogeneousCluster\",\"instance_type\":\"ml.g4dn.xlarge\"}}\u001b[0m\n",
      "\u001b[34mSM_DISTRIBUTION_INSTANCE_GROUPS=[]\u001b[0m\n",
      "\u001b[34mSM_IS_HETERO=false\u001b[0m\n",
      "\u001b[34mSM_MODULE_NAME=train\u001b[0m\n",
      "\u001b[34mSM_LOG_LEVEL=20\u001b[0m\n",
      "\u001b[34mSM_FRAMEWORK_MODULE=\u001b[0m\n",
      "\u001b[34mSM_INPUT_DIR=/opt/ml/input\u001b[0m\n",
      "\u001b[34mSM_INPUT_CONFIG_DIR=/opt/ml/input/config\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_DIR=/opt/ml/output\u001b[0m\n",
      "\u001b[34mSM_NUM_CPUS=4\u001b[0m\n",
      "\u001b[34mSM_NUM_GPUS=1\u001b[0m\n",
      "\u001b[34mSM_NUM_NEURONS=0\u001b[0m\n",
      "\u001b[34mSM_MODEL_DIR=/opt/ml/model\u001b[0m\n",
      "\u001b[34mSM_MODULE_DIR=s3://sagemaker-us-east-1-843263297212/sagemaker-merlin-tensorflow-2022-11-09-10-18-29-376/source/sourcedir.tar.gz\u001b[0m\n",
      "\u001b[34mSM_TRAINING_ENV={\"additional_framework_parameters\":{},\"channel_input_dirs\":{\"train\":\"/opt/ml/input/data/train\",\"valid\":\"/opt/ml/input/data/valid\"},\"current_host\":\"algo-1\",\"current_instance_group\":\"homogeneousCluster\",\"current_instance_group_hosts\":[\"algo-1\"],\"current_instance_type\":\"ml.g4dn.xlarge\",\"distribution_hosts\":[],\"distribution_instance_groups\":[],\"framework_module\":null,\"hosts\":[\"algo-1\"],\"hyperparameters\":{\"batch_size\":1024,\"epoch\":10},\"input_config_dir\":\"/opt/ml/input/config\",\"input_data_config\":{\"train\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"},\"valid\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"}},\"input_dir\":\"/opt/ml/input\",\"instance_groups\":[\"homogeneousCluster\"],\"instance_groups_dict\":{\"homogeneousCluster\":{\"hosts\":[\"algo-1\"],\"instance_group_name\":\"homogeneousCluster\",\"instance_type\":\"ml.g4dn.xlarge\"}},\"is_hetero\":false,\"is_master\":true,\"is_modelparallel_enabled\":null,\"is_smddpmprun_installed\":false,\"job_name\":\"sagemaker-merlin-tensorflow-2022-11-09-10-18-29-376\",\"log_level\":20,\"master_hostname\":\"algo-1\",\"model_dir\":\"/opt/ml/model\",\"module_dir\":\"s3://sagemaker-us-east-1-843263297212/sagemaker-merlin-tensorflow-2022-11-09-10-18-29-376/source/sourcedir.tar.gz\",\"module_name\":\"train\",\"network_interface_name\":\"eth0\",\"num_cpus\":4,\"num_gpus\":1,\"num_neurons\":0,\"output_data_dir\":\"/opt/ml/output/data\",\"output_dir\":\"/opt/ml/output\",\"output_intermediate_dir\":\"/opt/ml/output/intermediate\",\"resource_config\":{\"current_group_name\":\"homogeneousCluster\",\"current_host\":\"algo-1\",\"current_instance_type\":\"ml.g4dn.xlarge\",\"hosts\":[\"algo-1\"],\"instance_groups\":[{\"hosts\":[\"algo-1\"],\"instance_group_name\":\"homogeneousCluster\",\"instance_type\":\"ml.g4dn.xlarge\"}],\"network_interface_name\":\"eth0\"},\"user_entry_point\":\"train.py\"}\u001b[0m\n",
      "\u001b[34mSM_USER_ARGS=[\"--batch_size\",\"1024\",\"--epoch\",\"10\"]\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_INTERMEDIATE_DIR=/opt/ml/output/intermediate\u001b[0m\n",
      "\u001b[34mSM_CHANNEL_TRAIN=/opt/ml/input/data/train\u001b[0m\n",
      "\u001b[34mSM_CHANNEL_VALID=/opt/ml/input/data/valid\u001b[0m\n",
      "\u001b[34mSM_HP_BATCH_SIZE=1024\u001b[0m\n",
      "\u001b[34mSM_HP_EPOCH=10\u001b[0m\n",
      "\u001b[34mPYTHONPATH=/opt/ml/code:/usr/local/bin:/opt/tritonserver:/usr/local/lib/python3.8/dist-packages:/usr/lib/python38.zip:/usr/lib/python3.8:/usr/lib/python3.8/lib-dynload:/usr/local/lib/python3.8/dist-packages/faiss-1.7.2-py3.8.egg:/usr/local/lib/python3.8/dist-packages/merlin_sok-1.1.4-py3.8-linux-x86_64.egg:/usr/local/lib/python3.8/dist-packages/merlin_hps-1.0.0-py3.8-linux-x86_64.egg:/usr/lib/python3/dist-packages\u001b[0m\n",
      "\u001b[34mInvoking script with the following command:\u001b[0m\n",
      "\u001b[34m/usr/bin/python3 train.py --batch_size 1024 --epoch 10\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:03,486 sagemaker-training-toolkit INFO     Exceptions not imported for SageMaker Debugger as it is not installed.\u001b[0m\n",
      "\n",
      "2022-11-09 10:27:16 Training - Training image download completed. Training in progress.\u001b[34m2022-11-09 10:27:08.761711: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.818302: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.819693: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.819906: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.894084: I tensorflow/core/platform/cpu_feature_guard.cc:194] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE3 SSE4.1 SSE4.2 AVX\u001b[0m\n",
      "\u001b[34mTo enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.895367: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.895631: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:12.895807: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:16.651703: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:16.651981: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:16.652183: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\u001b[0m\n",
      "\u001b[34m2022-11-09 10:27:16.653025: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 10752 MB memory:  -> device: 0, name: Tesla T4, pci bus id: 0000:00:1e.0, compute capability: 7.5\u001b[0m\n",
      "\u001b[34mWorkflow saved to /tmp/tmp5fpdavsc/workflow.\u001b[0m\n",
      "\u001b[34mbatch_size = 1024, epochs = 10\u001b[0m\n",
      "\u001b[34mEpoch 1/10\u001b[0m\n",
      "\u001b[34m684/684 - 14s - loss: 0.6932 - auc: 0.4998 - regularization_loss: 0.0000e+00 - val_loss: 0.6931 - val_auc: 0.5000 - val_regularization_loss: 0.0000e+00 - 14s/epoch - 20ms/step\u001b[0m\n",
      "\u001b[34mEpoch 2/10\u001b[0m\n",
      "\u001b[34m684/684 - 8s - loss: 0.6931 - auc: 0.5026 - regularization_loss: 0.0000e+00 - val_loss: 0.6932 - val_auc: 0.4990 - val_regularization_loss: 0.0000e+00 - 8s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 3/10\u001b[0m\n",
      "\u001b[34m684/684 - 7s - loss: 0.6922 - auc: 0.5222 - regularization_loss: 0.0000e+00 - val_loss: 0.6941 - val_auc: 0.4989 - val_regularization_loss: 0.0000e+00 - 7s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 4/10\u001b[0m\n",
      "\u001b[34m684/684 - 7s - loss: 0.6858 - auc: 0.5509 - regularization_loss: 0.0000e+00 - val_loss: 0.6991 - val_auc: 0.4994 - val_regularization_loss: 0.0000e+00 - 7s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 5/10\u001b[0m\n",
      "\u001b[34m684/684 - 7s - loss: 0.6790 - auc: 0.5660 - regularization_loss: 0.0000e+00 - val_loss: 0.7052 - val_auc: 0.4993 - val_regularization_loss: 0.0000e+00 - 7s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 6/10\u001b[0m\n",
      "\u001b[34m684/684 - 8s - loss: 0.6751 - auc: 0.5722 - regularization_loss: 0.0000e+00 - val_loss: 0.7096 - val_auc: 0.4994 - val_regularization_loss: 0.0000e+00 - 8s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 7/10\u001b[0m\n",
      "\u001b[34m684/684 - 7s - loss: 0.6722 - auc: 0.5755 - regularization_loss: 0.0000e+00 - val_loss: 0.7184 - val_auc: 0.4991 - val_regularization_loss: 0.0000e+00 - 7s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 8/10\u001b[0m\n",
      "\u001b[34m684/684 - 7s - loss: 0.6700 - auc: 0.5777 - regularization_loss: 0.0000e+00 - val_loss: 0.7289 - val_auc: 0.4990 - val_regularization_loss: 0.0000e+00 - 7s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 9/10\u001b[0m\n",
      "\u001b[34m684/684 - 8s - loss: 0.6687 - auc: 0.5792 - regularization_loss: 0.0000e+00 - val_loss: 0.7404 - val_auc: 0.4994 - val_regularization_loss: 0.0000e+00 - 8s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34mEpoch 10/10\u001b[0m\n",
      "\u001b[34m684/684 - 8s - loss: 0.6678 - auc: 0.5801 - regularization_loss: 0.0000e+00 - val_loss: 0.7393 - val_auc: 0.4988 - val_regularization_loss: 0.0000e+00 - 8s/epoch - 11ms/step\u001b[0m\n",
      "\u001b[34m/usr/lib/python3/dist-packages/requests/__init__.py:89: RequestsDependencyWarning: urllib3 (1.26.12) or chardet (3.0.4) doesn't match a supported version!\n",
      "  warnings.warn(\"urllib3 ({}) or chardet ({}) doesn't match a supported \"\u001b[0m\n",
      "\u001b[34m/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.USER_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.USER: 'user'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\u001b[0m\n",
      "\u001b[34m/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.ITEM_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.ITEM: 'item'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\u001b[0m\n",
      "\u001b[34mWARNING:absl:Found untraced functions such as train_compute_metrics, model_context_layer_call_fn, model_context_layer_call_and_return_conditional_losses, output_layer_layer_call_fn, output_layer_layer_call_and_return_conditional_losses while saving (showing 5 of 97). These functions will not be directly callable after loading.\u001b[0m\n",
      "\u001b[34mINFO:__main__:Model saved to /tmp/tmp5fpdavsc/dlrm.\u001b[0m\n",
      "\u001b[34mModel saved to /tmp/tmp5fpdavsc/dlrm.\u001b[0m\n",
      "\u001b[34mWARNING:absl:Found untraced functions such as train_compute_metrics, model_context_layer_call_fn, model_context_layer_call_and_return_conditional_losses, output_layer_layer_call_fn, output_layer_layer_call_and_return_conditional_losses while saving (showing 5 of 97). These functions will not be directly callable after loading.\u001b[0m\n",
      "\u001b[34m/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.USER_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.USER: 'user'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\u001b[0m\n",
      "\u001b[34m/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.ITEM_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.ITEM: 'item'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\u001b[0m\n",
      "\u001b[34mWARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.\u001b[0m\n",
      "\u001b[34mWARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.\u001b[0m\n",
      "\u001b[34mWARNING:absl:Found untraced functions such as train_compute_metrics, model_context_layer_call_fn, model_context_layer_call_and_return_conditional_losses, output_layer_layer_call_fn, output_layer_layer_call_and_return_conditional_losses while saving (showing 5 of 97). These functions will not be directly callable after loading.\u001b[0m\n",
      "\u001b[34mEnsemble graph saved to /opt/ml/model.\u001b[0m\n",
      "\u001b[34mINFO:__main__:Ensemble graph saved to /opt/ml/model.\u001b[0m\n",
      "\u001b[34m2022-11-09 10:29:21,498 sagemaker-training-toolkit INFO     Reporting training SUCCESS\u001b[0m\n",
      "\n",
      "2022-11-09 10:29:41 Uploading - Uploading generated training model\n",
      "2022-11-09 10:29:41 Completed - Training job completed\n",
      "Training seconds: 589\n",
      "Billable seconds: 589\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "\n",
    "training_instance_type = \"ml.g4dn.xlarge\"  # GPU instance, T4\n",
    "\n",
    "estimator = Estimator(\n",
    "    role=role,\n",
    "    instance_count=1,\n",
    "    instance_type=training_instance_type,\n",
    "    image_uri=ecr_image,\n",
    "    entry_point=\"train.py\",\n",
    "    hyperparameters={\n",
    "        \"batch_size\": 1_024,\n",
    "        \"epoch\": 10,\n",
    "    },\n",
    ")\n",
    "\n",
    "estimator.fit(\n",
    "    {\n",
    "        \"train\": f\"{data_location}/train/\",\n",
    "        \"valid\": f\"{data_location}/valid/\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d44a96bd-a6ab-4ff7-9232-9e0f0c330ecb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-east-1-843263297212/sagemaker-merlin-tensorflow-2022-11-09-10-18-29-376/output/model.tar.gz\n"
     ]
    }
   ],
   "source": [
    "print(estimator.model_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c5d6d979-1976-46dd-bf88-669bbad39ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.s3 import S3Downloader as s3down\n",
    "\n",
    "s3down.download(estimator.model_data, \"/tmp/ensemble/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d90341dc-8c5a-4500-acf8-1491664cd426",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_predicttensorflow/\n",
      "1_predicttensorflow/config.pbtxt\n",
      "1_predicttensorflow/1/\n",
      "1_predicttensorflow/1/model.savedmodel/\n",
      "1_predicttensorflow/1/model.savedmodel/assets/\n",
      "1_predicttensorflow/1/model.savedmodel/variables/\n",
      "1_predicttensorflow/1/model.savedmodel/variables/variables.index\n",
      "1_predicttensorflow/1/model.savedmodel/variables/variables.data-00000-of-00001\n",
      "1_predicttensorflow/1/model.savedmodel/saved_model.pb\n",
      "1_predicttensorflow/1/model.savedmodel/keras_metadata.pb\n",
      "executor_model/\n",
      "executor_model/config.pbtxt\n",
      "executor_model/1/\n",
      "0_transformworkflow/\n",
      "0_transformworkflow/config.pbtxt\n",
      "0_transformworkflow/1/\n",
      "0_transformworkflow/1/model.py\n",
      "0_transformworkflow/1/workflow/\n",
      "0_transformworkflow/1/workflow/categories/\n",
      "0_transformworkflow/1/workflow/categories/unique.user_profile.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_age.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_group.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_intentions.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.item_brand.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_geography.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_is_occupied.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_id.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_gender.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_shops.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.item_category.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_brands.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_consumption_2.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.item_id.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.item_shop.parquet\n",
      "0_transformworkflow/1/workflow/categories/unique.user_categories.parquet\n",
      "0_transformworkflow/1/workflow/workflow.pkl\n",
      "0_transformworkflow/1/workflow/metadata.json\n"
     ]
    }
   ],
   "source": [
    "! cd /tmp/ensemble && tar xvzf model.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e0c0a7-0815-4aa9-8536-5a2085406f7a",
   "metadata": {},
   "source": [
    "## Part 3: Retrieving Recommendations from Triton Inference Server\n",
    "\n",
    "Although we use the Sagemaker Python SDK to train our model, here we will use `boto3` to launch our inference endpoint as it offers more low-level control than the Python SDK.\n",
    "\n",
    "The model artificat `model.tar.gz` uploaded to S3 from the Sagemaker training job contained three directories: `0_transformworkflow` for the NVTabular workflow, `1_predicttensorflow` for the Tensorflow model, and `executor_model` for the ensemble graph that we can use in Triton.\n",
    "\n",
    "```shell\n",
    "/tmp/ensemble/\n",
    "├── 0_transformworkflow\n",
    "│   ├── 1\n",
    "│   │   ├── model.py\n",
    "│   │   └── workflow\n",
    "│   │       ├── categories\n",
    "│   │       │   ├── unique.item_brand.parquet\n",
    "│   │       │   ├── unique.item_category.parquet\n",
    "│   │       │   ├── unique.item_id.parquet\n",
    "│   │       │   ├── unique.item_shop.parquet\n",
    "│   │       │   ├── unique.user_age.parquet\n",
    "│   │       │   ├── unique.user_brands.parquet\n",
    "│   │       │   ├── unique.user_categories.parquet\n",
    "│   │       │   ├── unique.user_consumption_2.parquet\n",
    "│   │       │   ├── unique.user_gender.parquet\n",
    "│   │       │   ├── unique.user_geography.parquet\n",
    "│   │       │   ├── unique.user_group.parquet\n",
    "│   │       │   ├── unique.user_id.parquet\n",
    "│   │       │   ├── unique.user_intentions.parquet\n",
    "│   │       │   ├── unique.user_is_occupied.parquet\n",
    "│   │       │   ├── unique.user_profile.parquet\n",
    "│   │       │   └── unique.user_shops.parquet\n",
    "│   │       ├── metadata.json\n",
    "│   │       └── workflow.pkl\n",
    "│   └── config.pbtxt\n",
    "├── 1_predicttensorflow\n",
    "│   ├── 1\n",
    "│   │   └── model.savedmodel\n",
    "│   │       ├── assets\n",
    "│   │       ├── keras_metadata.pb\n",
    "│   │       ├── saved_model.pb\n",
    "│   │       └── variables\n",
    "│   │           ├── variables.data-00000-of-00001\n",
    "│   │           └── variables.index\n",
    "│   └── config.pbtxt\n",
    "├── executor_model\n",
    "│   ├── 1\n",
    "│   └── config.pbtxt\n",
    "└── model.tar.gz\n",
    "```\n",
    "\n",
    "We specify that we only want to use `executor_model` in Triton by passing the environment variable `SAGEMAKER_TRITON_DEFAULT_MODEL_NAME`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2aaa86d7-093e-4bdc-a606-a730c1bf9a5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model Arn: arn:aws:sagemaker:us-east-1:843263297212:model/model-triton-merlin-ensemble-2022-11-09-10-29-57\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "import boto3\n",
    "\n",
    "sm_client = boto3.client(service_name=\"sagemaker\")\n",
    "\n",
    "container = {\n",
    "    \"Image\": ecr_image,\n",
    "    \"ModelDataUrl\": estimator.model_data,\n",
    "    \"Environment\": {\n",
    "        \"SAGEMAKER_TRITON_TENSORFLOW_VERSION\": \"2\",\n",
    "        \"SAGEMAKER_TRITON_DEFAULT_MODEL_NAME\": \"executor_model\",\n",
    "    },\n",
    "}\n",
    "\n",
    "model_name = \"model-triton-merlin-ensemble-\" + time.strftime(\n",
    "    \"%Y-%m-%d-%H-%M-%S\", time.gmtime()\n",
    ")\n",
    "\n",
    "create_model_response = sm_client.create_model(\n",
    "    ModelName=model_name, ExecutionRoleArn=role, PrimaryContainer=container\n",
    ")\n",
    "\n",
    "model_arn = create_model_response[\"ModelArn\"]\n",
    "\n",
    "print(f\"Model Arn: {model_arn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d97661e2-e54b-435c-96e3-63507418afe1",
   "metadata": {},
   "source": [
    "We again use the `g4dn` GPU instance that are equipped with NVIDIA T4 GPUs for launching the Triton inference server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9cade896-0e92-42ff-b507-96a82a248814",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Endpoint Config Arn: arn:aws:sagemaker:us-east-1:843263297212:endpoint-config/endpoint-config-triton-merlin-ensemble-2022-11-09-10-29-58\n"
     ]
    }
   ],
   "source": [
    "endpoint_instance_type = \"ml.g4dn.xlarge\"\n",
    "\n",
    "endpoint_config_name = \"endpoint-config-triton-merlin-ensemble-\" + time.strftime(\n",
    "    \"%Y-%m-%d-%H-%M-%S\", time.gmtime()\n",
    ")\n",
    "\n",
    "create_endpoint_config_response = sm_client.create_endpoint_config(\n",
    "    EndpointConfigName=endpoint_config_name,\n",
    "    ProductionVariants=[\n",
    "        {\n",
    "            \"InstanceType\": endpoint_instance_type,\n",
    "            \"InitialVariantWeight\": 1,\n",
    "            \"InitialInstanceCount\": 1,\n",
    "            \"ModelName\": model_name,\n",
    "            \"VariantName\": \"AllTraffic\",\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "endpoint_config_arn = create_endpoint_config_response[\"EndpointConfigArn\"]\n",
    "\n",
    "print(f\"Endpoint Config Arn: {endpoint_config_arn}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "004eb730-0eaa-4945-b8be-6c820ed96485",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Endpoint Arn: arn:aws:sagemaker:us-east-1:843263297212:endpoint/endpoint-triton-merlin-ensemble-2022-11-09-10-29-58\n"
     ]
    }
   ],
   "source": [
    "endpoint_name = \"endpoint-triton-merlin-ensemble-\" + time.strftime(\n",
    "    \"%Y-%m-%d-%H-%M-%S\", time.gmtime()\n",
    ")\n",
    "\n",
    "create_endpoint_response = sm_client.create_endpoint(\n",
    "    EndpointName=endpoint_name, EndpointConfigName=endpoint_config_name\n",
    ")\n",
    "\n",
    "endpoint_arn = create_endpoint_response[\"EndpointArn\"]\n",
    "\n",
    "print(f\"Endpoint Arn: {endpoint_arn}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bd230b47-6a95-46ad-9555-ed8f1ee29be0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: Creating\n",
      "Endpoint Creation Status: InService\n",
      "Endpoint Arn: arn:aws:sagemaker:us-east-1:843263297212:endpoint/endpoint-triton-merlin-ensemble-2022-11-09-10-29-58\n",
      "Endpoint Status: InService\n"
     ]
    }
   ],
   "source": [
    "status = sm_client.describe_endpoint(EndpointName=endpoint_name)[\"EndpointStatus\"]\n",
    "print(f\"Endpoint Creation Status: {status}\")\n",
    "\n",
    "while status == \"Creating\":\n",
    "    time.sleep(60)\n",
    "    rv = sm_client.describe_endpoint(EndpointName=endpoint_name)\n",
    "    status = rv[\"EndpointStatus\"]\n",
    "    print(f\"Endpoint Creation Status: {status}\")\n",
    "\n",
    "endpoint_arn = rv[\"EndpointArn\"]\n",
    "\n",
    "print(f\"Endpoint Arn: {endpoint_arn}\")\n",
    "print(f\"Endpoint Status: {status}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acb458d8-3169-40cc-8cae-6b72a300394c",
   "metadata": {},
   "source": [
    "### Send a Request to Triton Inference Server to Transform a Raw Dataset\n",
    "\n",
    "Once we have an endpoint running, we can test it by sending requests.\n",
    "Here, we use the raw validation set and transform it using the saved NVTabular workflow we have downloaded from S3 in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9524fc1d-e613-4c24-b733-f7dcf673f584",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                     user_id  item_id  item_category  item_shop  item_brand  \\\n",
      "__null_dask_index__                                                           \n",
      "700000                    12        2              3        194          67   \n",
      "700001                    12       30             80       5621        1936   \n",
      "700002                    18        5             12        776         267   \n",
      "700003                    35        6             14        970         334   \n",
      "700004                    51       11             28       1939         668   \n",
      "700005                    22       83            226      15893        5474   \n",
      "700006                    13       38            102       7172        2470   \n",
      "700007                    10        7             17       1163         401   \n",
      "700008                     4        4              9        582         201   \n",
      "700009                     4       24             64       4458        1536   \n",
      "\n",
      "                     user_shops  user_profile  user_group  user_gender  \\\n",
      "__null_dask_index__                                                      \n",
      "700000                      636             1           1            1   \n",
      "700001                      636             1           1            1   \n",
      "700002                      983             1           1            1   \n",
      "700003                     1965             2           1            1   \n",
      "700004                     2890             3           1            1   \n",
      "700005                     1214             2           1            1   \n",
      "700006                      694             1           1            1   \n",
      "700007                      521             1           1            1   \n",
      "700008                      174             1           1            1   \n",
      "700009                      174             1           1            1   \n",
      "\n",
      "                     user_age  user_consumption_2  user_is_occupied  \\\n",
      "__null_dask_index__                                                   \n",
      "700000                      1                   1                 1   \n",
      "700001                      1                   1                 1   \n",
      "700002                      1                   1                 1   \n",
      "700003                      1                   1                 1   \n",
      "700004                      1                   1                 1   \n",
      "700005                      1                   1                 1   \n",
      "700006                      1                   1                 1   \n",
      "700007                      1                   1                 1   \n",
      "700008                      1                   1                 1   \n",
      "700009                      1                   1                 1   \n",
      "\n",
      "                     user_geography  user_intentions  user_brands  \\\n",
      "__null_dask_index__                                                 \n",
      "700000                            1              184          316   \n",
      "700001                            1              184          316   \n",
      "700002                            1              285          489   \n",
      "700003                            1              569          977   \n",
      "700004                            1              837         1436   \n",
      "700005                            1              352          604   \n",
      "700006                            1              201          345   \n",
      "700007                            1              151          259   \n",
      "700008                            1               51           87   \n",
      "700009                            1               51           87   \n",
      "\n",
      "                     user_categories  \n",
      "__null_dask_index__                   \n",
      "700000                            34  \n",
      "700001                            34  \n",
      "700002                            52  \n",
      "700003                           103  \n",
      "700004                           151  \n",
      "700005                            64  \n",
      "700006                            37  \n",
      "700007                            28  \n",
      "700008                            10  \n",
      "700009                            10  \n"
     ]
    }
   ],
   "source": [
    "from merlin.schema.tags import Tags\n",
    "from merlin.core.dispatch import get_lib\n",
    "from nvtabular.workflow import Workflow\n",
    "\n",
    "df_lib = get_lib()\n",
    "\n",
    "workflow = Workflow.load(\"/tmp/ensemble/0_transformworkflow/1/workflow/\")\n",
    "\n",
    "label_columns = workflow.output_schema.select_by_tag(Tags.TARGET).column_names\n",
    "workflow.remove_inputs(label_columns)\n",
    "\n",
    "# read in data for request\n",
    "batch = df_lib.read_parquet(\n",
    "    os.path.join(DATA_FOLDER, \"valid\", \"part.0.parquet\"),\n",
    "    columns=workflow.input_schema.column_names,\n",
    ")[:10]\n",
    "print(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "577d86d6-a17a-47e8-9966-8738f5bdcbd4",
   "metadata": {},
   "source": [
    "In the following code cell, we use a utility function provided in [Merlin Systems](https://github.com/NVIDIA-Merlin/systems) to convert our dataframe to the payload format that can be used as inference request format for Triton."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6dcae463-6d80-4a7a-981a-a9d034d49fa9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'{\"inputs\":[{\"name\":\"user_id\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"item_id\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"item_category\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"item_shop\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"item_brand\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_shops\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_profile\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_group\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_gender\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_age\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_consumption_2\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_is_occupied\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_geography\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_intentions\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_brands\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}},{\"name\":\"user_categories\",\"shape\":[10,1],\"datatype\":\"INT32\",\"parameters\":{\"binary_data_size\":40}}],\"parameters\":{\"binary_data_output\":true}}\\x0c\\x00\\x00\\x00\\x0c\\x00\\x00\\x00\\x12\\x00\\x00\\x00#\\x00\\x00\\x003\\x00\\x00\\x00\\x16\\x00\\x00\\x00\\r\\x00\\x00\\x00\\n\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x1e\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x0b\\x00\\x00\\x00S\\x00\\x00\\x00&\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x18\\x00\\x00\\x00\\x03\\x00\\x00\\x00P\\x00\\x00\\x00\\x0c\\x00\\x00\\x00\\x0e\\x00\\x00\\x00\\x1c\\x00\\x00\\x00\\xe2\\x00\\x00\\x00f\\x00\\x00\\x00\\x11\\x00\\x00\\x00\\t\\x00\\x00\\x00@\\x00\\x00\\x00\\xc2\\x00\\x00\\x00\\xf5\\x15\\x00\\x00\\x08\\x03\\x00\\x00\\xca\\x03\\x00\\x00\\x93\\x07\\x00\\x00\\x15>\\x00\\x00\\x04\\x1c\\x00\\x00\\x8b\\x04\\x00\\x00F\\x02\\x00\\x00j\\x11\\x00\\x00C\\x00\\x00\\x00\\x90\\x07\\x00\\x00\\x0b\\x01\\x00\\x00N\\x01\\x00\\x00\\x9c\\x02\\x00\\x00b\\x15\\x00\\x00\\xa6\\t\\x00\\x00\\x91\\x01\\x00\\x00\\xc9\\x00\\x00\\x00\\x00\\x06\\x00\\x00|\\x02\\x00\\x00|\\x02\\x00\\x00\\xd7\\x03\\x00\\x00\\xad\\x07\\x00\\x00J\\x0b\\x00\\x00\\xbe\\x04\\x00\\x00\\xb6\\x02\\x00\\x00\\t\\x02\\x00\\x00\\xae\\x00\\x00\\x00\\xae\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\xb8\\x00\\x00\\x00\\xb8\\x00\\x00\\x00\\x1d\\x01\\x00\\x009\\x02\\x00\\x00E\\x03\\x00\\x00`\\x01\\x00\\x00\\xc9\\x00\\x00\\x00\\x97\\x00\\x00\\x003\\x00\\x00\\x003\\x00\\x00\\x00<\\x01\\x00\\x00<\\x01\\x00\\x00\\xe9\\x01\\x00\\x00\\xd1\\x03\\x00\\x00\\x9c\\x05\\x00\\x00\\\\\\x02\\x00\\x00Y\\x01\\x00\\x00\\x03\\x01\\x00\\x00W\\x00\\x00\\x00W\\x00\\x00\\x00\"\\x00\\x00\\x00\"\\x00\\x00\\x004\\x00\\x00\\x00g\\x00\\x00\\x00\\x97\\x00\\x00\\x00@\\x00\\x00\\x00%\\x00\\x00\\x00\\x1c\\x00\\x00\\x00\\n\\x00\\x00\\x00\\n\\x00\\x00\\x00'\n"
     ]
    }
   ],
   "source": [
    "from merlin.systems.triton import convert_df_to_triton_input\n",
    "import tritonclient.http as httpclient\n",
    "\n",
    "inputs = convert_df_to_triton_input(workflow.input_schema, batch, httpclient.InferInput)\n",
    "\n",
    "request_body, header_length = httpclient.InferenceServerClient.generate_request_body(\n",
    "    inputs\n",
    ")\n",
    "\n",
    "print(request_body)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21cdb596-0cd8-415c-bab5-f750a213a938",
   "metadata": {},
   "source": [
    "Triton uses the [KServe community standard inference protocols](https://github.com/triton-inference-server/server/blob/main/docs/protocol/README.md).\n",
    "Here, we use the [binary+json format](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_binary_data.md) for optimal performance in the inference request.\n",
    "\n",
    "In order for Triton to correctly parse the binary payload, we have to specify the length of the request metadata in the header `json-header-size`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2484510b-4b50-410a-8b93-93ef8c546cb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicted sigmoid result:\n",
      " [[0.48595208]\n",
      " [0.4647554 ]\n",
      " [0.50048226]\n",
      " [0.53553176]\n",
      " [0.5209902 ]\n",
      " [0.54944164]\n",
      " [0.5032344 ]\n",
      " [0.475241  ]\n",
      " [0.5077254 ]\n",
      " [0.5009623 ]]\n"
     ]
    }
   ],
   "source": [
    "runtime_sm_client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "response = runtime_sm_client.invoke_endpoint(\n",
    "    EndpointName=endpoint_name,\n",
    "    ContentType=f\"application/vnd.sagemaker-triton.binary+json;json-header-size={header_length}\",\n",
    "    Body=request_body,\n",
    ")\n",
    "\n",
    "# Parse json header size length from the response\n",
    "header_length_prefix = \"application/vnd.sagemaker-triton.binary+json;json-header-size=\"\n",
    "header_length_str = response[\"ContentType\"][len(header_length_prefix):]\n",
    "\n",
    "# Read response body\n",
    "result = httpclient.InferenceServerClient.parse_response_body(\n",
    "    response[\"Body\"].read(), header_length=int(header_length_str)\n",
    ")\n",
    "output_data = result.as_numpy(\"click/binary_classification_task\")\n",
    "print(\"predicted sigmoid result:\\n\", output_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b660ca-6376-4615-aef4-474c66fa2bec",
   "metadata": {},
   "source": [
    "## Terminate endpoint and clean up artifacts\n",
    "\n",
    "Don't forget to clean up artifacts and terminate the endpoint, or the endpoint will continue to incur costs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2ce1630c-96aa-49ee-ac79-cf44120adb78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ResponseMetadata': {'RequestId': '6ad24616-5c7c-4525-a63c-62d1b06ee8ad',\n",
       "  'HTTPStatusCode': 200,\n",
       "  'HTTPHeaders': {'x-amzn-requestid': '6ad24616-5c7c-4525-a63c-62d1b06ee8ad',\n",
       "   'content-type': 'application/x-amz-json-1.1',\n",
       "   'content-length': '0',\n",
       "   'date': 'Wed, 09 Nov 2022 10:38:12 GMT'},\n",
       "  'RetryAttempts': 0}}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sm_client.delete_model(ModelName=model_name)\n",
    "sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)\n",
    "sm_client.delete_endpoint(EndpointName=endpoint_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}