aws/rapids_studio_hpo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:22pt\"> **Demo Overview** </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Automated Model Tuning (AMT) also known as 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 can feel prohibitive. \n",
    "\n",
    "\n",
    "In the notebook demo below, we show how [SageMaker Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio.html) and RAPIDS working together can tackle model tuning by accelerating compute parallelism within a node's GPUs; and simultaneously accelerating the search by leveraging sets of cloud nodes running parallel experiments. \n",
    "\n",
    "For more check out our [AWS blog](https://aws.amazon.com/blogs/machine-learning/rapids-and-amazon-sagemaker-scale-up-and-scale-out-to-tackle-ml-challenges/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:22pt\"> **0. Preamble** </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook was tested in an Amazon SageMaker Studio notebook, on a ml.t3.medium instance with Python 3 (Data Science) kernel.\n",
    "\n",
    "To get things rolling let's make sure we can query our AWS SageMaker execution role and session as well as our account ID and AWS region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from helper_functions import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "execution_role = sagemaker.get_execution_role()\n",
    "session = sagemaker.Session()\n",
    "\n",
    "account=!(aws sts get-caller-identity --query Account --output text)\n",
    "region = [session.boto_region_name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "account, region"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:22pt\"> **1. Key Choices** </span>\n",
    "\n",
    "<span style=\"color:#8735fb; font-size:18pt\"> 1.1 - HPO Configurations </span>\n",
    "\n",
    "Let's go ahead and choose the configuration options for our HPO run. We will be using the default algorithm configurations in [code/Dockerfile](code/Dockerfile) (three-fold cross validation with XGBoost on a single GPU), which are explained in detail in our [extended notebook example](rapids_sagemaker_hpo_extended.ipynb). If you are using your own workflow and training scripts, be sure to write your Dockerfile accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# please choose dataset S3 bucket and directory\n",
    "data_bucket = 'sagemaker-rapids-hpo-' + region[0]\n",
    "dataset_directory = '3_year' # '1_year', '3_year', '10_year', 'NYC_taxi'\n",
    "\n",
    "# please choose output bucket for trained model(s)\n",
    "model_output_bucket = session.default_bucket()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "s3_data_input = f\"s3://{data_bucket}/{dataset_directory}\"\n",
    "s3_model_output = f\"s3://{model_output_bucket}/trained-models\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# please choose HPO search ranges\n",
    "hyperparameter_ranges = {\n",
    "    'max_depth'    : sagemaker.parameter.IntegerParameter        ( 5, 15 ),\n",
    "    'num_boost_round' : sagemaker.parameter.IntegerParameter        ( 100, 500 ),\n",
    "    'max_features' : sagemaker.parameter.ContinuousParameter     ( 0.1, 1.0 ),    \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:18pt\"> 1.2 - Experiment Scale </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to decide how may total experiments to run, and how many should run in parallel. Below we have a very conservative number of maximum jobs to run so that you don't accidently spawn large computations when starting out, however for meaningful HPO searches this number should be much higher (e.g., in our experiments we often run 100 max_jobs). Note that you may need to request a [quota limit increase](https://docs.aws.amazon.com/general/latest/gr/sagemaker.html) for additional  `max_parallel_jobs` parallel workers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# please choose total number of HPO experiments[ we have set this number very low to allow for automated CI testing ]\n",
    "max_jobs = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# please choose number of experiments that can run in parallel\n",
    "max_parallel_jobs = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also set the max duration for an individual job to 24 hours so we don't have run-away compute jobs taking too long."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "max_duration_of_experiment_seconds = 60 * 60 * 24"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:18pt\"> 1.3 - Compute Platform </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on the workflow you have chosen, your instance should reflect the specifications needed. For example, for the singleGPU workflow, you should choose an instance with a GPU, such as the p3.2xlarge instance. You can [read about Amazon EC2 Instance Types here](https://aws.amazon.com/ec2/instance-types/). \n",
    "> e.g., For the 10_year dataset option, we suggest ml.p3.8xlarge instances (4 GPUs) and ml.m5.24xlarge CPU instances ( we will need upwards of 200GB CPU RAM during model training)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# we will recommend a compute instance type, feel free to modify \n",
    "instance_type = 'ml.p3.2xlarge'  # recommend_instance_type(ml_workflow_choice, dataset_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to choosing our instance type, we can also enable significant savings by leveraging [AWS EC2 Spot Instances](https://aws.amazon.com/ec2/spot/).\n",
    "\n",
    "We **highly recommend** that you set this flag to `True` as it typically leads to 60-70% cost savings. Note, however that you may need to request a [quota limit increase](https://docs.aws.amazon.com/general/latest/gr/sagemaker.html) to enable Spot instances in SageMaker.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# please choose whether spot instances should be used\n",
    "use_spot_instances_flag = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:22pt\"> **2. Build Estimator** </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To build a RAPIDS enabled SageMaker HPO we first need to build a [SageMaker Estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html). An Estimator is a container image that captures all the software needed to run an HPO experiment. The container is augmented with entrypoint code that will be trggered at runtime by each worker. The entrypoint code enables us to write custom models and hook them up to data. \n",
    "\n",
    "In order to work with SageMaker HPO, the entrypoint logic should parse hyperparameters (supplied by AWS SageMaker), load 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. We've already built multiple variations of this code.\n",
    "\n",
    "If you would like to make changes by adding your custom model logic feel free to modify the **train.py** and/or the specific workflow files in the `code/workflows` directory.\n",
    "\n",
    "First, let's switch our working directory to the location of the Estimator entrypoint and library code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%cd code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 2.1 - Containerize and Push to ECR </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's turn to building our container so that it can integrate with the AWS SageMaker HPO API.\n",
    "\n",
    "Let's first decide on the full name of our container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rapids_base_container = 'rapidsai/rapidsai-cloud-ml:latest'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "image_base = 'cloud-ml-sagemaker'\n",
    "image_tag  = rapids_base_container.split(':')[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ecr_fullname = f\"{account[0]}.dkr.ecr.{region[0]}.amazonaws.com/{image_base}:{image_tag}\"\n",
    "ecr_fullname"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, let's ensure that our Dockerfile correctly captured our base image selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "validate_dockerfile(rapids_base_container)\n",
    "!cat Dockerfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:18pt\"> 2.1.1 Build and Publish to ECR</span>\n",
    "\n",
    "In order to build and push to the ECR from SageMaker Studio, we must first install sm-docker: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install sagemaker-studio-image-build"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we’re ready to start taking advantage of the new CLI to easily build our custom bring-your-own Docker image from Amazon SageMaker Studio without worrying about the underlying setup and configuration of build services. Once in ECR, AWS SageMaker will be able to leverage our image to build Estimators and run experiments. We are able to build and publish to the ECR with: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "!sm-docker build . --repository cloud-ml-sagemaker:latest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 2.2 - Create Estimator </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having built our container [ +custom logic] and pushed it to ECR, we can finally compile all of efforts into an Estimator instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 'volume_size' - EBS volume size in GB, default = 30\n",
    "estimator_params = {\n",
    "    'image_uri': ecr_fullname,\n",
    "    'role': execution_role,    \n",
    "    \n",
    "    'instance_type': instance_type,\n",
    "    'instance_count': 1,\n",
    "    \n",
    "    'input_mode': 'File',\n",
    "    'output_path': s3_model_output,\n",
    "    \n",
    "    'use_spot_instances': use_spot_instances_flag,\n",
    "    \n",
    "    'max_run': max_duration_of_experiment_seconds, # 24 hours \n",
    "    'sagemaker_session': session,\n",
    "}\n",
    "\n",
    "if use_spot_instances_flag == True:\n",
    "    estimator_params.update({'max_wait' : max_duration_of_experiment_seconds + 1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "estimator = sagemaker.estimator.Estimator(**estimator_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 2.3 - Test Estimator </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to test by asking SageMaker to run the BYOContainer logic inside our Estimator. This is a useful step if you've made changes to your custom logic and are interested in making sure everything works before launching a large HPO search. \n",
    "\n",
    "> Note: This verification step will use the default hyperparameter values declared in our custom train code, as SageMaker HPO will not be orchestrating a search for this single run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "estimator.fit(inputs = s3_data_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"display: block; text-align: center; color:#8735fb; font-size:30pt\"> **3. Run HPO** </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With a working SageMaker Estimator in hand, the hardest part is behind us. In the key choices section we <a href='#strategy-and-param-ranges'>already defined our search strategy and hyperparameter ranges</a>, so all that remains is to choose a metric to evaluate performance on. For more documentation check out the AWS SageMaker [Hyperparameter Tuner documentation](https://sagemaker.readthedocs.io/en/stable/tuner.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 3.1 - Define Metric </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We only focus on a single metric, which we call 'final-score', that captures the accuracy of our model on the test data unseen during training. You are of course welcome to add aditional metrics, see [AWS SageMaker documentation on Metrics](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-metrics.html). When defining a metric we provide a regular expression (i.e., string parsing rule) to extract the key metric from the output of each Estimator/worker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "metric_definitions = [{'Name': 'final-score', 'Regex': 'final-score: (.*);'}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "objective_metric_name = 'final-score'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 3.2 - Define Tuner </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we put all of the elements we've been building up together into a HyperparameterTuner declaration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hpo = sagemaker.tuner.HyperparameterTuner(estimator=estimator,\n",
    "                                          metric_definitions=metric_definitions, \n",
    "                                          objective_metric_name=objective_metric_name,\n",
    "                                          objective_type='Maximize',\n",
    "                                          hyperparameter_ranges=hyperparameter_ranges,\n",
    "                                          strategy='Random',  \n",
    "                                          max_jobs=max_jobs,\n",
    "                                          max_parallel_jobs=max_parallel_jobs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 3.3 - Run HPO </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's be sure we take a moment to confirm before launching all of our HPO experiments. Depending on your configuration options running this cell can kick off a massive amount of computation!\n",
    "> Once this process begins, we recommend that you use the SageMaker UI to keep track of the <a href='../img/gpu_hpo_100x10.png'>health of the HPO process and the individual workers</a>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import string\n",
    "\n",
    "tuning_job_name = 'unified-hpo-' + ''.join(random.choices(string.digits, k = 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hpo.fit( inputs=s3_data_input,\n",
    "         job_name=tuning_job_name,\n",
    "         wait=True,\n",
    "         logs='All')\n",
    "\n",
    "hpo.wait()  # block until the .fit call above is completed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:20pt\"> 3.4 - Results and Summary </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your job is complete there are multiple ways to analyze the results. Below we display the performance of the best job, as well printing each HPO trial/job as a row of a dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hpo_results = summarize_hpo_results(tuning_job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name).dataframe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:30pt\">RAPIDS References</span>\n",
    "\n",
    "> [cloud-ml-examples](http://github.com/rapidsai/cloud-ml-examples)\n",
    "\n",
    "> [RAPIDS HPO](https://rapids.ai/hpo)\n",
    "\n",
    "> [cuML Documentation](https://docs.rapids.ai/api/cuml/stable/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:#8735fb; font-size:30pt\">SageMaker References</span>\n",
    "\n",
    "> [SageMaker Training Toolkit](https://github.com/aws/sagemaker-training-toolkit)\n",
    "\n",
    "> [Getting Started with SageMaker Studio](https://sagemaker-examples.readthedocs.io/en/latest/aws_sagemaker_studio/index.html)\n",
    "\n",
    "> [Using the Amazon SageMaker Studio Image Build CLI to build container images from your Studio notebooks](https://aws.amazon.com/blogs/machine-learning/using-the-amazon-sagemaker-studio-image-build-cli-to-build-container-images-from-your-studio-notebooks/)\n",
    "\n",
    "> [Docker containers with SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/docker-basics.html)\n",
    "\n",
    "> [Estimator Parameters](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html)\n",
    "\n",
    "> Spot Instances [docs](https://docs.aws.amazon.com/sagemaker/latest/dg/model-managed-spot-training.html), and [blog]()"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}