Repository: kubeflow/pipelines

Issue:
allow output artifact store configuration (vs hard coded)
it seems like the output artifacts are always stored in a specific minio service, port, namespace, bucket, secrets, etc (`minio-service.kubeflow:9000`). 

see: https://github.com/kubeflow/pipelines/blob/f40a22a3f4a8e06d20cf3e3f425b5058d5c87e0b/sdk/python/kfp/compiler/_op_to_template.py#L148

it would be great to make it flexible, e.g. allow using S3, or change namespace or bucket names.
i suggest making it configurable, i can do such PR if we agree its needed. 
flexible pipeline service (host) path in client SDK 
when creating an SDK `Client()` the path to `ml-pipeline` API service is loaded from a hard coded value (`ml-pipeline.kubeflow.svc.cluster.local:8888`) which indicate a specific k8s namespace. it can be valuable to load that default value from an env variable, i.e. changing the line in `_client.py` from:

`config.host = host if host else Client.IN_CLUSTER_DNS_NAME`

to:

`config.host = host or os.environ.get('ML_PIPELINE_DNS_NAME',Client.IN_CLUSTER_DNS_NAME)`

also note that when a user provide the `host` parameter, the ipython output points to the API server and not to the UI service (see the logic in `_get_url_prefix()`), it seems like a potential bug

if its acceptable i can submit a PR for the line change above
 


Discussion:
There is already a PR for this. But we'd like to solve this in a different way - by having an in-cluster configMap with the storage configuration. Argo already supports that, but we're hesitating to use that mechanism since KF Pipelines might not be the exclusive user of the Argo controller.

Fortunately, the problem will be solved in Argo 2.4 (per-pipeline configMap reference)
Also https://github.com/kubeflow/pipelines/issues/363 https://github.com/kubeflow/pipelines/issues/1195 https://github.com/kubeflow/pipelines/issues/362
@Ark-kun a mechanism like config map sounds better, it will also make the YAMLs cleaner, currently u get the same 27 lines added to every container, see below, would be nice to hide it:

```yaml
    outputs:
      artifacts:
      - name: mlpipeline-ui-metadata
        path: /mlpipeline-ui-metadata.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-ui-metadata.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
      - name: mlpipeline-metrics
        path: /mlpipeline-metrics.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-metrics.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
```
Fixed by https://github.com/kubeflow/pipelines/pull/1297
>it can be valuable to load that default value from an env variable

I'm not sure about the user scenarios, but I guess it might be useful for some cases.

>ML_PIPELINE_DNS_NAME

Maybe, `KF_PIPELINES_ENDPOINT`?

>it seems like a potential bug

It probably is. The Client class can definitely use some improvements, especially for out-of-cluster usage.

> if its acceptable i can submit a PR

We always appreciate PRs.
Fixed in https://github.com/kubeflow/pipelines/pull/1293

Resolving patch:
diff --git a/components/aws/sagemaker/common/_utils.py b/components/aws/sagemaker/common/_utils.py
--- a/components/aws/sagemaker/common/_utils.py
+++ b/components/aws/sagemaker/common/_utils.py
@@ -22,6 +22,7 @@
 import re
 import json
 from pathlib2 import Path
+from enum import Enum, auto
 
 import boto3
 from boto3.session import Session
@@ -36,7 +37,7 @@
 from botocore.exceptions import ClientError
 from botocore.session import Session as BotocoreSession
 
-from sagemaker.amazon.amazon_estimator import get_image_uri
+from sagemaker.image_uris import retrieve
 
 import logging
 logging.getLogger().setLevel(logging.INFO)
@@ -99,6 +100,9 @@ def get_component_version():
     return component_version
 
 
+def print_log_header(header_len, title=""):
+    logging.info(f"{title:*^{header_len}}")
+
 def print_logs_for_job(cw_client, log_grp, job_name):
     """Gets the CloudWatch logs for SageMaker jobs"""
     try:
@@ -206,12 +210,12 @@ def create_training_job_request(args):
         # TODO: Adjust this implementation to account for custom algorithm resources names that are the same as built-in algorithm names
         algo_name = args['algorithm_name'].lower().strip()
         if algo_name in built_in_algos.keys():
-            request['AlgorithmSpecification']['TrainingImage'] = get_image_uri(args['region'], built_in_algos[algo_name])
+            request['AlgorithmSpecification']['TrainingImage'] = retrieve(built_in_algos[algo_name], args['region'])
             request['AlgorithmSpecification'].pop('AlgorithmName')
             logging.warning('Algorithm name is found as an Amazon built-in algorithm. Using built-in algorithm.')
         # Just to give the user more leeway for built-in algorithm name inputs
         elif algo_name in built_in_algos.values():
-            request['AlgorithmSpecification']['TrainingImage'] = get_image_uri(args['region'], algo_name)
+            request['AlgorithmSpecification']['TrainingImage'] = retrieve(algo_name, args['region'])
             request['AlgorithmSpecification'].pop('AlgorithmName')
             logging.warning('Algorithm name is found as an Amazon built-in algorithm. Using built-in algorithm.')
         else:
@@ -258,6 +262,17 @@ def create_training_job_request(args):
 
     enable_spot_instance_support(request, args)
 
+    ### Update DebugHookConfig and DebugRuleConfigurations
+    if args['debug_hook_config']:
+        request['DebugHookConfig'] = args['debug_hook_config']
+    else:
+        request.pop('DebugHookConfig')
+
+    if args['debug_rule_config']:
+        request['DebugRuleConfigurations'] = args['debug_rule_config']
+    else:
+        request.pop('DebugRuleConfigurations')
+
     ### Update tags
     for key, val in args['tags'].items():
         request['Tags'].append({'Key': key, 'Value': val})
@@ -282,18 +297,94 @@ def create_training_job(client, args):
 
 
 def wait_for_training_job(client, training_job_name, poll_interval=30):
-  while(True):
-    response = client.describe_training_job(TrainingJobName=training_job_name)
-    status = response['TrainingJobStatus']
-    if status == 'Completed':
-      logging.info("Training job ended with status: " + status)
-      break
-    if status == 'Failed':
-      message = response['FailureReason']
-      logging.info('Training failed with the following error: {}'.format(message))
-      raise Exception('Training job failed')
-    logging.info("Training job is still in status: " + status)
-    time.sleep(poll_interval)
+    while(True):
+        response = client.describe_training_job(TrainingJobName=training_job_name)
+        status = response['TrainingJobStatus']
+        if status == 'Completed':
+            logging.info("Training job ended with status: " + status)
+            break
+        if status == 'Failed':
+            message = response['FailureReason']
+            logging.info(f'Training failed with the following error: {message}')
+            raise Exception('Training job failed')
+        logging.info("Training job is still in status: " + status)
+        time.sleep(poll_interval)
+
+
+def wait_for_debug_rules(client, training_job_name, poll_interval=30):
+    first_poll = True
+    while(True):
+        response = client.describe_training_job(TrainingJobName=training_job_name)
+        if 'DebugRuleEvaluationStatuses' not in response:
+            break
+        if first_poll:
+            logging.info("Polling for status of all debug rules:")
+            first_poll = False
+        if DebugRulesStatus.from_describe(response) != DebugRulesStatus.INPROGRESS:
+            logging.info("Rules have ended with status:\n")
+            print_debug_rule_status(response, True)
+            break
+        print_debug_rule_status(response)
+        time.sleep(poll_interval)
+
+
+class DebugRulesStatus(Enum):
+    COMPLETED = auto()
+    ERRORED = auto()
+    INPROGRESS = auto()
+
+    @classmethod
+    def from_describe(self, response):
+        has_error = False
+        for debug_rule in response['DebugRuleEvaluationStatuses']:
+            if debug_rule['RuleEvaluationStatus'] == "Error":
+                has_error = True
+            if debug_rule['RuleEvaluationStatus'] == "InProgress":
+                return DebugRulesStatus.INPROGRESS
+        if has_error:
+            return DebugRulesStatus.ERRORED
+        else:
+            return DebugRulesStatus.COMPLETED
+
+
+def print_debug_rule_status(response, last_print=False):
+    """
+    Example of DebugRuleEvaluationStatuses:
+    response['DebugRuleEvaluationStatuses'] =
+        [{
+            "RuleConfigurationName": "VanishingGradient",
+            "RuleEvaluationStatus": "IssuesFound",
+            "StatusDetails": "There was an issue."
+        }]
+
+    If last_print is False:
+    INFO:root: - LossNotDecreasing: InProgress
+    INFO:root: - Overtraining: NoIssuesFound
+    ERROR:root:- CustomGradientRule: Error
+
+    If last_print is True:
+    INFO:root: - LossNotDecreasing: IssuesFound
+    INFO:root:   - RuleEvaluationConditionMet: Evaluation of the rule LossNotDecreasing at step 10 resulted in the condition being met
+    """
+    for debug_rule in response['DebugRuleEvaluationStatuses']:
+        line_ending = "\n" if last_print else ""
+        if 'StatusDetails' in debug_rule:
+            status_details = f"- {debug_rule['StatusDetails'].rstrip()}{line_ending}"
+            line_ending = ""
+        else:
+            status_details = ""
+        rule_status = f"- {debug_rule['RuleConfigurationName']}: {debug_rule['RuleEvaluationStatus']}{line_ending}"
+        if debug_rule['RuleEvaluationStatus'] == "Error":
+            log = logging.error
+            status_padding = 1
+        else:
+            log = logging.info
+            status_padding = 2
+
+        log(f"{status_padding * ' '}{rule_status}")
+        if last_print and status_details:
+            log(f"{(status_padding + 2) * ' '}{status_details}")
+    print_log_header(50)
 
 
 def get_model_artifacts_from_job(client, job_name):
@@ -314,10 +405,13 @@ def get_image_from_job(client, job_name):
 
 
 def stop_training_job(client, job_name):
-    try:
-        client.stop_training_job(TrainingJobName=job_name)
-    except ClientError as e:
-        raise Exception(e.response['Error']['Message'])
+    response = client.describe_training_job(TrainingJobName=job_name)
+    if response["TrainingJobStatus"] == "InProgress":
+        try:
+            client.stop_training_job(TrainingJobName=job_name)
+            return job_name
+        except ClientError as e:
+            raise Exception(e.response['Error']['Message'])
 
 
 def create_model(client, args):
@@ -611,12 +705,12 @@ def create_hyperparameter_tuning_job_request(args):
         # TODO: Adjust this implementation to account for custom algorithm resources names that are the same as built-in algorithm names
         algo_name = args['algorithm_name'].lower().strip()
         if algo_name in built_in_algos.keys():
-            request['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'] = get_image_uri(args['region'], built_in_algos[algo_name])
+            request['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'] = retrieve(built_in_algos[algo_name], args['region'])
             request['TrainingJobDefinition']['AlgorithmSpecification'].pop('AlgorithmName')
             logging.warning('Algorithm name is found as an Amazon built-in algorithm. Using built-in algorithm.')
         # To give the user more leeway for built-in algorithm name inputs
         elif algo_name in built_in_algos.values():
-            request['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'] = get_image_uri(args['region'], algo_name)
+            request['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'] = retrieve(algo_name, args['region'])
             request['TrainingJobDefinition']['AlgorithmSpecification'].pop('AlgorithmName')
             logging.warning('Algorithm name is found as an Amazon built-in algorithm. Using built-in algorithm.')
         else:
@@ -1135,4 +1229,4 @@ def write_output(output_path, output_value, json_encode=False):
     write_value = json.dumps(output_value) if json_encode else output_value 
 
     Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-    Path(output_path).write_text(write_value)
\ No newline at end of file
+    Path(output_path).write_text(write_value)
diff --git a/components/aws/sagemaker/train/src/train.py b/components/aws/sagemaker/train/src/train.py
--- a/components/aws/sagemaker/train/src/train.py
+++ b/components/aws/sagemaker/train/src/train.py
@@ -40,6 +40,8 @@ def create_parser():
   parser.add_argument('--vpc_subnets', type=str, required=False, help='The ID of the subnets in the VPC to which you want to connect your hpo job.')
   parser.add_argument('--network_isolation', type=_utils.str_to_bool, required=False, help='Isolates the training container.', default=True)
   parser.add_argument('--traffic_encryption', type=_utils.str_to_bool, required=False, help='Encrypts all communications between ML compute instances in distributed training.', default=False)
+  parser.add_argument('--debug_hook_config', type=_utils.yaml_or_json_str, required=False, help='Configuration information for the debug hook parameters, collection configuration, and storage paths.', default={})
+  parser.add_argument('--debug_rule_config', type=_utils.yaml_or_json_str, required=False, help='Configuration information for debugging rules.', default=[])
 
   ### Start spot instance support
   parser.add_argument('--spot_instance', type=_utils.str_to_bool, required=False, help='Use managed spot training.', default=False)
@@ -68,13 +70,15 @@ def main(argv=None):
   job_name = _utils.create_training_job(client, vars(args))
 
   def signal_term_handler(signalNumber, frame):
-    _utils.stop_training_job(client, job_name)
-    logging.info(f"Training Job: {job_name} request submitted to Stop")
+    job_stopped = _utils.stop_training_job(client, job_name)
+    if job_stopped:
+        logging.info(f"Training Job: {job_stopped} request submitted to Stop")
   signal.signal(signal.SIGTERM, signal_term_handler)
 
   logging.info('Job request submitted. Waiting for completion...')
   try:
     _utils.wait_for_training_job(client, job_name)
+    _utils.wait_for_debug_rules(client, job_name)
   except:
     raise
   finally:
diff --git a/samples/contrib/aws-samples/sagemaker_debugger_demo/debugger-training-pipeline.py b/samples/contrib/aws-samples/sagemaker_debugger_demo/debugger-training-pipeline.py
new file mode 100644
--- /dev/null
+++ b/samples/contrib/aws-samples/sagemaker_debugger_demo/debugger-training-pipeline.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+
+import kfp
+import json
+import os
+import copy
+from kfp import components
+from kfp import dsl
+
+
+cur_file_dir = os.path.dirname(__file__)
+components_dir = os.path.join(cur_file_dir, '../../../../components/aws/sagemaker/')
+
+sagemaker_train_op = components.load_component_from_file(components_dir + '/train/component.yaml')
+
+def training_input(input_name, s3_uri, content_type):
+    return {
+        "ChannelName": input_name,
+        "DataSource": {"S3DataSource": {"S3Uri": s3_uri, "S3DataType": "S3Prefix"}},
+        "ContentType": content_type
+    }
+
+
+def training_debug_hook(s3_uri, collection_dict):
+    return {
+        'S3OutputPath': s3_uri,
+        'CollectionConfigurations': format_collection_config(collection_dict)
+    }
+
+
+def format_collection_config(collection_dict):
+    output = []
+    for key, val in collection_dict.items():
+        output.append({'CollectionName': key, 'CollectionParameters': val})
+    return output
+
+
+def training_debug_rules(rule_name, parameters):
+    return {
+        'RuleConfigurationName': rule_name,
+        'RuleEvaluatorImage': '503895931360.dkr.ecr.us-east-1.amazonaws.com/sagemaker-debugger-rules:latest',
+        'RuleParameters': parameters
+    }
+
+
+collections = {
+    'feature_importance' : {
+        'save_interval': '5'
+    },
+    'losses' : {
+        'save_interval': '10'
+    },
+    'average_shap': {
+        'save_interval': '5'
+    },
+    'metrics': {
+        'save_interval': '3'
+    }
+}
+
+
+bad_hyperparameters = {
+    'max_depth': '5',
+    'eta': '0',
+    'gamma': '4',
+    'min_child_weight': '6',
+    'silent': '0',
+    'subsample': '0.7',
+    'num_round': '50'
+}
+
+
+@dsl.pipeline(
+    name='XGBoost Training Pipeline with bad hyperparameters',
+    description='SageMaker training job test with debugger'
+)
+def training(role_arn="", bucket_name="my-bucket"):
+    train_channels = [
+        training_input("train", f"s3://{bucket_name}/mnist_kmeans_example/input/valid_data.csv", 'text/csv')
+    ]
+    train_debug_rules = [
+        training_debug_rules("LossNotDecreasing", {"rule_to_invoke": "LossNotDecreasing", "tensor_regex": ".*"}),
+        training_debug_rules("Overtraining", {'rule_to_invoke': 'Overtraining', 'patience_train': '10', 'patience_validation': '20'}),
+    ]
+    training = sagemaker_train_op(
+        region='us-east-1',
+        image='683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-xgboost:0.90-2-cpu-py3',
+        hyperparameters=bad_hyperparameters,
+        channels=train_channels,
+        instance_type='ml.m5.2xlarge',
+        model_artifact_path=f's3://{bucket_name}/mnist_kmeans_example/output/model',
+        debug_hook_config=training_debug_hook(f's3://{bucket_name}/mnist_kmeans_example/hook_config', collections),
+        debug_rule_config=train_debug_rules,
+        role=role_arn,
+    )
+
+
+if __name__ == '__main__':
+    kfp.compiler.Compiler().compile(training, __file__ + '.zip')


Test patch:
diff --git a/components/aws/sagemaker/tests/integration_tests/README.md b/components/aws/sagemaker/tests/integration_tests/README.md
--- a/components/aws/sagemaker/tests/integration_tests/README.md
+++ b/components/aws/sagemaker/tests/integration_tests/README.md
@@ -9,7 +9,7 @@
 
 1. In the following Python script, change the bucket name and run the [`s3_sample_data_creator.py`](https://github.com/kubeflow/pipelines/tree/master/samples/contrib/aws-samples/mnist-kmeans-sagemaker#the-sample-dataset) to create an S3 bucket with the sample mnist dataset in the region where you want to run the tests.
 2. To prepare the dataset for the SageMaker GroundTruth Component test, follow the steps in the `[GroundTruth Sample README](https://github.com/kubeflow/pipelines/tree/master/samples/contrib/aws-samples/ground_truth_pipeline_demo#prep-the-dataset-label-categories-and-ui-template)`.
-3. To prepare the processing script for the SageMaker Processing Component tests, upload the `scripts/kmeans_preprocessing.py` script to your bucket. This can be done by replacing `<my-bucket> with your bucket name and running `aws s3 cp scripts/kmeans_preprocessing.py s3://<my-bucket>/mnist_kmeans_example/processing_code/kmeans_preprocessing.py`
+3. To prepare the processing script for the SageMaker Processing Component tests, upload the `scripts/kmeans_preprocessing.py` script to your bucket. This can be done by replacing `<my-bucket>` with your bucket name and running `aws s3 cp scripts/kmeans_preprocessing.py s3://<my-bucket>/mnist_kmeans_example/processing_code/kmeans_preprocessing.py`
 
 
 ## Step to run integration tests
@@ -22,4 +22,4 @@
     1. Navigate to the root of this github directory.
     1. Run `docker build . -f components/aws/sagemaker/tests/integration_tests/Dockerfile -t amazon/integration_test`
 1. Run the image, injecting your environment variable files:
-    1. Run `docker run --env-file components/aws/sagemaker/tests/integration_tests/.env amazon/integration_test`
\ No newline at end of file
+    1. Run `docker run --env-file components/aws/sagemaker/tests/integration_tests/.env amazon/integration_test`
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py
@@ -13,9 +13,13 @@
         pytest.param(
             "resources/config/simple-mnist-training", marks=pytest.mark.canary_test
         ),
-        pytest.param("resources/config/fsx-mnist-training", marks=pytest.mark.fsx_test),
+        pytest.param(
+            "resources/config/fsx-mnist-training",
+            marks=pytest.mark.fsx_test
+        ),
         "resources/config/spot-sample-pipeline-training",
         "resources/config/assume-role-training",
+        "resources/config/xgboost-mnist-trainingjob-debugger"
     ],
 )
 def test_trainingjob(
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py
@@ -44,7 +44,12 @@ def create_workteamjob(
 
 @pytest.mark.parametrize(
     "test_file_dir",
-    [pytest.param("resources/config/create-workteam", marks=pytest.mark.canary_test)],
+    [
+        pytest.param(
+            "resources/config/create-workteam",
+            marks=pytest.mark.canary_test
+        )
+    ],
 )
 def test_workteamjob(
     kfp_client, experiment_id, region, sagemaker_client, test_file_dir
diff --git a/components/aws/sagemaker/tests/integration_tests/environment.yml b/components/aws/sagemaker/tests/integration_tests/environment.yml
--- a/components/aws/sagemaker/tests/integration_tests/environment.yml
+++ b/components/aws/sagemaker/tests/integration_tests/environment.yml
@@ -17,7 +17,7 @@ dependencies:
     - kubernetes==11.0.*
     - kfp==0.5.*
     - minio==5.0.10
-    - sagemaker==1.56.*
+    - sagemaker==2.1.*
     - ruamel.yaml==0.16.*
       
       
diff --git a/components/aws/sagemaker/tests/integration_tests/resources/config/xgboost-mnist-trainingjob-debugger/config.yaml b/components/aws/sagemaker/tests/integration_tests/resources/config/xgboost-mnist-trainingjob-debugger/config.yaml
new file mode 100644
--- /dev/null
+++ b/components/aws/sagemaker/tests/integration_tests/resources/config/xgboost-mnist-trainingjob-debugger/config.yaml
@@ -0,0 +1,7 @@
+PipelineDefinition: ../../../../../samples/contrib/aws-samples/sagemaker_debugger_demo/debugger-training-pipeline.py
+TestName: xgboost-mnist-trainingjob-debugger
+Timeout: 3600
+ExpectedTrainingImage: 683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-xgboost:0.90-2-cpu-py3
+Arguments:
+  bucket_name: ((DATA_BUCKET))
+  role_arn: ((ROLE_ARN))
diff --git a/components/aws/sagemaker/tests/integration_tests/resources/definition/training_pipeline.py b/components/aws/sagemaker/tests/integration_tests/resources/definition/training_pipeline.py
--- a/components/aws/sagemaker/tests/integration_tests/resources/definition/training_pipeline.py
+++ b/components/aws/sagemaker/tests/integration_tests/resources/definition/training_pipeline.py
@@ -24,7 +24,7 @@ def training_pipeline(
     traffic_encryption="",
     spot_instance="",
     max_wait_time="",
-    checkpoint_config="{}",
+    checkpoint_config="",
     vpc_security_group_ids="",
     vpc_subnets="",
     assume_role="",
diff --git a/components/aws/sagemaker/tests/integration_tests/utils/__init__.py b/components/aws/sagemaker/tests/integration_tests/utils/__init__.py
--- a/components/aws/sagemaker/tests/integration_tests/utils/__init__.py
+++ b/components/aws/sagemaker/tests/integration_tests/utils/__init__.py
@@ -7,7 +7,7 @@
 import string
 import shutil
 
-from sagemaker.amazon.amazon_estimator import get_image_uri
+from sagemaker.image_uris import retrieve
 
 
 def get_region():
@@ -42,12 +42,12 @@ def get_fsx_id():
     return os.environ.get("FSX_ID")
 
 
-def get_assume_role_arn():
-    return os.environ.get("ASSUME_ROLE_ARN")
+def get_algorithm_image_registry(framework, region, version=None):
+    return retrieve(framework, region, version).split(".")[0]
 
 
-def get_algorithm_image_registry(region, algorithm):
-    return get_image_uri(region, algorithm).split(".")[0]
+def get_assume_role_arn():
+    return os.environ.get("ASSUME_ROLE_ARN")
 
 
 def run_command(cmd, *popenargs, **kwargs):
@@ -84,7 +84,9 @@ def replace_placeholders(input_filename, output_filename):
         "((REGION))": region,
         "((ROLE_ARN))": get_role_arn(),
         "((DATA_BUCKET))": get_s3_data_bucket(),
-        "((KMEANS_REGISTRY))": get_algorithm_image_registry(region, "kmeans"),
+        "((KMEANS_REGISTRY))": get_algorithm_image_registry("kmeans", region, "1"),
+        "((XGBOOST_REGISTRY))": get_algorithm_image_registry("xgboost", region, "1.0-1"),
+        "((BUILTIN_RULE_IMAGE))": get_algorithm_image_registry("debugger", region),
         "((FSX_ID))": get_fsx_id(),
         "((FSX_SUBNET))": get_fsx_subnet(),
         "((FSX_SECURITY_GROUP))": get_fsx_security_group(),
diff --git a/components/aws/sagemaker/tests/unit_tests/requirements.txt b/components/aws/sagemaker/tests/unit_tests/requirements.txt
--- a/components/aws/sagemaker/tests/unit_tests/requirements.txt
+++ b/components/aws/sagemaker/tests/unit_tests/requirements.txt
@@ -1,6 +1,6 @@
-boto3==1.12.33
+boto3==1.14.12
 coverage==5.1
 pathlib2==2.3.5
 pytest==5.4.1
 pyyaml==5.3.1
-sagemaker==1.56.1
\ No newline at end of file
+sagemaker==2.1.0
diff --git a/components/aws/sagemaker/tests/unit_tests/tests/test_hpo.py b/components/aws/sagemaker/tests/unit_tests/tests/test_hpo.py
--- a/components/aws/sagemaker/tests/unit_tests/tests/test_hpo.py
+++ b/components/aws/sagemaker/tests/unit_tests/tests/test_hpo.py
@@ -287,13 +287,13 @@ def test_known_algorithm_key(self):
 
     parsed_args = self.parser.parse_args(known_algorithm_args)
 
-    # Patch get_image_uri
-    _utils.get_image_uri = MagicMock()
-    _utils.get_image_uri.return_value = "seq2seq-url"
+    # Patch retrieve
+    _utils.retrieve = MagicMock()
+    _utils.retrieve.return_value = "seq2seq-url"
 
     response = _utils.create_hyperparameter_tuning_job_request(vars(parsed_args))
 
-    _utils.get_image_uri.assert_called_with('us-west-2', 'seq2seq')
+    _utils.retrieve.assert_called_with('seq2seq', 'us-west-2')
     self.assertEqual(response['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'], "seq2seq-url")
 
 
@@ -306,13 +306,13 @@ def test_known_algorithm_value(self):
 
     parsed_args = self.parser.parse_args(known_algorithm_args)
 
-    # Patch get_image_uri
-    _utils.get_image_uri = MagicMock()
-    _utils.get_image_uri.return_value = "seq2seq-url"
+    # Patch retrieve
+    _utils.retrieve = MagicMock()
+    _utils.retrieve.return_value = "seq2seq-url"
 
     response = _utils.create_hyperparameter_tuning_job_request(vars(parsed_args))
 
-    _utils.get_image_uri.assert_called_with('us-west-2', 'seq2seq')
+    _utils.retrieve.assert_called_with('seq2seq', 'us-west-2')
     self.assertEqual(response['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'], "seq2seq-url")
 
 
@@ -324,14 +324,14 @@ def test_unknown_algorithm(self):
 
     parsed_args = self.parser.parse_args(known_algorithm_args)
 
-    # Patch get_image_uri
-    _utils.get_image_uri = MagicMock()
-    _utils.get_image_uri.return_value = "unknown-url"
+    # Patch retrieve
+    _utils.retrieve = MagicMock()
+    _utils.retrieve.return_value = "unknown-url"
 
     response = _utils.create_hyperparameter_tuning_job_request(vars(parsed_args))
 
     # Should just place the algorithm name in regardless
-    _utils.get_image_uri.assert_not_called()
+    _utils.retrieve.assert_not_called()
     self.assertEqual(response['TrainingJobDefinition']['AlgorithmSpecification']['AlgorithmName'], "unknown algorithm")
 
   def test_no_channels(self):
diff --git a/components/aws/sagemaker/tests/unit_tests/tests/test_train.py b/components/aws/sagemaker/tests/unit_tests/tests/test_train.py
--- a/components/aws/sagemaker/tests/unit_tests/tests/test_train.py
+++ b/components/aws/sagemaker/tests/unit_tests/tests/test_train.py
@@ -112,14 +112,18 @@ def test_main_stop_training_job(self):
 
   def test_utils_stop_training_job(self):
     mock_sm_client = MagicMock()
-    mock_sm_client.stop_training_job.return_value = None
+    mock_sm_client.stop_training_job.return_value = 'FakeJobName'
+
+    mock_sm_client.describe_training_job.side_effect = [
+        {"TrainingJobStatus": "InProgress"}
+    ]
 
     response = _utils.stop_training_job(mock_sm_client, 'FakeJobName')
 
     mock_sm_client.stop_training_job.assert_called_once_with(
         TrainingJobName='FakeJobName'
     )
-    self.assertEqual(response, None)
+    self.assertEqual(response, 'FakeJobName')
 
   def test_sagemaker_exception_in_create_training_job(self):
     mock_client = MagicMock()
@@ -157,6 +161,28 @@ def test_wait_for_failed_job(self):
       _utils.wait_for_training_job(mock_client, 'training-job', 0)
 
     self.assertEqual(mock_client.describe_training_job.call_count, 4)
+ 
+  def test_wait_for_debug_rules(self):
+      mock_client = MagicMock()
+      mock_client.describe_training_job.side_effect = [
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "InProgress"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "InProgress"}]},
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "NoIssuesFound"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "InProgress"}]},
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "NoIssuesFound"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "IssuesFound"}]},
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "Should not be called"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "Should not be called"}]},
+      ]
+      _utils.wait_for_debug_rules(mock_client, 'training-job', 0)
+      self.assertEqual(mock_client.describe_training_job.call_count, 3)
+
+  def test_wait_for_errored_rule(self):
+      mock_client = MagicMock()
+      mock_client.describe_training_job.side_effect = [
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "InProgress"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "InProgress"}]},
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "Error"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "InProgress"}]},
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "Error"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "NoIssuesFound"}]},
+        {"DebugRuleEvaluationStatuses": [{"RuleConfigurationName": "rule1", "RuleEvaluationStatus": "Should not be called"}, {"RuleConfigurationName": "rule2", "RuleEvaluationStatus": "Should not be called"}]},
+      ]
+      _utils.wait_for_debug_rules(mock_client, 'training-job', 0)
+      self.assertEqual(mock_client.describe_training_job.call_count, 3)
 
   def test_get_model_artifacts_from_job(self):
     mock_client = MagicMock()
@@ -233,13 +259,13 @@ def test_known_algorithm_key(self):
 
     parsed_args = self.parser.parse_args(known_algorithm_args)
 
-    # Patch get_image_uri
-    _utils.get_image_uri = MagicMock()
-    _utils.get_image_uri.return_value = "seq2seq-url"
+    # Patch retrieve
+    _utils.retrieve = MagicMock()
+    _utils.retrieve.return_value = "seq2seq-url"
 
     response = _utils.create_training_job_request(vars(parsed_args))
 
-    _utils.get_image_uri.assert_called_with('us-west-2', 'seq2seq')
+    _utils.retrieve.assert_called_with('seq2seq', 'us-west-2')
     self.assertEqual(response['AlgorithmSpecification']['TrainingImage'], "seq2seq-url")
 
   def test_known_algorithm_value(self):
@@ -251,13 +277,13 @@ def test_known_algorithm_value(self):
 
     parsed_args = self.parser.parse_args(known_algorithm_args)
 
-    # Patch get_image_uri
-    _utils.get_image_uri = MagicMock()
-    _utils.get_image_uri.return_value = "seq2seq-url"
+    # Patch retrieve
+    _utils.retrieve = MagicMock()
+    _utils.retrieve.return_value = "seq2seq-url"
 
     response = _utils.create_training_job_request(vars(parsed_args))
 
-    _utils.get_image_uri.assert_called_with('us-west-2', 'seq2seq')
+    _utils.retrieve.assert_called_with('seq2seq', 'us-west-2')
     self.assertEqual(response['AlgorithmSpecification']['TrainingImage'], "seq2seq-url")
 
   def test_unknown_algorithm(self):
@@ -268,14 +294,14 @@ def test_unknown_algorithm(self):
 
     parsed_args = self.parser.parse_args(known_algorithm_args)
 
-    # Patch get_image_uri
-    _utils.get_image_uri = MagicMock()
-    _utils.get_image_uri.return_value = "unknown-url"
+    # Patch retrieve
+    _utils.retrieve = MagicMock()
+    _utils.retrieve.return_value = "unknown-url"
 
     response = _utils.create_training_job_request(vars(parsed_args))
 
     # Should just place the algorithm name in regardless
-    _utils.get_image_uri.assert_not_called()
+    _utils.retrieve.assert_not_called()
     self.assertEqual(response['AlgorithmSpecification']['AlgorithmName'], "unknown algorithm")
 
   def test_no_channels(self):
@@ -342,6 +368,49 @@ def test_spot_bad_args(self):
       with self.assertRaises(Exception):
         _utils.create_training_job_request(vars(arg))
 
+  def test_hook_min_args(self):
+    good_args = self.parser.parse_args(required_args + ['--debug_hook_config', '{"S3OutputPath": "s3://fake-uri/"}'])
+    response = _utils.create_training_job_request(vars(good_args))
+    self.assertEqual(response['DebugHookConfig']['S3OutputPath'], "s3://fake-uri/")
+
+  def test_hook_max_args(self):
+    good_args = self.parser.parse_args(required_args + ['--debug_hook_config', '{"S3OutputPath": "s3://fake-uri/", "LocalPath": "/local/path/", "HookParameters": {"key": "value"}, "CollectionConfigurations": [{"CollectionName": "collection1", "CollectionParameters": {"key1": "value1"}}, {"CollectionName": "collection2", "CollectionParameters": {"key2": "value2", "key3": "value3"}}]}'])
+    response = _utils.create_training_job_request(vars(good_args))
+    self.assertEqual(response['DebugHookConfig']['S3OutputPath'], "s3://fake-uri/")
+    self.assertEqual(response['DebugHookConfig']['LocalPath'], "/local/path/")
+    self.assertEqual(response['DebugHookConfig']['HookParameters'], {"key": "value"})
+    self.assertEqual(response['DebugHookConfig']['CollectionConfigurations'], [
+      {
+        "CollectionName": "collection1",
+        "CollectionParameters": {
+          "key1": "value1"
+        }
+      }, {
+        "CollectionName": "collection2", 
+        "CollectionParameters": {
+          "key2": "value2",
+          "key3": "value3"
+        }
+      }
+    ])
+
+  def test_rule_max_args(self):
+    good_args = self.parser.parse_args(required_args + ['--debug_rule_config', '[{"InstanceType": "ml.m4.xlarge", "LocalPath": "/local/path/", "RuleConfigurationName": "rule_name", "RuleEvaluatorImage": "test-image", "RuleParameters": {"key1": "value1"}, "S3OutputPath": "s3://fake-uri/", "VolumeSizeInGB": 1}]'])
+    response = _utils.create_training_job_request(vars(good_args))
+    self.assertEqual(response['DebugRuleConfigurations'][0]['InstanceType'], 'ml.m4.xlarge')
+    self.assertEqual(response['DebugRuleConfigurations'][0]['LocalPath'], '/local/path/')
+    self.assertEqual(response['DebugRuleConfigurations'][0]['RuleConfigurationName'], 'rule_name')
+    self.assertEqual(response['DebugRuleConfigurations'][0]['RuleEvaluatorImage'], 'test-image')
+    self.assertEqual(response['DebugRuleConfigurations'][0]['RuleParameters'], {"key1": "value1"})
+    self.assertEqual(response['DebugRuleConfigurations'][0]['S3OutputPath'], 's3://fake-uri/')
+    self.assertEqual(response['DebugRuleConfigurations'][0]['VolumeSizeInGB'], 1)
+
+  def test_rule_min_good_args(self):
+    good_args = self.parser.parse_args(required_args + ['--debug_rule_config', '[{"RuleConfigurationName": "rule_name", "RuleEvaluatorImage": "test-image"}]'])
+    response = _utils.create_training_job_request(vars(good_args))
+    self.assertEqual(response['DebugRuleConfigurations'][0]['RuleConfigurationName'], 'rule_name')
+    self.assertEqual(response['DebugRuleConfigurations'][0]['RuleEvaluatorImage'], 'test-image')
+
   def test_spot_lesser_wait_time(self):
     args = self.parser.parse_args(required_args + ['--spot_instance', 'True', '--max_wait_time', '3599', '--checkpoint_config', '{"S3Uri": "s3://fake-uri/", "LocalPath": "local-path"}'])
     with self.assertRaises(Exception):