{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "cloud/azure/azure-vm-multi",
     "tools/dask-cloudprovider",
     "library/cudf",
     "library/cuml",
     "library/xgboost",
     "library/dask",
     "library/fil",
     "data-storage/azure-data-lake",
     "dataset/nyc-taxi",
     "workflow/xgboost"
    ]
   },
   "source": [
    "# Multi-Node Multi-GPU XGBoost Example on Azure using dask-cloudprovider\n",
    "\n",
    "_November, 2023_\n",
    "\n",
    "[Dask Cloud Provider](https://cloudprovider.dask.org/en/latest/) is a native cloud integration library for Dask. It helps manage Dask clusters on different cloud platforms. In this notebook, we will look at how we can use this package to set-up an Azure cluster and run a multi-node multi-GPU (MNMG) example with [RAPIDS](https://rapids.ai/). RAPIDS provides a suite of libraries to accelerate data science pipelines on the GPU entirely. This can be scaled to multiple nodes using Dask as we will see in this notebook. \n",
    "\n",
    "For the purposes of this demo, we will use a part of the NYC Taxi Dataset (only the files of 2014 calendar year will be used here). The goal is to predict the fare amount for a given trip given the times and coordinates of the taxi trip. We will download the data from [Azure Open Datasets](https://docs.microsoft.com/en-us/azure/open-datasets/overview-what-are-open-datasets), where the dataset is publicly hosted by Microsoft.\n",
    "\n",
    "```{note}\n",
    "In this notebook, we will explore two possible ways to use `dask-cloudprovider` to run our workloads on Azure VM clusters:\n",
    "\n",
    "1. [Option 1](use-an-Azure-marketplace-VM-image): Using an [Azure Marketplace image](https://azuremarketplace.microsoft.com/en-us/marketplace/apps/nvidia.ngc_azure_17_11?tab=overview) made available for free from NVIDIA. The RAPIDS container will be subsequently downloaded once the VMs start up.\n",
    "2. [Option 2](set-up-an-Azure-Customized-VM): Using [`packer`](https://docs.microsoft.com/en-us/azure/virtual-machines/linux/build-image-with-packer) to create a custom VM image to be used in the cluster. This image will include the RAPIDS container, and having the container already inside the image should speed up the process of provisioning the cluster.\n",
    "\n",
    "**You can either use Option 1 or use Option 2**\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "library/xgboost"
    ]
   },
   "source": [
    "## Step 0: Set up Azure credentials and CLI\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running the notebook, run the following commands in the terminal to setup Azure CLI\n",
    "\n",
    "```\n",
    "curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash\n",
    "az login\n",
    "```\n",
    "\n",
    "Then, follow the instructions on the prompt to finish setting up the account. If you are running the notebook from inside a Docker container, you can remove `sudo`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[93mA web browser has been opened at https://login.microsoftonline.com/organizations/oauth2/v2.0/authorize. Please continue the login in the web browser. If no web browser is available or if the web browser fails to open, use device code flow with `az login --use-device-code`.\u001b[0m\n",
      "[\n",
      "  {\n",
      "    \"cloudName\": \"AzureCloud\",\n",
      "    \"homeTenantId\": \"43083d15-7273-40c1-b7db-39efd9ccc17a\",\n",
      "    \"id\": \"fc4f4a6b-4041-4b1c-8249-854d68edcf62\",\n",
      "    \"isDefault\": true,\n",
      "    \"managedByTenants\": [\n",
      "      {\n",
      "        \"tenantId\": \"2f4a9838-26b7-47ee-be60-ccc1fdec5953\"\n",
      "      }\n",
      "    ],\n",
      "    \"name\": \"NV-AI-Infra\",\n",
      "    \"state\": \"Enabled\",\n",
      "    \"tenantId\": \"43083d15-7273-40c1-b7db-39efd9ccc17a\",\n",
      "    \"user\": {\n",
      "      \"name\": \"skirui@nvidia.com\",\n",
      "      \"type\": \"user\"\n",
      "    }\n",
      "  }\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "!az login"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Step 1: Import necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # Uncomment the following and install some libraries at the beginning.\n",
    "# If adlfs is not present, install adlfs to read from Azure data lake.\n",
    "! pip install adlfs\n",
    "! pip install \"dask-cloudprovider[azure]\" --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from timeit import default_timer as timer\n",
    "\n",
    "import dask\n",
    "import dask_cudf\n",
    "import numpy as np\n",
    "import xgboost as xgb\n",
    "from cuml.metrics import mean_squared_error\n",
    "from dask.distributed import Client, wait\n",
    "from dask_cloudprovider.azure import AzureVMCluster\n",
    "from dask_ml.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Step 2: Set up the Azure VM Cluster\n",
    "\n",
    "We will now set up a Dask cluster on Azure Virtual machines using `AzureVMCluster` from Dask Cloud Provider following these [instructions](https://docs.rapids.ai/deployment/stable/cloud/azure/azure-vm-multi/).\n",
    "\n",
    "To do this, you will first need to set up a Resource Group, a Virtual Network and a Security Group on Azure. [Learn more about how you can set this up](https://cloudprovider.dask.org/en/latest/azure.html#resource-groups). Note that you can also set it up using the Azure portal.\n",
    "\n",
    "Once you have set it up, you can now plug in the names of the entities you have created in the cell below. \n",
    "\n",
    "We need to pass in the docker argument `docker_args = '--shm-size=256m'` to allow larger shared memory for successfully running multiple docker containers in the same VM. This is the case when each VM has more than one worker. Even if you don't have such a case, there is no harm in having a larger shared memory. Finally, note that we use the RAPIDS docker image to build the VM and use the `dask_cuda.CUDAWorker` to run within the VM. This will run the worker docker image with GPU capabilities instead of CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "location = \"West US 2\"\n",
    "resource_group = \"rapidsai-deployment\"\n",
    "vnet = \"rapidsai-deployment-vnet\"\n",
    "security_group = \"rapidsaiclouddeploymenttest-nsg\"\n",
    "vm_size = \"Standard_NC12s_v3\"  # or choose a different GPU enabled VM type\n",
    "\n",
    "docker_image = \"rapidsai/base:26.08a-cuda13-py3.13\"\n",
    "docker_args = \"--shm-size=256m\"\n",
    "worker_class = \"dask_cuda.CUDAWorker\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(use-an-Azure-marketplace-VM-image)=\n",
    "### Option 1: Use an Azure Marketplace VM image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this method, we can use an Azure marketplace VM provided by NVIDIA for free. These VM images contain all the necessary dependencies and NVIDIA drivers preinstalled. These images are made available by NVIDIA as an out-of-the-box solution to decrease the cluster setup time for data scientists. Fortunately for us, `dask-cloudprovider` has made it simple to pass in information of a marketplace VM, and it will use the selected VM image instead of a vanilla image. \n",
    "\n",
    "We will use the following image: [NVIDIA GPU-Optimized Image for AI and HPC](https://azuremarketplace.microsoft.com/en-us/marketplace/apps/nvidia.ngc_azure_17_11?tab=overview).\n",
    "\n",
    "```{note}\n",
    "Please make sure you have [dask-cloudprovider](https://cloudprovider.dask.org/en/latest/) version 2021.6.0 or above. Marketplace VMs in Azure is not supported in older versions.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set up Marketplace VM information and clear default dask config. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'vnet': None,\n",
       " 'security_group': None,\n",
       " 'public_ingress': True,\n",
       " 'vm_size': 'Standard_DS1_v2',\n",
       " 'disk_size': 50,\n",
       " 'scheduler_vm_size': None,\n",
       " 'docker_image': 'daskdev/dask:latest',\n",
       " 'vm_image': {'publisher': 'Canonical',\n",
       "  'offer': 'UbuntuServer',\n",
       "  'sku': '18.04-LTS',\n",
       "  'version': 'latest'},\n",
       " 'bootstrap': True,\n",
       " 'auto_shutdown': True,\n",
       " 'marketplace_plan': {'publisher': 'nvidia',\n",
       "  'name': 'ngc-base-version-23_03_0',\n",
       "  'product': 'ngc_azure_17_11',\n",
       "  'version': '23.03.0'}}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask.config.set(\n",
    "    {\n",
    "        \"logging.distributed\": \"info\",\n",
    "        \"cloudprovider.azure.azurevm.marketplace_plan\": {\n",
    "            \"publisher\": \"nvidia\",\n",
    "            \"name\": \"ngc-base-version-23_03_0\",\n",
    "            \"product\": \"ngc_azure_17_11\",\n",
    "            \"version\": \"23.03.0\",\n",
    "        },\n",
    "    }\n",
    ")\n",
    "vm_image = \"\"\n",
    "config = dask.config.get(\"cloudprovider.azure.azurevm\", {})\n",
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If necessary, you must uncomment and accept the Azure Marketplace image terms so that the image can be used to create VMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"accepted\": true,\n",
      "  \"id\": \"/subscriptions/fc4f4a6b-4041-4b1c-8249-854d68edcf62/providers/Microsoft.MarketplaceOrdering/offerTypes/Microsoft.MarketplaceOrdering/offertypes/publishers/nvidia/offers/ngc_azure_17_11/plans/ngc-base-version-23_03_0/agreements/current\",\n",
      "  \"licenseTextLink\": \"https://mpcprodsa.blob.core.windows.net/legalterms/3E5ED_legalterms_NVIDIA%253a24NGC%253a5FAZURE%253a5F17%253a5F11%253a24NGC%253a2DBASE%253a2DVERSION%253a2D23%253a5F03%253a5F0%253a24KJVKRIWKTRQ3CIEPNL6YTG4AVORBHHPZCDQDVWX7JPPDEF6UM7R4XO76VDRHXCNTQYATKLGYYW3KA7DSIKTYXBZ3HJ2FMWYCINEY4WQ.txt\",\n",
      "  \"marketplaceTermsLink\": \"https://mpcprodsa.blob.core.windows.net/marketplaceterms/3EDEF_marketplaceterms_VIRTUALMACHINE%253a24AAK2OAIZEAWW5H4MSP5KSTVB6NDKKRTUBAU23BRFTWN4YC2MQLJUB5ZEYUOUJBVF3YK34CIVPZL2HWYASPGDUY5O2FWEGRBYOXWZE5Y.txt\",\n",
      "  \"name\": \"ngc-base-version-23_03_0\",\n",
      "  \"plan\": \"ngc-base-version-23_03_0\",\n",
      "  \"privacyPolicyLink\": \"https://www.nvidia.com/en-us/about-nvidia/privacy-policy/\",\n",
      "  \"product\": \"ngc_azure_17_11\",\n",
      "  \"publisher\": \"nvidia\",\n",
      "  \"retrieveDatetime\": \"2023-10-02T08:17:40.3203275Z\",\n",
      "  \"signature\": \"SWCKS7PPTL3XIBGBE2IZCMF43KBRDLSIZ7XLXXTLI6SXDCPCXY53BAISH6DNIELVV63GPZ44AOMMMZ6RV2AL5ARNM6XWHXRJ4HDNTJI\",\n",
      "  \"systemData\": {\n",
      "    \"createdAt\": \"2023-10-02T08:17:43.219827+00:00\",\n",
      "    \"createdBy\": \"fc4f4a6b-4041-4b1c-8249-854d68edcf62\",\n",
      "    \"createdByType\": \"ManagedIdentity\",\n",
      "    \"lastModifiedAt\": \"2023-10-02T08:17:43.219827+00:00\",\n",
      "    \"lastModifiedBy\": \"fc4f4a6b-4041-4b1c-8249-854d68edcf62\",\n",
      "    \"lastModifiedByType\": \"ManagedIdentity\"\n",
      "  },\n",
      "  \"type\": \"Microsoft.MarketplaceOrdering/offertypes\"\n",
      "}\n",
      "\u001b[32mCommand ran in 7.879 seconds (init: 0.159, invoke: 7.720)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "! az vm image terms accept --urn \"nvidia:ngc_azure_17_11:ngc-base-version-23_03_0:23.03.0\" --verbose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have set up the necessary configurations to use the NVIDIA VM image, directly move to [Step 2.1](start-the-VM-Cluster-in-Azure) to start the AzureVMCluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "(set-up-an-Azure-Customized-VM)=\n",
    "### Option 2: Set up an Azure Customized VM\n",
    "\n",
    "If you already have a customized VM and you know its resource id, jump to [Step f. of Option 2](set-up-customized-VM-information-and-clear-default-dask-config) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, if we use a generic image to create a cluster, we would have to wait till the new VMs are provisioned fully with all dependencies. The provisioning step does several things such as set the VM up with required libraries, set up Docker, install the NVIDIA drivers and also pull and decompress the RAPIDS container etc. This usually takes around 10-15 minutes of time depending on the cloud provider. If the user wants to fire up a cluster quickly, setting up a VM from a generic image every time may not be optimal. \n",
    "\n",
    "Further, as detailed in Option 1, we can also choose to use a custom Marketplace VM from NVIDIA. However, we will still have to download and decompress the RAPIDS container. So the setup time to start the workers and the scheduler would still be around 8-10 minutes.\n",
    "\n",
    "Luckily we can improve on this. We can make our own customized VM bundled with all the necessary packages, drivers, containers and dependencies. This way, firing up the cluster using the customized VM will take minimal time. \n",
    "\n",
    "In this example, we will be using a tool called [packer](https://www.packer.io/) to create our customized virtual machine image. Packer automates the process of building and customizing VMs across all major cloud providers. \n",
    "\n",
    "Now, to create a customized VM image, follow steps *a.* to *f.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### a. Install `packer`\n",
    "\n",
    "Follow the [getting started guide](https://learn.hashicorp.com/tutorials/packer/get-started-install-cli?in=packer/azure-get-started) to download the necessary binary according to your platform and install it. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### b. Authenticate `packer` with Azure\n",
    "There are several ways to authenticate `packer` to work with Azure (details provided [here](https://learn.hashicorp.com/tutorials/packer/get-started-install-cli?in=packer/azure-get-started)). However, since we already have installed Azure cli (`az`) at the beginning of the notebook, authenticating `packer` with `az` cli is the easiest option. We will let `packer` use the Azure credentials from `az` cli, and so, you do not have to do anything further in this step. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### c. Generate the cloud init script for customizing the VM image\n",
    "\n",
    "`packer` can use a [cloud-init](https://cloudinit.readthedocs.io/en/latest/) script to initialize a VM. The cloud init script contains the set of commands that will set up the environment of our customized VM. We will pass this as an external file to the `packer` command via a configuration script.\n",
    "\n",
    "The cloud init file [cloud_init.yaml.j2](./configs/cloud_init.yaml.j2) file is present in the `configs` folder. In case you want to add/modify any configuration, edit the [cloud_init.yaml.j2](./configs/cloud_init.yaml.j2) before proceeding to the next steps. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### d. Write packer configuration to a configuration file\n",
    "We now need to provide `packer` with a build file with platform related and cloud-init configurations. `packer` will use this to create the customized VM. \n",
    "\n",
    "In this example, we are creating a single custom VM image that will be accessible by the user only. We will use a Ubuntu Server 18.04 base image and customize it. Later on, we will instantiate all our VMs from this customized VM image.\n",
    "\n",
    "If you are curious about what else you can configure, take a look at all the available [Azure build parameters for `packer`](https://www.packer.io/docs/builders/azure/arm).\n",
    "\n",
    "```{note}\n",
    "Our resource group already exists in this example. Hence we simply pass in our resource group name in the required parameters `managed_image_resource_group_name` and `build_resource_group_name`.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_vm_image_name = \"FILL-THIS-IN\"\n",
    "packer_config = {\n",
    "    \"builders\": [\n",
    "        {\n",
    "            \"type\": \"azure-arm\",\n",
    "            \"use_azure_cli_auth\": True,\n",
    "            \"managed_image_resource_group_name\": resource_group,\n",
    "            \"managed_image_name\": custom_vm_image_name,\n",
    "            \"custom_data_file\": \"./configs/cloud_init.yaml.j2\",\n",
    "            \"os_type\": \"Linux\",\n",
    "            \"image_publisher\": \"Canonical\",\n",
    "            \"image_offer\": \"UbuntuServer\",\n",
    "            \"image_sku\": \"18.04-LTS\",\n",
    "            \"azure_tags\": {\n",
    "                \"dept\": \"RAPIDS-CSP\",\n",
    "                \"task\": \"RAPIDS Custom Image deployment\",\n",
    "            },\n",
    "            \"build_resource_group_name\": resource_group,\n",
    "            \"vm_size\": vm_size,\n",
    "        }\n",
    "    ],\n",
    "    \"provisioners\": [\n",
    "        {\n",
    "            \"inline\": [\n",
    "                (\n",
    "                    \"echo 'Waiting for cloud-init'; \"\n",
    "                    \"while [ ! -f /var/lib/cloud/instance/boot-finished ]; \"\n",
    "                    \"do sleep 1; done; echo 'Done'\"\n",
    "                )\n",
    "            ],\n",
    "            \"type\": \"shell\",\n",
    "        }\n",
    "    ],\n",
    "}\n",
    "\n",
    "with open(\"packer_config.json\", \"w\") as fh:\n",
    "    fh.write(json.dumps(packer_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### e. Run `packer` build and create the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Uncomment the following line and run to create the custom image\n",
    "# ! packer build packer_config.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will take around 15 minutes. Grab a coffee or watch an episode of your favourite tv show and come back. But remember, you will only have to do this once, unless you want to update the packages in the VM. This means that you can make this custom image once, and then keep on using it for hundreds of times. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While packer is building the image, you will see an output similar to what is shown below. \n",
    "\n",
    "```console\n",
    "$ packer build packer_config.json\n",
    "azure-arm: output will be in this color.\n",
    "\n",
    "==> azure-arm: Running builder ...\n",
    "==> azure-arm: Getting tokens using Azure CLI\n",
    "==> azure-arm: Getting tokens using Azure CLI\n",
    "    azure-arm: Creating Azure Resource Manager (ARM) client ...\n",
    "==> azure-arm: Using existing resource group ...\n",
    "==> azure-arm:  -> ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> Location          : <some chosen location>\n",
    "==> azure-arm: Validating deployment template ...\n",
    "==> azure-arm:  -> ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> DeploymentName    : 'pkrdp04rrahxkg9'\n",
    "==> azure-arm: Deploying deployment template ...\n",
    "==> azure-arm:  -> ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> DeploymentName    : 'pkrdp04rrahxkg9'\n",
    "==> azure-arm:\n",
    "==> azure-arm: Getting the VM's IP address ...\n",
    "==> azure-arm:  -> ResourceGroupName   : <your resource group>\n",
    "==> azure-arm:  -> PublicIPAddressName : 'pkrip04rrahxkg9'\n",
    "==> azure-arm:  -> NicName             : 'pkrni04rrahxkg9'\n",
    "==> azure-arm:  -> Network Connection  : 'PublicEndpoint'\n",
    "==> azure-arm:  -> IP Address          : '40.77.62.118'\n",
    "==> azure-arm: Waiting for SSH to become available...\n",
    "==> azure-arm: Connected to SSH!\n",
    "==> azure-arm: Provisioning with shell script: /tmp/packer-shell614221056\n",
    "    azure-arm: Waiting for cloud-init\n",
    "    azure-arm: Done\n",
    "==> azure-arm: Querying the machine's properties ...\n",
    "==> azure-arm:  -> ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> ComputeName       : 'pkrvm04rrahxkg9'\n",
    "==> azure-arm:  -> Managed OS Disk   : '/subscriptions/<your subscription id>/resourceGroups/<your resource group>/providers/Microsoft.Compute/disks/pkros04rrahxkg9'\n",
    "==> azure-arm: Querying the machine's additional disks properties ...\n",
    "==> azure-arm:  -> ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> ComputeName       : 'pkrvm04rrahxkg9'\n",
    "==> azure-arm: Powering off machine ...\n",
    "==> azure-arm:  -> ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> ComputeName       : 'pkrvm04rrahxkg9'\n",
    "==> azure-arm: Capturing image ...\n",
    "==> azure-arm:  -> Compute ResourceGroupName : <your resource group>\n",
    "==> azure-arm:  -> Compute Name              : 'pkrvm04rrahxkg9'\n",
    "==> azure-arm:  -> Compute Location          : <some chosen location>\n",
    "==> azure-arm:  -> Image ResourceGroupName   : <your resource group>\n",
    "==> azure-arm:  -> Image Name                : <your chosen custom image name>\n",
    "==> azure-arm:  -> Image Location            : <some chosen location>\n",
    "==> azure-arm: \n",
    "==> azure-arm: Deleting individual resources ...\n",
    "==> azure-arm: Adding to deletion queue -> Microsoft.Compute/virtualMachines : 'pkrvm04rrahxkg9'\n",
    "==> azure-arm: Adding to deletion queue -> Microsoft.Network/networkInterfaces : 'pkrni04rrahxkg9'\n",
    "==> azure-arm: Adding to deletion queue -> Microsoft.Network/publicIPAddresses : 'pkrip04rrahxkg9'\n",
    "==> azure-arm: Adding to deletion queue -> Microsoft.Network/virtualNetworks : 'pkrvn04rrahxkg9'\n",
    "==> azure-arm: Attempting deletion -> Microsoft.Network/networkInterfaces : 'pkrni04rrahxkg9'\n",
    "==> azure-arm: Waiting for deletion of all resources...\n",
    "==> azure-arm: Attempting deletion -> Microsoft.Network/publicIPAddresses : 'pkrip04rrahxkg9'\n",
    "==> azure-arm: Attempting deletion -> Microsoft.Compute/virtualMachines : 'pkrvm04rrahxkg9'\n",
    "==> azure-arm: Attempting deletion -> Microsoft.Network/virtualNetworks : 'pkrvn04rrahxkg9'\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "    .\n",
    "==> azure-arm:  Deleting -> Microsoft.Compute/disks : '/subscriptions/<your subscription id>/resourceGroups/<your resource group>/providers/Microsoft.Compute/disks/pkros04rrahxkg9'\n",
    "==> azure-arm: Removing the created Deployment object: 'pkrdp04rrahxkg9'\n",
    "==> azure-arm: \n",
    "==> azure-arm: The resource group was not created by Packer, not deleting ...\n",
    "Build 'azure-arm' finished after 16 minutes 22 seconds.\n",
    "\n",
    "==> Wait completed after 16 minutes 22 seconds\n",
    "\n",
    "==> Builds finished. The artifacts of successful builds are:\n",
    "--> azure-arm: Azure.ResourceManagement.VMImage:\n",
    "\n",
    "OSType: Linux\n",
    "ManagedImageResourceGroupName: <your resource group>\n",
    "ManagedImageName: <your chosen custom image name>\n",
    "ManagedImageId: /subscriptions/<your subscription id>/resourceGroups/<your resource group>/providers/Microsoft.Compute/images/<your chosen custom image name>\n",
    "ManagedImageLocation: <some chosen location>\n",
    "\n",
    "```\n",
    "---\n",
    "    \n",
    "When `packer` finishes, at the bottom of the output, you will see something similar to the following:\n",
    "\n",
    "```\n",
    "ManagedImageResourceGroupName: <your resource group>\n",
    "ManagedImageName: <your chosen custom image name>\n",
    "ManagedImageId: /subscriptions/<your subscription id>/resourceGroups/<your resource group>/providers/Microsoft.Compute/images/<your chosen custom image name>\n",
    "ManagedImageLocation: <some chosen location>\n",
    "```\n",
    "\n",
    "Make note of the `ManagedImageId`. This is the resource id of the custom image we will use. \n",
    "\n",
    "As shown above the `ManagedImageId` will look something like : `/subscriptions/12345/resourceGroups/myown-rg/providers/Microsoft.Compute/images/myCustomImage`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(set-up-customized-VM-information-and-clear-default-dask-config)=\n",
    "#### f. Set up customized VM information and clear default dask config "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have the custom VM resource id, you should reset the default VM image information in `dask.config`. The default image value loaded in `dask.config` is that of a basic Ubuntu Server 18.04 LTS (the one that you already customized). If you do not reset it, `dask` will try to use that image instead of your custom made one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fill this in with the value from above\n",
    "# or the customized VM id if you already have resource id of the customized VM from a previous run.\n",
    "ManagedImageId = \"FILL-THIS-IN\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dask.config.set({\"cloudprovider.azure.azurevm.vm_image\": {}})\n",
    "config = dask.config.get(\"cloudprovider.azure.azurevm\", {})\n",
    "print(config)\n",
    "vm_image = {\"id\": ManagedImageId}\n",
    "print(vm_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "(start-the-VM-Cluster-in-Azure)=\n",
    "### Step 2.1: Start the VM Cluster in Azure\n",
    "\n",
    "Here, if you have used Option 1, i.e., the NVIDIA VM image, pass an empty string for `vm_image` information. \n",
    "\n",
    "For Option 2, pass the `vm_image` information that you got from the output of `packer` run as a parameter to `AzureVMCluster`. \n",
    "\n",
    "Also turn off the bootstrapping of the VM by passing `bootstrap=False`. This will turn off installation of the dependencies in the VM while instantiating, since we already have them on our custom VM in either cases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "The rest of the notebook should be the same irrespective of whether you chose Option 1 or Option 2.\n",
    "```\n",
    "\n",
    "```{note}\n",
    "The number of actual workers that our cluster would have is not always equal to the number of VMs spawned i.e. the value of $n\\_workers$ passed in. If the number of GPUs in the chosen `vm_size` is $G$ and number of VMs spawned is $n\\_workers$, then we have then number of actual workers $W = n\\_workers \\times G$. For example, for `Standard_NC12s_v3` VMs that have 2 V100 GPUs per VM, for $n\\_workers=2$, we have $W = 2 \\times 2=4$.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "cluster = AzureVMCluster(\n",
    "    location=location,\n",
    "    resource_group=resource_group,\n",
    "    vnet=vnet,\n",
    "    security_group=security_group,\n",
    "    vm_image=vm_image,\n",
    "    vm_size=vm_size,\n",
    "    disk_size=200,\n",
    "    docker_image=docker_image,\n",
    "    worker_class=worker_class,\n",
    "    n_workers=2,\n",
    "    security=True,\n",
    "    docker_args=docker_args,\n",
    "    debug=False,\n",
    "    bootstrap=False,  # This is to prevent the cloud init jinja2 script from running in the custom VM.\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "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-b8982284-60fe-11ee-a1e9-80e82cd32958</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_cloudprovider.AzureVMCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://4.155.2.188:8787/status\" target=\"_blank\">http://4.155.2.188: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;\">AzureVMCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">e0c95e38</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://4.155.2.188:8787/status\" target=\"_blank\">http://4.155.2.188:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 4\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 4\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 440.42 GiB\n",
       "                </td>\n",
       "            </tr>\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-3bae5a4d-29d1-4317-bbfc-931e97a077fb</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tls://10.5.0.42:8786\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 4\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://10.5.0.42:8787/status\" target=\"_blank\">http://10.5.0.42:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 4\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> 12 minutes ago\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 440.42 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: dask-92c5978e-worker-54f8d057-0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.5.0.43:38107\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://10.5.0.43:45421/status\" target=\"_blank\">http://10.5.0.43:45421/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 110.11 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.5.0.43:33657\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-scratch-space/worker-ix8y4_eg\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-PCIE-16GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 16.00 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: dask-92c5978e-worker-54f8d057-1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.5.0.43:36201\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://10.5.0.43:44817/status\" target=\"_blank\">http://10.5.0.43:44817/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 110.11 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.5.0.43:42265\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-scratch-space/worker-6bghw_yx\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-PCIE-16GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 16.00 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: dask-92c5978e-worker-9f9a9c9b-0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.5.0.44:37791\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://10.5.0.44:32965/status\" target=\"_blank\">http://10.5.0.44:32965/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 110.11 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.5.0.44:36779\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-scratch-space/worker-7y8g_hu7\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-PCIE-16GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 16.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks executing: </strong> \n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks in memory: </strong> \n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks ready: </strong> \n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks in flight: </strong>\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>CPU usage:</strong> 0.0%\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Last seen: </strong> 6 minutes ago\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory usage: </strong> 329.71 MiB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Spilled bytes: </strong> 0 B\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Read bytes: </strong> 0.0 B\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Write bytes: </strong> 0.0 B\n",
       "                        </td>\n",
       "                    </tr>\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: dask-92c5978e-worker-9f9a9c9b-1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.5.0.44:34087\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://10.5.0.44:36073/status\" target=\"_blank\">http://10.5.0.44:36073/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 110.11 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.5.0.44:37979\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-scratch-space/worker-1d7vbddw\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-PCIE-16GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 16.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks executing: </strong> \n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks in memory: </strong> \n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks ready: </strong> \n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Tasks in flight: </strong>\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>CPU usage:</strong> 0.0%\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Last seen: </strong> 6 minutes ago\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory usage: </strong> 329.50 MiB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Spilled bytes: </strong> 0 B\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Read bytes: </strong> 0.0 B\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Write bytes: </strong> 0.0 B\n",
       "                        </td>\n",
       "                    </tr>\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: 'tls://10.5.0.42:8786' processes=4 threads=4, memory=440.42 GiB>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 6.1 ms, total: 6.1 ms\n",
      "Wall time: 29 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "client.wait_for_workers(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment if you only have the scheduler with n_workers=0 and want to scale the workers separately.\n",
    "# %%time\n",
    "# client.cluster.scale(n_workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wait till all the workers are up. This will wait for `n_workers` number of VMs to be up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we start the training process, let us take a quick look at the details of the GPUs in the worker pods that we will be using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'address': 'tls://10.5.0.42:8786',\n",
      " 'id': 'Scheduler-3bae5a4d-29d1-4317-bbfc-931e97a077fb',\n",
      " 'services': {'dashboard': 8787},\n",
      " 'started': 1696235012.5914223,\n",
      " 'type': 'Scheduler',\n",
      " 'workers': {'tls://10.5.0.43:36201': {'gpu': {'memory-total': 17179869184,\n",
      "                                               'name': 'Tesla V100-PCIE-16GB'},\n",
      "                                       'host': '10.5.0.43',\n",
      "                                       'id': 'dask-92c5978e-worker-54f8d057-1',\n",
      "                                       'last_seen': 1696235778.2340653,\n",
      "                                       'local_directory': '/tmp/dask-scratch-space/worker-6bghw_yx',\n",
      "                                       'memory_limit': 118225670144,\n",
      "                                       'metrics': {'bandwidth': {'total': 100000000,\n",
      "                                                                 'types': {},\n",
      "                                                                 'workers': {}},\n",
      "                                                   'cpu': 4.0,\n",
      "                                                   'digests_total_since_heartbeat': {'latency': 0.004627227783203125,\n",
      "                                                                                     'tick-duration': 0.5006744861602783},\n",
      "                                                   'event_loop_interval': 0.019985613822937013,\n",
      "                                                   'gpu': {'memory-used': 598867968,\n",
      "                                                           'utilization': 0},\n",
      "                                                   'gpu_memory_used': 598867968,\n",
      "                                                   'gpu_utilization': 0,\n",
      "                                                   'host_disk_io': {'read_bps': 0.0,\n",
      "                                                                    'write_bps': 0.0},\n",
      "                                                   'host_net_io': {'read_bps': 612.42422993883,\n",
      "                                                                   'write_bps': 3346.3180145677247},\n",
      "                                                   'managed_bytes': 0,\n",
      "                                                   'memory': 623116288,\n",
      "                                                   'num_fds': 86,\n",
      "                                                   'rmm': {'rmm-total': 0,\n",
      "                                                           'rmm-used': 0},\n",
      "                                                   'spilled_bytes': {'disk': 0,\n",
      "                                                                     'memory': 0},\n",
      "                                                   'task_counts': {},\n",
      "                                                   'time': 1696235777.730071,\n",
      "                                                   'transfer': {'incoming_bytes': 0,\n",
      "                                                                'incoming_count': 0,\n",
      "                                                                'incoming_count_total': 0,\n",
      "                                                                'outgoing_bytes': 0,\n",
      "                                                                'outgoing_count': 0,\n",
      "                                                                'outgoing_count_total': 0}},\n",
      "                                       'name': 'dask-92c5978e-worker-54f8d057-1',\n",
      "                                       'nanny': 'tls://10.5.0.43:42265',\n",
      "                                       'nthreads': 1,\n",
      "                                       'resources': {},\n",
      "                                       'services': {'dashboard': 44817},\n",
      "                                       'status': 'running',\n",
      "                                       'type': 'Worker'},\n",
      "             'tls://10.5.0.43:38107': {'gpu': {'memory-total': 17179869184,\n",
      "                                               'name': 'Tesla V100-PCIE-16GB'},\n",
      "                                       'host': '10.5.0.43',\n",
      "                                       'id': 'dask-92c5978e-worker-54f8d057-0',\n",
      "                                       'last_seen': 1696235778.2329032,\n",
      "                                       'local_directory': '/tmp/dask-scratch-space/worker-ix8y4_eg',\n",
      "                                       'memory_limit': 118225670144,\n",
      "                                       'metrics': {'bandwidth': {'total': 100000000,\n",
      "                                                                 'types': {},\n",
      "                                                                 'workers': {}},\n",
      "                                                   'cpu': 2.0,\n",
      "                                                   'digests_total_since_heartbeat': {'latency': 0.004603147506713867,\n",
      "                                                                                     'tick-duration': 0.4996976852416992},\n",
      "                                                   'event_loop_interval': 0.019999494552612306,\n",
      "                                                   'gpu': {'memory-used': 598867968,\n",
      "                                                           'utilization': 0},\n",
      "                                                   'gpu_memory_used': 598867968,\n",
      "                                                   'gpu_utilization': 0,\n",
      "                                                   'host_disk_io': {'read_bps': 0.0,\n",
      "                                                                    'write_bps': 0.0},\n",
      "                                                   'host_net_io': {'read_bps': 611.5250712835996,\n",
      "                                                                   'write_bps': 3341.404964660714},\n",
      "                                                   'managed_bytes': 0,\n",
      "                                                   'memory': 623882240,\n",
      "                                                   'num_fds': 86,\n",
      "                                                   'rmm': {'rmm-total': 0,\n",
      "                                                           'rmm-used': 0},\n",
      "                                                   'spilled_bytes': {'disk': 0,\n",
      "                                                                     'memory': 0},\n",
      "                                                   'task_counts': {},\n",
      "                                                   'time': 1696235777.729443,\n",
      "                                                   'transfer': {'incoming_bytes': 0,\n",
      "                                                                'incoming_count': 0,\n",
      "                                                                'incoming_count_total': 0,\n",
      "                                                                'outgoing_bytes': 0,\n",
      "                                                                'outgoing_count': 0,\n",
      "                                                                'outgoing_count_total': 0}},\n",
      "                                       'name': 'dask-92c5978e-worker-54f8d057-0',\n",
      "                                       'nanny': 'tls://10.5.0.43:33657',\n",
      "                                       'nthreads': 1,\n",
      "                                       'resources': {},\n",
      "                                       'services': {'dashboard': 45421},\n",
      "                                       'status': 'running',\n",
      "                                       'type': 'Worker'},\n",
      "             'tls://10.5.0.44:34087': {'gpu': {'memory-total': 17179869184,\n",
      "                                               'name': 'Tesla V100-PCIE-16GB'},\n",
      "                                       'host': '10.5.0.44',\n",
      "                                       'id': 'dask-92c5978e-worker-9f9a9c9b-1',\n",
      "                                       'last_seen': 1696235778.5268767,\n",
      "                                       'local_directory': '/tmp/dask-scratch-space/worker-1d7vbddw',\n",
      "                                       'memory_limit': 118225670144,\n",
      "                                       'metrics': {'bandwidth': {'total': 100000000,\n",
      "                                                                 'types': {},\n",
      "                                                                 'workers': {}},\n",
      "                                                   'cpu': 0.0,\n",
      "                                                   'digests_total_since_heartbeat': {'latency': 0.004075765609741211,\n",
      "                                                                                     'tick-duration': 0.4998819828033447},\n",
      "                                                   'event_loop_interval': 0.02001068115234375,\n",
      "                                                   'gpu': {'memory-used': 598867968,\n",
      "                                                           'utilization': 0},\n",
      "                                                   'gpu_memory_used': 598867968,\n",
      "                                                   'gpu_utilization': 0,\n",
      "                                                   'host_disk_io': {'read_bps': 0.0,\n",
      "                                                                    'write_bps': 12597732.652975753},\n",
      "                                                   'host_net_io': {'read_bps': 612.7208378808626,\n",
      "                                                                   'write_bps': 3347.938695871903},\n",
      "                                                   'managed_bytes': 0,\n",
      "                                                   'memory': 624406528,\n",
      "                                                   'num_fds': 86,\n",
      "                                                   'rmm': {'rmm-total': 0,\n",
      "                                                           'rmm-used': 0},\n",
      "                                                   'spilled_bytes': {'disk': 0,\n",
      "                                                                     'memory': 0},\n",
      "                                                   'task_counts': {},\n",
      "                                                   'time': 1696235778.023989,\n",
      "                                                   'transfer': {'incoming_bytes': 0,\n",
      "                                                                'incoming_count': 0,\n",
      "                                                                'incoming_count_total': 0,\n",
      "                                                                'outgoing_bytes': 0,\n",
      "                                                                'outgoing_count': 0,\n",
      "                                                                'outgoing_count_total': 0}},\n",
      "                                       'name': 'dask-92c5978e-worker-9f9a9c9b-1',\n",
      "                                       'nanny': 'tls://10.5.0.44:37979',\n",
      "                                       'nthreads': 1,\n",
      "                                       'resources': {},\n",
      "                                       'services': {'dashboard': 36073},\n",
      "                                       'status': 'running',\n",
      "                                       'type': 'Worker'},\n",
      "             'tls://10.5.0.44:37791': {'gpu': {'memory-total': 17179869184,\n",
      "                                               'name': 'Tesla V100-PCIE-16GB'},\n",
      "                                       'host': '10.5.0.44',\n",
      "                                       'id': 'dask-92c5978e-worker-9f9a9c9b-0',\n",
      "                                       'last_seen': 1696235778.528408,\n",
      "                                       'local_directory': '/tmp/dask-scratch-space/worker-7y8g_hu7',\n",
      "                                       'memory_limit': 118225670144,\n",
      "                                       'metrics': {'bandwidth': {'total': 100000000,\n",
      "                                                                 'types': {},\n",
      "                                                                 'workers': {}},\n",
      "                                                   'cpu': 0.0,\n",
      "                                                   'digests_total_since_heartbeat': {'latency': 0.003975629806518555,\n",
      "                                                                                     'tick-duration': 0.4994323253631592},\n",
      "                                                   'event_loop_interval': 0.020001530647277832,\n",
      "                                                   'gpu': {'memory-used': 598867968,\n",
      "                                                           'utilization': 0},\n",
      "                                                   'gpu_memory_used': 598867968,\n",
      "                                                   'gpu_utilization': 0,\n",
      "                                                   'host_disk_io': {'read_bps': 0.0,\n",
      "                                                                    'write_bps': 12589746.67130889},\n",
      "                                                   'host_net_io': {'read_bps': 612.3324205749067,\n",
      "                                                                   'write_bps': 3345.8163634027583},\n",
      "                                                   'managed_bytes': 0,\n",
      "                                                   'memory': 623104000,\n",
      "                                                   'num_fds': 86,\n",
      "                                                   'rmm': {'rmm-total': 0,\n",
      "                                                           'rmm-used': 0},\n",
      "                                                   'spilled_bytes': {'disk': 0,\n",
      "                                                                     'memory': 0},\n",
      "                                                   'task_counts': {},\n",
      "                                                   'time': 1696235778.0250378,\n",
      "                                                   'transfer': {'incoming_bytes': 0,\n",
      "                                                                'incoming_count': 0,\n",
      "                                                                'incoming_count_total': 0,\n",
      "                                                                'outgoing_bytes': 0,\n",
      "                                                                'outgoing_count': 0,\n",
      "                                                                'outgoing_count_total': 0}},\n",
      "                                       'name': 'dask-92c5978e-worker-9f9a9c9b-0',\n",
      "                                       'nanny': 'tls://10.5.0.44:36779',\n",
      "                                       'nthreads': 1,\n",
      "                                       'resources': {},\n",
      "                                       'services': {'dashboard': 32965},\n",
      "                                       'status': 'running',\n",
      "                                       'type': 'Worker'}}}\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "pp = pprint.PrettyPrinter()\n",
    "\n",
    "pp.pprint(\n",
    "    client.scheduler_info()\n",
    ")  # will show some information of the GPUs of the workers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Step 3: Data Setup, Cleanup and Enhancement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Step 3.a: Set up the workers for reading parquet files from Azure Data Lake endpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now enable all the workers to read the `parquet` files directly from the Azure Data Lake endpoints. This requires the [`adlfs`](https://github.com/dask/adlfs) python library in the workers. We will pass in the simple function `installAdlfs` in `client.run` which will install the python package in all the workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tls://10.5.0.43:36201': {'status': 'OK'},\n",
       " 'tls://10.5.0.43:38107': {'status': 'OK'},\n",
       " 'tls://10.5.0.44:34087': {'status': 'OK'},\n",
       " 'tls://10.5.0.44:37791': {'status': 'OK'}}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dask.distributed import PipInstall\n",
    "\n",
    "client.register_worker_plugin(PipInstall(packages=[\"adlfs\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Step 3.b: Data Cleanup, Enhancement and Persisting Scripts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data needs to be cleaned up first. We remove some columns that we are not interested in. We also define the datatypes each of the columns need to be read as."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also add some new features to our dataframe via some custom functions, namely: \n",
    "1. Haversine distance: This is used for calculating the total trip distance.\n",
    "\n",
    "2. Day of the week: This can be useful information for determining the fare cost.\n",
    "\n",
    "`add_features` function combines the two to produce a new dataframe that has the added features.\n",
    "\n",
    "```{note}\n",
    "In the function `persist_train_infer_split`, We will also persist the test dataset in the workers. If the `X_infer` i.e. the test dataset is small enough, we can call `compute()` on it to bring the test dataset to the local machine and then perform predict on it. But in general, if the `X_infer` is large, it may not fit in the GPU(s) of the local machine. Moreover, moving around a large amount of data will also add to the prediction latency. Therefore it is better to persist the test dataset on the dask workers, and then call the predict functionality on the individual workers. Finally we collect the prediction results from the dask workers. \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adding features functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from math import asin, cos, pi, sin, sqrt\n",
    "\n",
    "\n",
    "def haversine_distance_kernel(\n",
    "    pickup_latitude_r,\n",
    "    pickup_longitude_r,\n",
    "    dropoff_latitude_r,\n",
    "    dropoff_longitude_r,\n",
    "    h_distance,\n",
    "    radius,\n",
    "):\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(\n",
    "        zip(\n",
    "            pickup_latitude_r,\n",
    "            pickup_longitude_r,\n",
    "            dropoff_latitude_r,\n",
    "            dropoff_longitude_r,\n",
    "            strict=False,\n",
    "        )\n",
    "    ):\n",
    "        x_1 = pi / 180 * x_1\n",
    "        y_1 = pi / 180 * y_1\n",
    "        x_2 = pi / 180 * x_2\n",
    "        y_2 = pi / 180 * y_2\n",
    "\n",
    "        dlon = y_2 - y_1\n",
    "        dlat = x_2 - x_1\n",
    "        a = sin(dlat / 2) ** 2 + cos(x_1) * cos(x_2) * sin(dlon / 2) ** 2\n",
    "\n",
    "        c = 2 * asin(sqrt(a))\n",
    "        # radius = 6371 # Radius of earth in kilometers # currently passed as input arguments\n",
    "\n",
    "        h_distance[i] = c * radius\n",
    "\n",
    "\n",
    "def day_of_the_week_kernel(day, month, year, day_of_week):\n",
    "    for i, (_, _, _) in enumerate(zip(day, month, year, strict=False)):\n",
    "        if month[i] < 3:\n",
    "            shift = month[i]\n",
    "        else:\n",
    "            shift = 0\n",
    "        Y = year[i] - (month[i] < 3)\n",
    "        y = Y - 2000\n",
    "        c = 20\n",
    "        d = day[i]\n",
    "        m = month[i] + shift + 1\n",
    "        day_of_week[i] = (d + math.floor(m * 2.6) + y + (y // 4) + (c // 4) - 2 * c) % 7\n",
    "\n",
    "\n",
    "def add_features(df):\n",
    "    df[\"hour\"] = df[\"tpepPickupDateTime\"].dt.hour\n",
    "    df[\"year\"] = df[\"tpepPickupDateTime\"].dt.year\n",
    "    df[\"month\"] = df[\"tpepPickupDateTime\"].dt.month\n",
    "    df[\"day\"] = df[\"tpepPickupDateTime\"].dt.day\n",
    "    df[\"diff\"] = (\n",
    "        df[\"tpepDropoffDateTime\"] - df[\"tpepPickupDateTime\"]\n",
    "    ).dt.seconds  # convert difference between pickup and dropoff into seconds\n",
    "\n",
    "    df[\"pickup_latitude_r\"] = df[\"startLat\"] // 0.01 * 0.01\n",
    "    df[\"pickup_longitude_r\"] = df[\"startLon\"] // 0.01 * 0.01\n",
    "    df[\"dropoff_latitude_r\"] = df[\"endLat\"] // 0.01 * 0.01\n",
    "    df[\"dropoff_longitude_r\"] = df[\"endLon\"] // 0.01 * 0.01\n",
    "\n",
    "    df = df.drop(\"tpepDropoffDateTime\", axis=1)\n",
    "    df = df.drop(\"tpepPickupDateTime\", axis=1)\n",
    "\n",
    "    df = df.apply_rows(\n",
    "        haversine_distance_kernel,\n",
    "        incols=[\n",
    "            \"pickup_latitude_r\",\n",
    "            \"pickup_longitude_r\",\n",
    "            \"dropoff_latitude_r\",\n",
    "            \"dropoff_longitude_r\",\n",
    "        ],\n",
    "        outcols=dict(h_distance=np.float32),\n",
    "        kwargs=dict(radius=6371),\n",
    "    )\n",
    "\n",
    "    df = df.apply_rows(\n",
    "        day_of_the_week_kernel,\n",
    "        incols=[\"day\", \"month\", \"year\"],\n",
    "        outcols=dict(day_of_week=np.float32),\n",
    "        kwargs=dict(),\n",
    "    )\n",
    "\n",
    "    df[\"is_weekend\"] = df[\"day_of_week\"] < 2\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Functions for cleaning and persisting the data in the workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def persist_train_infer_split(\n",
    "    client,\n",
    "    df,\n",
    "    response_dtype,\n",
    "    response_id,\n",
    "    infer_frac=1.0,\n",
    "    random_state=42,\n",
    "    shuffle=True,\n",
    "):\n",
    "    workers = client.has_what().keys()\n",
    "    X, y = df.drop([response_id], axis=1), df[response_id].astype(\"float32\")\n",
    "    infer_frac = max(0, min(infer_frac, 1.0))\n",
    "    X_train, X_infer, y_train, y_infer = train_test_split(\n",
    "        X, y, shuffle=True, random_state=random_state, test_size=infer_frac\n",
    "    )\n",
    "\n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        X_train, y_train = client.persist(collections=[X_train, y_train])\n",
    "\n",
    "    if infer_frac != 1.0:\n",
    "        with dask.annotate(workers=set(workers)):\n",
    "            X_infer, y_infer = client.persist(collections=[X_infer, y_infer])\n",
    "\n",
    "        wait([X_train, y_train, X_infer, y_infer])\n",
    "    else:\n",
    "        X_infer = X_train\n",
    "        y_infer = y_train\n",
    "\n",
    "        wait([X_train, y_train])\n",
    "\n",
    "    return X_train, y_train, X_infer, y_infer\n",
    "\n",
    "\n",
    "def clean(df_part, must_haves):\n",
    "    \"\"\"\n",
    "    This function performs the various clean up tasks for the data\n",
    "    and returns the cleaned dataframe.\n",
    "    \"\"\"\n",
    "    # iterate through columns in this df partition\n",
    "    for col in df_part.columns:\n",
    "        # drop anything not in our expected list\n",
    "        if col not in must_haves:\n",
    "            df_part = df_part.drop(col, axis=1)\n",
    "            continue\n",
    "\n",
    "        # fixes datetime error found by Ty Mckercher and fixed by Paul Mahler\n",
    "        if df_part[col].dtype == \"object\" and col in [\n",
    "            \"tpepPickupDateTime\",\n",
    "            \"tpepDropoffDateTime\",\n",
    "        ]:\n",
    "            df_part[col] = df_part[col].astype(\"datetime64[ms]\")\n",
    "            continue\n",
    "\n",
    "        # if column was read as a string, recast as float\n",
    "        if df_part[col].dtype == \"object\":\n",
    "            df_part[col] = df_part[col].str.fillna(\"-1\")\n",
    "            df_part[col] = df_part[col].astype(\"float32\")\n",
    "        else:\n",
    "            # downcast from 64bit to 32bit types\n",
    "            # Tesla T4 are faster on 32bit ops\n",
    "            if \"int\" in str(df_part[col].dtype):\n",
    "                df_part[col] = df_part[col].astype(\"int32\")\n",
    "            if \"float\" in str(df_part[col].dtype):\n",
    "                df_part[col] = df_part[col].astype(\"float32\")\n",
    "            df_part[col] = df_part[col].fillna(-1)\n",
    "\n",
    "    return df_part\n",
    "\n",
    "\n",
    "def taxi_data_loader(\n",
    "    client,\n",
    "    adlsaccount,\n",
    "    adlspath,\n",
    "    response_dtype=np.float32,\n",
    "    infer_frac=1.0,\n",
    "    random_state=0,\n",
    "):\n",
    "    # create a list of columns & dtypes the df must have\n",
    "    must_haves = {\n",
    "        \"tpepPickupDateTime\": \"datetime64[ms]\",\n",
    "        \"tpepDropoffDateTime\": \"datetime64[ms]\",\n",
    "        \"passengerCount\": \"int32\",\n",
    "        \"tripDistance\": \"float32\",\n",
    "        \"startLon\": \"float32\",\n",
    "        \"startLat\": \"float32\",\n",
    "        \"rateCodeId\": \"int32\",\n",
    "        \"endLon\": \"float32\",\n",
    "        \"endLat\": \"float32\",\n",
    "        \"fareAmount\": \"float32\",\n",
    "    }\n",
    "\n",
    "    workers = client.has_what().keys()\n",
    "    response_id = \"fareAmount\"\n",
    "    storage_options = {\"account_name\": adlsaccount}\n",
    "    taxi_data = dask_cudf.read_parquet(\n",
    "        adlspath,\n",
    "        storage_options=storage_options,\n",
    "        chunksize=25e6,\n",
    "        npartitions=len(workers),\n",
    "    )\n",
    "    taxi_data = clean(taxi_data, must_haves)\n",
    "    taxi_data = taxi_data.map_partitions(add_features)\n",
    "    # Drop NaN values and convert to float32\n",
    "    taxi_data = taxi_data.dropna()\n",
    "    fields = [\n",
    "        \"passengerCount\",\n",
    "        \"tripDistance\",\n",
    "        \"startLon\",\n",
    "        \"startLat\",\n",
    "        \"rateCodeId\",\n",
    "        \"endLon\",\n",
    "        \"endLat\",\n",
    "        \"fareAmount\",\n",
    "        \"diff\",\n",
    "        \"h_distance\",\n",
    "        \"day_of_week\",\n",
    "        \"is_weekend\",\n",
    "    ]\n",
    "    taxi_data = taxi_data.astype(\"float32\")\n",
    "    taxi_data = taxi_data[fields]\n",
    "    taxi_data = taxi_data.reset_index()\n",
    "\n",
    "    return persist_train_infer_split(\n",
    "        client, taxi_data, response_dtype, response_id, infer_frac, random_state\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Step 3.c: Get the split data and persist across workers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will make use of the data from November and December 2014 for the purposes of the demo. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tic = timer()\n",
    "X_train, y_train, X_infer, y_infer = taxi_data_loader(\n",
    "    client,\n",
    "    adlsaccount=\"azureopendatastorage\",\n",
    "    adlspath=\"az://nyctlc/yellow/puYear=2014/puMonth=1*/*.parquet\",\n",
    "    infer_frac=0.1,\n",
    "    random_state=42,\n",
    ")\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for ETL and persisting : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48817562"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train.shape[0].compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The size of our training dataset is around 49 million rows. Let's look at the data locally to see what we're dealing with. We see that there are columns for pickup and dropoff latitude and longitude, passenger count, trip distance, day of week etc. These are the information we'll use to estimate the trip fare amount. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>passengerCount</th>\n",
       "      <th>tripDistance</th>\n",
       "      <th>startLon</th>\n",
       "      <th>startLat</th>\n",
       "      <th>rateCodeId</th>\n",
       "      <th>endLon</th>\n",
       "      <th>endLat</th>\n",
       "      <th>diff</th>\n",
       "      <th>h_distance</th>\n",
       "      <th>day_of_week</th>\n",
       "      <th>is_weekend</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>300446</th>\n",
       "      <td>300446</td>\n",
       "      <td>1.0</td>\n",
       "      <td>4.00</td>\n",
       "      <td>-73.984955</td>\n",
       "      <td>40.768543</td>\n",
       "      <td>1.0</td>\n",
       "      <td>-74.008789</td>\n",
       "      <td>40.719330</td>\n",
       "      <td>1324.0</td>\n",
       "      <td>5.809536e+00</td>\n",
       "      <td>6.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>163817</th>\n",
       "      <td>163817</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.93</td>\n",
       "      <td>-74.008179</td>\n",
       "      <td>40.722198</td>\n",
       "      <td>1.0</td>\n",
       "      <td>-73.992989</td>\n",
       "      <td>40.739151</td>\n",
       "      <td>840.0</td>\n",
       "      <td>1.395489e+00</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>236958</th>\n",
       "      <td>236958</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1.10</td>\n",
       "      <td>-73.987595</td>\n",
       "      <td>40.775360</td>\n",
       "      <td>1.0</td>\n",
       "      <td>-73.976494</td>\n",
       "      <td>40.785755</td>\n",
       "      <td>360.0</td>\n",
       "      <td>1.394768e+00</td>\n",
       "      <td>5.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>73461</th>\n",
       "      <td>73461</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.76</td>\n",
       "      <td>-73.994698</td>\n",
       "      <td>40.725929</td>\n",
       "      <td>1.0</td>\n",
       "      <td>-73.994698</td>\n",
       "      <td>40.725929</td>\n",
       "      <td>180.0</td>\n",
       "      <td>1.005159e-13</td>\n",
       "      <td>5.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>294464</th>\n",
       "      <td>294464</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.60</td>\n",
       "      <td>-73.974342</td>\n",
       "      <td>40.748165</td>\n",
       "      <td>1.0</td>\n",
       "      <td>-73.982536</td>\n",
       "      <td>40.750767</td>\n",
       "      <td>229.0</td>\n",
       "      <td>1.395336e+00</td>\n",
       "      <td>6.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         index  passengerCount  tripDistance   startLon   startLat  \\\n",
       "300446  300446             1.0          4.00 -73.984955  40.768543   \n",
       "163817  163817             3.0          1.93 -74.008179  40.722198   \n",
       "236958  236958             5.0          1.10 -73.987595  40.775360   \n",
       "73461    73461             1.0          0.76 -73.994698  40.725929   \n",
       "294464  294464             1.0          0.60 -73.974342  40.748165   \n",
       "\n",
       "        rateCodeId     endLon     endLat    diff    h_distance  day_of_week  \\\n",
       "300446         1.0 -74.008789  40.719330  1324.0  5.809536e+00          6.0   \n",
       "163817         1.0 -73.992989  40.739151   840.0  1.395489e+00          2.0   \n",
       "236958         1.0 -73.976494  40.785755   360.0  1.394768e+00          5.0   \n",
       "73461          1.0 -73.994698  40.725929   180.0  1.005159e-13          5.0   \n",
       "294464         1.0 -73.982536  40.750767   229.0  1.395336e+00          6.0   \n",
       "\n",
       "        is_weekend  \n",
       "300446         0.0  \n",
       "163817         0.0  \n",
       "236958         0.0  \n",
       "73461          0.0  \n",
       "294464         0.0  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>passengerCount</th>\n",
       "      <th>tripDistance</th>\n",
       "      <th>startLon</th>\n",
       "      <th>startLat</th>\n",
       "      <th>rateCodeId</th>\n",
       "      <th>endLon</th>\n",
       "      <th>endLat</th>\n",
       "      <th>diff</th>\n",
       "      <th>h_distance</th>\n",
       "      <th>day_of_week</th>\n",
       "      <th>is_weekend</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=84</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>int64</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: split, 1 graph layer</div>"
      ],
      "text/plain": [
       "<dask_cudf.DataFrame | 84 tasks | 84 npartitions>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_infer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Step 4: Train a XGBoost Model\n",
    "\n",
    "We are now ready to train a XGBoost model on the data and then predict the fare for each trip."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Step 4.a: Set training Parameters\n",
    "\n",
    "In this training example, we will use RMSE as the evaluation metric. It is also worth noting that performing HPO will lead to a set of more optimal hyperparameters.\n",
    "\n",
    "Refer to the notebook [HPO-RAPIDS](../rapids-azureml-hpo/notebook.ipynb) in this repository for how to perform HPO on Azure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {\n",
    "    \"learning_rate\": 0.15,\n",
    "    \"max_depth\": 8,\n",
    "    \"objective\": \"reg:squarederror\",\n",
    "    \"subsample\": 0.7,\n",
    "    \"colsample_bytree\": 0.7,\n",
    "    \"min_child_weight\": 1,\n",
    "    \"gamma\": 1,\n",
    "    \"silent\": True,\n",
    "    \"verbose_eval\": True,\n",
    "    \"booster\": \"gbtree\",  # 'gblinear' not implemented in dask\n",
    "    \"debug_synchronize\": True,\n",
    "    \"eval_metric\": \"rmse\",\n",
    "    \"tree_method\": \"gpu_hist\",\n",
    "    \"num_boost_rounds\": 100,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Step 4.b: Train XGBoost Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the data is already persisted in the dask workers in the cluster, the next steps should not take a lot of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall clock time taken for this cell : 9.483002611901611 s\n"
     ]
    }
   ],
   "source": [
    "data_train = xgb.dask.DaskDMatrix(client, X_train, y_train)\n",
    "tic = timer()\n",
    "xgboost_output = xgb.dask.train(\n",
    "    client, params, data_train, num_boost_round=params[\"num_boost_rounds\"]\n",
    ")\n",
    "xgb_gpu_model = xgboost_output[\"booster\"]\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for this cell : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Step 4.c: Save the trained model to disk locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xgboost.core.Booster at 0x7fcee8055bd0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xgb_gpu_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_filename = \"trained-model_nyctaxi.xgb\"\n",
    "xgb_gpu_model.save_model(model_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Step 5: Predict & Score using vanilla XGBoost Predict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Here we will use the `predict` and `inplace_predict` methods provided by the `xgboost.dask` library, out of the box. Later we will also use [Forest Inference Library (FIL)](https://docs.rapids.ai/api/cuml/~~~rapids_api_docs_version~~~/api.html?highlight=forestinference#cuml.ForestInference) to perform prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DoneAndNotDoneFutures(done=set(), not_done=set())"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_y_test = y_infer.compute()\n",
    "wait(_y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall clock time taken for xgb.dask.predict : 1.5550181320868433 s\n"
     ]
    }
   ],
   "source": [
    "d_test = xgb.dask.DaskDMatrix(client, X_infer)\n",
    "tic = timer()\n",
    "y_pred = xgb.dask.predict(client, xgb_gpu_model, d_test)\n",
    "y_pred = y_pred.compute()\n",
    "wait(y_pred)\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for xgb.dask.predict : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "#### Inference with the inplace predict method of dask XGBoost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall clock time taken for inplace inference : 1.8849179210374132 s\n"
     ]
    }
   ],
   "source": [
    "tic = timer()\n",
    "y_pred = xgb.dask.inplace_predict(client, xgb_gpu_model, X_infer)\n",
    "y_pred = y_pred.compute()\n",
    "wait(y_pred)\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for inplace inference : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating MSE\n",
      "Workflow Complete - RMSE:  2.2968235\n",
      "Wall clock time taken for this cell : 0.009336891933344305 s\n"
     ]
    }
   ],
   "source": [
    "tic = timer()\n",
    "print(\"Calculating MSE\")\n",
    "score = mean_squared_error(y_pred, _y_test)\n",
    "print(\"Workflow Complete - RMSE: \", np.sqrt(score))\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for this cell : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Step 6: Predict & Score using FIL or Forest Inference Library\n",
    "\n",
    "[Forest Inference Library (FIL)](https://docs.rapids.ai/api/cuml/~~~rapids_api_docs_version~~~/api.html?highlight=forestinference#cuml.ForestInference) provides GPU accelerated inference capabilities for tree models. We will import the FIL functionality from [cuML](https://github.com/rapidsai/cuml) library.\n",
    "\n",
    "It accepts a **trained** tree model in a treelite format (currently LightGBM, XGBoost and SKLearn GBDT and random forest models\n",
    "are supported). In general, using FIL allows for faster inference while using a large number of workers, and the latency benefits are more pronounced as the size of the dataset grows large. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6.a: Predict using `compute` on a single worker in case the test dataset is small.\n",
    "\n",
    "As noted in *Step 3.b*,  in case the test dataset is huge, it makes sense to call predict individually on the dask workers instead of bringing the entire test dataset to the local machine.\n",
    "\n",
    "To perform prediction individually on the dask workers, each dask worker needs to load the XGB model using FIL. However, the dask workers are remote and do not have access to the locally saved model. Hence we need to send the locally saved XGB model to the dask workers. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Persist the local model in the remote dask workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the code below will read the locally saved xgboost model\n",
    "# in binary format and write a copy of it to all dask workers\n",
    "def read_model(path):\n",
    "    \"\"\"Read model file into memory.\"\"\"\n",
    "    with open(path, \"rb\") as fh:\n",
    "        return fh.read()\n",
    "\n",
    "\n",
    "def write_model(path, data):\n",
    "    \"\"\"Write model file to disk.\"\"\"\n",
    "    with open(path, \"wb\") as fh:\n",
    "        fh.write(data)\n",
    "    return path\n",
    "\n",
    "\n",
    "model_data = read_model(\"trained-model_nyctaxi.xgb\")\n",
    "\n",
    "# Tell all the workers to write the model to disk\n",
    "client.run(write_model, \"/tmp/model.dat\", model_data)\n",
    "\n",
    "\n",
    "# this code reads the binary file in worker directory\n",
    "# and loads the model via FIL for prediction\n",
    "def predict_model(input_df):\n",
    "    from cuml import ForestInference\n",
    "\n",
    "    # load xgboost model using FIL and make prediction\n",
    "    fm = ForestInference.load(\"/tmp/model.dat\", model_type=\"xgboost\")\n",
    "    print(fm)\n",
    "    pred = fm.predict(input_df)\n",
    "\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inference with distributed predict with FIL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tic = timer()\n",
    "predictions = X_infer.map_partitions(\n",
    "    predict_model, meta=\"float\"\n",
    ")  # this is like MPI reduce\n",
    "y_pred = predictions.compute()\n",
    "wait(y_pred)\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for this cell : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "It took 5.638823717948981 seconds to predict on 5426301 rows using FIL distributedly on each worker\n"
     ]
    }
   ],
   "source": [
    "rows_csv = X_infer.iloc[:, 0].shape[0].compute()\n",
    "print(\n",
    "    f\"It took {toc-tic} seconds to predict on {rows_csv} rows using FIL distributedly on each worker\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final - RMSE:  2.2968235\n"
     ]
    }
   ],
   "source": [
    "tic = timer()\n",
    "score = mean_squared_error(y_pred, _y_test)\n",
    "toc = timer()\n",
    "print(\"Final - RMSE: \", np.sqrt(score))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Step 7: Clean up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Terminated VM dask-92c5978e-worker-54f8d057\n",
      "Terminated VM dask-92c5978e-worker-9f9a9c9b\n",
      "Removed disks for VM dask-92c5978e-worker-54f8d057\n",
      "Removed disks for VM dask-92c5978e-worker-9f9a9c9b\n",
      "Deleted network interface\n",
      "Deleted network interface\n",
      "Terminated VM dask-92c5978e-scheduler\n",
      "Removed disks for VM dask-92c5978e-scheduler\n",
      "Deleted network interface\n",
      "Unassigned public IP\n"
     ]
    }
   ],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
