{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3dc40b34",
   "metadata": {
    "tags": [
     "cloud/aws/ec2-multi",
     "library/cuml",
     "library/dask",
     "library/numpy",
     "library/dask-ml",
     "library/cudf",
     "workflow/randomforest",
     "tools/dask-cloudprovider",
     "data-format/csv",
     "data-storage/gcs"
    ]
   },
   "source": [
    "# Multi-node Multi-GPU Example on AWS using dask-cloudprovider\n",
    "\n",
    "_February, 2023_\n",
    "\n",
    "[Dask Cloud Provider](https://cloudprovider.dask.org/en/latest/) is a native cloud integration for dask. It helps manage Dask clusters on different cloud platforms. In this notebook, we will look at how we can use the package to set-up a AWS 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 through this notebook. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98cf9dae",
   "metadata": {},
   "source": [
    "## Create Your Cluster\n",
    "\n",
    "\n",
    "```{note}\n",
    "First follow the [full instructions](../../cloud/aws/ec2-multi) on launching a multi-node GPU cluster with Dask Cloud Provider.\n",
    "```\n",
    "\n",
    "Once you have a `cluster` object up and running head back here and continue.\n",
    "\n",
    "```python\n",
    "from dask_cloudprovider.aws import EC2Cluster\n",
    "\n",
    "cluster = EC2Cluster(...)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af8c63a",
   "metadata": {},
   "source": [
    "## Client Set Up\n",
    "\n",
    "Now we can create a [Dask Client](https://distributed.dask.org/en/latest/client.html) with the cluster we just defined. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e471a136",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f933be1c",
   "metadata": {},
   "source": [
    "````{dropdown} Optionally: We can wait for all workers to be up and running.\n",
    ":color: info\n",
    ":icon: info\n",
    "\n",
    "We do so by adding:\n",
    "\n",
    "```python\n",
    "# n_workers is the number of GPUs your cluster will have\n",
    "client.wait_for_workers(n_workers)  \n",
    "```\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d3e8c42",
   "metadata": {},
   "source": [
    "## Machine Learning Workflow\n",
    "\n",
    "Once workers become available, we can now run the rest of our workflow:\n",
    "\n",
    "- read and clean the data\n",
    "- add features\n",
    "- split into training and validation sets\n",
    "- fit a Random Forest model\n",
    "- predict on the validation set\n",
    "- compute RMSE\n",
    "\n",
    "Let's import the rest of our dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "796b203d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask_cudf\n",
    "import numpy as np\n",
    "from cuml.dask.common import utils as dask_utils\n",
    "from cuml.dask.ensemble import RandomForestRegressor\n",
    "from cuml.metrics import mean_squared_error\n",
    "from dask_ml.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46967a4d",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### 1. Read and Clean Data\n",
    "\n",
    "The data needs to be cleaned up before it can be used in a meaningful way. We verify the columns have appropriate datatypes to make it ready for computation using cuML."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909a9714",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create a list of all columns & dtypes the df must have for reading\n",
    "col_dtype = {\n",
    "    \"VendorID\": \"int32\",\n",
    "    \"tpep_pickup_datetime\": \"datetime64[ms]\",\n",
    "    \"tpep_dropoff_datetime\": \"datetime64[ms]\",\n",
    "    \"passenger_count\": \"int32\",\n",
    "    \"trip_distance\": \"float32\",\n",
    "    \"pickup_longitude\": \"float32\",\n",
    "    \"pickup_latitude\": \"float32\",\n",
    "    \"RatecodeID\": \"int32\",\n",
    "    \"store_and_fwd_flag\": \"int32\",\n",
    "    \"dropoff_longitude\": \"float32\",\n",
    "    \"dropoff_latitude\": \"float32\",\n",
    "    \"payment_type\": \"int32\",\n",
    "    \"fare_amount\": \"float32\",\n",
    "    \"extra\": \"float32\",\n",
    "    \"mta_tax\": \"float32\",\n",
    "    \"tip_amount\": \"float32\",\n",
    "    \"total_amount\": \"float32\",\n",
    "    \"tolls_amount\": \"float32\",\n",
    "    \"improvement_surcharge\": \"float32\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53342476",
   "metadata": {},
   "outputs": [],
   "source": [
    "taxi_df = dask_cudf.read_csv(\n",
    "    \"https://storage.googleapis.com/anaconda-public-data/nyc-taxi/csv/2016/yellow_tripdata_2016-02.csv\",\n",
    "    dtype=col_dtype,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ce9c1dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary of required columns and their datatypes\n",
    "must_haves = {\n",
    "    \"pickup_datetime\": \"datetime64[ms]\",\n",
    "    \"dropoff_datetime\": \"datetime64[ms]\",\n",
    "    \"passenger_count\": \"int32\",\n",
    "    \"trip_distance\": \"float32\",\n",
    "    \"pickup_longitude\": \"float32\",\n",
    "    \"pickup_latitude\": \"float32\",\n",
    "    \"rate_code\": \"int32\",\n",
    "    \"dropoff_longitude\": \"float32\",\n",
    "    \"dropoff_latitude\": \"float32\",\n",
    "    \"fare_amount\": \"float32\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4266d311",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean(ddf, must_haves):\n",
    "    # replace the extraneous spaces in column names and lower the font type\n",
    "    tmp = {col: col.strip().lower() for col in list(ddf.columns)}\n",
    "    ddf = ddf.rename(columns=tmp)\n",
    "\n",
    "    ddf = ddf.rename(\n",
    "        columns={\n",
    "            \"tpep_pickup_datetime\": \"pickup_datetime\",\n",
    "            \"tpep_dropoff_datetime\": \"dropoff_datetime\",\n",
    "            \"ratecodeid\": \"rate_code\",\n",
    "        }\n",
    "    )\n",
    "\n",
    "    ddf[\"pickup_datetime\"] = ddf[\"pickup_datetime\"].astype(\"datetime64[ms]\")\n",
    "    ddf[\"dropoff_datetime\"] = ddf[\"dropoff_datetime\"].astype(\"datetime64[ms]\")\n",
    "\n",
    "    for col in ddf.columns:\n",
    "        if col not in must_haves:\n",
    "            ddf = ddf.drop(columns=col)\n",
    "            continue\n",
    "        if ddf[col].dtype == \"object\":\n",
    "            # Fixing error: could not convert arg to str\n",
    "            ddf = ddf.drop(columns=col)\n",
    "        else:\n",
    "            # downcast from 64bit to 32bit types\n",
    "            # Tesla T4 are faster on 32bit ops\n",
    "            if \"int\" in str(ddf[col].dtype):\n",
    "                ddf[col] = ddf[col].astype(\"int32\")\n",
    "            if \"float\" in str(ddf[col].dtype):\n",
    "                ddf[col] = ddf[col].astype(\"float32\")\n",
    "            ddf[col] = ddf[col].fillna(-1)\n",
    "\n",
    "    return ddf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45814cc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "taxi_df = taxi_df.map_partitions(clean, must_haves, meta=must_haves)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c49c543a",
   "metadata": {},
   "source": [
    "#### 2. Add Features\n",
    "\n",
    "We'll add new features to the dataframe:\n",
    "\n",
    "1. We can split the datetime column to retrieve year, month, day, hour, day_of_week columns. Find the difference between pickup time and drop off time. \n",
    "2. Haversine Distance between the pick-up and drop-off coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "708fb4f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## add features\n",
    "\n",
    "taxi_df[\"hour\"] = taxi_df[\"pickup_datetime\"].dt.hour.astype(\"int32\")\n",
    "taxi_df[\"year\"] = taxi_df[\"pickup_datetime\"].dt.year.astype(\"int32\")\n",
    "taxi_df[\"month\"] = taxi_df[\"pickup_datetime\"].dt.month.astype(\"int32\")\n",
    "taxi_df[\"day\"] = taxi_df[\"pickup_datetime\"].dt.day.astype(\"int32\")\n",
    "taxi_df[\"day_of_week\"] = taxi_df[\"pickup_datetime\"].dt.weekday.astype(\"int32\")\n",
    "taxi_df[\"is_weekend\"] = (taxi_df[\"day_of_week\"] >= 5).astype(\"int32\")\n",
    "\n",
    "# calculate the time difference between dropoff and pickup.\n",
    "taxi_df[\"diff\"] = taxi_df[\"dropoff_datetime\"].astype(\"int32\") - taxi_df[\n",
    "    \"pickup_datetime\"\n",
    "].astype(\"int32\")\n",
    "taxi_df[\"diff\"] = (taxi_df[\"diff\"] / 1000).astype(\"int32\")\n",
    "\n",
    "taxi_df[\"pickup_latitude_r\"] = taxi_df[\"pickup_latitude\"] // 0.01 * 0.01\n",
    "taxi_df[\"pickup_longitude_r\"] = taxi_df[\"pickup_longitude\"] // 0.01 * 0.01\n",
    "taxi_df[\"dropoff_latitude_r\"] = taxi_df[\"dropoff_latitude\"] // 0.01 * 0.01\n",
    "taxi_df[\"dropoff_longitude_r\"] = taxi_df[\"dropoff_longitude\"] // 0.01 * 0.01\n",
    "\n",
    "taxi_df = taxi_df.drop(\"pickup_datetime\", axis=1)\n",
    "taxi_df = taxi_df.drop(\"dropoff_datetime\", axis=1)\n",
    "\n",
    "\n",
    "def haversine_dist(df):\n",
    "    import cuspatial\n",
    "\n",
    "    pickup = cuspatial.GeoSeries.from_points_xy(\n",
    "        df[[\"pickup_longitude\", \"pickup_latitude\"]].interleave_columns()\n",
    "    )\n",
    "    dropoff = cuspatial.GeoSeries.from_points_xy(\n",
    "        df[[\"dropoff_longitude\", \"dropoff_latitude\"]].interleave_columns()\n",
    "    )\n",
    "    df[\"h_distance\"] = cuspatial.haversine_distance(pickup, dropoff)\n",
    "    df[\"h_distance\"] = df[\"h_distance\"].astype(\"float32\")\n",
    "    return df\n",
    "\n",
    "\n",
    "taxi_df = taxi_df.map_partitions(haversine_dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac02e714",
   "metadata": {},
   "source": [
    "#### 3. Split Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81c737f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split into training and validation sets\n",
    "X, y = taxi_df.drop([\"fare_amount\"], axis=1).astype(\"float32\"), taxi_df[\n",
    "    \"fare_amount\"\n",
    "].astype(\"float32\")\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb21a06",
   "metadata": {},
   "outputs": [],
   "source": [
    "workers = client.has_what().keys()\n",
    "X_train, X_test, y_train, y_test = dask_utils.persist_across_workers(\n",
    "    client, [X_train, X_test, y_train, y_test], workers=workers\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "882b2c3e",
   "metadata": {},
   "source": [
    "#### 4. Create and fit a Random Forest Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ce41ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create cuml.dask RF regressor\n",
    "cu_dask_rf = RandomForestRegressor(ignore_empty_partitions=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "727d7d5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit RF model\n",
    "cu_dask_rf = cu_dask_rf.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd362d4",
   "metadata": {},
   "source": [
    "#### 5. Predict on validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72128d71",
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict on validation set\n",
    "y_pred = cu_dask_rf.predict(X_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba599603",
   "metadata": {},
   "source": [
    "#### 6. Compute RMSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1bff0d6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# compute RMSE\n",
    "score = mean_squared_error(y_pred.compute().to_numpy(), y_test.compute().to_numpy())\n",
    "print(\"Workflow Complete - RMSE: \", np.sqrt(score))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44f0d5d2",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Resource Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f484866",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up resources\n",
    "client.close()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d18d8e",
   "metadata": {},
   "source": [
    "#### Learn More\n",
    "\n",
    "- [Dask Cloud Provider](https://cloudprovider.dask.org/en/latest/)"
   ]
  }
 ],
 "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.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "f7a54d993f849a0f97fda357a1a3bac7e25a43aff77e618e8d69a4ad36661dba"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}