Modeling 4 - Parallel HPO of XGBoost with Optuna and Dask (single cluster).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0652f264-b2c3-43b2-9310-7c00e50e7cdc",
   "metadata": {},
   "source": [
    "# Parallelize HPO of XGBoost with Optuna and Dask on a single large Coiled cluster\n",
    "\n",
    "In this workflow we use advanced techniques to perform the same work as the last notebook (many runs in many clusters) but now using a single cluster."
   ]
  },
  {
   "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 functools import partial\n",
    "from datetime import datetime\n",
    "\n",
    "import dask.array as da\n",
    "import dask.dataframe as dd\n",
    "import coiled\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": "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",
    "# Number of converging serial trials to run in each job.\n",
    "# NOTE: this is different from the `Modeling 3` notebook!\n",
    "N_TRIALS = 5\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",
    "# The number of training exercises that will be run in total is\n",
    "# N_JOBS * N_TRIALS * N_FOLDS\n",
    "\n",
    "# Keep the number of parallel training exercises to the bare minimum\n",
    "# to ensure pipelining. More than this would just overwhelm the scheduler.\n",
    "N_PARALLEL = 2\n",
    "\n",
    "# Dask worker instance type and number\n",
    "WORKER_INSTANCE_TYPE = \"r6i.large\"\n",
    "N_WORKERS = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32ee232d",
   "metadata": {},
   "source": [
    "### Start coiled cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5686386-60ea-4277-9116-daabd3ca5138",
   "metadata": {},
   "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\": \"30 minutes\"},\n",
    ")\n",
    "client = distributed.Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77d715d5",
   "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.\n",
    "# We'll also need to retrieve it from the workers.\n",
    "if \"ddf\" in client.datasets:\n",
    "    client.unpublish_dataset(\"ddf\")\n",
    "client.publish_dataset(ddf=ddf)\n",
    "\n",
    "ddf.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dd461c3-e551-420a-a3e1-1ca2515bb613",
   "metadata": {},
   "source": [
    "### Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e93a5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we subset data for cross-validation\n",
    "def make_cv_splits(\n",
    "    ddf: dd.DataFrame, n_folds: int = 5\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": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_once(\n",
    "    train: dd.DataFrame,\n",
    "    test: dd.DataFrame,\n",
    "    sem: distributed.Semaphore,\n",
    "    study_params: dict[str, float],\n",
    "):\n",
    "    distributed.secede()\n",
    "    # Block until there are less than N_PARALLEL train_once\n",
    "    # critical sections running\n",
    "    with sem:\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",
    "            None, X_train, y_train, enable_categorical=True\n",
    "        )\n",
    "        # This has its own internal semaphore with a limit of 1\n",
    "        model = xgboost.dask.train(\n",
    "            None,\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",
    "\n",
    "        # Note: in Modeling 3, score was a scalar dask collection wrapping a Future\n",
    "        # (due to the compute=False flag). It would be a very big mistake to do\n",
    "        # the same here, as returning a Future from a dask task will cause the\n",
    "        # scheduler to forget your data as soon as there isn't a client holding\n",
    "        # a reference to the future anymore.\n",
    "        score = mean_squared_error(\n",
    "            y_test.to_dask_array(),\n",
    "            predictions.to_dask_array(),\n",
    "            squared=False,\n",
    "        )\n",
    "        return score\n",
    "\n",
    "\n",
    "def train_model(\n",
    "    ddf_name: str,\n",
    "    n_folds: int,\n",
    "    sem: distributed.Semaphore,\n",
    "    study_params: dict[str, float],\n",
    "):\n",
    "    client = distributed.get_client()\n",
    "    ddf = client.get_dataset(ddf_name)\n",
    "\n",
    "    futures = [\n",
    "        client.submit(train_once, train, test, sem, study_params, pure=False)\n",
    "        for train, test in make_cv_splits(ddf, n_folds)\n",
    "    ]\n",
    "\n",
    "    try:\n",
    "        distributed.secede()\n",
    "    except KeyError:\n",
    "        # Already seceded in a previous iteration of study.optimize()\n",
    "        pass\n",
    "\n",
    "    scores = client.gather(futures)\n",
    "    return sum(scores) / len(scores)\n",
    "\n",
    "\n",
    "def objective(trial, ddf_name: str, n_folds: int, sem: distributed.Semaphore) -> float:\n",
    "    study_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(ddf_name, n_folds, sem, study_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d70a951",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a single study and run some trials\n",
    "start = datetime.now()\n",
    "storage = optuna.integration.DaskStorage()\n",
    "study = optuna.create_study(storage=storage, study_name=\"nyc-travel-time-model\")\n",
    "\n",
    "\n",
    "# Run N_JOBS in parallel\n",
    "# each job will run N_TRIALS converging trials in series\n",
    "# each trial will start N_FOLDS training exercises in parallel\n",
    "# but only up to N_PARALLEL training exercise will actually submit\n",
    "# tasks to the scheduler at the same time.\n",
    "\n",
    "# In pure optuna, we would have used\n",
    "# study.optimize(objective, n_trials=N_TRIALS, n_jobs=N_JOBS)\n",
    "sem = distributed.Semaphore(N_PARALLEL)\n",
    "futures = [\n",
    "    client.submit(\n",
    "        study.optimize,\n",
    "        partial(objective, ddf_name=\"ddf\", n_folds=N_FOLDS, sem=sem),\n",
    "        n_trials=N_TRIALS,\n",
    "        pure=False,\n",
    "    )\n",
    "    for _ in range(N_JOBS)\n",
    "]\n",
    "client.gather(futures)\n",
    "\n",
    "print(f\"Total time:  {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ebf98b6-9f2b-432e-8caa-9c7e1970d1c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.shutdown()"
   ]
  },
  {
   "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": "e4928342-4749-4e29-bb65-ccabd8190a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.best_trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4c41f7-afa1-44a7-ab36-d1ed15a4bd48",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save the results of your study to examine later\n",
    "joblib.dump(study, \"study_single_cluster.pickle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71bc6ef-a34a-45b8-bb04-ce583ac29623",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = optuna.visualization.matplotlib.plot_optimization_history(study)\n",
    "fig.legend(loc=\"upper right\")\n",
    "plt.savefig(\"optimization_history_study_1.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7afef92b",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ae808d",
   "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
}