{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f70ae5d7",
   "metadata": {},
   "source": [
    "# 02 - Continuous stirred tank reactor (CSTR)\n",
    "\n",
    "**Amazon SageMaker kernel**: conda_pytorch_p36\n",
    "\n",
    "## Sending insights to the edge\n",
    "\n",
    "In this notebook we will learn how to take the model we built in the previous notebook and deploy it to an edge device that is running AWS IoT Greengrass to allow process engineers to leverage machine learning model in their daily operations. This could be either by running the process model to predict likely outcomes for a set of process inputs (what-if analysis) or by solving an optimization problem to recommend process inputs for a desired output (prescriptive analytics)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "413ef03d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import time\n",
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "# Insert the name of your S3 bucket from CloudFormation outputs\n",
    "bucket = \"\" \n",
    "role = sagemaker.get_execution_role()\n",
    "sess = sagemaker.Session()\n",
    "region = \"us-east-1\"\n",
    "prefix = \"models/pytorch\"\n",
    "\n",
    "assert bucket, \"ERROR: Insert the name of your S3 bucket before moving on\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a856fb",
   "metadata": {},
   "source": [
    "### Compile using SageMaker Neo\n",
    "\n",
    "Neo is a capability of Amazon SageMaker that allows machine learning models to train once and run anywhere in the cloud and at the edge. SageMaker Neo optimizes the machine learning model for inference. The first cell will load the required libraries and will get the Amazon SageMaker session and execution role to be passed to the compilation jobs later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "676c74c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "from sagemaker.predictor import Predictor\n",
    "\n",
    "sm_client = boto3.client('sagemaker')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faff9dcb",
   "metadata": {},
   "source": [
    "The next cell contains the configuration parameters for the compilation job. We need to define the framework in which the model is built (PyTorch), its version, and the shape of the input tensor. We also need to define the paths to the S3 buckets where we will store: (1) the compiled model (machine lerarning model optimized for execution with SageMaker Neo), and (2) packaged model (package that contains the model and all the utilities required for it to run on AWS IoT Greengrass).\n",
    "\n",
    "The following cells contain the instructions to compile and export the process model `cstr-model`, but you can follow the same approach to compile and export the optimization surrogate model `cstr-optimizer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f0007bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "prefix = \"models/pytorch\"\n",
    "data_shape = '{\"input0\":[1,2]}'\n",
    "key = \"workshop\" # Use this key to facilitate complete rebuilds of all of the compilation/packaging\n",
    "\n",
    "# Target definition\n",
    "target_device = \"lambda\"\n",
    "framework = \"PYTORCH\"\n",
    "framework_version = \"1.7\"\n",
    "model_name = \"cstr-model\"\n",
    "model_version = \"1.0.0\"\n",
    "compiled_model_path = \"s3://{:s}/compiled\".format(bucket)\n",
    "packaged_model_path = \"s3://{:s}/packages\".format(bucket)\n",
    "model_s3_uri = (\"s3://{:s}/{:s}/\" + model_name + \".tar.gz\").format(bucket, prefix)\n",
    "\n",
    "# Greengrass V2 Component Config\n",
    "component_config = \"{\\\"ComponentName\\\":\\\"\" + model_name + \"-\" + target_device + \"-component\\\", \\\"ComponentVersion\\\":\\\"\" + model_version + \"\\\"}\"\n",
    "\n",
    "# Specify a name for your edge packaging job.\n",
    "edge_packaging_job_name = model_name + \"NeoPackage\" + key\n",
    "\n",
    "# Specify a name for the model compile... \n",
    "compilation_job_name = name_from_base(model_name + \"NeoCompile\" + key)\n",
    "\n",
    "print(\"Model name: \" + model_name + \" Version: \" + model_version)\n",
    "print(\"Target Device: \" + target_device)\n",
    "print(\"Framework: \" + framework + \" Version: \" + framework_version)\n",
    "print(\"Model S3 URI: \" + model_s3_uri)\n",
    "print(\"Component config: \" + component_config)\n",
    "print(\"Edge package job name: \" + edge_packaging_job_name)\n",
    "print(\"Compile job name: \" + compilation_job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b4798e2",
   "metadata": {},
   "source": [
    "#### Create a PyTorch model object in Amazon SageMaker\n",
    "\n",
    "The first step will take the model we built in the previous notebook and put it in a `PyTorchModel` object that also indicates the framework, and the SageMaker session that will be used to compile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969a23de",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Creating PyTorch model from content: \" + model_s3_uri + \"...\")\n",
    "sagemaker_model = PyTorchModel(\n",
    "    model_data = model_s3_uri,\n",
    "    predictor_cls = Predictor,\n",
    "    framework_version = framework_version,\n",
    "    role = role,\n",
    "    sagemaker_session = sess, \n",
    "    entry_point = None,\n",
    "    py_version = \"py3\",\n",
    "    env = {\"MMS_DEFAULT_RESPONSE_TIMEOUT\": \"900\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9052432f",
   "metadata": {},
   "source": [
    "#### Compile model for device\n",
    "\n",
    "Once the model has been defined, we will run a compilation job for a Lambda target device. This choice allows for portability as the machine learning model can be executed in the cloud (as a Lambda function) or on the edge (as an AWS IoT Greengrass component).\n",
    "\n",
    "This step can take around 4 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "affb95c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Compiling model for device: \" + target_device + \"...\")\n",
    "compiled_model = sagemaker_model.compile(\n",
    "    target_instance_family = target_device,\n",
    "    job_name = compilation_job_name,\n",
    "    input_shape = data_shape,\n",
    "    role = role,\n",
    "    framework = framework.lower(),\n",
    "    framework_version = framework_version,\n",
    "    output_path = compiled_model_path,\n",
    ")\n",
    "compiled_model.name = \"compiled-cstr-model-\" + target_device + \"-\" + key\n",
    "print(\"\")\n",
    "compiled_model_file = compiled_model_path + \"/\" + model_name + \"-\" + target_device + \".tar.gz\"\n",
    "print(\"Compiled model: \" + compiled_model_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f110931b",
   "metadata": {},
   "source": [
    "### (OPTIONAL) Test the SageMaker Neo optimized model prior to deployment\n",
    "\n",
    "The machine learning model now has been optimized to run as a Lambda. Are we sure that the model is still working as it used to? Let's test it!\n",
    "\n",
    "For testing we will use the Python [DLR library](https://pypi.org/project/dlr/), which is a compact runtime for deep learning models compiled by AWS SageMaker Neo, TVM, or Treelite.\n",
    "\n",
    "We will run the following steps:\n",
    "\n",
    "- Download the compiled model from S3 as a TGZ file\n",
    "\n",
    "- Unzip the content of the TGZ file in the `local_dirname` directory\n",
    "\n",
    "- Instantiate a model from the configuration in `local_dirname`\n",
    "\n",
    "- Evaluate the model for the same configuration we tested in the previous notebook: `F` is 12 [l/h] and `Q_dot` is -10 [kW]\n",
    "\n",
    "- Compare the results to the ones obtained in the previous notebook. They will be the same"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c0db308",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_dirname = \"dlr_compiled\"\n",
    "\n",
    "print(\"Pulling compiled model from S3: \" + compiled_model_file)\n",
    "! aws s3 cp {compiled_model_file} .\n",
    "    \n",
    "print(\"Expanding compiled model into our notebook...\")\n",
    "! rm -rf ./{local_dirname} ; mkdir {local_dirname} ; cd {local_dirname} ; tar xzpf ../{model_name}-{target_device}.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d8ff5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import dlr\n",
    "\n",
    "model = dlr.DLRModel(local_dirname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bab50c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_dlr_model(model, manipulated_vars):\n",
    "    \"\"\"\n",
    "    Return predictions from machine learning model compiled\n",
    "    with SageMaker Neo that are loaded in Python with the DLR\n",
    "    library\n",
    "    \"\"\"\n",
    "    X_test = manipulated_vars[[\"F\", \"Q_dot\"]].values\n",
    "    X_test[:,0] = (X_test[:,0]-5.0)/95.0\n",
    "    X_test[:,1] = (X_test[:,1]+5000.0)/5000.0\n",
    "    X_test = X_test.astype(np.float32)\n",
    "    Y_test = model.run(X_test)[0]\n",
    "    Y_test[:,2] = (Y_test[:,2]*25.0)+125.0\n",
    "    prediction = pd.DataFrame({\"F\" : manipulated_vars[\"F\"].values.round(2),\n",
    "                              \"Q_dot\": manipulated_vars[\"Q_dot\"].values.round(1),\n",
    "                              \"C_a\": Y_test[:,0].round(4),\n",
    "                              \"C_b\": Y_test[:,1].round(4),\n",
    "                              \"T_K\": Y_test[:,2].round(1)})\n",
    "    \n",
    "    return prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f524a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate predictions\n",
    "manipulated_vars = pd.DataFrame({\"F\": [12.0], \"Q_dot\": [-10.0]})\n",
    "\n",
    "measurements = evaluate_dlr_model(model, manipulated_vars)\n",
    "\n",
    "print(\"Measurements from CSTR:\")\n",
    "print(\"Flow rate: {:.2f} l/h\".format(measurements[\"F\"][0]))\n",
    "print(\"Concentration of reactant A: {:.4f} mol A/l\".format(measurements[\"C_a\"][0]))\n",
    "print(\"Concentration of reactant B: {:.4f} mol B/l\".format(measurements[\"C_b\"][0]))\n",
    "print(\"Temperature in the cooling jacket: {:.2f} C\".format(measurements[\"T_K\"][0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0edc255",
   "metadata": {},
   "source": [
    "### Package model to deploy as AWS IoT Greengrass component\n",
    "\n",
    "This cell will create a Greengrass v2 component. Make sure to use the `GreengrassV2Component` option for the `PresetDeploymentType`. Once this task is completed you should see a brand new component available in AWS IoT Greengrass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae29f5d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Packaging model for Greengrass Edge (\" + target_device + \")...\")\n",
    "print(\"Greengrass V2 Component Config: \" + component_config)\n",
    "sm_client.create_edge_packaging_job(\n",
    "    EdgePackagingJobName=edge_packaging_job_name,\n",
    "    CompilationJobName=compilation_job_name,\n",
    "    RoleArn=role,\n",
    "    ModelName=model_name,\n",
    "    ModelVersion=model_version,\n",
    "    OutputConfig={\n",
    "        \"S3OutputLocation\": packaged_model_path,\n",
    "        \"PresetDeploymentType\": \"GreengrassV2Component\",\n",
    "        \"PresetDeploymentConfig\": component_config\n",
    "        }\n",
    "    )\n",
    "\n",
    "print(\"DONE. GGV2 Component Created: \" + component_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2e43522",
   "metadata": {},
   "source": [
    "## Custom components: Using the machine learning model at the edge\n",
    "\n",
    "Once the model is ready to be deployed to the edge, we need to define the interfaces and modes of operation of the model. We have defined scripts that establish MQTT interfaces for both the prediction and the optimization cases. These will be added to your Greengrass deployment as components in the next section of the workshop but, for now, we will upload these scripts to S3 to they can used in the definition of the Greengrass components\n",
    "\n",
    "### Prediction\n",
    "\n",
    "We need to copy 5 files from our SageMaker notebook to our S3 bucket. The script that contains the logic for running the prediction is `cstr_predictor.py`. Feel free to read it to learn how to interact with your ML model at the edge with gRPC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8ce8427",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload scripts to run machine learning inference (prediction)\n",
    "artifact_dir = \"s3://\" + bucket + \"/artifacts/CSTRModelPredictor/1.0.0/\"\n",
    "\n",
    "# Copy the files from our notebook to S3\n",
    "print(\"Copying files from the notebook to S3\")\n",
    "! aws s3 cp ./predictor/agent_pb2.py {artifact_dir}\n",
    "! aws s3 cp ./predictor/agent_pb2_grpc.py {artifact_dir}\n",
    "! aws s3 cp ./predictor/cstr_predictor.py {artifact_dir}\n",
    "! aws s3 cp ./predictor/launch_predictor.sh {artifact_dir}\n",
    "! aws s3 cp ./predictor/stop.sh {artifact_dir}\n",
    "\n",
    "# Note the S3 file names - we'll need them later\n",
    "print(\"\")\n",
    "print(\"S3 Files for custom inference component artifacts:\")\n",
    "print(artifact_dir + \"agent_pb2.py\")\n",
    "print(artifact_dir + \"agent_pb2_grpc.py\")\n",
    "print(artifact_dir + \"cstr_predictor.py\")\n",
    "print(artifact_dir + \"launch_predictor.sh\")\n",
    "print(artifact_dir + \"stop.sh\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "820ec6b3",
   "metadata": {},
   "source": [
    "### Optimization\n",
    "\n",
    "We also need 5 files to run optimization. See that the `cstr_optimizer.py` runs the same `optimal_manipulated_vars` function we designed in the previous notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f63454",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload scripts to run optimization with machine learning model\n",
    "artifact_dir = \"s3://\" + bucket + \"/artifacts/CSTROptimizer/1.0.0/\"\n",
    "\n",
    "# Copy the files from our notebook to S3\n",
    "print(\"Copying files from the notebook to S3\")\n",
    "! aws s3 cp ./optimizer/agent_pb2.py {artifact_dir}\n",
    "! aws s3 cp ./optimizer/agent_pb2_grpc.py {artifact_dir}\n",
    "! aws s3 cp ./optimizer/cstr_optimizer.py {artifact_dir}\n",
    "! aws s3 cp ./optimizer/launch_optimizer.sh {artifact_dir}\n",
    "! aws s3 cp ./optimizer/stop.sh {artifact_dir}\n",
    "\n",
    "# Note the S3 file names - we'll need them later\n",
    "print(\"\")\n",
    "print(\"S3 Files for custom inference component artifacts:\")\n",
    "print(artifact_dir + \"agent_pb2.py\")\n",
    "print(artifact_dir + \"agent_pb2_grpc.py\")\n",
    "print(artifact_dir + \"cstr_optimizer.py\")\n",
    "print(artifact_dir + \"launch_optimizer.sh\")\n",
    "print(artifact_dir + \"stop.sh\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0324c4",
   "metadata": {},
   "source": [
    "We will also note the address of of our IoTCore endpoint because we will need it when configuring our custom Greengrass components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3799a2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note our IoTCore endpoint - we'll need this later\n",
    "print(\"\")\n",
    "print(\"IoTCore Endpoint: \")\n",
    "! aws iot describe-endpoint --endpoint-type iot:Data-ATS | jq \".endpointAddress\" | sed 's/\"//g'"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p36",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}