Sagemaker debugger stalled training rule example (aws#1301)

* Sagemaker debugger stalled training rule example

Author: Vikas-kum

sagemaker-debugger/tensorflow_action_on_rule/detect_stalled_training_job_and_stop.ipynb

@@ -0,0 +1,256 @@
+{
+ "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
+}

sagemaker-debugger/tensorflow_action_on_rule/src/simple_stalled_training.py

@@ -0,0 +1,96 @@
+"""
+This script is a simple training script which uses Tensorflow's MonitoredSession interface.
+It is designed to be used with SageMaker Debugger in an official SageMaker Framework container (i.e. AWS Deep Learning Container).
+Here we create the hook object which loads configuration from the json file that SageMaker will
+put in the training container based on the configuration provided using the SageMaker python SDK when creating a job.
+We use this hook object here to add our custom loss to the losses collection and set the mode.
+For more information, please refer to https://github.com/awslabs/sagemaker-debugger/blob/master/docs/
+"""
+
+# Standard Library
+import argparse
+import random
+
+# Third Party
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+# First Party
+import smdebug.tensorflow as smd
+
+
+def str2bool(v):
+ if isinstance(v, bool):
+ return v
+ if v.lower() in ("yes", "true", "t", "y", "1"):
+ return True
+ elif v.lower() in ("no", "false", "f", "n", "0"):
+ return False
+ else:
+ raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_dir", type=str, help="S3 path for the model")
+parser.add_argument("--lr", type=float, help="Learning Rate", default=0.001)
+parser.add_argument("--steps", type=int, help="Number of steps to run", default=100)
+parser.add_argument("--scale", type=float, help="Scaling factor for inputs", default=1.0)
+parser.add_argument("--random_seed", type=bool, default=False)
+args = parser.parse_args()
+
+# these random seeds are only intended for test purpose.
+# for now, 2,2,12 could promise no assert failure when running tests
+# if you wish to change the number, notice that certain steps' tensor value may be capable of variation
+if args.random_seed:
+ tf.set_random_seed(2)
+ np.random.seed(2)
+ random.seed(12)
+
+# Network definition
+# Note the use of name scopes
+with tf.name_scope("foobar"):
+ x = tf.placeholder(shape=(None, 2), dtype=tf.float32)
+ w = tf.Variable(initial_value=[[10.0], [10.0]], name="weight1")
+with tf.name_scope("foobaz"):
+ w0 = [[1], [1.0]]
+ y = tf.matmul(x, w0)
+loss = tf.reduce_mean((tf.matmul(x, w) - y) ** 2, name="loss")
+
+hook = smd.SessionHook.create_from_json_file()
+hook.add_to_collection("losses", loss)
+
+global_step = tf.Variable(17, name="global_step", trainable=False)
+increment_global_step_op = tf.assign(global_step, global_step + 1)
+
+optimizer = tf.train.AdamOptimizer(args.lr)
+
+# Do not need to wrap the optimizer if in a zero script change environment
+# i.e. SageMaker/AWS Deep Learning Containers
+# as the framework will automatically do that there if the hook exists
+optimizer = hook.wrap_optimizer(optimizer)
+
+# use this wrapped optimizer to minimize loss
+optimizer_op = optimizer.minimize(loss, global_step=increment_global_step_op)
+
+# Do not need to pass the hook to the session if in a zero script change environment
+# i.e. SageMaker/AWS Deep Learning Containers
+# as the framework will automatically do that there if the hook exists
+sess = tf.train.MonitoredSession()
+
+# use this session for running the tensorflow model
+hook.set_mode(smd.modes.TRAIN)
+for i in range(args.steps):
+ x_ = np.random.random((10, 2)) * args.scale
+ _loss, opt, gstep = sess.run([loss, optimizer_op, increment_global_step_op], {x: x_})
+ print(f"Step={i}, Loss={_loss}")
+
+# set the mode for monitored session based runs
+# so smdebug can separate out steps by mode
+hook.set_mode(smd.modes.EVAL)
+for i in range(args.steps):
+ x_ = np.random.random((10, 2)) * args.scale
+ sess.run([loss, increment_global_step_op], {x: x_})
+import time
+print("Sleeping for 10 minutes")
+time.sleep(10*60)
+print("Waking up and exiting")