Modeling 3 - Parallel HPO of XGBoost with Optuna and Dask (multi cluster).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0652f264-b2c3-43b2-9310-7c00e50e7cdc",
   "metadata": {},
   "source": [
    "# Parallel HPO with XGBoost/Dask/Optuna with multiple clusters\n",
    "\n",
    "Training one model can be slow.  Parallelism can help!\n",
    "\n",
    "Training many models during hyper-parameter optimization can be even slower.  Even more parallelism can help!\n",
    "\n",
    "In this example we extend our previous notebook to run many model trainings in parallel, each model running in a separate Dask cluster.  This allows us to accelerate our search for a good model by using more hardware.  It's important to note here that there are two levels of parallelism:\n",
    "\n",
    "1.  Each model runs in parallel using Dask\n",
    "2.  Trigger many such runs in different threads locally\n",
    "\n",
    "Each local thread does very little work, it just asks Dask to manage a large remote job.\n",
    "\n",
    "![high level diagram](Modeling_3.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cb987d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import threading\n",
    "from collections.abc import Iterator\n",
    "from datetime import datetime\n",
    "\n",
    "import coiled\n",
    "import dask\n",
    "import dask.array as da\n",
    "import dask.dataframe as dd\n",
    "import distributed\n",
    "import joblib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import optuna\n",
    "import xgboost\n",
    "from dask_ml.metrics import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d546fe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is only for xgboost 1.7.1.\n",
    "# Not necessary with xgboost 1.7.3.\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=RuntimeWarning)"
   ]
  },
  {
   "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 parallel optuna jobs to run\n",
    "N_JOBS = 10\n",
    "# Total number of converging trials to run across the various jobs\n",
    "N_TRIALS = 50\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 (per cluster)\n",
    "# Total number of EC2 instances spun up = N_JOBS * N_WORKERS\n",
    "WORKER_INSTANCE_TYPE = \"r6i.large\"\n",
    "N_WORKERS = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a31b9e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "clusters: dict[int, tuple[distributed.Client, dd.DataFrame]] = {}\n",
    "\n",
    "\n",
    "def get_ddf() -> tuple[distributed.Client, dd.DataFrame]:\n",
    "    thread_id = threading.get_ident()\n",
    "    try:\n",
    "        return clusters[thread_id]\n",
    "    except KeyError:\n",
    "        pass\n",
    "\n",
    "    cluster = coiled.Cluster(\n",
    "        name=f\"xgb-nyc-taxi-gbh-{thread_id}\",\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",
    "\n",
    "    client = distributed.Client(cluster, set_as_default=False)\n",
    "    print(\"Started cluster at\", client.dashboard_link)\n",
    "\n",
    "    with client.as_current():\n",
    "        # 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",
    "        # FIXME https://github.com/dask/dask/issues/9901\n",
    "        ddf = ddf.categorize(columns=categorical_vars, scheduler=client)\n",
    "\n",
    "        # We will need to access this multiple times. Let's persist it.\n",
    "        ddf = ddf.persist()\n",
    "\n",
    "        clusters[thread_id] = client, ddf\n",
    "        return client, ddf"
   ]
  },
  {
   "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",
    "    ddf: dd.DataFrame, 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": [
    "def train_model(study_params: dict[str, float]) -> float:\n",
    "    scores = []\n",
    "    client, ddf = get_ddf()\n",
    "\n",
    "    with client.as_current():\n",
    "        for train, test in make_cv_splits(ddf):\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",
    "            d_train = xgboost.dask.DaskDMatrix(\n",
    "                client, X_train, y_train, enable_categorical=True\n",
    "            )\n",
    "            model = xgboost.dask.train(\n",
    "                client,\n",
    "                {\"tree_method\": \"hist\", **study_params},\n",
    "                d_train,\n",
    "                num_boost_round=4,\n",
    "                evals=[(d_train, \"train\")],\n",
    "            )\n",
    "            predictions = xgboost.dask.predict(None, model, X_test)\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",
    "            del d_train, model, predictions, score\n",
    "\n",
    "        scores = da.concatenate(scores).compute()\n",
    "        return scores.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0250d23a-4d6c-4758-9200-f47a0a614bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    params = {\n",
    "        \"n_estimators\": trial.suggest_int(\"n_estimators\", 75, 125),\n",
    "        \"learning_rate\": trial.suggest_float(\"learning_rate\", 0.5, 0.7),\n",
    "        \"colsample_bytree\": trial.suggest_float(\"colsample_bytree\", 0.5, 1),\n",
    "        \"colsample_bynode\": trial.suggest_float(\"colsample_bynode\", 0.5, 1),\n",
    "        \"colsample_bylevel\": trial.suggest_float(\"colsample_bylevel\", 0.5, 1),\n",
    "        \"reg_lambda\": trial.suggest_float(\"reg_lambda\", 0, 1),\n",
    "        \"max_depth\": trial.suggest_int(\"max_depth\", 1, 6),\n",
    "        \"max_leaves\": trial.suggest_int(\"max_leaves\", 0, 2),\n",
    "        \"max_cat_to_onehot\": trial.suggest_int(\"max_cat_to_onehot\", 1, 10),\n",
    "    }\n",
    "    return train_model(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8f3b355-7e76-423f-8b4c-8e35eb1eb022",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a single study\n",
    "start = datetime.now()\n",
    "study = optuna.create_study(study_name=\"parallel-nyc-travel-time-model\")\n",
    "study.optimize(objective, n_trials=N_TRIALS, n_jobs=N_JOBS)\n",
    "print(f\"Total time:  {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64485e2-eeb7-4eb2-8661-f8aa4a2e519e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Tear down running clusters\n",
    "for client, _ in clusters.values():\n",
    "    client.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f913de4-9988-49d9-9b23-f5a90ecb9973",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(study.trials)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ede5a2bf-c88c-4266-9738-73cb6fcb9d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22a6da34-582c-4e2d-9b58-89a5369d7536",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4c41f7-afa1-44a7-ab36-d1ed15a4bd48",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Uncomment this if you want to save the results of your study to examine later.\n",
    "\n",
    "joblib.dump(study, \"study_many_threads.pickle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f65e308a-9569-40b1-9557-c6e16234be43",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = optuna.visualization.matplotlib.plot_optimization_history(study)\n",
    "fig.legend(loc=\"upper right\")\n",
    "plt.savefig(\"optimization_history_study_2.png\")"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}