sagemaker-debugger/tensorflow_action_on_rule/tf-mnist-stop-training-job.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amazon SageMaker Debugger - Reacting to Cloudwatch Events from Rules\n",
    "[Amazon SageMaker](https://aws.amazon.com/sagemaker/) is managed platform to build, train and host maching learning models. Amazon SageMaker Debugger is a new feature which offers the capability to debug machine learning models during training by identifying and detecting problems with the models in near real time. \n",
    "\n",
    "In this notebook, we'll show you how you can react off rule triggers and take some action, e.g. stop the training job through CloudWatch Events.\n",
    "\n",
    "## How does Amazon SageMaker Debugger work?\n",
    "\n",
    "Amazon SageMaker Debugger lets you go beyond just looking at scalars like losses and accuracies during training and gives you full visibility into all tensors 'flowing through the graph' during training. Furthermore, it helps you monitor your training in near real time using rules and provides you alerts, once it has detected inconsistency in training flow.\n",
    "\n",
    "### Concepts\n",
    "* **Tensors**: These represent the state of the training network at intermediate points during its execution\n",
    "* **Debug Hook**: Hook is the construct with which Amazon SageMaker Debugger looks into the training process and captures the tensors requested at the desired step intervals\n",
    "* **Rule**: A logical construct, implemented as Python code, which helps analyze the tensors captured by the hook and report anomalies, if at all\n",
    "\n",
    "With these concepts in mind, let's understand the overall flow of things that Amazon SageMaker Debugger uses to orchestrate debugging.\n",
    "\n",
    "### Saving tensors during training\n",
    "\n",
    "The tensors captured by the debug hook are stored in the S3 location specified by you. There are two ways you can configure Amazon SageMaker Debugger to save tensors:\n",
    "\n",
    "#### With no changes to your training script\n",
    "If you use one of the SageMaker provided [Deep Learning Containers](https://docs.aws.amazon.com/sagemaker/latest/dg/pre-built-containers-frameworks-deep-learning.html) for 1.15, then you don't need to make any changes to your training script for the tensors to be stored. SageMaker Debugger will use the configuration you provide through the SageMaker SDK's Tensorflow `Estimator` when creating your job to save the tensors in the fashion you specify. You can review the script we are going to use at [src/mnist_zerocodechange.py](src/mnist_zerocodechange.py). You will note that this is an untouched TensorFlow script which uses the Estimator interface. Please note that SageMaker Debugger only supports `tf.keras`, `tf.Estimator` and `tf.MonitoredSession` interfaces. Full description of support is available at [SageMaker Debugger with TensorFlow ](https://github.com/awslabs/sagemaker-debugger/tree/master/docs/tensorflow.md)\n",
    "\n",
    "#### Orchestrating your script to store tensors\n",
    "For other containers, you need to make couple of lines of changes to your training script. SageMaker Debugger exposes a library `smdebug` which allows you to capture these tensors and save them for analysis. It's highly customizable and allows to save the specific tensors you want at different frequencies and possibly with other configurations. Refer [DeveloperGuide](https://github.com/awslabs/sagemaker-debugger/tree/master/docs) for details on how to use SageMaker Debugger library with your choice of framework in your training script. Here we have an example script orchestrated at [src/mnist_byoc](src/mnist_byoc.py). You also need to ensure that your container has the `smdebug` library installed.\n",
    "\n",
    "### Analysis of tensors\n",
    "\n",
    "Once the tensors are saved, Amazon SageMaker Debugger can be configured to run debugging ***Rules*** on them. At a very broad level, a rule is Python code used to detect certain conditions during training. Some of the conditions that a data scientist training an algorithm may care about are monitoring for gradients getting too large or too small, detecting overfitting, and so on. Sagemaker Debugger comes pre-packaged with certain built-in rules. Users can write their own rules using the Sagemaker Debugger APIs. You can also analyze raw tensor data outside of the Rules construct in say, a Sagemaker notebook, using Amazon Sagemaker Debugger's full set of APIs. Please refer [Analysis Developer Guide](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/api.md) for more on these APIs.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cloudwatch Events for Rules\n",
    "Rule status changes in a training job trigger CloudWatch Events. These events can be acted upon by configuring a CloudWatch Rule (different from Amazon SageMaker Debugger Rule) to trigger each time a Debugger Rule changes status. In this notebook we'll go through how you can create a CloudWatch Rule to direct Training Job State change events to a lambda function that stops the training job in case a rule triggers and has status `\"IssuesFound\"`\n",
    "\n",
    "#### Lambda Function\n",
    "\n",
    "* In your AWS console, go to Lambda Management Console,\n",
    "* Create a new function by hitting Create Function,\n",
    "* Choose the language as Python 3.7 (or higher) and put in the following sample code for stopping the training job if one of the Rule statuses is `\"IssuesFound\"`:\n",
    "\n",
    "```python\n",
    "import json\n",
    "import boto3\n",
    "import logging\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "\n",
    "def lambda_handler(event, context):\n",
    "    training_job_name = event.get(\"detail\").get(\"TrainingJobName\")\n",
    "    logging.info(f'Evaluating Debugger rules for training job: {training_job_name}')\n",
    "\n",
    "    eval_statuses = event.get(\"detail\").get(\"DebugRuleEvaluationStatuses\", None)\n",
    "\n",
    "    if eval_statuses is None or len(eval_statuses) == 0:\n",
    "        logging.info(\"Couldn't find any debug rule statuses, skipping...\")\n",
    "        return {\n",
    "            'statusCode': 200,\n",
    "            'body': json.dumps('Nothing to do')\n",
    "        }\n",
    "\n",
    "    # should only attempt stopping jobs with InProgress status\n",
    "    training_job_status = event.get(\"detail\").get(\"TrainingJobStatus\", None)\n",
    "    if training_job_status != 'InProgress':\n",
    "        logging.debug(f\"Current Training job status({training_job_status}) is not 'InProgress'. Exiting\")\n",
    "        return {\n",
    "            'statusCode': 200,\n",
    "            'body': json.dumps('Nothing to do')\n",
    "        }\n",
    "\n",
    "    client = boto3.client('sagemaker')\n",
    "\n",
    "    for status in eval_statuses:\n",
    "        logging.info(status.get(\"RuleEvaluationStatus\") + ', RuleEvaluationStatus=' + str(status))\n",
    "        if status.get(\"RuleEvaluationStatus\") == \"IssuesFound\":\n",
    "            secondary_status = event.get(\"detail\").get(\"SecondaryStatus\", None)\n",
    "            logging.info(\n",
    "                f'About to stop training job, since evaluation of rule configuration {status.get(\"RuleConfigurationName\")} resulted in \"IssuesFound\". ' +\n",
    "                f'\\ntraining job \"{training_job_name}\" status is \"{training_job_status}\", secondary status is \"{secondary_status}\"' +\n",
    "                f'\\nAttempting to stop training job \"{training_job_name}\"'\n",
    "            )\n",
    "            try:\n",
    "                client.stop_training_job(\n",
    "                    TrainingJobName=training_job_name\n",
    "                )\n",
    "            except Exception as e:\n",
    "                logging.error(\n",
    "                    \"Encountered error while trying to \"\n",
    "                    \"stop training job {}: {}\".format(\n",
    "                        training_job_name, str(e)\n",
    "                    )\n",
    "                )\n",
    "                raise e\n",
    "    return None\n",
    "```\n",
    "* Create a new execution role for the Lambda, and\n",
    "* In your IAM console, search for the role and attach \"AmazonSageMakerFullAccess\" policy to the role. This is needed for the code in your Lambda function to stop the training job.\n",
    "* Basic settings > set Timeout to 30 seconds instead of 3 seconds. \n",
    "\n",
    "#### Create a CloudWatch Rule\n",
    "\n",
    "* In your AWS Console, go to CloudWatch and select Rule from the left column,\n",
    "* Hit Create Rule. The console will redirect you to the Rule creation page,\n",
    " * For the Service Name, select \"SageMaker\".\n",
    " * For the Event Type, select \"SageMaker Training Job State Change\".\n",
    "* In the Targets select the Lambda function you created above, and\n",
    "* For this example notebook, we'll leave everything as is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import os\n",
    "import sagemaker\n",
    "from sagemaker.tensorflow import TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.debugger import Rule, rule_configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the entrypoint script\n",
    "entrypoint_script = \"src/mnist_zerocodechange.py\"\n",
    "\n",
    "# these hyperparameters ensure that vanishing gradient will trigger for our tensorflow mnist script\n",
    "hyperparameters = {\"num_epochs\": \"10\", \"lr\": \"10.00\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "rules = [\n",
    "    Rule.sagemaker(rule_configs.vanishing_gradient()),\n",
    "    Rule.sagemaker(rule_configs.loss_not_decreasing()),\n",
    "]\n",
    "\n",
    "estimator = TensorFlow(\n",
    "    role=sagemaker.get_execution_role(),\n",
    "    base_job_name=\"smdebugger-demo-mnist-tensorflow\",\n",
    "    train_instance_count=1,\n",
    "    train_instance_type=\"ml.m4.xlarge\",\n",
    "    entry_point=entrypoint_script,\n",
    "    framework_version=\"1.15\",\n",
    "    train_volume_size=400,\n",
    "    py_version=\"py3\",\n",
    "    train_max_run=3600,\n",
    "    script_mode=True,\n",
    "    hyperparameters=hyperparameters,\n",
    "    ## New parameter\n",
    "    rules=rules,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [],
   "source": [
    "# After calling fit, SageMaker will spin off 1 training job and 1 rule job for you\n",
    "# The rule evaluation status(es) will be visible in the training logs\n",
    "# at regular intervals\n",
    "# wait=False makes this a fire and forget function. To stream the logs in the notebook leave this out\n",
    "\n",
    "estimator.fit(wait=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monitoring\n",
    "\n",
    "SageMaker kicked off rule evaluation jobs, one for each of the SageMaker rules - `VanishingGradient` and `LossNotDecreasing` as specified in the estimator. \n",
    "Given that we've tweaked the hyperparameters of our training script such that `VanishingGradient` is bound to fire, we should expect to see the `TrainingJobStatus` as\n",
    "`Stopped` once the `RuleEvaluationStatus` for `VanishingGradient` changes to `IssuesFound`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'RuleConfigurationName': 'VanishingGradient',\n",
       "  'RuleEvaluationJobArn': 'arn:aws:sagemaker:us-east-2:072677473360:processing-job/smdebugger-demo-mnist-tens-vanishinggradient-e23301a8',\n",
       "  'RuleEvaluationStatus': 'IssuesFound',\n",
       "  'StatusDetails': 'RuleEvaluationConditionMet: Evaluation of the rule VanishingGradient at step 500 resulted in the condition being met\\n',\n",
       "  'LastModifiedTime': datetime.datetime(2019, 12, 1, 7, 20, 55, 495000, tzinfo=tzlocal())},\n",
       " {'RuleConfigurationName': 'LossNotDecreasing',\n",
       "  'RuleEvaluationJobArn': 'arn:aws:sagemaker:us-east-2:072677473360:processing-job/smdebugger-demo-mnist-tens-lossnotdecreasing-27ee2da1',\n",
       "  'RuleEvaluationStatus': 'InProgress',\n",
       "  'LastModifiedTime': datetime.datetime(2019, 12, 1, 7, 20, 55, 495000, tzinfo=tzlocal())}]"
      ]
     },
     "execution_count": 191,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# rule job summary gives you the summary of the rule evaluations. You might have to run it over\n",
    "# a few times before you start to see all values populated/changing\n",
    "estimator.latest_training_job.rule_job_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'VanishingGradient': 'https://us-east-2.console.aws.amazon.com/cloudwatch/home?region=us-east-2#logStream:group=/aws/sagemaker/ProcessingJobs;prefix=smdebugger-demo-mnist-tens-VanishingGradient-e23301a8;streamFilter=typeLogStreamPrefix',\n",
       " 'LossNotDecreasing': 'https://us-east-2.console.aws.amazon.com/cloudwatch/home?region=us-east-2#logStream:group=/aws/sagemaker/ProcessingJobs;prefix=smdebugger-demo-mnist-tens-LossNotDecreasing-27ee2da1;streamFilter=typeLogStreamPrefix'}"
      ]
     },
     "execution_count": 194,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This utility gives the link to monitor the CW event\n",
    "def _get_rule_job_name(training_job_name, rule_configuration_name, rule_job_arn):\n",
    "    \"\"\"Helper function to get the rule job name\"\"\"\n",
    "    return \"{}-{}-{}\".format(\n",
    "        training_job_name[:26], rule_configuration_name[:26], rule_job_arn[-8:]\n",
    "    )\n",
    "\n",
    "\n",
    "def _get_cw_url_for_rule_job(rule_job_name, region):\n",
    "    return \"https://{}.console.aws.amazon.com/cloudwatch/home?region={}#logStream:group=/aws/sagemaker/ProcessingJobs;prefix={};streamFilter=typeLogStreamPrefix\".format(\n",
    "        region, region, rule_job_name\n",
    "    )\n",
    "\n",
    "\n",
    "def get_rule_jobs_cw_urls(estimator):\n",
    "    region = boto3.Session().region_name\n",
    "    training_job = estimator.latest_training_job\n",
    "    training_job_name = training_job.describe()[\"TrainingJobName\"]\n",
    "    rule_eval_statuses = training_job.describe()[\"DebugRuleEvaluationStatuses\"]\n",
    "\n",
    "    result = {}\n",
    "    for status in rule_eval_statuses:\n",
    "        if status.get(\"RuleEvaluationJobArn\", None) is not None:\n",
    "            rule_job_name = _get_rule_job_name(\n",
    "                training_job_name, status[\"RuleConfigurationName\"], status[\"RuleEvaluationJobArn\"]\n",
    "            )\n",
    "            result[status[\"RuleConfigurationName\"]] = _get_cw_url_for_rule_job(\n",
    "                rule_job_name, region\n",
    "            )\n",
    "    return result\n",
    "\n",
    "\n",
    "get_rule_jobs_cw_urls(estimator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running the last two cells over and until `VanishingGradient` reports `IssuesFound`, we'll attempt to describe the `TrainingJobStatus` for our training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Stopped'"
      ]
     },
     "execution_count": 193,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimator.latest_training_job.describe()[\"TrainingJobStatus\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Result\n",
    "\n",
    "This notebook attempted to show a very simple setup of how you can use CloudWatch events for your training job to take action on rule evaluation status changes. Learn more about Amazon SageMaker Debugger in the [GitHub Documentation](https://github.com/awslabs/sagemaker-debugger)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_tensorflow_p36",
   "language": "python",
   "name": "conda_tensorflow_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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}