examples/scaling-criteo/02-ETL-with-NVTabular.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_merlin_scaling-criteo-02-etl-with-nvtabular/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Scaling Criteo: ETL with NVTabular\n",
    "\n",
    "This notebook is created using the latest stable [merlin-hugectr](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-hugectr/tags), [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags), or [merlin-pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-pytorch/tags) container.\n",
    "\n",
    "## Overview\n",
    "\n",
    "NVTabular is a feature engineering and preprocessing library for tabular data designed to quickly and easily manipulate terabyte scale datasets used to train deep learning based recommender systems. It provides a high level abstraction to simplify code and accelerates computation on the GPU using the RAPIDS cuDF library.<br><br>\n",
    "\n",
    "**In this notebook, we will use NVTabular with multiple GPUs. The notebook can run on single GPU, as well**. Prerequisite is to be familiar with NVTabular and its API. You can read more NVTabular and its API in our [Documentation](https://nvidia-merlin.github.io/NVTabular/stable/Introduction.html) or check out our [NVTabular Examples](https://nvidia-merlin.github.io/NVTabular/stable/examples/index.html).<br><br>\n",
    "\n",
    "The full [Criteo 1TB Click Logs dataset](https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/) contains ~1.3 TB of uncompressed click logs containing over four billion samples spanning 24 days. In our benchmarks, we are able to preprocess and engineer features in **13.8min with 1x NVIDIA A100 GPU and 1.9min with 8x NVIDIA A100 GPUs**. This is a **speed-up of 100x-10000x** in comparison to different CPU versions, You can read more in our [blog](https://developer.nvidia.com/blog/announcing-the-nvtabular-open-beta-with-multi-gpu-support-and-new-data-loaders/).\n",
    "\n",
    "Our pipeline will be representative with most common preprocessing transformation for deep learning recommender models.\n",
    "\n",
    "* Categorical input features are `Categorified` to be continuous integers (0, ..., |C|) for the embedding layers\n",
    "* Missing values of continuous input features are filled with 0. Afterwards the continuous features are clipped and normalized.\n",
    "\n",
    "### Learning objectives\n",
    "In this notebook, we learn how to to scale ETLs with NVTabular\n",
    "\n",
    "- Learn to use larger than GPU/host memory datasets\n",
    "- Use multi-GPU or multi node for ETL\n",
    "- Apply common deep learning ETL workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ETL with NVTabular\n",
    "Here we'll show how to use NVTabular first as a preprocessing library to prepare the [Criteo 1TB Click Logs dataset](https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/) dataset. The following notebooks can use the output to train a deep learning model.\n",
    "\n",
    "### Data Prep\n",
    "The previous notebook [01-Download-Convert](./01-Download-Convert.ipynb) converted the tsv data published by Criteo into the parquet format that our accelerated readers prefer. Accelerating these pipelines on new hardware like GPUs may require us to make new choices about the representations we use to store that data, and parquet represents a strong alternative."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load the required libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Standard Libraries\n",
    "import os\n",
    "import re\n",
    "import shutil\n",
    "import warnings\n",
    "\n",
    "# External Dependencies\n",
    "import numpy as np\n",
    "import numba\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.distributed import Client\n",
    "\n",
    "# NVTabular\n",
    "import nvtabular as nvt\n",
    "from nvtabular.ops import (\n",
    "    Categorify,\n",
    "    Clip,\n",
    "    FillMissing,\n",
    "    Normalize,\n",
    "    AddMetadata\n",
    ")\n",
    "from merlin.core.compat import pynvml_mem_size, device_mem_size\n",
    "from merlin.schema.tags import Tags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once our data is ready, we'll define some high level parameters to describe where our data is and what it \"looks like\" at a high level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# define some information about where to get our data\n",
    "BASE_DIR = os.environ.get(\"BASE_DIR\", \"/raid/data/criteo\")\n",
    "INPUT_DATA_DIR = os.environ.get(\"INPUT_DATA_DIR\", BASE_DIR + \"/converted/criteo\")\n",
    "OUTPUT_DATA_DIR = os.environ.get(\"OUTPUT_DATA_DIR\", BASE_DIR + \"/test_dask/output\")\n",
    "USE_HUGECTR = bool(os.environ.get(\"USE_HUGECTR\", \"\"))\n",
    "stats_path = os.path.join(OUTPUT_DATA_DIR, \"test_dask/stats\")\n",
    "dask_workdir = os.path.join(OUTPUT_DATA_DIR, \"test_dask/workdir\")\n",
    "\n",
    "# Make sure we have a clean worker space for Dask\n",
    "if os.path.isdir(dask_workdir):\n",
    "    shutil.rmtree(dask_workdir)\n",
    "os.makedirs(dask_workdir)\n",
    "\n",
    "# Make sure we have a clean stats space for Dask\n",
    "if os.path.isdir(stats_path):\n",
    "    shutil.rmtree(stats_path)\n",
    "os.mkdir(stats_path)\n",
    "\n",
    "# Make sure we have a clean output path\n",
    "if os.path.isdir(OUTPUT_DATA_DIR):\n",
    "    shutil.rmtree(OUTPUT_DATA_DIR)\n",
    "os.mkdir(OUTPUT_DATA_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the last day as validation dataset and the remaining days as training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"day_{}.parquet\"\n",
    "num_days = len(\n",
    "    [i for i in os.listdir(INPUT_DATA_DIR) if re.match(fname.format(\"[0-9]{1,2}\"), i) is not None]\n",
    ")\n",
    "train_paths = [os.path.join(INPUT_DATA_DIR, fname.format(day)) for day in range(num_days - 1)]\n",
    "valid_paths = [\n",
    "    os.path.join(INPUT_DATA_DIR, fname.format(day)) for day in range(num_days - 1, num_days)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/raid/data/criteo/converted/criteo/day_0.parquet']\n",
      "['/raid/data/criteo/converted/criteo/day_1.parquet']\n"
     ]
    }
   ],
   "source": [
    "print(train_paths)\n",
    "print(valid_paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy a Distributed-Dask Cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we configure and deploy a Dask Cluster. Please, [read this document](https://github.com/NVIDIA-Merlin/NVTabular/blob/d419a4da29cf372f1547edc536729b0733560a44/bench/examples/MultiGPUBench.md) to know how to set the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-12-06 10:52:32,108 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2022-12-06 10:52:32,242 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-16739757-7554-11ed-88cc-2a33bb9638f6</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_cuda.LocalCUDACluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">LocalCUDACluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">c4f394c8</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 2\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 2\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 100.00 GiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "            <tr>\n",
       "    <td style=\"text-align: left;\"><strong>Status:</strong> running</td>\n",
       "    <td style=\"text-align: left;\"><strong>Using processes:</strong> True</td>\n",
       "</tr>\n",
       "\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-a605479f-02de-4ca6-8661-fd02b84655ac</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://127.0.0.1:44059\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 2\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 2\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 100.00 GiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:45833\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:38541/status\" target=\"_blank\">http://127.0.0.1:38541/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 50.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:43849\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /raid/data/criteo/test_dask/output/test_dask/workdir/dask-worker-space/worker-9hxldmq1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB-LS\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 31.75 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:35261\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:46685/status\" target=\"_blank\">http://127.0.0.1:46685/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 50.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:37015\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /raid/data/criteo/test_dask/output/test_dask/workdir/dask-worker-space/worker-nsltf61z\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB-LS\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 31.75 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:44059' processes=2 threads=2, memory=100.00 GiB>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dask dashboard\n",
    "dashboard_port = \"8787\"\n",
    "\n",
    "# Deploy a Single-Machine Multi-GPU Cluster\n",
    "protocol = \"tcp\"  # \"tcp\" or \"ucx\"\n",
    "if numba.cuda.is_available():\n",
    "    NUM_GPUS = list(range(len(numba.cuda.gpus)))\n",
    "else:\n",
    "    NUM_GPUS = []\n",
    "visible_devices = \",\".join([str(n) for n in NUM_GPUS])  # Select devices to place workers\n",
    "device_limit_frac = 0.7  # Spill GPU-Worker memory to host at this limit.\n",
    "device_pool_frac = 0.8\n",
    "part_mem_frac = 0.15\n",
    "\n",
    "# Use total device size to calculate args.device_limit_frac\n",
    "device_size = device_mem_size(kind=\"total\")\n",
    "device_limit = int(device_limit_frac * device_size)\n",
    "device_pool_size = int(device_pool_frac * device_size)\n",
    "part_size = int(part_mem_frac * device_size)\n",
    "\n",
    "# Check if any device memory is already occupied\n",
    "for dev in visible_devices.split(\",\"):\n",
    "    fmem = pynvml_mem_size(kind=\"free\", index=int(dev))\n",
    "    used = (device_size - fmem) / 1e9\n",
    "    if used > 1.0:\n",
    "        warnings.warn(f\"BEWARE - {used} GB is already occupied on device {int(dev)}!\")\n",
    "\n",
    "cluster = None  # (Optional) Specify existing scheduler port\n",
    "if cluster is None:\n",
    "    cluster = LocalCUDACluster(\n",
    "        protocol=protocol,\n",
    "        n_workers=len(visible_devices.split(\",\")),\n",
    "        CUDA_VISIBLE_DEVICES=visible_devices,\n",
    "        device_memory_limit=device_limit,\n",
    "        local_directory=dask_workdir,\n",
    "        dashboard_address=\":\" + dashboard_port,\n",
    "        rmm_pool_size=(device_pool_size // 256) * 256\n",
    "    )\n",
    "\n",
    "# Create the distributed client\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it. We initialized our Dask cluster and NVTabular will execute the workflow on multiple GPUs. Similar, we could define a cluster with multiple nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining our Preprocessing Pipeline\n",
    "At this point, our data still isn't in a form that's ideal for consumption by neural networks. The most pressing issues are missing values and the fact that our categorical variables are still represented by random, discrete identifiers, and need to be transformed into contiguous indices that can be leveraged by a learned embedding. Less pressing, but still important for learning dynamics, are the distributions of our continuous variables, which are distributed across multiple orders of magnitude and are uncentered (i.e. E[x] != 0).\n",
    "\n",
    "We can fix these issues in a concise and GPU-accelerated manner with an NVTabular `Workflow`. More information about NVTabular's API can be found in its [Documentation](https://nvidia-merlin.github.io/NVTabular/stable/Introduction.html).\n",
    "\n",
    "#### Frequency Thresholding\n",
    "One interesting thing worth pointing out is that we're using _frequency thresholding_ in our `Categorify` op. This handy functionality will map all categories which occur in the dataset with some threshold level of infrequency (which we've set here to be 15 occurrences throughout the dataset) to the _same_ index, keeping the model from overfitting to sparse signals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# define our dataset schema\n",
    "CONTINUOUS_COLUMNS = [\"I\" + str(x) for x in range(1, 14)]\n",
    "CATEGORICAL_COLUMNS = [\"C\" + str(x) for x in range(1, 27)]\n",
    "LABEL_COLUMNS = [\"label\"]\n",
    "COLUMNS = CONTINUOUS_COLUMNS + CATEGORICAL_COLUMNS + LABEL_COLUMNS\n",
    "\n",
    "num_buckets = 10000000\n",
    "categorify_op = Categorify(out_path=stats_path, max_size=num_buckets, dtype='int32')\n",
    "cat_features = CATEGORICAL_COLUMNS >> categorify_op\n",
    "cont_features = CONTINUOUS_COLUMNS >> FillMissing() >> Clip(min_value=0) >> Normalize(out_dtype='float32')\n",
    "label_features = LABEL_COLUMNS >> AddMetadata(\n",
    "    tags=[str(Tags.BINARY_CLASSIFICATION), \"target\"]\n",
    ")\n",
    "\n",
    "features = cat_features + cont_features + label_features\n",
    "\n",
    "workflow = nvt.Workflow(features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to enforce the required HugeCTR data types, so we set them in a dictionary and give as an argument when creating our dataset. The dictionary defines the output datatypes of our datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "dict_dtypes = {}\n",
    "\n",
    "# The environment variable USE_HUGECTR defines, if we want to use the output for HugeCTR or another framework\n",
    "for col in CATEGORICAL_COLUMNS:\n",
    "    dict_dtypes[col] = np.int64 if USE_HUGECTR else np.int32\n",
    "\n",
    "for col in CONTINUOUS_COLUMNS:\n",
    "    dict_dtypes[col] = np.float32\n",
    "\n",
    "for col in LABEL_COLUMNS:\n",
    "    dict_dtypes[col] = np.int32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now instantiate dataset iterators to loop through our dataset (which we couldn't fit into GPU memory)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "train_dataset = nvt.Dataset(train_paths, engine=\"parquet\", part_size=part_size)\n",
    "valid_dataset = nvt.Dataset(valid_paths, engine=\"parquet\", part_size=part_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now run them through our workflows to collect statistics on the train set, then transform and save to parquet files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "output_train_dir = os.path.join(OUTPUT_DATA_DIR, \"train/\")\n",
    "output_valid_dir = os.path.join(OUTPUT_DATA_DIR, \"valid/\")\n",
    "! mkdir -p $output_train_dir\n",
    "! mkdir -p $output_valid_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For reference, let's time it to see how long it takes..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.38 s, sys: 260 ms, total: 2.64 s\n",
      "Wall time: 12.6 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<nvtabular.workflow.workflow.Workflow at 0x7fdacec4fdc0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "workflow.fit(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 980 ms, sys: 202 ms, total: 1.18 s\n",
      "Wall time: 19.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "workflow.transform(train_dataset).to_parquet(\n",
    "    output_files=len(NUM_GPUS),\n",
    "    output_path=output_train_dir,\n",
    "    shuffle=nvt.io.Shuffle.PER_PARTITION,\n",
    "    dtypes=dict_dtypes,\n",
    "    cats=CATEGORICAL_COLUMNS,\n",
    "    conts=CONTINUOUS_COLUMNS,\n",
    "    labels=LABEL_COLUMNS,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 801 ms, sys: 203 ms, total: 1 s\n",
      "Wall time: 15.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "workflow.transform(valid_dataset).to_parquet(\n",
    "    output_path=output_valid_dir,\n",
    "    dtypes=dict_dtypes,\n",
    "    cats=CATEGORICAL_COLUMNS,\n",
    "    conts=CONTINUOUS_COLUMNS,\n",
    "    labels=LABEL_COLUMNS,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next notebooks, we will train a deep learning model. Our training pipeline requires information about the data schema to define the neural network architecture. We will save the NVTabular workflow to disk so that we can restore it in the next notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow.save(os.path.join(OUTPUT_DATA_DIR, \"workflow\"))"
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "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"
  },
  "merlin": {
   "containers": [
    "nvcr.io/nvidia/merlin/merlin-hugectr:latest",
    "nvcr.io/nvidia/merlin/merlin-tensorflow:latest",
    "nvcr.io/nvidia/merlin/merlin-pytorch:latest"
   ]
  },
  "mimetype": "text/x-python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}