azure/notebooks/Azure-MNMG-XGBoost.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <span style=\"color:#8735fb; font-size:24pt\"> Multi-Node Multi-GPU XGBoost example on Azure using dask-cloudprovider </span>\n",
    "\n",
    "[Dask Cloud Provider](https://cloudprovider.dask.org/en/latest/) is a native cloud intergration 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",
    "---\n",
    "\n",
    "**NOTE:**  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](#-Option-1:-Use-an-Azure-marketplace-VM.-): 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](#-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": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 0: Set up Azure credentials and CLI </span>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running the notebook, run the following commands in the terminal to setup Azure CLI\n",
    "```\n",
    "curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash\n",
    "az login\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": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 1: Import necessary packages </span>"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, wait, get_worker\n",
    "from dask_cloudprovider.azure import AzureVMCluster\n",
    "import dask_cudf\n",
    "from dask_ml.model_selection import train_test_split\n",
    "from cuml.dask.common import utils as dask_utils\n",
    "from cuml.metrics import mean_squared_error\n",
    "from cuml import ForestInference\n",
    "import cudf\n",
    "import xgboost as xgb\n",
    "from datetime import datetime\n",
    "from dateutil import parser\n",
    "import numpy as np\n",
    "from timeit import default_timer as timer\n",
    "import dask\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 2: Set up the Azure VM Cluster</span>\n",
    "\n",
    "We will now set up a Dask cluster on Azure Virtual machines using [`AzureVMCluster`](https://cloudprovider.dask.org/en/latest/azure.htm) from Dask Cloud Provider. 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# location = <your chosen location where the resource group and vnet exists>\n",
    "# resource_group = <your resource group>\n",
    "# vnet = <vnet in your resource group>\n",
    "# security_group = <security group in your resource group>\n",
    "# vm_size = \"Standard_NC12s_v3\" # or choose a different GPU enabled VM type\n",
    "\n",
    "docker_image = \"rapidsai/rapidsai:cuda11.2-runtime-ubuntu18.04-py3.8\"\n",
    "docker_args = '--shm-size=256m'\n",
    "worker_class = \"dask_cuda.CUDAWorker\"\n",
    "worker_options = {'rmm-managed-memory':True}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### <span style=\"color:#8735fb; font-size:20pt\"> Option 1: Use an Azure Marketplace VM image </span>\n",
    "\n",
    "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**: 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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set up Marketplace VM information and clear default dask config. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dask.config.set({\"logging.distributed\": \"info\",\n",
    "                \"cloudprovider.azure.azurevm.marketplace_plan\":{\n",
    "                        \"publisher\": \"nvidia\",\n",
    "                        \"name\": \"ngc-base-version-21-02-2\",\n",
    "                        \"product\": \"ngc_azure_17_11\",\n",
    "                        \"version\": \"21.02.2\"\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": [],
   "source": [
    "# ! az vm image terms accept --urn \"nvidia:ngc_azure_17_11:ngc-base-version-21-02-2:21.02.2\" --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](#-Step-2.1:-Start-the-VM-Cluster-in-Azure-) to start the AzureVMCluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:20pt\"> Option 2: Set up an Azure Customized VM </span>\n",
    "\n",
    "#### If you already have a customized VM and you know its resource id, jump to [Step f. of Option 2](#f.-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 guidelines in 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": {},
   "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:** 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`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "packer_config = {\n",
    "  \"builders\": [{\n",
    "    \"type\": \"azure-arm\", \n",
    "    \"use_azure_cli_auth\": True,\n",
    "    \"managed_image_resource_group_name\": resource_group,\n",
    "    \"managed_image_name\": <the name of the customized VM image>,\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",
    "  \"provisioners\": [{\n",
    "    \"inline\": [\n",
    "      \"echo 'Waiting for cloud-init'; while [ ! -f /var/lib/cloud/instance/boot-finished ]; do sleep 1; done; echo 'Done'\",\n",
    "    ],\n",
    "    \"type\": \"shell\"\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",
    "---\n",
    "\n",
    "```bash\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",
    "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",
    "As shown above the `ManagedImageId` will look something like : `/subscriptions/12345/resourceGroups/myown-rg/providers/Microsoft.Compute/images/myCustomImage`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 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": [
    "ManagedImageId = <value from the output above> # or the customized VM id if you already have resource id of the customized VM from a previous run. "
   ]
  },
  {
   "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": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:20pt\"> Step 2.1: Start the VM Cluster in Azure </span>\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:** The rest of the notebook should be the same irrespective of whether you chose Option 1 or Option 2.\n",
    "\n",
    "**NOTE on $n\\_workers$ as a parameter:** 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$."
   ]
  },
  {
   "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",
    "    docker_image=docker_image,\n",
    "    worker_class=worker_class,\n",
    "    n_workers=2,\n",
    "    security=True,\n",
    "    docker_args=docker_args,\n",
    "    worker_options=worker_options,\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster) \n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale_workers(client, n_workers, n_gpus_per_worker, timeout=300):\n",
    "    import time\n",
    "    client.cluster.scale(n_workers)\n",
    "    m = len(client.has_what().keys())    \n",
    "    start = end = time.perf_counter_ns()\n",
    "    while ((m != n_workers*n_gpus_per_worker) and (((end - start) / 1e9) < timeout) ):\n",
    "        time.sleep(5)\n",
    "        m = len(client.has_what().keys())\n",
    "        \n",
    "        end = time.perf_counter_ns()\n",
    "        \n",
    "    if (((end - start) / 1e9) >= timeout):\n",
    "        raise RuntimeError(f\"Failed to rescale cluster in {timeout} sec.\"\n",
    "              \"Try increasing timeout for very large containers, and verify available compute resources.\")"
   ]
  },
  {
   "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",
    "# scale_workers(client, 2, 2, timeout=600)"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "client.wait_for_workers(2)"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pretty_print(scheduler_dict):\n",
    "    print(f\"All workers for scheduler id: {scheduler_dict['id']}, address: {scheduler_dict['address']}\")\n",
    "    for worker in scheduler_dict['workers']:\n",
    "        print(f\"Worker: {worker} , gpu_machines: {scheduler_dict['workers'][worker]['gpu']}\")\n",
    "\n",
    "pretty_print(client.scheduler_info()) # will show some information of the GPUs of the workers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 3: Data Setup, Cleanup and Enhancement </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 3.a: Set up the workers for reading parquet files from Azure Data Lake endpoints </span>"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def installAdlfs():\n",
    "    import subprocess\n",
    "    subprocess.run([\"pip\", \"install\", \"adlfs\"])\n",
    "    return \"done\"\n",
    "\n",
    "results=client.run(installAdlfs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 3.b: Data Cleanup, Enhancement and Persisting Scripts </span>"
   ]
  },
  {
   "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:** 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. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adding features functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from math import cos, sin, asin, sqrt, pi\n",
    "\n",
    "def haversine_distance_kernel(pickup_latitude_r, pickup_longitude_r, dropoff_latitude_r, dropoff_longitude_r, h_distance, radius):\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(pickup_latitude_r, pickup_longitude_r, dropoff_latitude_r, dropoff_longitude_r,)):\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",
    "def day_of_the_week_kernel(day, month, year, day_of_week):\n",
    "    for i, (d_1, m_1, y_1) in enumerate(zip(day, month, year)):\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",
    "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'] = (df['tpepPickupDateTime'] - df['tpepPickupDateTime']).dt.seconds #convert difference between pickup and dropoff into seconds\n",
    "    \n",
    "    df['pickup_latitude_r'] = df['startLat']//.01*.01\n",
    "    df['pickup_longitude_r'] = df['startLon']//.01*.01\n",
    "    df['dropoff_latitude_r'] = df['endLat']//.01*.01\n",
    "    df['dropoff_longitude_r'] = df['endLon']//.01*.01\n",
    "    \n",
    "    df = df.drop('tpepDropoffDateTime', axis=1)\n",
    "    df = df.drop('tpepPickupDateTime', axis =1)\n",
    "    \n",
    "    \n",
    "    df = df.apply_rows(haversine_distance_kernel,\n",
    "                   incols=['pickup_latitude_r', 'pickup_longitude_r', 'dropoff_latitude_r', 'dropoff_longitude_r'],\n",
    "                   outcols=dict(h_distance=np.float32),\n",
    "                   kwargs=dict(radius=6371))\n",
    "    \n",
    "    \n",
    "    df = df.apply_rows(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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def persist_train_infer_split(client, df, response_dtype, response_id, infer_frac=1.0, random_state=42, shuffle=True):\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(X, y, shuffle=True, random_state=random_state, test_size=infer_frac)\n",
    "    \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        X_train, y_train = client.persist(\n",
    "            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(\n",
    "                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 ['tpepPickupDateTime', 'tpepDropoffDateTime']:\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(client, adlsaccount, adlspath, response_dtype=np.float32, infer_frac=1.0, random_state=0):\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(adlspath, storage_options=storage_options, chunksize=25e6, npartitions=len(workers))\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 = ['passengerCount', 'tripDistance', 'startLon', 'startLat', 'rateCodeId',\n",
    "                 'endLon', 'endLat', 'fareAmount', 'diff', 'h_distance', 'day_of_week', 'is_weekend']\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(client, taxi_data, response_dtype, response_id, infer_frac, random_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 3.c: Get the split data and persist across workers </span>"
   ]
  },
  {
   "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(client, \n",
    "                                                      adlsaccount=\"azureopendatastorage\", \n",
    "                                                      adlspath=\"az://nyctlc/yellow/puYear=2014/puMonth=1*/*.parquet\",\n",
    "                                                      infer_frac=0.1, random_state=42)\n",
    "toc = timer()\n",
    "print(f\"Wall clock time taken for ETL and persisting : {toc-tic} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 4: Train a XGBoost Model </span>\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": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 4.a: Set training Parameters </span>\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](./HPO-RAPIDS.ipynb) in this repository for how to perform HPO on Azure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 4.b: Train XGBoost Model </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the data is already persisted in the dask workers in the Kubernetes Cluster, the next steps should not take a lot of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_train = xgb.dask.DaskDMatrix(client, X_train, y_train)\n",
    "tic = timer()\n",
    "xgboost_output = xgb.dask.train(client, params,data_train, \n",
    "                                    num_boost_round=params['num_boost_rounds'])\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": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 4.c: Save the trained model to disk locally </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_filename = 'trained-model_nyctaxi.xgb'\n",
    "xgb_gpu_model.save_model(model_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 5: Predict & Score using vanilla XGBoost Predict </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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/stable/api.html?highlight=forestinference#cuml.ForestInference) to perform prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_y_test = y_infer.compute()\n",
    "wait(_y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": {},
   "source": [
    "#### Inference with the inplace predict method of dask XGBoost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 6: Predict & Score using FIL or Forest Inference Library </span>\n",
    "\n",
    "[Forest Inference Library (FIL)](https://docs.rapids.ai/api/cuml/stable/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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cuml import ForestInference\n",
    "from dask.distributed import get_worker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:#8735fb; font-size:18pt\"> Step 6.a: Predict using `compute` on a single worker in case the test dataset is small. </span>\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workers = client.has_what().keys()\n",
    "print(workers)\n",
    "n_workers = len(workers)\n",
    "n_partitions = n_workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unzipFile(zipname):\n",
    "    worker = get_worker()\n",
    "    import zipfile\n",
    "    import os\n",
    "    with zipfile.ZipFile(os.path.join(worker.local_directory, zipname)) as zf:\n",
    "        zf.extractall(worker.local_directory)\n",
    "\n",
    "def checkOrMakeLocalDir():\n",
    "    worker = get_worker()\n",
    "    import os\n",
    "    if not os.path.exists(worker.local_directory):\n",
    "        os.makedirs(worker.local_directory)\n",
    "    \n",
    "def workerModelInit(model_file):   \n",
    "    # this function will run in each worker and initialize the worker \n",
    "    import os\n",
    "    worker = get_worker()\n",
    "    worker.data[\"fil_model\"] = ForestInference.load(filename=os.path.join(worker.local_directory, model_file),model_type='xgboost')\n",
    "    \n",
    "def predict(input_df):\n",
    "    # this function will run in each worker and predict \n",
    "    worker = get_worker()\n",
    "    return worker.data[\"fil_model\"].predict(input_df)\n",
    "\n",
    "def persistModelonWorkers(client, zip_file_name, model_file_name):\n",
    "    import zipfile\n",
    "    zf = zipfile.ZipFile(zip_file_name, mode='w')\n",
    "    zf.write(f\"./{model_file_name}\")\n",
    "    zf.close()\n",
    "    # check to see if local directory present in workers\n",
    "    # if not present make it\n",
    "    fut = client.run(checkOrMakeLocalDir)\n",
    "    wait(fut)\n",
    "    # upload the zip file in workers\n",
    "    fut = client.upload_file(f\"./{zip_file_name}\")\n",
    "    wait(fut)\n",
    "    # unzip file in the workers\n",
    "    fut = client.run(unzipFile, zip_file_name)\n",
    "    wait(fut)\n",
    "    # load model using FIL in workers\n",
    "    fut = client.run(workerModelInit, model_file_name)\n",
    "    wait(fut)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Persist the local model in the remote dask workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "persistModelonWorkers(client, \"zipfile_write.zip\", \"trained-model_nyctaxi.xgb\")"
   ]
  },
  {
   "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(predict, meta=\"float\") # 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rows_csv = X_infer.iloc[:,0].shape[0].compute()\n",
    "print(f\"It took {toc-tic} seconds to predict on {rows_csv} rows using FIL distributedly on each worker\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": {},
   "source": [
    "## <span style=\"color:#8735fb; font-size:22pt\"> Step 7: Clean up </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}