notebooks/hugectr_e2e_demo_with_nvtabular.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdfec37b",
   "metadata": {},
   "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",
   "id": "71a30b86-9963-4ee7-b91d-5e3c77a47996",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_hugectr_training-with-hdfs/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# HugeCTR End-end Example with NVTabular"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2e738caa-67fa-4d02-95c8-4f3066932a9e",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this sample notebook, we are going to:\n",
    "1. Preprocess data using **NVTabular**\n",
    "2. Training model with **HugeCTR**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "655ef37c-eab7-43e5-9948-96d3845c655e",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To setup the environment, refer to [HugeCTR Example Notebooks](../notebooks) and follow the instructions there before running the following."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53252cc0-dbee-4b23-818a-e67c6c5d85c7",
   "metadata": {},
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e2be5dc5-c75e-4e0a-a8a9-30287bc8a6fe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "004808a1-698e-43b3-af14-b906288e240c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mkdir -p /hugectr_e2e\n",
    "!mkdir -p /hugectr_e2e/criteo/train\n",
    "!mkdir -p /hugectr_e2e/criteo/val\n",
    "!mkdir -p /hugectr_e2e/model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e899f72f-acba-4647-aba7-3f8bbba92e9e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "BASE_DIR = os.environ.get(\"BASE_DIR\", \"/hugectr_e2e\")\n",
    "DATA_DIR = os.environ.get(\"DATA_DIR\", BASE_DIR + \"/criteo\")\n",
    "TRAIN_DIR = os.environ.get(\"TRAIN_DIR\", DATA_DIR +\"/train\")\n",
    "VAL_DIR = os.environ.get(\"VAL_DIR\", DATA_DIR +\"/val\")\n",
    "MODEL_DIR = os.environ.get(\"MODEL_DIR\", BASE_DIR + \"/model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4f1b32-cd56-4acb-b644-e711e985b09d",
   "metadata": {},
   "source": [
    "**Download the Criteo data for 1 day:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5226a424-1e9b-424d-b343-146550174302",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ln: failed to create symbolic link '/hugectr_e2e/criteo/day_0': File exists\n"
     ]
    }
   ],
   "source": [
    "#!wget -P $DATA_DIR https://storage.googleapis.com/criteo-cail-datasets/day_0.gz  #decomment this line to download, otherwise soft link the data.\n",
    "#!gzip -d -c $DATA_DIR/day_0.gz > $DATA_DIR/day_0\n",
    "INPUT_DATA = os.environ.get(\"INPUT_DATA\", DATA_DIR + \"/day_0\")\n",
    "!ln -s $INPUT_DATA $DATA_DIR/day_0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aee65599-7f49-4672-b29e-ef10e89edc8f",
   "metadata": {},
   "source": [
    "**Unzip and split data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e212a66e-417e-4b82-9db1-86a301264ac3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!head -n 10000000 $DATA_DIR/day_0 > $DATA_DIR/train/train.txt\n",
    "!tail -n 2000000 $DATA_DIR/day_0 > $DATA_DIR/val/test.txt "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392fa5f5-e62b-4942-8abd-847b15e52873",
   "metadata": {},
   "source": [
    "## Data Preprocessing using NVTabular"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "368e7ab2-94bd-4946-8033-311db26bcf14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "warnings.simplefilter(\"ignore\", UserWarning)\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import argparse\n",
    "import glob\n",
    "import time\n",
    "import numpy as np\n",
    "import shutil\n",
    "import numba\n",
    "\n",
    "import dask_cudf\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.distributed import Client\n",
    "\n",
    "import nvtabular as nvt\n",
    "from merlin.core.compat import device_mem_size, pynvml_mem_size\n",
    "from nvtabular.ops import (\n",
    "    Categorify,\n",
    "    Clip,\n",
    "    FillMissing,\n",
    "    Normalize,\n",
    "    get_embedding_sizes,\n",
    ")\n",
    "\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(format=\"%(asctime)s %(message)s\")\n",
    "logging.root.setLevel(logging.NOTSET)\n",
    "logging.getLogger(\"numba\").setLevel(logging.WARNING)\n",
    "logging.getLogger(\"asyncio\").setLevel(logging.WARNING)\n",
    "\n",
    "# define dataset schema\n",
    "CATEGORICAL_COLUMNS=[\"C\" + str(x) for x in range(1, 27)]\n",
    "CONTINUOUS_COLUMNS=[\"I\" + str(x) for x in range(1, 14)]\n",
    "LABEL_COLUMNS = ['label']\n",
    "COLUMNS =  LABEL_COLUMNS + CONTINUOUS_COLUMNS +  CATEGORICAL_COLUMNS\n",
    "#/samples/criteo mode doesn't have dense features\n",
    "criteo_COLUMN=LABEL_COLUMNS +  CATEGORICAL_COLUMNS\n",
    "#For new feature cross columns\n",
    "CROSS_COLUMNS = [\"C1_C2\", \"C3_C4\"]\n",
    "\n",
    "NUM_INTEGER_COLUMNS = 13\n",
    "NUM_CATEGORICAL_COLUMNS = 26\n",
    "NUM_TOTAL_COLUMNS = 1 + NUM_INTEGER_COLUMNS + NUM_CATEGORICAL_COLUMNS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c31b4c02-e344-4c05-9c56-24268ea83c5c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-26 03:09:04,061 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,061 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,062 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,062 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,063 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,063 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,064 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,064 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,072 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,072 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,085 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,085 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,087 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,087 - distributed.preloading - INFO - Import preload module: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,093 - distributed.preloading - INFO - Creating preload: dask_cuda.initialize\n",
      "2023-05-26 03:09:04,093 - 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-acc90f7f-fb72-11ed-808f-54ab3adac0a5</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",
       "\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;\">c0d46f34</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> 8\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 8\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 503.79 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-789d4132-7d07-451f-aced-1867dfa9d7b3</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:33423\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 8\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> 8\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> 503.79 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:40925\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:43851/status\" target=\"_blank\">http://127.0.0.1:43851/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:40143\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> /tmp/dask-worker-space/worker-6pb36hck\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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:40769\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:44353/status\" target=\"_blank\">http://127.0.0.1:44353/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:43979\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> /tmp/dask-worker-space/worker-rv8itza6\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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: 2</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:45255\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:37165/status\" target=\"_blank\">http://127.0.0.1:37165/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:46241\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> /tmp/dask-worker-space/worker-co2ru8ea\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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: 3</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:40555\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:37399/status\" target=\"_blank\">http://127.0.0.1:37399/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:36117\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> /tmp/dask-worker-space/worker-04qjh_rt\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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: 4</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:39951\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:42631/status\" target=\"_blank\">http://127.0.0.1:42631/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:46323\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> /tmp/dask-worker-space/worker-9zzs6cz6\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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: 5</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:40815\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:37909/status\" target=\"_blank\">http://127.0.0.1:37909/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:40533\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> /tmp/dask-worker-space/worker-hq437puc\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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: 6</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:42963\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:41947/status\" target=\"_blank\">http://127.0.0.1:41947/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:42201\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> /tmp/dask-worker-space/worker-kuwua5fi\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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: 7</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:39607\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:45371/status\" target=\"_blank\">http://127.0.0.1:45371/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 62.97 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:39667\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> /tmp/dask-worker-space/worker-ouyeimq6\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\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:33423' processes=8 threads=8, memory=503.79 GiB>"
      ]
     },
     "execution_count": 8,
     "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])  # Delect 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",
    "        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": "code",
   "execution_count": 10,
   "id": "3f00da1c-6974-4fb9-8b33-c64aee10ee1b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training output data: /hugectr_e2e/criteo/train\n",
      "Validation output data: /hugectr_e2e/criteo/val\n",
      "Training dataset: /hugectr_e2e/criteo/train/train.txt\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-26 03:09:49,967 Preprocessing\n",
      "2023-05-26 03:09:50,513 Train Datasets Preprocessing.....\n",
      "2023-05-26 03:09:57,544 Valid Datasets Preprocessing.....\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1234907, 19683, 13780, 6867, 18490, 4, 6264, 1235, 50, 854680, 114026, 75736, 11, 2159, 7533, 61, 4, 919, 15, 1307783, 404742, 1105613, 87714, 9032, 77, 34, 1577645, 1093030]\n",
      "[1234907, 19683, 13780, 6867, 18490, 4, 6264, 1235, 50, 854680, 114026, 75736, 11, 2159, 7533, 61, 4, 919, 15, 1307783, 404742, 1105613, 87714, 9032, 77, 34, 1577645, 1093030]\n",
      "\n",
      "Dask-NVTabular Criteo Preprocessing Done!\n",
      "Runtime[s]         | 11.187256813049316\n",
      "======================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_output = os.path.join(DATA_DIR, \"train\")\n",
    "print(\"Training output data: \"+train_output)\n",
    "val_output = os.path.join(DATA_DIR, \"val\")\n",
    "print(\"Validation output data: \"+val_output)\n",
    "train_input = os.path.join(DATA_DIR, \"train/train.txt\")\n",
    "print(\"Training dataset: \"+train_input)\n",
    "val_input = os.path.join(DATA_DIR, \"val/test.txt\")\n",
    "PREPROCESS_DIR_temp_train = os.path.join(DATA_DIR, 'train/temp-parquet-after-conversion')  \n",
    "PREPROCESS_DIR_temp_val = os.path.join(DATA_DIR, \"val/temp-parquet-after-conversion\")\n",
    "if not os.path.exists(PREPROCESS_DIR_temp_train):\n",
    "    os.makedirs(PREPROCESS_DIR_temp_train)\n",
    "\n",
    "if not os.path.exists(PREPROCESS_DIR_temp_val):\n",
    "    os.makedirs(PREPROCESS_DIR_temp_val)\n",
    "\n",
    "PREPROCESS_DIR_temp = [PREPROCESS_DIR_temp_train, PREPROCESS_DIR_temp_val]\n",
    "\n",
    "# Make sure we have a clean parquet space for cudf conversion\n",
    "for one_path in PREPROCESS_DIR_temp:\n",
    "    if os.path.exists(one_path):\n",
    "        shutil.rmtree(one_path)\n",
    "    os.mkdir(one_path)\n",
    "\n",
    "#calculate the total processing time\n",
    "runtime = time.time()\n",
    "\n",
    "## train/valid txt to parquet\n",
    "train_valid_paths = [(train_input,PREPROCESS_DIR_temp_train),(val_input,PREPROCESS_DIR_temp_val)]\n",
    "\n",
    "for input, temp_output in train_valid_paths:\n",
    "\n",
    "    ddf = dask_cudf.read_csv(input,sep='\\t',names=LABEL_COLUMNS + CONTINUOUS_COLUMNS + CATEGORICAL_COLUMNS)\n",
    "    \n",
    "    if CROSS_COLUMNS:\n",
    "        for pair in CROSS_COLUMNS:\n",
    "            feature_pair = pair.split(\"_\")\n",
    "            ddf[pair] = ddf[feature_pair[0]] + ddf[feature_pair[1]]\n",
    "\n",
    "    ddf[\"label\"] = ddf['label'].astype('float32')\n",
    "    ddf[CONTINUOUS_COLUMNS] = ddf[CONTINUOUS_COLUMNS].astype('float32')\n",
    "\n",
    "    # Save it as parquet format for better memory usage\n",
    "    ddf.to_parquet(temp_output,header=True)\n",
    "    ##-----------------------------------##\n",
    "\n",
    "COLUMNS =  LABEL_COLUMNS + CONTINUOUS_COLUMNS + CROSS_COLUMNS + CATEGORICAL_COLUMNS\n",
    "train_paths = glob.glob(os.path.join(PREPROCESS_DIR_temp_train, \"*.parquet\"))\n",
    "valid_paths = glob.glob(os.path.join(PREPROCESS_DIR_temp_val, \"*.parquet\"))\n",
    "\n",
    "categorify_op = Categorify()\n",
    "cat_features = CATEGORICAL_COLUMNS >> categorify_op\n",
    "cont_features = CONTINUOUS_COLUMNS >> FillMissing() >> Clip(min_value=0) >> Normalize()\n",
    "cross_cat_op = Categorify()\n",
    "\n",
    "features = LABEL_COLUMNS\n",
    "\n",
    "features += cont_features\n",
    "if CROSS_COLUMNS:\n",
    "    for pair in CROSS_COLUMNS:\n",
    "        features += [pair] >> cross_cat_op\n",
    "\n",
    "features += cat_features\n",
    "\n",
    "workflow = nvt.Workflow(features, client=client)\n",
    "\n",
    "logging.info(\"Preprocessing\")\n",
    "\n",
    "output_format = 'parquet'\n",
    "\n",
    "# just for /samples/criteo model\n",
    "train_ds_iterator = nvt.Dataset(train_paths, engine='parquet')\n",
    "valid_ds_iterator = nvt.Dataset(valid_paths, engine='parquet')\n",
    "\n",
    "shuffle = nvt.io.Shuffle.PER_PARTITION\n",
    "\n",
    "logging.info('Train Datasets Preprocessing.....')\n",
    "\n",
    "dict_dtypes = {}\n",
    "for col in CATEGORICAL_COLUMNS:\n",
    "    dict_dtypes[col] = np.int64\n",
    "for col in CONTINUOUS_COLUMNS:\n",
    "    dict_dtypes[col] = np.float32\n",
    "for col in CROSS_COLUMNS:\n",
    "    dict_dtypes[col] = np.int64\n",
    "for col in LABEL_COLUMNS:\n",
    "    dict_dtypes[col] = np.float32\n",
    "\n",
    "conts = CONTINUOUS_COLUMNS\n",
    "\n",
    "workflow.fit(train_ds_iterator)\n",
    "\n",
    "if output_format == 'hugectr':\n",
    "    workflow.transform(train_ds_iterator).to_hugectr(\n",
    "            cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "            conts=conts,\n",
    "            labels=LABEL_COLUMNS,\n",
    "            output_path=train_output,\n",
    "            shuffle=shuffle)\n",
    "else:\n",
    "    workflow.transform(train_ds_iterator).to_parquet(\n",
    "            output_path=train_output,\n",
    "            dtypes=dict_dtypes,\n",
    "            cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "            conts=conts,\n",
    "            labels=LABEL_COLUMNS,\n",
    "            shuffle=shuffle)\n",
    "\n",
    "###Getting slot size###    \n",
    "#--------------------##\n",
    "embeddings_dict_cat = categorify_op.get_embedding_sizes(CATEGORICAL_COLUMNS)\n",
    "embeddings_dict_cross = cross_cat_op.get_embedding_sizes(CROSS_COLUMNS)\n",
    "embeddings = [embeddings_dict_cat[c][0] for c in CATEGORICAL_COLUMNS] + [embeddings_dict_cross[c][0] for c in CROSS_COLUMNS]\n",
    "\n",
    "print(embeddings)\n",
    "##--------------------##\n",
    "\n",
    "logging.info('Valid Datasets Preprocessing.....')\n",
    "\n",
    "if output_format == 'hugectr':\n",
    "    workflow.transform(valid_ds_iterator).to_hugectr(\n",
    "            cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "            conts=conts,\n",
    "            labels=LABEL_COLUMNS,\n",
    "            output_path=val_output,\n",
    "            shuffle=shuffle)\n",
    "else:\n",
    "    workflow.transform(valid_ds_iterator).to_parquet(\n",
    "            output_path=val_output,\n",
    "            dtypes=dict_dtypes,\n",
    "            cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "            conts=conts,\n",
    "            labels=LABEL_COLUMNS,\n",
    "            shuffle=shuffle)\n",
    "\n",
    "embeddings_dict_cat = categorify_op.get_embedding_sizes(CATEGORICAL_COLUMNS)\n",
    "embeddings_dict_cross = cross_cat_op.get_embedding_sizes(CROSS_COLUMNS)\n",
    "embeddings = [embeddings_dict_cat[c][0] for c in CATEGORICAL_COLUMNS] + [embeddings_dict_cross[c][0] for c in CROSS_COLUMNS]\n",
    "\n",
    "print(embeddings)\n",
    "##--------------------##\n",
    "\n",
    "## Shutdown clusters\n",
    "client.shutdown()\n",
    "\n",
    "runtime = time.time() - runtime\n",
    "\n",
    "print(\"\\nDask-NVTabular Criteo Preprocessing Done!\")\n",
    "print(f\"Runtime[s]         | {runtime}\")\n",
    "print(\"======================================\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fee364eb-6052-4459-90b3-cf5c5c253eab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Record the slot size array\n",
    "SLOT_SIZE_ARRAY = embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb98899c-d3b6-4e4e-9a25-f7ff94c82eaa",
   "metadata": {},
   "source": [
    "## Training a WDL model with HugeCTR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6ba18c4e-b33c-464e-9a74-579fbc6d6809",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing ./train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile './train.py'\n",
    "import hugectr\n",
    "import os\n",
    "import argparse\n",
    "from mpi4py import MPI\n",
    "parser = argparse.ArgumentParser(description=(\"HugeCTR Training\"))\n",
    "parser.add_argument(\"--data_path\", type=str, help=\"Input dataset path (Required)\")\n",
    "parser.add_argument(\"--model_path\", type=str, help=\"Directory path to write output (Required)\")\n",
    "args = parser.parse_args()\n",
    "SLOT_SIZE_ARRAY = [1234907, 19683, 13780, 6867, 18490, 4, 6264, 1235, 50, 854680, 114026, 75736, 11, 2159, 7533, 61, 4, 919, 15, 1307783, 404742, 1105613, 87714, 9032, 77, 34, 1581605, 1093030]\n",
    "\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 4000,\n",
    "                              batchsize_eval = 2720,\n",
    "                              batchsize = 2720,\n",
    "                              lr = 0.001,\n",
    "                              vvgpu = [[0]],\n",
    "                              repeat_dataset = True,\n",
    "                              i64_input_key = True)\n",
    "\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Parquet,\n",
    "                                  source = [os.path.join(args.data_path, \"train/_file_list.txt\")],\n",
    "                                  eval_source = os.path.join(args.data_path, \"val/_file_list.txt\"),\n",
    "                                  check_type = hugectr.Check_t.Non,\n",
    "                                  slot_size_array = SLOT_SIZE_ARRAY)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam,\n",
    "                                    update_type = hugectr.Update_t.Global,\n",
    "                                    beta1 = 0.9,\n",
    "                                    beta2 = 0.999,\n",
    "                                    epsilon = 0.0000001)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "\n",
    "model.add(hugectr.Input(label_dim = 1, label_name = \"label\",\n",
    "                        dense_dim = 13, dense_name = \"dense\",\n",
    "                        data_reader_sparse_param_array = \n",
    "                        [hugectr.DataReaderSparseParam(\"wide_data\", 1, True, 2),\n",
    "                        hugectr.DataReaderSparseParam(\"deep_data\", 2, False, 26)]))\n",
    "\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 80,\n",
    "                            embedding_vec_size = 1,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding2\",\n",
    "                            bottom_name = \"wide_data\",\n",
    "                            optimizer = optimizer))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 1350,\n",
    "                            embedding_vec_size = 16,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding1\",\n",
    "                            bottom_name = \"deep_data\",\n",
    "                            optimizer = optimizer))\n",
    "\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding1\"],\n",
    "                            top_names = [\"reshape1\"],\n",
    "                            leading_dim=416))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding2\"],\n",
    "                            top_names = [\"reshape2\"],\n",
    "                            leading_dim=2))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReduceSum,\n",
    "                            bottom_names = [\"reshape2\"],\n",
    "                            top_names = [\"wide_redn\"],\n",
    "                            axis = 1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,\n",
    "                            bottom_names = [\"reshape1\", \"dense\"],\n",
    "                            top_names = [\"concat1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"concat1\"],\n",
    "                            top_names = [\"fc1\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc1\"],\n",
    "                            top_names = [\"relu1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,\n",
    "                            bottom_names = [\"relu1\"],\n",
    "                            top_names = [\"dropout1\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout1\"],\n",
    "                            top_names = [\"fc2\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc2\"],\n",
    "                            top_names = [\"relu2\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,\n",
    "                            bottom_names = [\"relu2\"],\n",
    "                            top_names = [\"dropout2\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout2\"],\n",
    "                            top_names = [\"fc3\"],\n",
    "                            num_output=1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Add,\n",
    "                            bottom_names = [\"fc3\", \"wide_redn\"],\n",
    "                            top_names = [\"add1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.BinaryCrossEntropyLoss,\n",
    "                            bottom_names = [\"add1\", \"label\"],\n",
    "                            top_names = [\"loss\"]))\n",
    "model.compile()\n",
    "model.summary()\n",
    "model.fit(max_iter = 21000, display = 1000, eval_interval = 4000, snapshot = 20000, snapshot_prefix = os.path.join(args.model_path, \"wdl/\"))\n",
    "model.graph_to_json(graph_config_file = os.path.join(args.model_path, \"wdl.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2e42f382-14ff-485b-a7a2-dfcf41b2f5ac",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MpiInitService: MPI was already initialized by another (non-HugeCTR) mechanism.\n",
      "HugeCTR Version: 23.4\n",
      "====================================================Model Init=====================================================\n",
      "[HCTR][03:10:28.412][WARNING][RK0][main]: The model name is not specified when creating the solver.\n",
      "[HCTR][03:10:28.413][INFO][RK0][main]: Global seed is 4031005480\n",
      "[HCTR][03:10:29.069][INFO][RK0][main]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "[HCTR][03:10:32.353][WARNING][RK0][main]: Peer-to-peer access cannot be fully enabled.\n",
      "[HCTR][03:10:32.353][DEBUG][RK0][main]: [device 0] allocating 0.0000 GB, available 29.9792 \n",
      "[HCTR][03:10:32.353][INFO][RK0][main]: Start all2all warmup\n",
      "[HCTR][03:10:32.353][INFO][RK0][main]: End all2all warmup\n",
      "[HCTR][03:10:32.355][INFO][RK0][main]: Using All-reduce algorithm: NCCL\n",
      "[HCTR][03:10:32.361][INFO][RK0][main]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HCTR][03:10:32.362][INFO][RK0][main]: eval source /hugectr_e2e/criteo/val/_file_list.txt max_row_group_size 475000\n",
      "[HCTR][03:10:32.364][INFO][RK0][main]: train source /hugectr_e2e/criteo/train/_file_list.txt max_row_group_size 475000\n",
      "[HCTR][03:10:32.364][INFO][RK0][main]: num of DataReader workers for train: 1\n",
      "[HCTR][03:10:32.364][INFO][RK0][main]: num of DataReader workers for eval: 1\n",
      "[HCTR][03:10:32.365][DEBUG][RK0][main]: [device 0] allocating 0.0018 GB, available 29.7234 \n",
      "[HCTR][03:10:32.366][DEBUG][RK0][main]: [device 0] allocating 0.0018 GB, available 29.7175 \n",
      "[HCTR][03:10:32.379][INFO][RK0][main]: Vocabulary size: 7946054\n",
      "[HCTR][03:10:32.380][INFO][RK0][main]: max_vocabulary_size_per_gpu_=6990506\n",
      "[HCTR][03:10:32.391][DEBUG][RK0][main]: [device 0] allocating 0.0788 GB, available 28.3132 \n",
      "[HCTR][03:10:32.392][INFO][RK0][main]: max_vocabulary_size_per_gpu_=7372800\n",
      "[HCTR][03:10:32.396][DEBUG][RK0][main]: [device 0] allocating 1.3516 GB, available 26.5847 \n",
      "[HCTR][03:10:32.397][INFO][RK0][main]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HCTR][03:10:32.408][DEBUG][RK0][main]: [device 0] allocating 0.2162 GB, available 26.3523 \n",
      "[HCTR][03:10:32.409][DEBUG][RK0][main]: [device 0] allocating 0.0056 GB, available 26.3464 \n",
      "[HCTR][03:10:40.869][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][03:10:40.869][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][03:10:40.869][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][03:10:40.873][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][03:10:40.873][DEBUG][RK0][main]: [device 0] allocating 0.0001 GB, available 26.3464 \n",
      "[HCTR][03:10:40.874][INFO][RK0][main]: Starting AUC NCCL warm-up\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Model structure on each GPU\n",
      "Label                                   Dense                         Sparse                        \n",
      "label                                   dense                          wide_data,deep_data           \n",
      "(2720,1)                                (2720,13)                               \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "DistributedSlotSparseEmbeddingHash      wide_data                     sparse_embedding2             (2720,2,1)                    \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (2720,26,16)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding1             reshape1                      (2720,416)                    \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding2             reshape2                      (2720,2)                      \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReduceSum                               reshape2                      wide_redn                     (2720,1)                      \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Concat                                  reshape1                      concat1                       (2720,429)                    \n",
      "                                        dense                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            concat1                       fc1                           (2720,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc1                           relu1                         (2720,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu1                         dropout1                      (2720,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout1                      fc2                           (2720,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc2                           relu2                         (2720,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu2                         dropout2                      (2720,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout2                      fc3                           (2720,1)                      \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Add                                     fc3                           add1                          (2720,1)                      \n",
      "                                        wide_redn                                                                                 \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "BinaryCrossEntropyLoss                  add1                          loss                                                        \n",
      "                                        label                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "=====================================================Model Fit=====================================================\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Use non-epoch mode with number of iterations: 21000\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Training batchsize: 2720, evaluation batchsize: 2720\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Evaluation interval: 4000, snapshot interval: 20000\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Dense network trainable: True\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Sparse embedding sparse_embedding1 trainable: True\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Training source file: /hugectr_e2e/criteo/train/_file_list.txt\n",
      "[HCTR][03:10:40.879][INFO][RK0][main]: Evaluation source file: /hugectr_e2e/criteo/val/_file_list.txt\n",
      "[HCTR][03:10:49.588][INFO][RK0][main]: Iter: 1000 Time(1000 iters): 8.70458s Loss: 0.124098 lr:0.001\n",
      "[HCTR][03:10:58.211][INFO][RK0][main]: Iter: 2000 Time(1000 iters): 8.6176s Loss: 0.130088 lr:0.001\n",
      "[HCTR][03:11:06.835][INFO][RK0][main]: Iter: 3000 Time(1000 iters): 8.61959s Loss: 0.101731 lr:0.001\n",
      "[HCTR][03:11:15.449][INFO][RK0][main]: Iter: 4000 Time(1000 iters): 8.61009s Loss: 0.110557 lr:0.001\n",
      "[HCTR][03:11:19.929][INFO][RK0][main]: Evaluation, AUC: 0.738497\n",
      "[HCTR][03:11:19.929][INFO][RK0][main]: Eval Time for 4000 iters: 4.47924s\n",
      "[HCTR][03:11:28.559][INFO][RK0][main]: Iter: 5000 Time(1000 iters): 13.1046s Loss: 0.10236 lr:0.001\n",
      "[HCTR][03:11:37.182][INFO][RK0][main]: Iter: 6000 Time(1000 iters): 8.61852s Loss: 0.102157 lr:0.001\n",
      "[HCTR][03:11:45.771][INFO][RK0][main]: Iter: 7000 Time(1000 iters): 8.58452s Loss: 0.123451 lr:0.001\n",
      "[HCTR][03:11:54.385][INFO][RK0][main]: Iter: 8000 Time(1000 iters): 8.61023s Loss: 0.122763 lr:0.001\n",
      "[HCTR][03:11:58.867][INFO][RK0][main]: Evaluation, AUC: 0.698276\n",
      "[HCTR][03:11:58.867][INFO][RK0][main]: Eval Time for 4000 iters: 4.48087s\n",
      "[HCTR][03:12:07.487][INFO][RK0][main]: Iter: 9000 Time(1000 iters): 13.097s Loss: 0.0999177 lr:0.001\n",
      "[HCTR][03:12:16.103][INFO][RK0][main]: Iter: 10000 Time(1000 iters): 8.61106s Loss: 0.0999892 lr:0.001\n",
      "[HCTR][03:12:24.722][INFO][RK0][main]: Iter: 11000 Time(1000 iters): 8.61545s Loss: 0.0883301 lr:0.001\n",
      "[HCTR][03:12:33.348][INFO][RK0][main]: Iter: 12000 Time(1000 iters): 8.62134s Loss: 0.0828304 lr:0.001\n",
      "[HCTR][03:12:37.823][INFO][RK0][main]: Evaluation, AUC: 0.688598\n",
      "[HCTR][03:12:37.823][INFO][RK0][main]: Eval Time for 4000 iters: 4.4733s\n",
      "[HCTR][03:12:46.425][INFO][RK0][main]: Iter: 13000 Time(1000 iters): 13.0717s Loss: 0.108287 lr:0.001\n",
      "[HCTR][03:12:55.059][INFO][RK0][main]: Iter: 14000 Time(1000 iters): 8.62997s Loss: 0.0745141 lr:0.001\n",
      "[HCTR][03:13:03.671][INFO][RK0][main]: Iter: 15000 Time(1000 iters): 8.60764s Loss: 0.0720452 lr:0.001\n",
      "[HCTR][03:13:12.287][INFO][RK0][main]: Iter: 16000 Time(1000 iters): 8.61101s Loss: 0.0851126 lr:0.001\n",
      "[HCTR][03:13:16.758][INFO][RK0][main]: Evaluation, AUC: 0.685426\n",
      "[HCTR][03:13:16.758][INFO][RK0][main]: Eval Time for 4000 iters: 4.47088s\n",
      "[HCTR][03:13:25.378][INFO][RK0][main]: Iter: 17000 Time(1000 iters): 13.0865s Loss: 0.0632745 lr:0.001\n",
      "[HCTR][03:13:34.011][INFO][RK0][main]: Iter: 18000 Time(1000 iters): 8.62825s Loss: 0.0742994 lr:0.001\n",
      "[HCTR][03:13:42.626][INFO][RK0][main]: Iter: 19000 Time(1000 iters): 8.61035s Loss: 0.0679226 lr:0.001\n",
      "[HCTR][03:13:51.230][INFO][RK0][main]: Iter: 20000 Time(1000 iters): 8.59954s Loss: 0.0779185 lr:0.001\n",
      "[HCTR][03:13:55.704][INFO][RK0][main]: Evaluation, AUC: 0.684045\n",
      "[HCTR][03:13:55.704][INFO][RK0][main]: Eval Time for 4000 iters: 4.4736s\n",
      "[HCTR][03:13:55.733][INFO][RK0][main]: Rank0: Write hash table to file\n",
      "[HCTR][03:13:55.902][INFO][RK0][main]: Rank0: Write hash table to file\n",
      "[HCTR][03:13:56.075][INFO][RK0][main]: Dumping sparse weights to files, successful\n",
      "[HCTR][03:13:56.091][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][03:13:56.104][INFO][RK0][main]: Done\n",
      "[HCTR][03:13:56.119][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][03:13:56.133][INFO][RK0][main]: Done\n",
      "[HCTR][03:13:56.398][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][03:13:56.611][INFO][RK0][main]: Done\n",
      "[HCTR][03:13:56.903][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][03:13:57.152][INFO][RK0][main]: Done\n",
      "[HCTR][03:13:57.169][INFO][RK0][main]: Dumping sparse optimzer states to files, successful\n",
      "[HCTR][03:13:57.176][INFO][RK0][main]: Dumping dense weights to file, successful\n",
      "[HCTR][03:13:57.188][INFO][RK0][main]: Dumping dense optimizer states to file, successful\n",
      "[HCTR][03:14:05.788][INFO][RK0][main]: Iter: 21000 Time(1000 iters): 14.5538s Loss: 0.0770708 lr:0.001\n",
      "[HCTR][03:14:05.788][INFO][RK0][main]: Finish 21000 iterations with batchsize: 2720 in 204.91s.\n",
      "[HCTR][03:14:05.788][INFO][RK0][main]: Save the model graph to /hugectr_e2e/model/wdl.json successfully\n"
     ]
    }
   ],
   "source": [
    "!python train.py --data_path $DATA_DIR --model_path $MODEL_DIR"
   ]
  }
 ],
 "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
}