{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4424d2d5",
   "metadata": {},
   "source": [
    "# People counter model preparation and compilation\n",
    "# Section 1: Model preparation\n",
    "\n",
    "## Install and import necessary Python libraries\n",
    "In this first step, we install some useful Python libraries for us to easily prepare the model to be compiled by SageMaker NEO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f41615",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install gluoncv\n",
    "!pip3 install xtarfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7725b169",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mxnet as mx \n",
    "import gluoncv as gcv\n",
    "import xtarfile as tarfile\n",
    "import sagemaker\n",
    "import boto3\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78a11265",
   "metadata": {},
   "source": [
    "## Model download and hybridize\n",
    "Use the following code to download the pre-trained model from MXNet model zoo. If you want to use models from other frameworks, you can use check out the [model preparation page](https://docs.aws.amazon.com/sagemaker/latest/dg/neo-compilation-preparing-model.html) on developer's guide to find out more information.\n",
    "\n",
    "In the second line, we hybridize the model and run through it once with randomized input of the expected input size. For more information about hybridizing models, please checkout the [MXNet explanation on hybridize](https://mxnet.apache.org/versions/1.7.0/api/python/docs/tutorials/packages/gluon/blocks/hybridize.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d054d1cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = gcv.model_zoo.get_model('yolo3_mobilenet1.0_coco', pretrained=True)\n",
    "net.hybridize()\n",
    "out = net(mx.nd.ones((1,3,224,224)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36a2cc22",
   "metadata": {},
   "outputs": [],
   "source": [
    "net.export('yolo3_mobilenet1.0_coco')\n",
    "tar = tarfile.open(\"model.tar.gz\", \"w:gz\")\n",
    "for name in [\"yolo3_mobilenet1.0_coco-0000.params\", \"yolo3_mobilenet1.0_coco-symbol.json\"]:\n",
    "    tar.add(name)\n",
    "tar.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5554487b",
   "metadata": {},
   "source": [
    "## Tar the model, and upload our model to S3\n",
    "Since in this notebook instance, there is already a execution role assiciated with it, we do not need to provision all the credentials. And in this example, for the simplicity, we are going to get the default bucket. SageMaker has a default bucket for each account (starts with sagemaker-) in the same region of this notebook instance. We are going to store all of the artifacts and compiled models within that default bucket. Feel free to change the bucket location if needed otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59ca21f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "print(role)\n",
    "sess = sagemaker.Session()\n",
    "region = sess.boto_region_name\n",
    "bucket = sess.default_bucket()\n",
    "\n",
    "compilation_job_name = name_from_base(\"MXNet-yolo3-mobilenet10-Jetson-Nano\")\n",
    "prefix = compilation_job_name + \"/model\"\n",
    "model_path = sess.upload_data(path=\"model.tar.gz\", key_prefix=prefix)\n",
    "\n",
    "data_shape = '{\"data\":[1,3,224,224]}'\n",
    "target_platform = {'Os': 'LINUX', 'Arch': 'ARM64', 'Accelerator': 'NVIDIA'}\n",
    "compiler_options = '{\"trt-ver\": \"7.1.3\", \"cuda-ver\": \"10.2\", \"gpu-code\": \"sm_53\"}'\n",
    "framework = \"MXNET\"\n",
    "framework_version = \"1.7\"\n",
    "compiled_model_path = \"s3://{}/{}/output\".format(bucket, compilation_job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e028a4d3",
   "metadata": {},
   "source": [
    "## Start compilation job\n",
    "Start compilation job, and start a polling process to wait for the compilation job to succeed. It could take around 5 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2961d30f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a SageMaker client so you can submit a compilation job\n",
    "sagemaker_client = boto3.client('sagemaker')\n",
    "\n",
    "response = sagemaker_client.create_compilation_job(\n",
    "    CompilationJobName=compilation_job_name,\n",
    "    RoleArn=role,\n",
    "    InputConfig={\n",
    "        'S3Uri': model_path,\n",
    "        'DataInputConfig': data_shape,\n",
    "        'Framework': framework.upper()\n",
    "    },\n",
    "    OutputConfig={\n",
    "        'S3OutputLocation': compiled_model_path,\n",
    "        'TargetPlatform': target_platform,\n",
    "        'CompilerOptions': compiler_options\n",
    "    },\n",
    "    StoppingCondition={\n",
    "        'MaxRuntimeInSeconds': 900\n",
    "    }\n",
    ")\n",
    "while True:\n",
    "    response = sagemaker_client.describe_compilation_job(CompilationJobName=compilation_job_name)\n",
    "    if response['CompilationJobStatus'] == 'COMPLETED':\n",
    "        break\n",
    "    elif response['CompilationJobStatus'] == 'FAILED':\n",
    "        raise RuntimeError('Compilation failed')\n",
    "    print('Compiling ...')\n",
    "    time.sleep(30)\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "368ad464",
   "metadata": {},
   "source": [
    "## Repackage compiled model in zip file for Greengrass deployment\n",
    "Since NEO compilation output is in the format of .tar.gz, but AWS IoT Greengrass only accepts .zip. We need the following step to convert the compiled model to convert the format in order for the archived file to be downloaded and unpacked by Greengrass service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a4f7abf",
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_client = boto3.client('s3')\n",
    "neo_compiled_model = 'compiled-model.tar.gz'\n",
    "s3_client.download_file(bucket, compilation_job_name+'/output/model-LINUX_ARM64_NVIDIA.tar.gz', neo_compiled_model)\n",
    "!rm -rf model compiled-model.zip\n",
    "!mkdir model\n",
    "!tar zfxv compiled-model.tar.gz -C model/\n",
    "!zip compiled-model.zip model/*\n",
    "s3_client.upload_file('compiled-model.zip', bucket, '{}/model_zipped/model.zip'.format(compilation_job_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37ad0dc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled_model_s3_url = f's3://{bucket}/{compilation_job_name}/model_zipped/model.zip'\n",
    "compiled_model_s3_url\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4d1605d",
   "metadata": {},
   "source": [
    "# !!! PAUSE HERE !!!\n",
    "\n",
    "\n",
    "# Section 2: Please only proceed if you have a shared library uploaded on S3 as instructed in the workshop guide.\n",
    "\n",
    "### We need to pause here on this notebook and prepare the sample deepstream application before running the next steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49ddaede",
   "metadata": {},
   "outputs": [],
   "source": [
    "packaged_deepstream_app_artifacts_S3_url = f's3://{bucket}/deepstream-people-counter/artifact/app.zip'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a67ffdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "greengrass_recipe = {\n",
    "  \"RecipeFormatVersion\": \"2020-01-25\",\n",
    "  \"ComponentName\": \"iot-greengrass-managed-deepstream-application\",\n",
    "  \"ComponentVersion\": \"1.0.1\",\n",
    "  \"ComponentDescription\": \"Component that runs CV program at the edge\",\n",
    "  \"ComponentPublisher\": \"Amazon Web Services\",\n",
    "  \"ComponentConfiguration\": {\n",
    "    \"DefaultConfiguration\": {\n",
    "      \"accessControl\": {\n",
    "        \"aws.greengrass.ipc.pubsub\": {\n",
    "          \"$component_name:pubsub:1\": {\n",
    "            \"policyDescription\": \"Allows access to publish to all topics.\",\n",
    "            \"operations\": [\n",
    "              \"aws.greengrass#PublishToTopic\"\n",
    "            ],\n",
    "            \"resources\": [\n",
    "              \"*\"\n",
    "            ]\n",
    "          }\n",
    "        }\n",
    "      }\n",
    "    }\n",
    "  },\n",
    "  \"Manifests\": [\n",
    "    {\n",
    "      \"Platform\": {\n",
    "        \"os\": \"linux\"\n",
    "      },\n",
    "      \"Lifecycle\": {\n",
    "        \"Install\": {\n",
    "          \"Script\": \"chmod 777 {artifacts:decompressedPath}/app/deepstream-test5-analytics ; rm -rf ~/.cache/gstreamer-1.0/registry.aarch64.bin\"\n",
    "        },\n",
    "        \"Run\": {\n",
    "          \"Script\": \"sudo --preserve-env=COMPILED_MODEL_PATH {artifacts:decompressedPath}/app/deepstream-test5-analytics -c {artifacts:decompressedPath}/app/config/test5_config_file_src_infer_tlt_neo.txt\"\n",
    "        },\n",
    "        \"setenv\": {\n",
    "          \"COMPILED_MODEL_PATH\": \"{artifacts:decompressedPath}/model/model\"\n",
    "        }\n",
    "      },\n",
    "      \"Artifacts\": [\n",
    "        {\n",
    "          \"URI\": compiled_model_s3_url,\n",
    "          \"Unarchive\": \"ZIP\"\n",
    "        },\n",
    "        {\n",
    "          \"URI\": packaged_deepstream_app_artifacts_S3_url,\n",
    "          \"Unarchive\": \"ZIP\"\n",
    "        }\n",
    "      ]\n",
    "    }\n",
    "  ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b94daffa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "gg_client = boto3.client('greengrassv2')\n",
    "component_creation_ret = gg_client.create_component_version(inlineRecipe=json.dumps(greengrass_recipe).encode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd01d5d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "iot_client = boto3.client('iot')\n",
    "\n",
    "thing_name = \"jetsonDeepstreamDemo\"\n",
    "thing_group_name = \"jetsonDeepstreamDemoGroup\"\n",
    "deployment_name = \"deepstream-app-deployment\"\n",
    "response = gg_client.create_deployment(\n",
    "    targetArn=iot_client.describe_thing(thingName = thing_name)['thingArn'],\n",
    "    deploymentName=deployment_name,\n",
    "    components={\n",
    "        component_creation_ret['componentName']: {\n",
    "            'componentVersion': component_creation_ret['componentVersion']\n",
    "        }\n",
    "    },\n",
    "    deploymentPolicies={\n",
    "        'failureHandlingPolicy': 'DO_NOTHING'\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88d5cb3b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_latest_p37",
   "language": "python",
   "name": "conda_mxnet_latest_p37"
  },
  "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
}