sagemaker-debugger/tensorflow_action_on_rule/detect_stalled_training_job_and_stop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detect stalled training and stop training job using debugger rule\n",
    " \n",
    "\n",
    "In this notebook, we'll show you how you can use StalledTrainingRule rule which can take action like stopping your training job when it finds that there has been no update in training job for certain threshold duration.\n",
    "\n",
    "## How does StalledTrainingRule works?\n",
    "\n",
    "Amazon Sagemaker debugger automatically captures tensors from training job which use AWS DLC(tensorflow, pytorch, mxnet, xgboost)[refer doc for supported versions](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/sagemaker.md#zero-script-change). StalledTrainingRule keeps watching on emission of tensors like loss. The execution happens outside of training containers. It is evident that if training job is running good and is not stalled it is expected to emit loss and metrics tensors at frequent intervals. If Rule doesn't find new tensors being emitted from training job for threshold period of time, it takes automatic action to issue StopTrainingJob.\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). [Refer doc for supported framework versions](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/sagemaker.md#zero-script-change), then you don't need to make any changes to your training script for activating this rule. Loss tensors will automatically be captured and monitored by the rule.\n",
    "\n",
    "You can also emit tensors periodically by using [save scalar api of hook](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/api.md#common-hook-api) . \n",
    "\n",
    "Also look at example how to use save_scalar api [here](https://github.com/awslabs/sagemaker-debugger/blob/master/examples/tensorflow2/scripts/tf_keras_fit_non_eager.py#L42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import os\n",
    "import sagemaker\n",
    "from sagemaker.tensorflow import TensorFlow\n",
    "print(sagemaker.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.debugger import Rule, DebuggerHookConfig, TensorBoardOutputConfig, CollectionConfig\n",
    "import smdebug_rulesconfig as rule_configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the entrypoint script\n",
    "# Below script has 5 minutes sleep, we will create a stalledTrainingRule with 3 minutes of threshold.\n",
    "entrypoint_script='src/simple_stalled_training.py'\n",
    "\n",
    "# these hyperparameters ensure that vanishing gradient will trigger for our tensorflow mnist script\n",
    "hyperparameters = {\n",
    "    \"num_epochs\": \"10\",\n",
    "    \"lr\": \"10.00\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create unique training job prefix\n",
    "We will create unique training job name prefix. this prefix would be passed to StalledTrainingRule to identify which training job, rule should take action on once the stalled training rule condition is met.\n",
    "Note that, this prefix needs to be unique. If rule doesn't find exactly one job with provided prefix, it will fallback to safe mode and not take action of stop training job. Rule will still emit a cloudwatch event if the rule condition is met. To see details about cloud watch event, check [here](https://github.com/awslabs/amazon-sagemaker-examples/tree/master/sagemaker-debugger/tensorflow_action_on_rule/tf-mnist-stop-training-job.ipynb). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "print(int(time.time()))\n",
    "# Note that sagemaker appends date to your training job and truncates the provided name to 39 character. So, we will make \n",
    "# sure that we use less than 39 character in below prefix. Appending time is to provide a unique id\n",
    "base_job_name_prefix= 'smdebug-stalled-demo-' + str(int(time.time()))\n",
    "base_job_name_prefix = base_job_name_prefix[:34]\n",
    "print(base_job_name_prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stalled_training_job_rule = Rule.sagemaker(\n",
    "    base_config={\n",
    "                    'DebugRuleConfiguration': {\n",
    "                        'RuleConfigurationName': 'StalledTrainingRule', \n",
    "                        'RuleParameters': {'rule_to_invoke': 'StalledTrainingRule'}\n",
    "                    }\n",
    "                 },\n",
    "    rule_parameters={\n",
    "        'threshold': '120',\n",
    "        'training_job_name_prefix': base_job_name_prefix,\n",
    "        'stop_training_on_fire' : 'True'\n",
    "    },    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = TensorFlow(\n",
    "    role=sagemaker.get_execution_role(),\n",
    "    base_job_name=base_job_name_prefix,\n",
    "    train_instance_count=1,\n",
    "    train_instance_type='ml.m5.4xlarge',\n",
    "    entry_point=entrypoint_script,\n",
    "    #source_dir = 'src',\n",
    "    framework_version='1.15.0',\n",
    "    py_version='py3',\n",
    "    train_max_run=3600,\n",
    "    script_mode=True,\n",
    "    ## New parameter\n",
    "    rules = [stalled_training_job_rule]\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monitoring\n",
    "\n",
    "SageMaker kicked off rule evaluation job `StalledTrainingRule` as specified in the estimator. \n",
    "Given that we've stalled our training script for 10 minutes such that `StalledTrainingRule` is bound to fire and take action StopTrainingJob, we should expect to see the `TrainingJobStatus` as\n",
    "`Stopped` once the `RuleEvaluationStatus` for `StalledTrainingRule` changes to `IssuesFound`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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",
    "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(region, region, rule_job_name)\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(training_job_name, status[\"RuleConfigurationName\"], status[\"RuleEvaluationJobArn\"])\n",
    "            result[status[\"RuleConfigurationName\"]] = _get_cw_url_for_rule_job(rule_job_name, region)\n",
    "    return result\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": null,
   "metadata": {},
   "outputs": [],
   "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": "Python 3",
   "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}