{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a51c95d1-b447-4f1b-9571-cf597ca93ef4",
   "metadata": {
    "tags": [
     "cloud/aws/ec2",
     "data-storage/s3",
     "workflow/randomforest",
     "workflow/hpo",
     "workflow/xgboost",
     "library/dask",
     "library/dask-cuda",
     "library/xgboost",
     "library/optuna",
     "library/sklearn",
     "library/dask-ml"
    ]
   },
   "source": [
    "# HPO Benchmarking with RAPIDS and Dask\n",
    "\n",
    "_August, 2023_ \n",
    "\n",
    "Hyper-Parameter Optimization (HPO) helps to find the best version of a model by exploring the space of possible configurations. While generally desirable, this search is computationally expensive and time-consuming.\n",
    "\n",
    "In the notebook demo below, we compare benchmarking results to show how GPU can accelerate HPO tuning jobs relative to CPU.\n",
    "\n",
    "For instance, we find a 48x speedup in wall clock time (0.71 hrs vs 34.6 hrs) for XGBoost and 16x (3.86 hrs vs 63.2 hrs) for RandomForest when comparing between `p3.8xlarge` Tesla V100 GPUs  and `c5.24xlarge` CPU EC2 instances on 100 HPO trials of the 3-year Airline Dataset.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0d178f5-fc5d-471e-9898-544e5fdbc271",
   "metadata": {},
   "source": [
    "## Preamble"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c311bd4-76a1-4ee7-8841-5d44ca052566",
   "metadata": {},
   "source": [
    "You can set up local environment but it is recommended to launch a Virtual Machine service (Azure, AWS, GCP, etc).\n",
    "\n",
    "For the purposes of this notebook, we will be utilizing the [Amazon Machine Image (AMI)](https://aws.amazon.com/releasenotes/aws-deep-learning-ami-gpu-tensorflow-2-12-amazon-linux-2/) as the starting point.\n",
    "\n",
    "\n",
    "````{docref} /cloud/aws/ec2\n",
    "Please follow instructions in [AWS Elastic Cloud Compute)](../../cloud/aws/ec2) to launch an EC2 instance with GPUs, the NVIDIA Driver and the NVIDIA Container Runtime.\n",
    "\n",
    "```{note}\n",
    "When configuring your instance ensure you select the [Deep Learning AMI GPU TensorFlow or PyTorch](https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html) in the AMI selection box under **\"Amazon Machine Image (AMI)\"**\n",
    "\n",
    "![](../../_static/images/examples/xgboost-rf-gpu-cpu-benchmark/amazon-deeplearning-ami.png)\n",
    "```\n",
    "\n",
    "Once your instance is running and you have access to Jupyter save this notebook and run through the cells.\n",
    "\n",
    "````\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d57ba2-df8a-4757-a0df-44b3cd73b75c",
   "metadata": {},
   "source": [
    "## Python ML Workflow\n",
    "\n",
    "In order to work with RAPIDS container, the entrypoint logic should parse arguments, load, preprocess and split data, build and train a model, score/evaluate the trained model, and emit an output representing the final score for the given hyperparameter setting.\n",
    "\n",
    "Let's have a step-by-step look at each stage of the ML workflow:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84b3fbaf-754b-45a6-959c-7163edfe5c4f",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "\n",
    "We leverage the `Airline` dataset, which is a large public tracker of US domestic flight logs which we offer in various sizes (1 year, 3 year, and 10 year) and in [Parquet](https://parquet.apache.org/) (compressed column storage) format. The machine learning objective with this dataset is to predict whether flights will be more than 15 minutes late arriving to their destination.\n",
    "\n",
    "We host the demo dataset in public S3 demo buckets in both the `us-east-1` or `us-west-2`. To optimize performance, we recommend that you access the s3 bucket in the same region as your EC2 instance to reduce network latency and data transfer costs. \n",
    "\n",
    "For this demo, we are using the **`3_year`** dataset, which includes the following features to mention a few:\n",
    "\n",
    "* Date and distance ( Year, Month, Distance )\n",
    "* Airline / carrier ( Flight_Number_Reporting_Airline )\n",
    "* Actual departure and arrival times ( DepTime and ArrTime )\n",
    "* Difference between scheduled & actual times ( ArrDelay and DepDelay )\n",
    "* Binary encoded version of late, aka our target variable ( ArrDelay15 )\n",
    "\n",
    "Configure aws credentials for access to S3 storage\n",
    "\n",
    "```\n",
    "aws configure\n",
    "```\n",
    "\n",
    "Download dataset from S3 bucket to your current working dir\n",
    "\n",
    "```\n",
    "aws s3 cp --recursive s3://sagemaker-rapids-hpo-us-west-2/3_year/ ./data/\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4bc101b-3206-4f86-9b0d-d4313ad1ee43",
   "metadata": {},
   "source": [
    "### Algorithm\n",
    "\n",
    "From a ML/algorithm perspective, we offer `XGBoost` and `RandomForest`. You are free to switch between these algorithm choices and everything in the example will continue to work.\n",
    "\n",
    "```python\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\n",
    "    \"--model-type\", type=str, required=True, choices=[\"XGBoost\", \"RandomForest\"]\n",
    ")\n",
    "```\n",
    "\n",
    "We can also optionally increase robustness via reshuffles of the train-test split (i.e., cross-validation folds). Typical values are between 3 and 10 folds. We will use \n",
    "```python\n",
    "n_cv_folds = 5\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b1759de-98af-4628-a79b-a236a2dee5a2",
   "metadata": {},
   "source": [
    "### Dask Cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533be0b1-0d5e-46b3-9ff1-dd71751fe68f",
   "metadata": {},
   "source": [
    "To maximize on efficiency, we launch a Dask `LocalCluster` for cpu or `LocalCUDACluster` that utilizes GPUs for distributed computing. Then connect a Dask Client to submit and manage computations on the cluster. \n",
    "\n",
    "We can then ingest the data, and \"persist\" it in memory using dask as follows:\n",
    "\n",
    "```python\n",
    "if args.mode == \"gpu\":\n",
    "    cluster = LocalCUDACluster()\n",
    "else: # mode == \"cpu\"\n",
    "    cluster = LocalCluster(n_workers=os.cpu_count())\n",
    "\n",
    "with Client(cluster) as client:\n",
    "    dataset = ingest_data(mode=args.mode)\n",
    "    client.persist(dataset)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feef53b4-d7e0-43e2-b85b-372bd2d882f7",
   "metadata": {},
   "source": [
    "### Search Range\n",
    "\n",
    "One of the most important choices when running HPO is to choose the bounds of the hyperparameter search process. In this notebook, we leverage the power of `Optuna`, a widely used Python library for hyperparameter optimization.\n",
    "\n",
    "Here's the quick steps on getting started with Optuna:\n",
    "\n",
    "\n",
    "1) Define the Objective Function, which represents the model training and evaluation process. It takes hyperparameters as inputs and returns a metric to optimize (e.g, accuracy in our case,). Refer to `train_xgboost()` and `train_randomforest()` in `hpo.py`\n",
    "\n",
    "2. Specify the search space using the `Trial` object's methods to define the hyperparameters and their corresponding value ranges or distributions. For example:\n",
    "\n",
    "```python\n",
    "\"max_depth\": trial.suggest_int(\"max_depth\", 4, 8),\n",
    "\"max_features\": trial.suggest_float(\"max_features\", 0.1, 1.0),\n",
    "\"learning_rate\": trial.suggest_float(\"learning_rate\", 0.001, 0.1, log=True),\n",
    "\"min_samples_split\": trial.suggest_int(\"min_samples_split\", 2, 1000, log=True),\n",
    "```\n",
    "\n",
    "3. Create an Optuna study object to keep track of trials and their corresponding hyperparameter configurations and evaluation metrics.\n",
    "\n",
    "```python\n",
    "study = optuna.create_study(\n",
    "        sampler=RandomSampler(seed=args.seed), direction=\"maximize\"\n",
    "    )\n",
    "```\n",
    "\n",
    "4. Select an optimization algorithm to determine how Optuna explores and exploits the search space to find optimal configurations.  For instance, the `RandomSampler` is an algorithm provided by the Optuna library that samples hyperparameter configurations randomly from the search space.\n",
    "\n",
    "5. Run the Optimization by calling the Optuna's `optimize()` function on the study object. You can specify the number of trials or number of parallel jobs to run.\n",
    "\n",
    "```python\n",
    " study.optimize(lambda trial: train_xgboost(\n",
    "                    trial, dataset=dataset, client=client, mode=args.mode\n",
    "                ),\n",
    "                n_trials=100,\n",
    "                n_jobs=1,\n",
    "            )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a89edfea-ca14-4d26-94c6-0ef8eaf02d77",
   "metadata": {},
   "source": [
    "## Run HPO\n",
    "\n",
    "Let's try this out!\n",
    "\n",
    "The example file `hpo.py` included here implements the patterns described above.\n",
    "\n",
    "First make sure you have the correct CUDAtoolkit version by running `nvidia-smi`. See the RAPIDS installation docs ([link](https://docs.rapids.ai/install/#system-req)) for details on the supported range of GPUs and drivers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e4152ac-2712-4568-a807-8ef1267d8a2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baca52e2-09e7-42f3-bc98-5ee38f9e274f",
   "metadata": {},
   "source": [
    "Executing benchmark tests can be an arduous and time-consuming procedure that may extend over multiple days. By using  a tool like [tmux](https://www.redhat.com/sysadmin/introduction-tmux-linux), you can maintain active terminal sessions, ensuring that your tasks continue running even if the SSH connection is interrupted. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a87780-7a19-4c15-a942-477d8578e8ac",
   "metadata": {},
   "source": [
    "```\n",
    "tmux\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77df8ce3-39b8-41d9-a538-ae404be25b45",
   "metadata": {},
   "source": [
    "Run the following to run hyper-parameter optimization in a Docker container.\n",
    "\n",
    "If you don't yet have that image locally, the first time this runs it might take a few minutes to pull it.\n",
    "After that, startup should be very fast.\n",
    "\n",
    "Here's what the arguments in that command below are doing:\n",
    "\n",
    "* `--gpus all` = make all GPUs on the system available to processes in the container\n",
    "* `--env EXTRA_CONDA_PACKAGES` = install `optuna` and `optuna-integration` conda packages\n",
    "  - *the image already comes with all of the RAPIDS libraries and their dependencies installed*\n",
    "* `-p 8787:8787` = forward between port port 8787 on the host and 8787 on the container\n",
    "  - *navigate to `{public IP of box}:8787 to see the Dask dashboard!*\n",
    "* `-v / -w` = mount the current directory from the host machine into the container\n",
    "  - *this allows processes in the container to read the data you downloaded to the `./data` directory earlier*\n",
    "  - *it also means that any changes made to these files from inside the container will be reflected back on the host*\n",
    "\n",
    "Piping to a file called `xgboost_hpo_logs.txt` is helpful, as it preserves all the logs for later inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fcf8b8f-63a9-4981-b714-9fcc77a7b6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker run \\\n",
    "    --gpus all \\\n",
    "    --env EXTRA_CONDA_PACKAGES=\"optuna optuna-integration\" \\\n",
    "    -p 8787:8787 \\\n",
    "    -v $(pwd):/home/rapids/xgboost-hpo-example \\\n",
    "    -w /home/rapids/xgboost-hpo-example \\\n",
    "    -it rapidsai/base:26.08a-cuda13-py3.13 \\\n",
    "    /bin/bash -c \"python ./hpo.py --model-type 'XGBoost' --target 'gpu'\" \\\n",
    "> ./xgboost_hpo_logs.txt 2>&1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f881ee-ce7f-4810-a6b7-fa8aa72d91f3",
   "metadata": {},
   "source": [
    "## Try Some Modifications\n",
    "\n",
    "Now that you've run this example, try some modifications!\n",
    "\n",
    "For example:\n",
    "\n",
    "* use `--model-type \"RandomForest\"` to see how a random forest model compares to XGBoost\n",
    "* use `--target \"cpu\"` to estimate the speedup from GPU-accelerated training\n",
    "* modify the pipeline in `hpo.py` with other customizations"
   ]
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
