{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Local Mode training and inference in SageMaker Studio\n",
    "This Jupyter notebook is based on AWS  sagemaker local-mode example [xgboost_script_mode_local_training_and_serving](https://github.com/aws-samples/amazon-sagemaker-local-mode/blob/599b86851ffb4c081b8f6b5458c98e9bf7479260/xgboost_script_mode_local_training_and_serving/xgboost_script_mode_local_training_and_serving.py)\n",
    "The notebook has steps added to setup `local mode` on SageMaker Studio. This notebook is tested on `Python3 (Data Science)` kernel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intial setup\n",
    "### Run `setup.sh` to execute initial setup\n",
    "`setup.sh` will do the following:\n",
    "- Create `~/.sagemaker_studio_docker_cli` directory\n",
    "- Setup softlink for `sdocker` to make it possible to run it from anywhere from command line\n",
    "- Install `docker` and `docker-compose`\n",
    "- Create `~/temp` directory used in `local mode`\n",
    "- Create `config.yaml` to change temporay directory to `~/temp`\n",
    "- Install branch `remote_docker_host` from SageMaker Python SDK which introduces Remote Docker Host capability (see [PR 2864](https://github.com/aws/sagemaker-python-sdk/pull/2864))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd .. ;./setup.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fix dependencies\n",
    "`Python 3 (Data Science)` kernel comes with `pyyaml==6.0` and doesn't have `pgrep` or `procps`. `local mode` requires `pyyaml==5.4.1`, higher versions cause error. Also `pgrep` is required when deleting a local endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!conda update --force -y conda\n",
    "!conda install -y pyyaml==5.4.1\n",
    "!apt-get install -y procps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Important**\n",
    "Restart kernel for the above installations to take effect"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Docker Host\n",
    "`sdocker` help us provision an EC2 instance that we can use as a docker host to run docker commands remotely. `sdocker` does the following:\n",
    "- Setup networking and security groups between the instance and SageMaker Studio Apps and EFS\n",
    "- Provision EC2 instance\n",
    "- Mount SageMaker Studio EFS on EC2 instance\n",
    "- Create docker context to connect to docker host\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sdocker create-host --instance-type c5.xlarge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download dataset\n",
    "Run cell below to download dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data/train; cd data/train; wget https://raw.githubusercontent.com/aws-samples/amazon-sagemaker-local-mode/main/xgboost_script_mode_local_training_and_serving/data/train/abalone"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create XGBoost estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import TrainingInput\n",
    "from sagemaker.xgboost import XGBoost, XGBoostModel\n",
    "\n",
    "DUMMY_IAM_ROLE = 'arn:aws:iam::111111111111:role/service-role/AmazonSageMaker-ExecutionRole-20200101T000001'\n",
    "\n",
    "hyperparameters = {\n",
    "    \"max_depth\": \"5\",\n",
    "    \"eta\": \"0.2\",\n",
    "    \"gamma\": \"4\",\n",
    "    \"min_child_weight\": \"6\",\n",
    "    \"subsample\": \"0.7\",\n",
    "    \"objective\": \"reg:squarederror\",\n",
    "    \"num_round\": \"50\",\n",
    "    \"verbosity\": \"2\",\n",
    "}\n",
    "\n",
    "xgb_script_mode_estimator = XGBoost(\n",
    "    entry_point=\"./code/abalone.py\",\n",
    "    hyperparameters=hyperparameters,\n",
    "    role=DUMMY_IAM_ROLE,\n",
    "    instance_count=1,\n",
    "    instance_type='local',\n",
    "    framework_version=\"1.2-1\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Start local training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_input = TrainingInput(\"file://./data/train/abalone\", content_type=\"text/libsvm\")\n",
    "\n",
    "xgb_script_mode_estimator.fit({\"train\": train_input, \"validation\": train_input})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Deploy local endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_data = xgb_script_mode_estimator.model_data\n",
    "\n",
    "xgb_inference_model = XGBoostModel(\n",
    "    model_data=model_data,\n",
    "    role=DUMMY_IAM_ROLE,\n",
    "    entry_point=\"./code/inference.py\",\n",
    "    framework_version=\"1.2-1\",\n",
    ")\n",
    "\n",
    "predictor = xgb_inference_model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=\"local\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Invoke local endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_inference_on_local_endpoint(predictor, libsvm_str):\n",
    "    label, *features = libsvm_str.strip().split()\n",
    "    predictions = predictor.predict(\" \".join([\"-99\"] + features))  # use dummy label -99\n",
    "    print()\n",
    "    print(\"*********************************************************\")\n",
    "    print(\"Prediction: {}\".format(predictions))\n",
    "    print(\"*********************************************************\")\n",
    "    print()\n",
    "    \n",
    "a_young_abalone = \"6 1:3 2:0.37 3:0.29 4:0.095 5:0.249 6:0.1045 7:0.058 8:0.067\"\n",
    "do_inference_on_local_endpoint(predictor, a_young_abalone)\n",
    "\n",
    "an_old_abalone = \"15 1:1 2:0.655 3:0.53 4:0.175 5:1.2635 6:0.486 7:0.2635 8:0.415\"\n",
    "do_inference_on_local_endpoint(predictor, an_old_abalone)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Clean up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint(predictor.endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete Current Docker Host\n",
    "Once you finish with docker host, you can terminate the instance with the below command. Note this command will only work if a docker host was created successfully with `create-host` command. Otherwise, please make sure to terminate the instance manually to avoid incurring any extra costs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sdocker terminate-current-host"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}