advanced_functionality/inference_pipeline_sparkml_xgboost_abalone/inference_pipeline_sparkml_xgboost_abalone.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature processing with Spark, training with XGBoost and deploying as Inference Pipeline\n",
    "\n",
    "Typically a Machine Learning (ML) process consists of few steps: gathering data with various ETL jobs, 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",
    "In many cases, when the trained model is used for processing real time or batch prediction requests, the model receives data in a format which needs to pre-processed (e.g. featurized) before it can be passed to the algorithm. In the following notebook, we will demonstrate how you can build your ML Pipeline leveraging Spark Feature Transformers and SageMaker XGBoost algorithm & after the model is trained, deploy the Pipeline (Feature Transformer and XGBoost) as an Inference Pipeline behind a single Endpoint for real-time inference and for batch inferences using Amazon SageMaker Batch Transform.\n",
    "\n",
    "In this notebook, we use Amazon Glue to run serverless Spark. Though the notebook demonstrates the end-to-end flow on a small dataset, the setup can be seamlessly used to scale to larger datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective: predict the age of an Abalone from its physical measurement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset is available from [UCI Machine Learning](https://archive.ics.uci.edu/ml/datasets/abalone). The aim for this task is to determine age of an Abalone (a kind of shellfish) from its physical measurements. At the core, it's a regression problem. The dataset contains several features - `sex` (categorical), `length` (continuous), `diameter` (continuous), `height` (continuous), `whole_weight` (continuous), `shucked_weight` (continuous), `viscera_weight` (continuous), `shell_weight` (continuous) and `rings` (integer).Our goal is to predict the variable `rings` which is a good approximation for age (age is `rings` + 1.5). \n",
    "\n",
    "We'll use SparkML to process the dataset (apply one or many feature transformers) and upload the transformed dataset to S3 so that it can be used for training with XGBoost."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Methodologies\n",
    "The Notebook consists of a few high-level steps:\n",
    "\n",
    "* Using AWS Glue for executing the SparkML feature processing job.\n",
    "* Using SageMaker XGBoost to train on the processed dataset produced by SparkML job.\n",
    "* Building an Inference Pipeline consisting of SparkML & XGBoost models for a realtime inference endpoint.\n",
    "* Building an Inference Pipeline consisting of SparkML & XGBoost models for a single Batch Transform job."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using AWS Glue for executing the SparkML job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll be running the SparkML job using [AWS Glue](https://aws.amazon.com/glue). AWS Glue is a serverless ETL service which can be used to execute standard Spark/PySpark jobs. Glue currently only supports `Python 2.7`, hence we'll write the script in `Python 2.7`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Permission setup for invoking AWS Glue from this Notebook\n",
    "In order to enable this Notebook to run AWS Glue jobs, we need to add one additional permission to the default execution role of this notebook. We will be using SageMaker Python SDK to retrieve the default execution role and then you have to go to [IAM Dashboard](https://console.aws.amazon.com/iam/home) to edit the Role to add AWS Glue specific permission. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finding out the current execution role of the Notebook\n",
    "We are using SageMaker Python SDK to retrieve the current role for this Notebook which needs to be enhanced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import SageMaker Python SDK to get the Session and execution_role\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "sess = sagemaker.Session()\n",
    "role = get_execution_role()\n",
    "print(role[role.rfind('/') + 1:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding AWS Glue as an additional trusted entity to this role\n",
    "This step is needed if you want to pass the execution role of this Notebook while calling Glue APIs as well without creating an additional **Role**. If you have not used AWS Glue before, then this step is mandatory. \n",
    "\n",
    "If you have used AWS Glue previously, then you should have an already existing role that can be used to invoke Glue APIs. In that case, you can pass that role while calling Glue (later in this notebook) and skip this next step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the IAM dashboard, please click on **Roles** on the left sidenav and search for this Role. Once the Role appears, click on the Role to go to its **Summary** page. Click on the **Trust relationships** tab on the **Summary** page to add AWS Glue as an additional trusted entity. \n",
    "\n",
    "Click on **Edit trust relationship** and replace the JSON with this JSON.\n",
    "```\n",
    "{\n",
    "  \"Version\": \"2012-10-17\",\n",
    "  \"Statement\": [\n",
    "    {\n",
    "      \"Effect\": \"Allow\",\n",
    "      \"Principal\": {\n",
    "        \"Service\": [\n",
    "          \"sagemaker.amazonaws.com\",\n",
    "          \"glue.amazonaws.com\"\n",
    "        ]\n",
    "      },\n",
    "      \"Action\": \"sts:AssumeRole\"\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```\n",
    "Once this is complete, click on **Update Trust Policy** and you are done."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading dataset and uploading to S3\n",
    "SageMaker team has downloaded the dataset from UCI and uploaded to one of the S3 buckets in our account. In this Notebook, we will download from that bucket and upload to your bucket so that AWS Glue can access the data. The default AWS Glue permissions we just added expects the data to be present in a bucket with the string `aws-glue`. Hence, after we download the dataset, we will create an S3 bucket in your account with a valid name and then upload the data to S3. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://s3-us-west-2.amazonaws.com/sparkml-mleap/data/abalone/abalone.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating an S3 bucket and uploading this dataset\n",
    "Next we will create an S3 bucket with the `aws-glue` string in the name and upload this data to the S3 bucket. In case you want to use some existing bucket to run your Spark job via AWS Glue, you can use that bucket to upload your data provided the `Role` has access permission to upload and download from that bucket.\n",
    "\n",
    "Once the bucket is created, the following cell would also update the `abalone.csv` file downloaded locally to this bucket under the `input/abalone` prefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import botocore\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "boto_session = sess.boto_session\n",
    "s3 = boto_session.resource('s3')\n",
    "account = boto_session.client('sts').get_caller_identity()['Account']\n",
    "region = boto_session.region_name\n",
    "default_bucket = 'aws-glue-{}-{}'.format(account, region)\n",
    "\n",
    "try:\n",
    "    if region == 'us-east-1':\n",
    "        s3.create_bucket(Bucket=default_bucket)\n",
    "    else:\n",
    "        s3.create_bucket(Bucket=default_bucket, CreateBucketConfiguration={'LocationConstraint': region})\n",
    "except ClientError as e:\n",
    "    error_code = e.response['Error']['Code']\n",
    "    message = e.response['Error']['Message']\n",
    "    if error_code == 'BucketAlreadyOwnedByYou':\n",
    "        print ('A bucket with the same name already exists in your account - using the same bucket.')\n",
    "        pass        \n",
    "\n",
    "# Uploading the training data to S3\n",
    "sess.upload_data(path='abalone.csv', bucket=default_bucket, key_prefix='input/abalone')    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing the feature processing script using SparkML\n",
    "\n",
    "The code for feature transformation using SparkML can be found in `abalone_processing.py` file written in the same directory. You can go through the code itself to see how it is using standard SparkML constructs to define the Pipeline for featurizing the data.\n",
    "\n",
    "Once the Spark ML Pipeline `fit` and `transform` is done, we are splitting our dataset into 80-20 train & validation as part of the script and uploading to S3 so that it can be used with XGBoost for training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Serializing the trained Spark ML Model with [MLeap](https://github.com/combust/mleap)\n",
    "Apache Spark is best suited batch processing workloads. In order to use the Spark ML model we trained for low latency inference, we need to use the MLeap library to serialize it to an MLeap bundle and later use the [SageMaker SparkML Serving](https://github.com/aws/sagemaker-sparkml-serving-container) to perform realtime and batch inference. \n",
    "\n",
    "By using the `SerializeToBundle()` method from MLeap in the script, we are serializing the ML Pipeline into an MLeap bundle and uploading to S3 in `tar.gz` format as SageMaker expects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uploading the code and other dependencies to S3 for AWS Glue\n",
    "Unlike SageMaker, in order to run your code in AWS Glue, we do not need to prepare a Docker image. We can upload the code and dependencies directly to S3 and pass those locations while invoking the Glue job."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload the SparkML script to S3\n",
    "We will be uploading the `abalone_processing.py` script to S3 now so that Glue can use it to run the PySpark job. You can replace it with your own script if needed. If your code has multiple files, you need to zip those files and upload to S3 instead of uploading a single file like it's being done here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "script_location = sess.upload_data(path='abalone_processing.py', bucket=default_bucket, key_prefix='codes')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload MLeap dependencies to S3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our job, we will also have to pass MLeap dependencies to Glue. MLeap is an additional library we are using which does not come bundled with default Spark.\n",
    "\n",
    "Similar to most of the packages in the Spark ecosystem, MLeap is also implemented as a Scala package with a front-end wrapper written in Python so that it can be used from PySpark. We need to make sure that the MLeap Python library as well as the JAR is available within the Glue job environment. In the following cell, we will download the MLeap Python dependency & JAR from a SageMaker hosted bucket and upload to the S3 bucket we created above in your account. \n",
    "\n",
    "If you are using some other Python libraries like `nltk` in your code, you need to download the wheel file from PyPI and upload to S3 in the same way. At this point, Glue only supports passing pure Python libraries in this way (e.g. you can not pass `Pandas` or `OpenCV`). However you can use `NumPy` & `SciPy` without having to pass these as packages because these are pre-installed in the Glue environment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://s3-us-west-2.amazonaws.com/sparkml-mleap/0.9.6/python/python.zip\n",
    "!wget https://s3-us-west-2.amazonaws.com/sparkml-mleap/0.9.6/jar/mleap_spark_assembly.jar    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "python_dep_location = sess.upload_data(path='python.zip', bucket=default_bucket, key_prefix='dependencies/python')\n",
    "jar_dep_location = sess.upload_data(path='mleap_spark_assembly.jar', bucket=default_bucket, key_prefix='dependencies/jar')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining output locations for the data and model\n",
    "Next we define the output location where the transformed dataset should be uploaded. We are also specifying a model location where the MLeap serialized model would be updated. This locations should be consumed as part of the Spark script using `getResolvedOptions` method of AWS Glue library (see `abalone_processing.py` for details).\n",
    "\n",
    "By designing our code in that way, we can re-use these variables as part of other SageMaker operations from this Notebook (details below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import gmtime, strftime\n",
    "import time\n",
    "\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "\n",
    "# Input location of the data, We uploaded our train.csv file to input key previously\n",
    "s3_input_bucket = default_bucket\n",
    "s3_input_key_prefix = 'input/abalone'\n",
    "\n",
    "# Output location of the data. The input data will be split, transformed, and \n",
    "# uploaded to output/train and output/validation\n",
    "s3_output_bucket = default_bucket\n",
    "s3_output_key_prefix = timestamp_prefix + '/abalone'\n",
    "\n",
    "# the MLeap serialized SparkML model will be uploaded to output/mleap\n",
    "s3_model_bucket = default_bucket\n",
    "s3_model_key_prefix = s3_output_key_prefix + '/mleap'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calling Glue APIs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we'll be creating Glue client via Boto so that we can invoke the `create_job` API of Glue. `create_job` API will create a job definition which can be used to execute your jobs in Glue. The job definition created here is mutable. While creating the job, we are also passing the code location as well as the dependencies location to Glue.\n",
    "\n",
    "`AllocatedCapacity` parameter controls the hardware resources that Glue will use to execute this job. It is measures in units of `DPU`. For more information on `DPU`, please see [here](https://docs.aws.amazon.com/glue/latest/dg/add-job.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glue_client = boto_session.client('glue')\n",
    "job_name = 'sparkml-abalone-' + timestamp_prefix\n",
    "response = glue_client.create_job(\n",
    "    Name=job_name,\n",
    "    Description='PySpark job to featurize the Abalone dataset',\n",
    "    Role=role, # you can pass your existing AWS Glue role here if you have used Glue before\n",
    "    ExecutionProperty={\n",
    "        'MaxConcurrentRuns': 1\n",
    "    },\n",
    "    Command={\n",
    "        'Name': 'glueetl',\n",
    "        'ScriptLocation': script_location\n",
    "    },\n",
    "    DefaultArguments={\n",
    "        '--job-language': 'python',\n",
    "        '--extra-jars' : jar_dep_location,\n",
    "        '--extra-py-files': python_dep_location\n",
    "    },\n",
    "    AllocatedCapacity=5,\n",
    "    Timeout=60,\n",
    ")\n",
    "glue_job_name = response['Name']\n",
    "print(glue_job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The aforementioned job will be executed now by calling `start_job_run` API. This API creates an immutable run/execution corresponding to the job definition created above. We will require the `job_run_id` for the particular job execution to check for status. We'll pass the data and model locations as part of the job execution parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_run_id = glue_client.start_job_run(JobName=job_name,\n",
    "                                       Arguments = {\n",
    "                                        '--S3_INPUT_BUCKET': s3_input_bucket,\n",
    "                                        '--S3_INPUT_KEY_PREFIX': s3_input_key_prefix,\n",
    "                                        '--S3_OUTPUT_BUCKET': s3_output_bucket,\n",
    "                                        '--S3_OUTPUT_KEY_PREFIX': s3_output_key_prefix,\n",
    "                                        '--S3_MODEL_BUCKET': s3_model_bucket,\n",
    "                                        '--S3_MODEL_KEY_PREFIX': s3_model_key_prefix\n",
    "                                       })['JobRunId']\n",
    "print(job_run_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checking Glue job status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will check for the job status to see if it has `succeeded`, `failed` or `stopped`. Once the job is succeeded, we have the transformed data into S3 in CSV format which we can use with XGBoost for training. If the job fails, you can go to [AWS Glue console](https://us-west-2.console.aws.amazon.com/glue/home), click on **Jobs** tab on the left, and from the page, click on this particular job and you will be able to find the CloudWatch logs (the link under **Logs**) link for these jobs which can help you to see what exactly went wrong in the job execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_run_status = glue_client.get_job_run(JobName=job_name,RunId=job_run_id)['JobRun']['JobRunState']\n",
    "while job_run_status not in ('FAILED', 'SUCCEEDED', 'STOPPED'):\n",
    "    job_run_status = glue_client.get_job_run(JobName=job_name,RunId=job_run_id)['JobRun']['JobRunState']\n",
    "    print (job_run_status)\n",
    "    time.sleep(30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using SageMaker XGBoost to train on the processed dataset produced by SparkML job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will use SageMaker XGBoost algorithm to train on this dataset. We already know the S3 location\n",
    "where the preprocessed training data was uploaded as part of the Glue job."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We need to retrieve the XGBoost algorithm image\n",
    "We will retrieve the XGBoost built-in algorithm image so that it can leveraged for the training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
    "\n",
    "training_image = get_image_uri(sess.boto_region_name, 'xgboost', repo_version=\"latest\")\n",
    "print (training_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next XGBoost model parameters and dataset details will be set properly\n",
    "We have parameterized this Notebook so that the same data location which was used in the PySpark script can now be passed to XGBoost Estimator as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_train_data = 's3://{}/{}/{}'.format(s3_output_bucket, s3_output_key_prefix, 'train')\n",
    "s3_validation_data = 's3://{}/{}/{}'.format(s3_output_bucket, s3_output_key_prefix, 'validation')\n",
    "s3_output_location = 's3://{}/{}/{}'.format(s3_output_bucket, s3_output_key_prefix, 'xgboost_model')\n",
    "\n",
    "xgb_model = sagemaker.estimator.Estimator(training_image,\n",
    "                                         role, \n",
    "                                         train_instance_count=1, \n",
    "                                         train_instance_type='ml.m4.xlarge',\n",
    "                                         train_volume_size = 20,\n",
    "                                         train_max_run = 3600,\n",
    "                                         input_mode= 'File',\n",
    "                                         output_path=s3_output_location,\n",
    "                                         sagemaker_session=sess)\n",
    "\n",
    "xgb_model.set_hyperparameters(objective = \"reg:linear\",\n",
    "                              eta = .2,\n",
    "                              gamma = 4,\n",
    "                              max_depth = 5,\n",
    "                              num_round = 10,\n",
    "                              subsample = 0.7,\n",
    "                              silent = 0,\n",
    "                              min_child_weight = 6)\n",
    "\n",
    "train_data = sagemaker.session.s3_input(s3_train_data, distribution='FullyReplicated', \n",
    "                        content_type='text/csv', s3_data_type='S3Prefix')\n",
    "validation_data = sagemaker.session.s3_input(s3_validation_data, distribution='FullyReplicated', \n",
    "                             content_type='text/csv', s3_data_type='S3Prefix')\n",
    "\n",
    "data_channels = {'train': train_data, 'validation': validation_data}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finally XGBoost training will be performed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb_model.fit(inputs=data_channels, logs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building an Inference Pipeline consisting of SparkML & XGBoost models for a realtime inference endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will proceed with deploying the models in SageMaker to create an Inference Pipeline. You can create an Inference Pipeline with upto five containers.\n",
    "\n",
    "Deploying a model in SageMaker requires two components:\n",
    "\n",
    "* Docker image residing in ECR.\n",
    "* Model artifacts residing in S3.\n",
    "\n",
    "**SparkML**\n",
    "\n",
    "For SparkML, Docker image for MLeap based SparkML serving is provided by SageMaker team. For more information on this, please see [SageMaker SparkML Serving](https://github.com/aws/sagemaker-sparkml-serving-container). MLeap serialized SparkML model was uploaded to S3 as part of the SparkML job we executed in AWS Glue.\n",
    "\n",
    "**XGBoost**\n",
    "\n",
    "For XGBoost, we will use the same Docker image we used for training. The model artifacts for XGBoost was uploaded as part of the training job we just ran."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Passing the schema of the payload via environment variable\n",
    "SparkML serving container needs to know the schema of the request that'll be passed to it while calling the `predict` method. In order to alleviate the pain of not having to pass the schema with every request, `sagemaker-sparkml-serving` allows you to pass it via an environment variable while creating the model definitions. This schema definition will be required in our next step for creating a model.\n",
    "\n",
    "We will see later that you can overwrite this schema on a per request basis by passing it as part of the individual request payload as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "schema = {\n",
    "    \"input\": [\n",
    "        {\n",
    "            \"name\": \"sex\",\n",
    "            \"type\": \"string\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"length\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"diameter\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"height\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"whole_weight\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"shucked_weight\",\n",
    "            \"type\": \"double\"\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"viscera_weight\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"shell_weight\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "    ],\n",
    "    \"output\": \n",
    "        {\n",
    "            \"name\": \"features\",\n",
    "            \"type\": \"double\",\n",
    "            \"struct\": \"vector\"\n",
    "        }\n",
    "}\n",
    "schema_json = json.dumps(schema)\n",
    "print(schema_json)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a `PipelineModel` which comprises of the SparkML and XGBoost model in the right order\n",
    "\n",
    "Next we'll create a SageMaker `PipelineModel` with SparkML and XGBoost.The `PipelineModel` will ensure that both the containers get deployed behind a single API endpoint in the correct order. The same model would later be used for Batch Transform as well to ensure that a single job is sufficient to do prediction against the Pipeline. \n",
    "\n",
    "Here, during the `Model` creation for SparkML, we will pass the schema definition that we built in the previous cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.model import Model\n",
    "from sagemaker.pipeline import PipelineModel\n",
    "from sagemaker.sparkml.model import SparkMLModel\n",
    "\n",
    "sparkml_data = 's3://{}/{}/{}'.format(s3_model_bucket, s3_model_key_prefix, 'model.tar.gz')\n",
    "# passing the schema defined above by using an environment variable that sagemaker-sparkml-serving understands\n",
    "sparkml_model = SparkMLModel(model_data=sparkml_data, env={'SAGEMAKER_SPARKML_SCHEMA' : schema_json})\n",
    "xgb_model = Model(model_data=xgb_model.model_data, image=training_image)\n",
    "\n",
    "model_name = 'inference-pipeline-' + timestamp_prefix\n",
    "sm_model = PipelineModel(name=model_name, role=role, models=[sparkml_model, xgb_model])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploying the `PipelineModel` to an endpoint for realtime inference\n",
    "Next we will deploy the model we just created with the `deploy()` method to start an inference endpoint and we will send some requests to the endpoint to verify that it works as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = 'inference-pipeline-ep-' + timestamp_prefix\n",
    "sm_model.deploy(initial_instance_count=1, instance_type='ml.c4.xlarge', endpoint_name=endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoking the newly created inference endpoint with a payload to transform the data\n",
    "Now we will invoke the endpoint with a valid payload that SageMaker SparkML Serving can recognize. There are three ways in which input payload can be passed to the request:\n",
    "\n",
    "* Pass it as a valid CSV string. In this case, the schema passed via the environment variable will be used to determine the schema. For CSV format, every column in the input has to be a basic datatype (e.g. int, double, string) and it can not be a Spark `Array` or `Vector`.\n",
    "\n",
    "* Pass it as a valid JSON string. In this case as well, the schema passed via the environment variable will be used to infer the schema. With JSON format, every column in the input can be a basic datatype or a Spark `Vector` or `Array` provided that the corresponding entry in the schema mentions the correct value.\n",
    "\n",
    "* Pass the request in JSON format along with the schema and the data. In this case, the schema passed in the payload will take precedence over the one passed via the environment variable (if any)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Passing the payload in CSV format\n",
    "We will first see how the payload can be passed to the endpoint in CSV format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.predictor import json_serializer, csv_serializer, json_deserializer, RealTimePredictor\n",
    "from sagemaker.content_types import CONTENT_TYPE_CSV, CONTENT_TYPE_JSON\n",
    "payload = \"F,0.515,0.425,0.14,0.766,0.304,0.1725,0.255\"\n",
    "predictor = RealTimePredictor(endpoint=endpoint_name, sagemaker_session=sess, serializer=csv_serializer,\n",
    "                                content_type=CONTENT_TYPE_CSV, accept=CONTENT_TYPE_CSV)\n",
    "print(predictor.predict(payload))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Passing the payload in JSON format\n",
    "We will now pass a different payload in JSON format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = {\"data\": [\"F\",0.515,0.425,0.14,0.766,0.304,0.1725,0.255]}\n",
    "predictor = RealTimePredictor(endpoint=endpoint_name, sagemaker_session=sess, serializer=json_serializer,\n",
    "                                content_type=CONTENT_TYPE_JSON, accept=CONTENT_TYPE_CSV)\n",
    "\n",
    "print(predictor.predict(payload))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [Optional] Passing the payload with both schema and the data\n",
    "Next we will pass the input payload comprising of both the schema and the data. If you notice carefully, this schema will be slightly different than what we have passed via the environment variable. The locations of `length` and `sex` column have been swapped and so the data. The server now parses the payload with this schema and works properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = {\n",
    "    \"schema\": {\n",
    "        \"input\": [\n",
    "        {\n",
    "            \"name\": \"length\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"sex\",\n",
    "            \"type\": \"string\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"diameter\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"height\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"whole_weight\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"shucked_weight\",\n",
    "            \"type\": \"double\"\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"viscera_weight\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "        {\n",
    "            \"name\": \"shell_weight\",\n",
    "            \"type\": \"double\"\n",
    "        }, \n",
    "    ],\n",
    "    \"output\": \n",
    "        {\n",
    "            \"name\": \"features\",\n",
    "            \"type\": \"double\",\n",
    "            \"struct\": \"vector\"\n",
    "        }\n",
    "    },\n",
    "    \"data\": [0.515,\"F\",0.425,0.14,0.766,0.304,0.1725,0.255]\n",
    "}\n",
    "\n",
    "predictor = RealTimePredictor(endpoint=endpoint_name, sagemaker_session=sess, serializer=json_serializer,\n",
    "                                content_type=CONTENT_TYPE_JSON, accept=CONTENT_TYPE_CSV)\n",
    "\n",
    "print(predictor.predict(payload))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Optional] Deleting the Endpoint\n",
    "If you do not plan to use this endpoint, then it is a good practice to delete the endpoint so that you do not incur the cost of running it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client = boto_session.client('sagemaker')\n",
    "sm_client.delete_endpoint(EndpointName=endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building an Inference Pipeline consisting of SparkML & XGBoost models for a single Batch Transform job\n",
    "SageMaker Batch Transform also supports chaining multiple containers together when deploying an Inference Pipeline and performing a single batch transform jobs to transform your data for a batch use-case similar to the real-time use-case we have seen above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing data for Batch Transform\n",
    "Batch Transform requires data in the same format described above, with one CSV or JSON being per line. For this Notebook, SageMaker team has created a sample input in CSV format which Batch Transform can process. The input is basically a similar CSV file to the training file with only difference is that it does not contain the label (``rings``) field.\n",
    "\n",
    "Next we will download a sample of this data from one of the SageMaker buckets (named `batch_input_abalone.csv`) and upload to your S3 bucket. We will also inspect first five rows of the data post downloading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://s3-us-west-2.amazonaws.com/sparkml-mleap/data/batch_input_abalone.csv\n",
    "!printf \"\\n\\nShowing first five lines\\n\\n\"    \n",
    "!head -n 5 batch_input_abalone.csv \n",
    "!printf \"\\n\\nAs we can see, it is identical to the training file apart from the label being absent here.\\n\\n\"  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_input_loc = sess.upload_data(path='batch_input_abalone.csv', bucket=default_bucket, key_prefix='batch')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoking the Transform API to create a Batch Transform job\n",
    "Next we will create a Batch Transform job using the `Transformer` class from Python SDK to create a Batch Transform job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data_path = 's3://{}/{}/{}'.format(default_bucket, 'batch', 'batch_input_abalone.csv')\n",
    "output_data_path = 's3://{}/{}/{}'.format(default_bucket, 'batch_output/abalone', timestamp_prefix)\n",
    "job_name = 'serial-inference-batch-' + timestamp_prefix\n",
    "transformer = sagemaker.transformer.Transformer(\n",
    "    # This was the model created using PipelineModel and it contains feature processing and XGBoost\n",
    "    model_name = model_name,\n",
    "    instance_count = 1,\n",
    "    instance_type = 'ml.m4.xlarge',\n",
    "    strategy = 'SingleRecord',\n",
    "    assemble_with = 'Line',\n",
    "    output_path = output_data_path,\n",
    "    base_transform_job_name='serial-inference-batch',\n",
    "    sagemaker_session=sess,\n",
    "    accept = CONTENT_TYPE_CSV\n",
    ")\n",
    "transformer.transform(data = input_data_path,\n",
    "                      job_name = job_name,\n",
    "                      content_type = CONTENT_TYPE_CSV, \n",
    "                      split_type = 'Line')\n",
    "transformer.wait()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python2",
   "language": "python",
   "name": "conda_python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}