Refactor the Debugger detect_stalled_training_job_and_stop.ipynb notebook

sagemaker-debugger/tensorflow_action_on_rule/detect_stalled_training_job_and_stop.ipynb

@@ -4,30 +4,24 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Detect stalled training and stop training job using debugger rule\n",
+    "# Detect Stalled Training and Stop Training Job Using SageMaker Debugger Rule\n",
     " \n",
+    "This notebook guides you how to use the `StalledTrainingRule` built-in rule. This rule can take an action to stop your training job, when the rule detects an inactivity in your training job for a certain time period. This functionality helps you monitor the training job status and save redundant resource usage.\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",
+    "## How the StalledTrainingRule Built-in Rule Works\n",
     "\n",
-    "## How does StalledTrainingRule works?\n",
+    "Amazon Sagemaker Debugger captures tensors that you want to watch from training jobs on [AWS Deep Learning Containers](https://docs.aws.amazon.com/sagemaker/latest/dg/train-debugger.html#debugger-supported-aws-containers) or your local machine. If you use one of the Debugger-integrated Deep Learning Containers, you don't need to make any changes to your training script to use the functionality of built-in rules. For information about Debugger-supported SageMaker frameworks and versions, see [Debugger-supported framework versions for zero script change](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/sagemaker.md#zero-script-change). \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",
+    "If you want to run a training script that uses partially supported framework by Debugger or your own custom container, you need to manually register the Debugger hook to your training script. The `smdebug` library provides tools to help the hook registration, and the sample script provided in the `src` folder includes the hook registration code as comment lines. If you want to explore using the manual hook registration, see the training script at `./src/simple_stalled_training.py`, and documentation at [smdebug TensorFlow hook](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/tensorflow.md), [smdebug PyTorch hook](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/pytorch.md), [smdebug MXNet hook](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/mxnet.md), and [smdebug XGBoost hook](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/xgboost.md).\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)"
+    "The Debugger `StalledTrainingRule` watches tensor updates from your training job. If the rule doesn't find new tensors updated to the default S3 URI for a threshold period of time, it takes an action to trigger the `StopTrainingJob` API operation. The following code cells set up a SageMaker TensorFlow estimator with the Debugger `StalledTrainingRule` to watch the `losses` pre-built tensor collection."
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "! pip install -q sagemaker"
+    "### Import SageMaker Python SDK"
    ]
   },
   {
@@ -36,21 +30,16 @@
    "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,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "from sagemaker.debugger import Rule, DebuggerHookConfig, TensorBoardOutputConfig, CollectionConfig\n",
-    "import smdebug_rulesconfig as rule_configs"
+    "### Import SageMaker Debugger classes for rule configuration"
    ]
   },
   {
@@ -59,24 +48,23 @@
    "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",
-    "}"
+    "from sagemaker.debugger import Rule, CollectionConfig, rule_configs"
    ]
   },
   {
    "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). "
+    "### Create a unique training job prefix\n",
+    "The unique prefix must be specified for `StalledTrainingRule` to identify the exact training job name that you want to monitor and stop when the rule triggers the stalled training job issue.\n",
+    "If there are multiple training jobs sharing the same prefix, this rule may react to other training jobs. If the rule cannot find the exact training job name with a provided prefix, it will fallback to safe mode and not take action of stop the training job. Such rule evaluation process goes on in parallel while the training jobs are running. If you want to access the rule job logs, you will later find how to get the information at [Get a direct Amazon CloudWatch URL to find the current rule processing job log](#cw-url).\n",
+    "\n",
+    "The following code cell includes:\n",
+    "* a code line to create a unique `base_job_name_prefix`\n",
+    "* a stalled training job rule configuration object\n",
+    "* a SageMaker TensorFlow estimator configuration with the Debugger `rules` parameter to run the built-in rule\n",
+    "\n",
+    "**Note**: Debugger collects `loss` tensors by default every 500 steps."
    ]
   },
   {
@@ -85,13 +73,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Append current time to your training job name to generate a unique base_job_name_prefix\n",
     "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)"
+    "\n",
+    "# Configure a StalledTrainingRule rule parameter object\n",
+    "stalled_training_job_rule = [\n",
+    "    Rule.sagemaker(\n",
+    "        base_config=rule_configs.stalled_training_rule(),\n",
+    "        rule_parameters={\n",
+    "                \"threshold\": \"120\", \n",
+    "                \"stop_training_on_fire\": \"True\",\n",
+    "                \"training_job_name_prefix\": base_job_name_prefix\n",
+    "        }\n",
+    "    )\n",
+    "]\n",
+    "\n",
+    "# Configure a SageMaker TensorFlow estimator\n",
+    "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='src/simple_stalled_training.py', # This sample script forces the training job to sleep for 10 minutes\n",
+    "    framework_version='1.15.0',\n",
+    "    py_version='py3',\n",
+    "    train_max_run=3600,\n",
+    "    ## Debugger-specific parameter\n",
+    "    rules = stalled_training_job_rule\n",
+    ")"
    ]
   },
   {
@@ -100,41 +110,26 @@
    "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",
-    ")"
+    "estimator.fit(wait=False)"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "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"
+    "## Monitoring Training and Rule Evaluation Status\n",
+    "\n",
+    "Once you excute the `estimator.fit()` API, SageMaker initiates a trining job in the background, and Debugger initiates a `StalledTrainingRule` rule evaluation job in parallel.\n",
+    "Because the training scripts has a few lines of code at the end to force a sleep mode for 10 minutes, the `RuleEvaluationStatus` for `StalledTrainingRule` will change to `IssuesFound` in 2 minutes after the sleep mode is on and trigger the `StopTrainingJob` API."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Print the training job name\n",
+    "\n",
+    "The following cell outputs the training job name and its training status running in the background."
    ]
   },
   {
@@ -143,23 +138,21 @@
    "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",
+    "job_name = estimator.latest_training_job.name\n",
+    "print('Training job name: {}'.format(job_name))\n",
     "\n",
-    "estimator.fit(wait=True)"
+    "client = estimator.sagemaker_session.sagemaker_client\n",
+    "\n",
+    "description = client.describe_training_job(TrainingJobName=job_name)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Monitoring\n",
+    "### Output the current job status and the rule evaluation status\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`"
+    "The following cell tracks the status of training job until the `SecondaryStatus` changes to `Stopped` or `Completed`. While training, Debugger collects output tensors from the training job and monitors the training job with the rules. "
    ]
   },
   {
@@ -168,9 +161,30 @@
    "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()"
+    "import time\n",
+    "\n",
+    "if description['TrainingJobStatus'] != 'Completed':\n",
+    "    while description['SecondaryStatus'] not in {'Stopped', 'Completed'}:\n",
+    "        description = client.describe_training_job(TrainingJobName=job_name)\n",
+    "        primary_status = description['TrainingJobStatus']\n",
+    "        secondary_status = description['SecondaryStatus']\n",
+    "        print('Current job status: [PrimaryStatus: {}, SecondaryStatus: {}] | {} Rule Evaluation Status: {}'\n",
+    "            .format(primary_status, secondary_status, \n",
+    "                estimator.latest_training_job.rule_job_summary()[0][\"RuleConfigurationName\"],\n",
+    "                estimator.latest_training_job.rule_job_summary()[0][\"RuleEvaluationStatus\"]\n",
+    "            )\n",
+    "        )\n",
+    "        time.sleep(15)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "<a class=\"anchor\" id=\"cw-url\"></a>\n",
+    "### Get a direct Amazon CloudWatch URL to find the current rule processing job log\n",
+    "\n",
+    "The following script returns a CloudWatch URL. Copy the URL and Paste it to a browser. This will directly lead you to the rule job log page."
    ]
   },
   {
@@ -203,40 +217,34 @@
     "            result[status[\"RuleConfigurationName\"]] = _get_cw_url_for_rule_job(rule_job_name, region)\n",
     "    return result\n",
     "\n",
-    "get_rule_jobs_cw_urls(estimator)"
+    "print(\n",
+    "    \"The direct CloudWatch URL to the current rule job:\", \n",
+    "    get_rule_jobs_cw_urls(estimator)[estimator.latest_training_job.rule_job_summary()[0][\"RuleConfigurationName\"]]\n",
+    ")"
    ]
   },
   {
    "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."
+    "## Conclusion\n",
+    "\n",
+    "This notebook showed how you can use the Debugger `StalledTrainingRule` built-in rule for your training job to take action on rule evaluation status changes. To find more information about Debugger, see [Amazon SageMaker Debugger Developer Guide](https://integ-docs-aws.amazon.com/sagemaker/latest/dg/train-debugger.html) and the [smdebug GitHub documentation](https://github.com/awslabs/sagemaker-debugger)."
    ]
   },
   {
    "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)."
-   ]
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "conda_tensorflow_p36",
    "language": "python",
-   "name": "python3"
+   "name": "conda_tensorflow_p36"
   },
   "language_info": {
    "codemirror_mode": {