sagemaker_processing/feature_transformation_with_sagemaker_processing_dask/feature_transformation_with_sagemaker_processing_dask.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature transformation with Amazon SageMaker Processing and Dask\n",
    "\n",
    "Typically a machine learning (ML) process consists of few steps. First, gathering data with various ETL jobs, then pre-processing the data, featurizing the dataset by incorporating standard techniques or prior knowledge, and finally training an ML model using an algorithm.\n",
    "\n",
    "Often, distributed data processing frameworks such as Dask are used to pre-process data sets in order to prepare them for training. In this notebook we'll use Amazon SageMaker Processing, and leverage the power of Dask in a managed SageMaker environment to run our preprocessing workload.\n",
    "\n",
    "### What is Dask Distributed?\n",
    "Dask.distributed: is a lightweight and open source library for distributed computing in Python. It is also a centrally managed, distributed, dynamic task scheduler. It is also a centrally managed, distributed, dynamic task scheduler. Dask has three main components:\n",
    "\n",
    "**dask-scheduler process:** coordinates the actions of several workers. The scheduler is asynchronous and event-driven, simultaneously responding to requests for computation from multiple clients and tracking the progress of multiple workers.\n",
    "\n",
    "**dask-worker processes:** Which are spread across multiple machines and the concurrent requests of several clients.\n",
    "\n",
    "**dask-client process:** which is is the primary entry point for users of dask.distributed\n",
    "\n",
    "<img src=\"https://docs.dask.org/en/latest/_images/dask-overview.svg\">\n",
    "\n",
    "source: https://docs.dask.org/en/latest/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "1. [Objective](#Objective:-predict-the-age-of-an-Abalone-from-its-physical-measurement)\n",
    "1. [Setup](#Setup)\n",
    "1. [Using Amazon SageMaker Processing to execute a Dask Job](#Using-Amazon-SageMaker-Processing-to-execute-a-Dask-Job)\n",
    "  1. [Downloading dataset and uploading to S3](#Downloading-dataset-and-uploading-to-S3)\n",
    "  1. [Build a Dask container for running the preprocessing job](#Build-a-Dask-container-for-running-the-preprocessing-job)\n",
    "  1. [Run the preprocessing job using Amazon SageMaker Processing](#Run-the-preprocessing-job-using-Amazon-SageMaker-Processing)\n",
    "    1. [Inspect the preprocessed dataset](#Inspect-the-preprocessed-dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Let's start by specifying:\n",
    "* The S3 bucket and prefixes that you use for training and model data. Use the default bucket specified by the Amazon SageMaker session.\n",
    "* The IAM role ARN used to give processing and training access to the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import gmtime, strftime\n",
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "role = sagemaker.get_execution_role()\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "\n",
    "prefix = \"sagemaker/dask-preprocess-demo\"\n",
    "input_prefix = prefix + \"/input/raw/census\"\n",
    "input_preprocessed_prefix = prefix + \"/input/preprocessed/census\"\n",
    "model_prefix = prefix + \"/model\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Amazon SageMaker Processing to execute a Dask job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Downloading dataset and uploading to Amazon Simple Storage Service (Amazon S3)\n",
    "\n",
    "The dataset used here is the Census-Income KDD Dataset. The first step are to select features, clean the data, and turn the data into features that the training algorithm can use to train a binary classification model which can then be used to predict whether rows representing census responders have an income greater or less than $50,000. In this example, we will use Dask distributed to preprocess and transform the data to make it ready for the training process. In the next section, you download from the bucket below then upload to your own bucket so that Amazon SageMaker can access the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import pandas as pd\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "region = sagemaker_session.boto_region_name\n",
    "input_data = 's3://sagemaker-sample-data-{}/processing/census/census-income.csv'.format(region)\n",
    "!aws s3 cp $input_data .\n",
    "\n",
    "# Uploading the training data to S3\n",
    "sagemaker_session.upload_data(path='census-income.csv', bucket=bucket, key_prefix=input_prefix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a dask container for running the preprocessing job\n",
    "\n",
    "An example Dask container is included in the `./container` directory of this example. The container handles the bootstrapping of Dask Scheduler and mapping each instance to a Dask Worke. At a high level the container provides:\n",
    "\n",
    "* A set of default worker/scheduler configurations\n",
    "* A bootstrapping script for configuring and starting up  scheduler/worker nodes\n",
    "* Starting dask cluster from all the workers including the scheduler node\n",
    "\n",
    "\n",
    "After the container build and push process is complete, use the Amazon SageMaker Python SDK to submit a managed, distributed dask application that performs our dataset preprocessing.\n",
    "\n",
    "### Build the example Dask container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd container\n",
    "!docker build -t sagemaker-dask-example .\n",
    "%cd ../"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an Amazon Elastic Container Registry (Amazon ECR) repository for the Dask container and push the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "region = boto3.session.Session().region_name\n",
    "\n",
    "ecr_repository = 'sagemaker-dask-example'\n",
    "tag = ':latest'\n",
    "uri_suffix = 'amazonaws.com'\n",
    "if region in ['cn-north-1', 'cn-northwest-1']:\n",
    "    uri_suffix = 'amazonaws.com.cn'\n",
    "dask_repository_uri = '{}.dkr.ecr.{}.{}/{}'.format(account_id, region, uri_suffix, ecr_repository + tag)\n",
    "\n",
    "# Create ECR repository and push docker image\n",
    "!$(aws ecr get-login --region $region --registry-ids $account_id --no-include-email)\n",
    "!aws ecr create-repository --repository-name $ecr_repository\n",
    "!docker tag {ecr_repository + tag} $dask_repository_uri\n",
    "!docker push $dask_repository_uri"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the preprocessing job using Amazon SageMaker Processing on Dask Cluster\n",
    "\n",
    "Next, use the Amazon SageMaker Python SDK to submit a processing job. Use the the custom Dask container that was just built, and a Scikit Learn script for preprocessing in the job configuration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create the Dask preprocessing script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile preprocess.py\n",
    "from __future__ import print_function, unicode_literals\n",
    "import argparse\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import warnings\n",
    "import boto3\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tornado import gen\n",
    "import dask.dataframe as dd\n",
    "import joblib\n",
    "from dask.distributed import Client\n",
    "from sklearn.compose import make_column_transformer\n",
    "from sklearn.exceptions import DataConversionWarning\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import (\n",
    "    KBinsDiscretizer,\n",
    "    LabelBinarizer,\n",
    "    OneHotEncoder,\n",
    "    PolynomialFeatures,\n",
    "    StandardScaler,\n",
    ")\n",
    "\n",
    "warnings.filterwarnings(action=\"ignore\", category=DataConversionWarning)\n",
    "s3_client = boto3.resource(\"s3\")\n",
    "attempts_counter = 3\n",
    "attempts = 0\n",
    "\n",
    "\n",
    "def upload_objects(bucket, prefix, local_path):\n",
    "    try:\n",
    "        bucket_name = bucket  # s3 bucket name\n",
    "        root_path = local_path  # local folder for upload\n",
    "\n",
    "        s3_bucket = s3_client.Bucket(bucket_name)\n",
    "\n",
    "        for path, subdirs, files in os.walk(root_path):\n",
    "            for file in files:\n",
    "                s3_bucket.upload_file(\n",
    "                    os.path.join(path, file), \"{}/output/{}\".format(prefix, file)\n",
    "                )\n",
    "    except Exception as err:\n",
    "        logging.exception(err)\n",
    "\n",
    "\n",
    "def print_shape(df):\n",
    "    negative_examples, positive_examples = np.bincount(df[\"income\"])\n",
    "    print(\n",
    "        \"Data shape: {}, {} positive examples, {} negative examples\".format(\n",
    "            df.shape, positive_examples, negative_examples\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--train-test-split-ratio\", type=float, default=0.3)\n",
    "    args, _ = parser.parse_known_args()\n",
    "\n",
    "    # Get processor scrip arguments\n",
    "    args_iter = iter(sys.argv[1:])\n",
    "    script_args = dict(zip(args_iter, args_iter))\n",
    "    scheduler_ip = sys.argv[-1]\n",
    "\n",
    "    # Start the Dask cluster client\n",
    "    try:\n",
    "        client = Client(\"tcp://{ip}:8786\".format(ip=scheduler_ip))\n",
    "        logging.info(\"Printing cluster information: {}\".format(client))\n",
    "    except Exception as err:\n",
    "        logging.exception(err)\n",
    "\n",
    "    columns = [\n",
    "        \"age\",\n",
    "        \"education\",\n",
    "        \"major industry code\",\n",
    "        \"class of worker\",\n",
    "        \"num persons worked for employer\",\n",
    "        \"capital gains\",\n",
    "        \"capital losses\",\n",
    "        \"dividends from stocks\",\n",
    "        \"income\",\n",
    "    ]\n",
    "    class_labels = [\" - 50000.\", \" 50000+.\"]\n",
    "    input_data_path = \"s3://{}\".format(os.path.join(\n",
    "        script_args[\"s3_input_bucket\"],\n",
    "        script_args[\"s3_input_key_prefix\"],\n",
    "        \"census-income.csv\",\n",
    "    ))\n",
    "    \n",
    "    # Creating the necessary paths to save the output files\n",
    "    if not os.path.exists(\"/opt/ml/processing/train\"):\n",
    "        os.makedirs(\"/opt/ml/processing/train\")\n",
    "\n",
    "    if not os.path.exists(\"/opt/ml/processing/test\"):\n",
    "        os.makedirs(\"/opt/ml/processing/test\")\n",
    "\n",
    "    print(\"Reading input data from {}\".format(input_data_path))\n",
    "    df = pd.read_csv(input_data_path)\n",
    "    df = pd.DataFrame(data=df, columns=columns)\n",
    "    df.dropna(inplace=True)\n",
    "    df.drop_duplicates(inplace=True)\n",
    "    df.replace(class_labels, [0, 1], inplace=True)\n",
    "\n",
    "    negative_examples, positive_examples = np.bincount(df[\"income\"])\n",
    "    print(\n",
    "        \"Data after cleaning: {}, {} positive examples, {} negative examples\".format(\n",
    "            df.shape, positive_examples, negative_examples\n",
    "        )\n",
    "    )\n",
    "\n",
    "    split_ratio = args.train_test_split_ratio\n",
    "    print(\"Splitting data into train and test sets with ratio {}\".format(split_ratio))\n",
    "    X_train, X_test, y_train, y_test = train_test_split(\n",
    "        df.drop(\"income\", axis=1), df[\"income\"], test_size=split_ratio, random_state=0\n",
    "    )\n",
    "\n",
    "    preprocess = make_column_transformer(\n",
    "        (\n",
    "            KBinsDiscretizer(encode=\"onehot-dense\", n_bins=2),\n",
    "            [\"age\", \"num persons worked for employer\"],\n",
    "        ),\n",
    "        (\n",
    "            StandardScaler(),\n",
    "            [\"capital gains\", \"capital losses\", \"dividends from stocks\"],\n",
    "        ),\n",
    "        (\n",
    "            OneHotEncoder(sparse=False),\n",
    "            [\"education\", \"major industry code\", \"class of worker\"],\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    print(\"Running preprocessing and feature engineering transformations in Dask\")\n",
    "    with joblib.parallel_backend(\"dask\"):\n",
    "        train_features = preprocess.fit_transform(X_train)\n",
    "        test_features = preprocess.transform(X_test)\n",
    "\n",
    "    print(\"Train data shape after preprocessing: {}\".format(train_features.shape))\n",
    "    print(\"Test data shape after preprocessing: {}\".format(test_features.shape))\n",
    "\n",
    "    train_features_output_path = os.path.join(\n",
    "        \"/opt/ml/processing/train\", \"train_features.csv\"\n",
    "    )\n",
    "    train_labels_output_path = os.path.join(\n",
    "        \"/opt/ml/processing/train\", \"train_labels.csv\"\n",
    "    )\n",
    "\n",
    "    test_features_output_path = os.path.join(\n",
    "        \"/opt/ml/processing/test\", \"test_features.csv\"\n",
    "    )\n",
    "    test_labels_output_path = os.path.join(\"/opt/ml/processing/test\", \"test_labels.csv\")\n",
    "\n",
    "    print(\"Saving training features to {}\".format(train_features_output_path))\n",
    "    pd.DataFrame(train_features).to_csv(\n",
    "        train_features_output_path, header=False, index=False\n",
    "    )\n",
    "\n",
    "    print(\"Saving test features to {}\".format(test_features_output_path))\n",
    "    pd.DataFrame(test_features).to_csv(\n",
    "        test_features_output_path, header=False, index=False\n",
    "    )\n",
    "\n",
    "    print(\"Saving training labels to {}\".format(train_labels_output_path))\n",
    "    y_train.to_csv(train_labels_output_path, header=False, index=False)\n",
    "\n",
    "    print(\"Saving test labels to {}\".format(test_labels_output_path))\n",
    "    y_test.to_csv(test_labels_output_path, header=False, index=False)\n",
    "    upload_objects(\n",
    "        script_args[\"s3_output_bucket\"],\n",
    "        script_args[\"s3_output_key_prefix\"],\n",
    "        \"/opt/ml/processing/train/\",\n",
    "    )\n",
    "    upload_objects(\n",
    "        script_args[\"s3_output_bucket\"],\n",
    "        script_args[\"s3_output_key_prefix\"],\n",
    "        \"/opt/ml/processing/test/\",\n",
    "    )\n",
    "\n",
    "    # wait for the file creation\n",
    "    while attempts < attempts_counter:\n",
    "        if os.path.exists(train_features_output_path) and os.path.isfile(\n",
    "            train_features_output_path\n",
    "        ):\n",
    "            try:\n",
    "                # Calculate the processed dataset baseline statistics on the Dask cluster\n",
    "                dask_df = dd.read_csv(train_features_output_path)\n",
    "                dask_df = client.persist(dask_df)\n",
    "                baseline = dask_df.describe().compute()\n",
    "                print(baseline)\n",
    "                break\n",
    "\n",
    "            except:\n",
    "                time.sleep(2)\n",
    "    if attempts == attempts_counter:\n",
    "        raise Exception(\n",
    "            \"Output file {} couldn't be found\".format(train_features_output_path)\n",
    "        )\n",
    "\n",
    "    print(client)\n",
    "    sys.exit(os.EX_OK)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run a processing job using the Docker image and preprocessing script you just created. When invoking the `dask_processor.run()` function, pass the Amazon S3 input and output paths as arguments that are required by our preprocessing script to determine input and output location in Amazon S3. Here, you also specify the number of instances and instance type that will be used for the distributed Spark job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.processing import ProcessingInput, ScriptProcessor\n",
    "\n",
    "dask_processor = ScriptProcessor(\n",
    "    base_job_name=\"dask-preprocessor\",\n",
    "    image_uri=dask_repository_uri,\n",
    "    command=[\"/opt/program/bootstrap.py\"],\n",
    "    role=role,\n",
    "    instance_count=2,\n",
    "    instance_type=\"ml.m5.large\",\n",
    "    max_runtime_in_seconds=1200,\n",
    ")\n",
    "\n",
    "dask_processor.run(\n",
    "    code=\"preprocess.py\",\n",
    "    arguments=[\n",
    "        \"s3_input_bucket\",\n",
    "        bucket,\n",
    "        \"s3_input_key_prefix\",\n",
    "        input_prefix,\n",
    "        \"s3_output_bucket\",\n",
    "        bucket,\n",
    "        \"s3_output_key_prefix\",\n",
    "        input_preprocessed_prefix,\n",
    "    ],\n",
    "    logs=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inspect the preprocessed dataset\n",
    "Take a look at a few rows of the transformed dataset to make sure the preprocessing was successful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Top 5 rows from s3://{}/{}/train/'.format(bucket, input_preprocessed_prefix))\n",
    "!aws s3 cp --quiet s3://$bucket/$input_preprocessed_prefix/output/train_features.csv - | head -n5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you can use the output files of the transformation process as input to a training job and train a regression model."
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}