databricks/notebooks/rapids_airline_hyperopt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyper-parameter Optimization with RAPIDS + MLflow + Hyperopt\n",
    "\n",
    "### Introduction\n",
    "\n",
    "Hyperparameter optimization is the task of picking the values for the hyperparameters of the model that provide the optimal results for the problem, as measured on a test dataset. This is often a crucial step and can help boost the model performance when done correctly. Despite its theoretical importance, HPO has been difficult to implement in practical applications because of the resources needed to run so many distinct training jobs. In this notebook, we explore the combination of RAPIDS, MLflow and Hyperopt to perform HPO on GPU and compare them with the performance of the CPU runs. We want to illustrate that HPO can be performed in an efficient manner with RAPIDS libraries on GPU.\n",
    "      \n",
    "### MLFlow\n",
    "\n",
    "[MLflow](https://mlflow.org/) is used for managing the machine learning lifecycle. It provides a way to track experiements and store data about them. MLflow also provides a way to deploy ML models. We'll make use of MLflow to store information about the models and use the built-in integration of MLflow into Databricks to register and store models.\n",
    "\n",
    "### Hyperopt\n",
    "\n",
    "[Hyperopt](http://hyperopt.github.io/hyperopt/) is a library for finding the best hyperparameters for a given objective function. It provides a way to choose the objective function, the search algorithm, the database in which to store the eval points. We'll use this to define the parameter space to search over. \n",
    "\n",
    "We'll use FAA flight history data for this demo - the aim is to predict if a given flight will be delayed or not using a target variable `ArrDelayBinary`. We'll use a Random Forest Classifier as the model for the learning task. \n",
    "\n",
    "We will compare the performance between the CPU version with scikit-learn to the GPU version with RAPIDS libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
    "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
    "\n",
    "import cudf\n",
    "import cuml\n",
    "\n",
    "import mlflow\n",
    "import hyperopt\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import mlflow.sklearn\n",
    "from mlflow.tracking.client import MlflowClient\n",
    "from hyperopt import fmin, tpe, hp, Trials, STATUS_OK\n",
    "\n",
    "import cuml.ensemble\n",
    "import cuml.metrics\n",
    "import cuml.preprocessing.model_selection\n",
    "\n",
    "import sklearn.ensemble\n",
    "import sklearn.metrics\n",
    "import sklearn.model_selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Small utility that times a block of code and prints how long it took to execute\n",
    "\n",
    "from contextlib import contextmanager\n",
    "import time\n",
    "\n",
    "@contextmanager\n",
    "def timed(name):\n",
    "    t0 = time.time()\n",
    "    yield\n",
    "    t1 = time.time()\n",
    "    print(\"..%-24s:  %8.4f\" % (name, t1 - t0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment set-up\n",
    "Here we'll define the input max evaluation runs and parallelism. \n",
    "\n",
    "Note: To use `MAX_PARALLEL` on DataBricks, make sure your cluster has at least `MAX_PARALLEL` nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_EVALS = 20\n",
    "MAX_PARALLEL = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the data and take a look\n",
    "\n",
    "### Data Acquisition\n",
    "\n",
    "We'll use the cell below to download the data. The `file_name` specifies which of the two available files - `airline_small.parquet` (smaller file) and `airline_20000000.parquet` we want to use. By default, we use the smaller file, but the benchmarks were run with the larger file. You are free to change it for experimentation.\n",
    "\n",
    "Run the cell below just once in the cluster to acquire the data and the comment it out for future runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Read above instructions - RUN ONLY ONCE\n",
    "\n",
    "# from urllib.request import urlretrieve\n",
    "# import os\n",
    "\n",
    "# file_name = 'airline_small.parquet' # NOTE: Change to airline_20000000.parquet to use a larger dataset\n",
    "\n",
    "# data_dir = \"/_dbfs_p8ath/\" # NOTE: Change to DBFS path where you want to save the file\n",
    "# INPUT_FILE = os.path.join(data_dir, file_name)\n",
    "\n",
    "# if os.path.isfile(INPUT_FILE):\n",
    "#         print(f\" > File already exists. Ready to load at {INPUT_FILE}\")\n",
    "# else:\n",
    "#     # Ensure folder exists\n",
    "#     os.makedirs(data_dir, exist_ok=True)\n",
    "        \n",
    "# url = \"https://rapidsai-cloud-ml-sample-data.s3-us-west-2.amazonaws.com/\" + file_name\n",
    "\n",
    "# urlretrieve(url= url,filename=INPUT_FILE)\n",
    "\n",
    "# print(\"Completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a quick peek into the data we will be using. As described before, this is a binary classification problem with the target variable as `ArrDelayBinary`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = cudf.read_parquet(INPUT_FILE)\n",
    "print(\"Data shape: \", df.shape)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting up MLFlow runs\n",
    "\n",
    "For the experiment to use MLFlow, we will define a training function (one each for CPU and GPU). In these functions, we will see a few mlflow log statements, these help in keep track of the experiment set-up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_cpu(params, test_set_frac=0.2, registered_model_name=None):\n",
    "    \"\"\"\n",
    "    Train scikit-learn model on the data, and calculate the accuracy on the model.\n",
    "    This method will be passed to `hyperopt.fmin()`.\n",
    "\n",
    "    Params:\n",
    "    params - dict; The range of the HPO space for different parameters (max_depth, max_features, n_estimators)\n",
    "                   in that order.\n",
    "    test_set_frac - float; Value between (0,1) for the size of the test set to be used for validation split\n",
    "    registered_model_name - string; Name under which the best model should be registered with MLFlow.\n",
    "\n",
    "    Returns:\n",
    "    dict with fields 'loss' (scalar loss) and 'status' (success/failure status of run)\n",
    "    \"\"\"\n",
    "    max_depth, max_features, n_estimators = params\n",
    "\n",
    "    with timed(\"load\"):\n",
    "        df = pd.read_parquet(INPUT_FILE)\n",
    "\n",
    "    with timed(\"etl\"):\n",
    "        X = df.drop([\"ArrDelayBinary\"], axis=1)\n",
    "        y = df[\"ArrDelayBinary\"].astype('int32')\n",
    "\n",
    "        X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(X, y,\n",
    "                                                                test_size=test_set_frac,\n",
    "                                                                random_state=123)\n",
    "\n",
    "    with timed(\"fit\"):\n",
    "        mod = sklearn.ensemble.RandomForestClassifier(max_depth=max_depth,\n",
    "                                             max_features=max_features,\n",
    "                                             n_estimators=n_estimators,\n",
    "                                             n_jobs=-1) # Use all available CPUs\n",
    "        mod.fit(X_train, y_train)\n",
    "\n",
    "    mlflow.sklearn.log_model(mod, \"RF_model_cpu_large_\",\n",
    "                           registered_model_name=registered_model_name)\n",
    "\n",
    "    with timed(\"predict\"):\n",
    "        if test_set_frac > 0.0:\n",
    "            preds = mod.predict(X_test)\n",
    "            acc = sklearn.metrics.accuracy_score(y_test, preds)\n",
    "            mlflow.log_metric(\"accuracy\", acc)\n",
    "        else:\n",
    "            acc = np.nan\n",
    "\n",
    "    # Returning -1 * acc because fmin minimizes the \"loss\" and we want to maximize accuracy.\n",
    "    return {'loss': -acc, 'status': STATUS_OK}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example run with a sample parameter value\n",
    "\n",
    "with timed(\"sample train skl\"):\n",
    "    result = train_cpu((8, 1.0, 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_rapids(params, test_set_frac=0.2, registered_model_name=None):\n",
    "    \"\"\"\n",
    "    Train RAPIDS cuml model on the data, and calculate the accuracy on the model.\n",
    "    This method will be passed to `hyperopt.fmin()`.\n",
    "\n",
    "    Params and Return values same as train_cpu\n",
    "    \"\"\"\n",
    "\n",
    "    max_depth, max_features, n_estimators = params\n",
    "\n",
    "    # Read using cudf\n",
    "    with timed(\"read_raw\"):\n",
    "        df = cudf.read_parquet(INPUT_FILE)\n",
    "\n",
    "    # Converting to dtypes expected by cuml model\n",
    "    with timed(\"etl\"):\n",
    "        X = df.drop([\"ArrDelayBinary\"], axis=1)\n",
    "        y = df[\"ArrDelayBinary\"].astype('int32')\n",
    "\n",
    "        # Splitting the data into 80/20 for training and validation\n",
    "        X_train, X_test, y_train, y_test = cuml.preprocessing.model_selection.train_test_split(\n",
    "                                                    X, y,\n",
    "                                                    test_size=test_set_frac,\n",
    "                                                    random_state=123)\n",
    "\n",
    "    with timed(\"fit\"):\n",
    "        n_bins = 16\n",
    "        mod = cuml.ensemble.RandomForestClassifier(max_depth=max_depth,\n",
    "                                        max_features=max_features,\n",
    "                                        n_bins=n_bins,\n",
    "                                        n_estimators=n_estimators)\n",
    "        mod.fit(X_train, y_train)\n",
    "\n",
    "\n",
    "    mlflow.sklearn.log_model(mod, \"RF_model_GPU_\",\n",
    "                             registered_model_name=registered_model_name)\n",
    "\n",
    "    with timed(\"predict\"):\n",
    "        if test_set_frac > 0.0:\n",
    "            preds = mod.predict(X_test)\n",
    "            acc = cuml.metrics.accuracy_score(y_test, preds)\n",
    "            mlflow.log_metric(\"accuracy\", acc)\n",
    "        else:\n",
    "            acc = np.nan\n",
    "\n",
    "    # Returning -1 * acc because fmin minimizes the \"loss\" and we want to maximize accuracy.\n",
    "    return {'loss': -acc, 'status': STATUS_OK}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example run with a sample parameter value\n",
    "with timed(\"sample train rapids\"):\n",
    "    result = train_rapids((8, 1.0, 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure hyperopt parameter search\n",
    "\n",
    "Let's define the ranges of the hyperparameter space using HyperOpt. Generally, the larger the ranges the better although, it is useful to keep in mind the limitations posed by the system or cluster you're running on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Shared search parameters\n",
    "from hyperopt.pyll import scope\n",
    "\n",
    "search_space = [\n",
    "        scope.int(hp.quniform('max_depth', 5, 15, 1)),\n",
    "        hp.uniform('max_features', 0., 1.0),\n",
    "        scope.int(hp.quniform('n_estimators', 100, 500, 100))\n",
    "    ]\n",
    "\n",
    "algo = tpe.suggest\n",
    "spark.conf.set('spark.task.maxFailures', '1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperopt with CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark_trials = hyperopt.SparkTrials(parallelism=MAX_PARALLEL)\n",
    "mlflow.end_run() # Close out any run in progress\n",
    "\n",
    "with mlflow.start_run() as run: \n",
    "    mlflow.set_tag(\"mlflow.runName\", \"CPU_run\")\n",
    "    best = fmin(\n",
    "      fn=train_cpu,\n",
    "      space=search_space,\n",
    "      algo=algo,\n",
    "      trials = spark_trials,\n",
    "      max_evals=MAX_EVALS)\n",
    "    mlflow.set_tag(\"best params\", str(best))\n",
    "\n",
    "    # Re-fit the best model on ALL of the data (no test set)\n",
    "    print(best)\n",
    "    train_cpu((int(best[\"max_depth\"]), best[\"max_features\"], int(best[\"n_estimators\"])),\n",
    "            test_set_frac=0.0, registered_model_name=\"MLFlow_Airline_CPU_large_\")\n",
    "\n",
    "mlflow.end_run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperopt with GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SparkTrials object will automatically log runs to MLFlow in DataBricks\n",
    "spark_trials = hyperopt.SparkTrials(parallelism=MAX_PARALLEL)\n",
    "mlflow.end_run() # Close out any run in progress\n",
    "\n",
    "with mlflow.start_run() as run:\n",
    "    mlflow.set_tag(\"mlflow.runName\", \"GPU_run\")\n",
    "    best = fmin(fn=train_rapids,\n",
    "      space=search_space,\n",
    "      algo=algo,\n",
    "      trials = spark_trials,\n",
    "      max_evals=MAX_EVALS)\n",
    "    mlflow.set_tag(\"best params\", str(best))\n",
    "\n",
    "    # Re-fit the best model on ALL of the data (no test set)\n",
    "    train_rapids((int(best[\"max_depth\"]), best[\"max_features\"], int(best[\"n_estimators\"])),\n",
    "               test_set_frac=0.0, registered_model_name=\"MLFlow_Airline_RAPIDS\")\n",
    "mlflow.end_run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "On a `p3.2xlarge` cluster on Databricks, the runtimes observed were 24 minutes for GPU and nearly 14 hours for CPU. We illustrate that RAPIDS on GPU can give up to 35x speedups. Hopefully this will make it easier to integrate hyperparameter optimization into your workflow if it can run in a coffee break rather than running overnight!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.10"
  },
  "name": "rapids_airline_hyperopt_large (1)",
  "notebookId": 2710846968050540
 },
 "nbformat": 4,
 "nbformat_minor": 4
}