Modeling 1 - Train an XGBoost Model with Dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0652f264-b2c3-43b2-9310-7c00e50e7cdc",
   "metadata": {},
   "source": [
    "# Train an XGBoost Model with Dask\n",
    "\n",
    "We start by training a single XGBoost model with Dask using the [`xgboost.dask`](https://xgboost.readthedocs.io/en/stable/tutorials/dask.html) module built into XGBoost. In this notebook we ...\n",
    "\n",
    "-  Load the data\n",
    "-  Perform basic feature engineering (date type optimization, categorization)\n",
    "-  Train a single model with XGBoost, using custom cross-validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b929987-5089-4c13-be35-c1812d65fbec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from collections.abc import Iterator\n",
    "from datetime import datetime\n",
    "\n",
    "import coiled\n",
    "import dask.array as da\n",
    "import dask.dataframe as dd\n",
    "import distributed\n",
    "import numpy as np\n",
    "import xgboost\n",
    "from dask_ml.metrics import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7179b003-e075-4be5-bde1-4ce90b7d726d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Location of feature table\n",
    "FILEPATH = \"s3://coiled-datasets/prefect-dask/nyc-uber-lyft/feature_table.parquet\"\n",
    "\n",
    "# Number of folds in each trial. This also determines the train/test split\n",
    "# (e.g. N_FOLDS=5 -> train=4/5 of the total data, test=1/5)\n",
    "N_FOLDS = 5\n",
    "\n",
    "# Dask worker instance type and number\n",
    "WORKER_INSTANCE_TYPE = \"r6i.large\"\n",
    "N_WORKERS = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e349a0e8-842c-4c27-a9dd-1a4686e6a339",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cluster = coiled.Cluster(\n",
    "    worker_vm_types=[WORKER_INSTANCE_TYPE],\n",
    "    scheduler_vm_types=[\"m6i.large\"],\n",
    "    package_sync=True,  # align remote packages to local ones\n",
    "    n_workers=N_WORKERS,\n",
    "    backend_options={\n",
    "        \"region\": \"us-east-2\",\n",
    "        \"multizone\": True,\n",
    "    },\n",
    "    scheduler_options={\"idle_timeout\": \"15 minutes\"},\n",
    ")\n",
    "client = distributed.Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5fa0dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load feature table generated by Feature Engineering.ipynb\n",
    "ddf = dd.read_parquet(FILEPATH)\n",
    "\n",
    "# Reduce dataset size. Uncomment to speed up the exercise.\n",
    "# ddf = ddf.partitions[:20]\n",
    "\n",
    "# Under the hood, XGBoost converts floats to `float32`.\n",
    "# Let's do it only once here.\n",
    "float_cols = ddf.select_dtypes(include=\"float\").columns.tolist()\n",
    "ddf = ddf.astype({c: np.float32 for c in float_cols})\n",
    "\n",
    "# We need the categories to be known\n",
    "categorical_vars = ddf.select_dtypes(include=\"category\").columns.tolist()\n",
    "\n",
    "# categorize() reads the whole input and then discards it.\n",
    "# Let's read from disk only once.\n",
    "ddf = ddf.persist()\n",
    "ddf = ddf.categorize(columns=categorical_vars)\n",
    "\n",
    "# We will need to access this multiple times. Let's persist it.\n",
    "ddf = ddf.persist()\n",
    "\n",
    "ddf.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dd461c3-e551-420a-a3e1-1ca2515bb613",
   "metadata": {},
   "source": [
    "### Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccf70461-6ab3-429b-a0a1-29f6b162894e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we subset data for cross-validation\n",
    "def make_cv_splits(\n",
    "    n_folds: int = N_FOLDS,\n",
    ") -> Iterator[tuple[dd.DataFrame, dd.DataFrame]]:\n",
    "    frac = [1 / n_folds] * n_folds\n",
    "    splits = ddf.random_split(frac, shuffle=True)\n",
    "    for i in range(n_folds):\n",
    "        train = [splits[j] for j in range(n_folds) if j != i]\n",
    "        test = splits[i]\n",
    "        yield dd.concat(train), test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132568f4-ac1b-43ff-85fd-e5a5324a38a0",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "start = datetime.now()\n",
    "scores = []\n",
    "\n",
    "for i, (train, test) in enumerate(make_cv_splits()):\n",
    "    print(f\"Training/Test split #{i}\")\n",
    "    y_train = train[\"trip_time\"]\n",
    "    X_train = train.drop(columns=[\"trip_time\"])\n",
    "    y_test = test[\"trip_time\"]\n",
    "    X_test = test.drop(columns=[\"trip_time\"])\n",
    "\n",
    "    print(\"Building DMatrix...\")\n",
    "    d_train = xgboost.dask.DaskDMatrix(\n",
    "        None, X_train, y_train, enable_categorical=True\n",
    "    )\n",
    "\n",
    "    print(\"Training model...\")\n",
    "    model = xgboost.dask.train(\n",
    "        None,\n",
    "        {\"tree_method\": \"hist\"},\n",
    "        d_train,\n",
    "        num_boost_round=4,\n",
    "        evals=[(d_train, \"train\")],\n",
    "    )\n",
    "\n",
    "    print(\"Running model on test data...\")\n",
    "    predictions = xgboost.dask.predict(None, model, X_test)\n",
    "\n",
    "    print(\"Measuring accuracy of model vs. ground truth...\")\n",
    "    score = mean_squared_error(\n",
    "        y_test.to_dask_array(),\n",
    "        predictions.to_dask_array(),\n",
    "        squared=False,\n",
    "        compute=False,\n",
    "    )\n",
    "    # Compute predictions and mean squared error for this iteration\n",
    "    # while we start the next one\n",
    "    scores.append(score.reshape(1).persist())\n",
    "    print(\"-\" * 80)\n",
    "\n",
    "scores = da.concatenate(scores).compute()\n",
    "print(f\"RSME={scores.mean()} +/- {scores.std()}\")\n",
    "print(f\"Total time:  {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af1248ad-7147-473c-801a-ca0a2f2ed282",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6debe738-0e8d-4919-9efe-7e11cb6e9264",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66bf4890-b1a8-47c0-aeda-b24e55b5e7e8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}