<jupyter_start><jupyter_text>Multi-Node Multi-GPU XGBoost example on Azure using dask-cloudprovider [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. 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.---**NOTE:**  In this notebook, we will explore two possible ways to use `dask-cloudprovider` to run our workloads on Azure VM clusters: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.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. You can either use Option 1 or use Option 2.---   Step 0: Set up Azure credentials and CLI  Before running the notebook, run the following commands in the terminal to setup Azure CLI```curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bashaz login```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`.   Step 1: Import necessary packages<jupyter_code># # Uncomment the following and install some libraries at the beginning. 
# # If adlfs is not present, install adlfs to read from Azure data lake. 
# ! pip install adlfs
# ! pip install "dask-cloudprovider[azure]" --upgrade
from dask.distributed import Client, wait, get_worker
from dask_cloudprovider.azure import AzureVMCluster
import dask_cudf
from dask_ml.model_selection import train_test_split
from cuml.dask.common import utils as dask_utils
from cuml.metrics import mean_squared_error
from cuml import ForestInference
import cudf
import xgboost as xgb
from datetime import datetime
from dateutil import parser
import numpy as np
from timeit import default_timer as timer
import dask
import json<jupyter_output><empty_output><jupyter_text>Step 2: Set up the Azure VM ClusterWe 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.htmlresource-groups). Note that you can also set it up using the Azure portal.Once you have set it up, you can now plug in the names of the entities you have created in the cell below. 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.<jupyter_code># location = <your chosen location where the resource group and vnet exists>
# resource_group = <your resource group>
# vnet = <vnet in your resource group>
# security_group = <security group in your resource group>
# vm_size = "Standard_NC12s_v3" # or choose a different GPU enabled VM type

docker_image = "rapidsai/rapidsai:cuda11.2-runtime-ubuntu18.04-py3.8"
docker_args = '--shm-size=256m'
worker_class = "dask_cuda.CUDAWorker"
worker_options = {'rmm-managed-memory':True}<jupyter_output><empty_output><jupyter_text>---  Option 1: Use an Azure Marketplace VM image 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. 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).**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.  Set up Marketplace VM information and clear default dask config.<jupyter_code>dask.config.set({"logging.distributed": "info",
                "cloudprovider.azure.azurevm.marketplace_plan":{
                        "publisher": "nvidia",
                        "name": "ngc-base-version-21-02-2",
                        "product": "ngc_azure_17_11",
                        "version": "21.02.2"
                }})
vm_image = ""
config = dask.config.get("cloudprovider.azure.azurevm", {})
config<jupyter_output><empty_output><jupyter_text>If necessary, you must uncomment and accept the Azure Marketplace image terms so that the image can be used to create VMs<jupyter_code># ! az vm image terms accept --urn "nvidia:ngc_azure_17_11:ngc-base-version-21-02-2:21.02.2" --verbose<jupyter_output><empty_output><jupyter_text>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. ---   Option 2: Set up an Azure Customized VM  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.)  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. 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.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. 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.  Now, to create a customized VM image, follow steps *a.* to *f.*  a. Install `packer`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.   b. Authenticate `packer` with AzureThere 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.   c. Generate the cloud init script for customizing the VM image`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.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.   d. Write packer configuration to a configuration fileWe 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. 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.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).**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`.<jupyter_code>packer_config = {
  "builders": [{
    "type": "azure-arm", 
    "use_azure_cli_auth": True,
    "managed_image_resource_group_name": resource_group,
    "managed_image_name": <the name of the customized VM image>,
    "custom_data_file": "./configs/cloud_init.yaml.j2",
    "os_type": "Linux",
    "image_publisher": "Canonical", 
    "image_offer": "UbuntuServer",
    "image_sku": "18.04-LTS",
    "azure_tags": {
        "dept": "RAPIDS-CSP",
        "task": "RAPIDS Custom Image deployment"
    },
    "build_resource_group_name": resource_group,
    "vm_size": vm_size
  }],
  "provisioners": [{
    "inline": [
      "echo 'Waiting for cloud-init'; while [ ! -f /var/lib/cloud/instance/boot-finished ]; do sleep 1; done; echo 'Done'",
    ],
    "type": "shell"
  }]
}

with open("packer_config.json", "w") as fh:
    fh.write(json.dumps(packer_config))<jupyter_output><empty_output><jupyter_text>e. Run `packer` build and create the image<jupyter_code># # Uncomment the following line and run to create the custom image
# ! packer build packer_config.json<jupyter_output><empty_output><jupyter_text>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.   While packer is building the image, you will see an output similar to what is shown below. ---```bashazure-arm: output will be in this color.==> azure-arm: Running builder ...==> azure-arm: Getting tokens using Azure CLI==> azure-arm: Getting tokens using Azure CLI    azure-arm: Creating Azure Resource Manager (ARM) client ...==> azure-arm: Using existing resource group ...==> azure-arm:  -> ResourceGroupName : ==> azure-arm:  -> Location          : ==> azure-arm: Validating deployment template ...==> azure-arm:  -> ResourceGroupName : ==> azure-arm:  -> DeploymentName    : 'pkrdp04rrahxkg9'==> azure-arm: Deploying deployment template ...==> azure-arm:  -> ResourceGroupName : ==> azure-arm:  -> DeploymentName    : 'pkrdp04rrahxkg9'==> azure-arm:==> azure-arm: Getting the VM's IP address ...==> azure-arm:  -> ResourceGroupName   : ==> azure-arm:  -> PublicIPAddressName : 'pkrip04rrahxkg9'==> azure-arm:  -> NicName             : 'pkrni04rrahxkg9'==> azure-arm:  -> Network Connection  : 'PublicEndpoint'==> azure-arm:  -> IP Address          : '192.168.3.11'==> azure-arm: Waiting for SSH to become available...==> azure-arm: Connected to SSH!==> azure-arm: Provisioning with shell script: /tmp/packer-shell614221056    azure-arm: Waiting for cloud-init    azure-arm: Done==> azure-arm: Querying the machine's properties ...==> azure-arm:  -> ResourceGroupName : ==> azure-arm:  -> ComputeName       : 'pkrvm04rrahxkg9'==> azure-arm:  -> Managed OS Disk   : '/subscriptions//resourceGroups//providers/Microsoft.Compute/disks/pkros04rrahxkg9'==> azure-arm: Querying the machine's additional disks properties ...==> azure-arm:  -> ResourceGroupName : ==> azure-arm:  -> ComputeName       : 'pkrvm04rrahxkg9'==> azure-arm: Powering off machine ...==> azure-arm:  -> ResourceGroupName : ==> azure-arm:  -> ComputeName       : 'pkrvm04rrahxkg9'==> azure-arm: Capturing image ...==> azure-arm:  -> Compute ResourceGroupName : ==> azure-arm:  -> Compute Name              : 'pkrvm04rrahxkg9'==> azure-arm:  -> Compute Location          : ==> azure-arm:  -> Image ResourceGroupName   : ==> azure-arm:  -> Image Name                : ==> azure-arm:  -> Image Location            : ==> azure-arm: ==> azure-arm: Deleting individual resources ...==> azure-arm: Adding to deletion queue -> Microsoft.Compute/virtualMachines : 'pkrvm04rrahxkg9'==> azure-arm: Adding to deletion queue -> Microsoft.Network/networkInterfaces : 'pkrni04rrahxkg9'==> azure-arm: Adding to deletion queue -> Microsoft.Network/publicIPAddresses : 'pkrip04rrahxkg9'==> azure-arm: Adding to deletion queue -> Microsoft.Network/virtualNetworks : 'pkrvn04rrahxkg9'==> azure-arm: Attempting deletion -> Microsoft.Network/networkInterfaces : 'pkrni04rrahxkg9'==> azure-arm: Waiting for deletion of all resources...==> azure-arm: Attempting deletion -> Microsoft.Network/publicIPAddresses : 'pkrip04rrahxkg9'==> azure-arm: Attempting deletion -> Microsoft.Compute/virtualMachines : 'pkrvm04rrahxkg9'==> azure-arm: Attempting deletion -> Microsoft.Network/virtualNetworks : 'pkrvn04rrahxkg9'    .    .    .    .    .    .    .    .==> azure-arm:  Deleting -> Microsoft.Compute/disks : '/subscriptions//resourceGroups//providers/Microsoft.Compute/disks/pkros04rrahxkg9'==> azure-arm: Removing the created Deployment object: 'pkrdp04rrahxkg9'==> azure-arm: ==> azure-arm: The resource group was not created by Packer, not deleting ...Build 'azure-arm' finished after 16 minutes 22 seconds.==> Wait completed after 16 minutes 22 seconds==> Builds finished. The artifacts of successful builds are:--> azure-arm: Azure.ResourceManagement.VMImage:OSType: LinuxManagedImageResourceGroupName: ManagedImageName: ManagedImageId: /subscriptions//resourceGroups//providers/Microsoft.Compute/images/ManagedImageLocation: ```---     When `packer` finishes, at the bottom of the output, you will see something similar to the following:```ManagedImageResourceGroupName: ManagedImageName: ManagedImageId: /subscriptions//resourceGroups//providers/Microsoft.Compute/images/ManagedImageLocation: ``` Make note of the `ManagedImageId`. This is the resource id of the custom image we will use. As shown above the `ManagedImageId` will look something like : `/subscriptions/12345/resourceGroups/myown-rg/providers/Microsoft.Compute/images/myCustomImage`  f. Set up customized VM information and clear default dask config  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.<jupyter_code>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. 
dask.config.set({"cloudprovider.azure.azurevm.vm_image":{}})
config = dask.config.get("cloudprovider.azure.azurevm", {})
print(config)
vm_image = {"id": ManagedImageId}
print(vm_image)<jupyter_output><empty_output><jupyter_text>Step 2.1: Start the VM Cluster in Azure Here, if you have used Option 1, i.e., the NVIDIA VM image, pass an empty string for `vm_image` information. For Option 2, pass the `vm_image` information that you got from the output of `packer` run as a parameter to `AzureVMCluster`. 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. **NOTE:** The rest of the notebook should be the same irrespective of whether you chose Option 1 or Option 2.**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$.<jupyter_code>%%time

cluster = AzureVMCluster(
    location=location,
    resource_group=resource_group,
    vnet=vnet,
    security_group=security_group,
    vm_image=vm_image,
    vm_size=vm_size,
    docker_image=docker_image,
    worker_class=worker_class,
    n_workers=2,
    security=True,
    docker_args=docker_args,
    worker_options=worker_options,
    debug=False,
    bootstrap=False, # This is to prevent the cloud init jinja2 script from running in the custom VM.
)
client = Client(cluster) 
client
def scale_workers(client, n_workers, n_gpus_per_worker, timeout=300):
    import time
    client.cluster.scale(n_workers)
    m = len(client.has_what().keys())    
    start = end = time.perf_counter_ns()
    while ((m != n_workers*n_gpus_per_worker) and (((end - start) / 1e9) < timeout) ):
        time.sleep(5)
        m = len(client.has_what().keys())
        
        end = time.perf_counter_ns()
        
    if (((end - start) / 1e9) >= timeout):
        raise RuntimeError(f"Failed to rescale cluster in {timeout} sec."
              "Try increasing timeout for very large containers, and verify available compute resources.")
# # Uncomment if you only have the scheduler with n_workers=0 and want to scale the workers separately.
# %%time
# scale_workers(client, 2, 2, timeout=600)<jupyter_output><empty_output><jupyter_text>Wait till all the workers are up. This will wait for `n_workers` number of VMs to be up.<jupyter_code>%%time
client.wait_for_workers(2)<jupyter_output><empty_output><jupyter_text>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.<jupyter_code>def pretty_print(scheduler_dict):
    print(f"All workers for scheduler id: {scheduler_dict['id']}, address: {scheduler_dict['address']}")
    for worker in scheduler_dict['workers']:
        print(f"Worker: {worker} , gpu_machines: {scheduler_dict['workers'][worker]['gpu']}")

pretty_print(client.scheduler_info()) # will show some information of the GPUs of the workers<jupyter_output><empty_output><jupyter_text>Step 3: Data Setup, Cleanup and Enhancement    Step 3.a: Set up the workers for reading parquet files from Azure Data Lake endpoints  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.<jupyter_code>def installAdlfs():
    import subprocess
    subprocess.run(["pip", "install", "adlfs"])
    return "done"

results=client.run(installAdlfs)<jupyter_output><empty_output><jupyter_text>Step 3.b: Data Cleanup, Enhancement and Persisting Scripts  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. We also add some new features to our dataframe via some custom functions, namely: 1. Haversine distance: This is used for calculating the total trip distance.2. Day of the week: This can be useful information for determining the fare cost.`add_features` function combines the two to produce a new dataframe that has the added features.**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.   Adding features functions<jupyter_code>import math
from math import cos, sin, asin, sqrt, pi

def haversine_distance_kernel(pickup_latitude_r, pickup_longitude_r, dropoff_latitude_r, dropoff_longitude_r, h_distance, radius):
    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,)):
        x_1 = pi/180 * x_1
        y_1 = pi/180 * y_1
        x_2 = pi/180 * x_2
        y_2 = pi/180 * y_2
        
        dlon = y_2 - y_1
        dlat = x_2 - x_1
        a = sin(dlat/2)**2 + cos(x_1) * cos(x_2) * sin(dlon/2)**2
        
        c = 2 * asin(sqrt(a)) 
        # radius = 6371 # Radius of earth in kilometers # currently passed as input arguments
        
        h_distance[i] = c * radius

def day_of_the_week_kernel(day, month, year, day_of_week):
    for i, (d_1, m_1, y_1) in enumerate(zip(day, month, year)):
        if month[i] <3:
            shift = month[i]
        else:
            shift = 0
        Y = year[i] - (month[i] < 3)
        y = Y - 2000
        c = 20
        d = day[i]
        m = month[i] + shift + 1
        day_of_week[i] = (d + math.floor(m*2.6) + y + (y//4) + (c//4) -2*c)%7
        
def add_features(df):
    df['hour'] = df['tpepPickupDateTime'].dt.hour
    df['year'] = df['tpepPickupDateTime'].dt.year
    df['month'] = df['tpepPickupDateTime'].dt.month
    df['day'] = df['tpepPickupDateTime'].dt.day
    df['diff'] = (df['tpepPickupDateTime'] - df['tpepPickupDateTime']).dt.seconds #convert difference between pickup and dropoff into seconds
    
    df['pickup_latitude_r'] = df['startLat']//.01*.01
    df['pickup_longitude_r'] = df['startLon']//.01*.01
    df['dropoff_latitude_r'] = df['endLat']//.01*.01
    df['dropoff_longitude_r'] = df['endLon']//.01*.01
    
    df = df.drop('tpepDropoffDateTime', axis=1)
    df = df.drop('tpepPickupDateTime', axis =1)
    
    
    df = df.apply_rows(haversine_distance_kernel,
                   incols=['pickup_latitude_r', 'pickup_longitude_r', 'dropoff_latitude_r', 'dropoff_longitude_r'],
                   outcols=dict(h_distance=np.float32),
                   kwargs=dict(radius=6371))
    
    
    df = df.apply_rows(day_of_the_week_kernel,
                      incols=['day', 'month', 'year'],
                      outcols=dict(day_of_week=np.float32),
                      kwargs=dict())
    
    
    df['is_weekend'] = (df['day_of_week']<2)
    return df<jupyter_output><empty_output><jupyter_text>Functions for cleaning and persisting the data in the workers.<jupyter_code>def persist_train_infer_split(client, df, response_dtype, response_id, infer_frac=1.0, random_state=42, shuffle=True):
    workers = client.has_what().keys()
    X, y = df.drop([response_id], axis=1), df[response_id].astype('float32')
    infer_frac = max(0, min(infer_frac, 1.0))
    X_train, X_infer, y_train, y_infer = train_test_split(X, y, shuffle=True, random_state=random_state, test_size=infer_frac)
    
    with dask.annotate(workers=set(workers)):
        X_train, y_train = client.persist(
            collections=[X_train, y_train]) 
    
    if (infer_frac != 1.0):
        with dask.annotate(workers=set(workers)):
            X_infer, y_infer = client.persist(
                collections=[X_infer, y_infer])

        wait([X_train, y_train, X_infer, y_infer])
    else:
        X_infer = X_train
        y_infer = y_train

        wait([X_train, y_train])
    
    return X_train, y_train, X_infer, y_infer


def clean(df_part, must_haves):
    """
    This function performs the various clean up tasks for the data
    and returns the cleaned dataframe.
    """    
    # iterate through columns in this df partition
    for col in df_part.columns:
        # drop anything not in our expected list
        if col not in must_haves:
            df_part = df_part.drop(col, axis=1)
            continue

        # fixes datetime error found by <NAME> and fixed by <NAME>
        if df_part[col].dtype == 'object' and col in ['tpepPickupDateTime', 'tpepDropoffDateTime']:
            df_part[col] = df_part[col].astype('datetime64[ms]')
            continue

        # if column was read as a string, recast as float
        if df_part[col].dtype == 'object':
            df_part[col] = df_part[col].str.fillna('-1')
            df_part[col] = df_part[col].astype('float32')
        else:
            # downcast from 64bit to 32bit types
            # Tesla T4 are faster on 32bit ops
            if 'int' in str(df_part[col].dtype):
                df_part[col] = df_part[col].astype('int32')
            if 'float' in str(df_part[col].dtype):
                df_part[col] = df_part[col].astype('float32')
            df_part[col] = df_part[col].fillna(-1)
            
    return df_part

    
def taxi_data_loader(client, adlsaccount, adlspath, response_dtype=np.float32, infer_frac=1.0, random_state=0):
    #create a list of columns & dtypes the df must have
    must_haves = {
     'tpepPickupDateTime': 'datetime64[ms]',
     'tpepDropoffDateTime': 'datetime64[ms]',
     'passengerCount': 'int32',
     'tripDistance': 'float32',
     'startLon': 'float32',
     'startLat': 'float32',
     'rateCodeId': 'int32',
     'endLon': 'float32',
     'endLat': 'float32',
     'fareAmount': 'float32'
    }

    workers = client.has_what().keys()
    response_id = 'fareAmount'
    storage_options = {'account_name': adlsaccount}
    taxi_data = dask_cudf.read_parquet(adlspath, storage_options=storage_options, chunksize=25e6, npartitions=len(workers))
    taxi_data = clean(taxi_data, must_haves)
    taxi_data = taxi_data.map_partitions(add_features)
    # Drop NaN values and convert to float32
    taxi_data = taxi_data.dropna()
    fields = ['passengerCount', 'tripDistance', 'startLon', 'startLat', 'rateCodeId',
                 'endLon', 'endLat', 'fareAmount', 'diff', 'h_distance', 'day_of_week', 'is_weekend']
    taxi_data = taxi_data.astype("float32")
    taxi_data = taxi_data[fields]
    taxi_data = taxi_data.reset_index()
    
    return persist_train_infer_split(client, taxi_data, response_dtype, response_id, infer_frac, random_state)<jupyter_output><empty_output><jupyter_text>Step 3.c: Get the split data and persist across workers  We will make use of the data from November and December 2014 for the purposes of the demo.<jupyter_code>tic = timer()
X_train, y_train, X_infer, y_infer = taxi_data_loader(client, 
                                                      adlsaccount="azureopendatastorage", 
                                                      adlspath="az://nyctlc/yellow/puYear=2014/puMonth=1*/*.parquet",
                                                      infer_frac=0.1, random_state=42)
toc = timer()
print(f"Wall clock time taken for ETL and persisting : {toc-tic} s")
X_train.shape[0].compute()<jupyter_output><empty_output><jupyter_text>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.<jupyter_code>X_train.head()<jupyter_output><empty_output><jupyter_text>Step 4: Train a XGBoost Model We are now ready to train a XGBoost model on the data and then predict the fare for each trip.   Step 4.a: Set training Parameters 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.Refer to the notebook [HPO-RAPIDS](./HPO-RAPIDS.ipynb) in this repository for how to perform HPO on Azure.<jupyter_code>params = {
    'learning_rate': 0.15,
    'max_depth': 8,
    'objective': 'reg:squarederror',
    'subsample': 0.7,
    'colsample_bytree': 0.7,
    'min_child_weight': 1,
    'gamma': 1,
    'silent': True,
    'verbose_eval': True,
    'booster' : 'gbtree', # 'gblinear' not implemented in dask
    'debug_synchronize': True,
    'eval_metric': 'rmse',
    'tree_method':'gpu_hist',
    'num_boost_rounds': 100,
}<jupyter_output><empty_output><jupyter_text>Step 4.b: Train XGBoost Model  Since the data is already persisted in the dask workers in the Kubernetes Cluster, the next steps should not take a lot of time.<jupyter_code>data_train = xgb.dask.DaskDMatrix(client, X_train, y_train)
tic = timer()
xgboost_output = xgb.dask.train(client, params,data_train, 
                                    num_boost_round=params['num_boost_rounds'])
xgb_gpu_model = xgboost_output['booster']
toc = timer()
print(f"Wall clock time taken for this cell : {toc-tic} s")<jupyter_output><empty_output><jupyter_text>Step 4.c: Save the trained model to disk locally<jupyter_code>model_filename = 'trained-model_nyctaxi.xgb'
xgb_gpu_model.save_model(model_filename)<jupyter_output><empty_output><jupyter_text>Step 5: Predict & Score using vanilla XGBoost Predict  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=forestinferencecuml.ForestInference) to perform prediction.<jupyter_code>_y_test = y_infer.compute()
wait(_y_test)
d_test = xgb.dask.DaskDMatrix(client, X_infer)
tic = timer()
y_pred = xgb.dask.predict(client, xgb_gpu_model, d_test)
y_pred= y_pred.compute()
wait(y_pred)
toc = timer()
print(f"Wall clock time taken for xgb.dask.predict : {toc-tic} s")<jupyter_output><empty_output><jupyter_text>Inference with the inplace predict method of dask XGBoost<jupyter_code>tic = timer()
y_pred = xgb.dask.inplace_predict(client, xgb_gpu_model, X_infer)
y_pred = y_pred.compute()
wait(y_pred)
toc = timer()
print(f"Wall clock time taken for inplace inference : {toc-tic} s")
tic = timer()
print("Calculating MSE")
score = mean_squared_error(y_pred, _y_test)
print("Workflow Complete - RMSE: ", np.sqrt(score))
toc = timer()
print(f"Wall clock time taken for this cell : {toc-tic} s")<jupyter_output><empty_output><jupyter_text>Step 6: Predict & Score using FIL or Forest Inference Library [Forest Inference Library (FIL)](https://docs.rapids.ai/api/cuml/stable/api.html?highlight=forestinferencecuml.ForestInference) provides GPU accelerated inference capabilities for tree models. We will import the FIL functionality from [cuML](https://github.com/rapidsai/cuml) library.It accepts a **trained** tree model in a treelite format (currently LightGBM, XGBoost and SKLearn GBDT and random forest modelsare 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.<jupyter_code>from cuml import ForestInference
from dask.distributed import get_worker<jupyter_output><empty_output><jupyter_text>Step 6.a: Predict using `compute` on a single worker in case the test dataset is small. 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.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.<jupyter_code>workers = client.has_what().keys()
print(workers)
n_workers = len(workers)
n_partitions = n_workers
def unzipFile(zipname):
    worker = get_worker()
    import zipfile
    import os
    with zipfile.ZipFile(os.path.join(worker.local_directory, zipname)) as zf:
        zf.extractall(worker.local_directory)

def checkOrMakeLocalDir():
    worker = get_worker()
    import os
    if not os.path.exists(worker.local_directory):
        os.makedirs(worker.local_directory)
    
def workerModelInit(model_file):   
    # this function will run in each worker and initialize the worker 
    import os
    worker = get_worker()
    worker.data["fil_model"] = ForestInference.load(filename=os.path.join(worker.local_directory, model_file),model_type='xgboost')
    
def predict(input_df):
    # this function will run in each worker and predict 
    worker = get_worker()
    return worker.data["fil_model"].predict(input_df)

def persistModelonWorkers(client, zip_file_name, model_file_name):
    import zipfile
    zf = zipfile.ZipFile(zip_file_name, mode='w')
    zf.write(f"./{model_file_name}")
    zf.close()
    # check to see if local directory present in workers
    # if not present make it
    fut = client.run(checkOrMakeLocalDir)
    wait(fut)
    # upload the zip file in workers
    fut = client.upload_file(f"./{zip_file_name}")
    wait(fut)
    # unzip file in the workers
    fut = client.run(unzipFile, zip_file_name)
    wait(fut)
    # load model using FIL in workers
    fut = client.run(workerModelInit, model_file_name)
    wait(fut)<jupyter_output><empty_output><jupyter_text>Persist the local model in the remote dask workers<jupyter_code>%%time
persistModelonWorkers(client, "zipfile_write.zip", "trained-model_nyctaxi.xgb")<jupyter_output><empty_output><jupyter_text>Inference with distributed predict with FIL<jupyter_code>tic = timer()
predictions = X_infer.map_partitions(predict, meta="float") # this is like MPI reduce
y_pred = predictions.compute()
wait(y_pred)
toc = timer()
print(f"Wall clock time taken for this cell : {toc-tic} s")
rows_csv = X_infer.iloc[:,0].shape[0].compute()
print(f"It took {toc-tic} seconds to predict on {rows_csv} rows using FIL distributedly on each worker")
tic = timer()
score = mean_squared_error(y_pred, _y_test)
toc = timer()
print("Final - RMSE: ", np.sqrt(score))<jupyter_output><empty_output><jupyter_text>Step 7: Clean up<jupyter_code>client.close()
cluster.close()<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup