SageMaker Debugger revision (#253)

* Add get_rule_container_image_uri and rename container_local_path parameter
* Fix bug with merging rule parameters
* Remove syntax incompatible with Python 2.7
* Add more detail to RuntimeError message
* Remove rule evaluator ECR repo from mapping
* Do not pass instance type and volume size for 1P rules
* Default CollectionConfig.parameters to None
* Temporarily comment out Debugger integ tests until API changes are in Prod
* Fix linting violations
* Fix unit tests to account for changes in S3OutputPath
* Minor change to trigger a new build
* Remove sagemaker_gamma_session
* Remove instance type and volume size from Sagemaker.rule classmethod
* Add fixtures to test_processing to fix cross-region errors

Author: ajaykarpur

src/sagemaker/debugger.py

@@ -23,6 +23,44 @@
 import smdebug_rulesconfig as rule_configs  # noqa: F401 # pylint: disable=unused-import
 
 
+RULES_ECR_REPO_NAME = "sagemaker-debugger-rules"
+
+SAGEMAKER_RULE_CONTAINERS_ACCOUNTS_MAP = {
+    "eu-north-1": {RULES_ECR_REPO_NAME: "314864569078"},
+    "me-south-1": {RULES_ECR_REPO_NAME: "986000313247"},
+    "ap-south-1": {RULES_ECR_REPO_NAME: "904829902805"},
+    "eu-west-3": {RULES_ECR_REPO_NAME: "447278800020"},
+    "us-east-2": {RULES_ECR_REPO_NAME: "915447279597"},
+    "eu-west-1": {RULES_ECR_REPO_NAME: "929884845733"},
+    "eu-central-1": {RULES_ECR_REPO_NAME: "482524230118"},
+    "sa-east-1": {RULES_ECR_REPO_NAME: "818342061345"},
+    "ap-east-1": {RULES_ECR_REPO_NAME: "199566480951"},
+    "us-east-1": {RULES_ECR_REPO_NAME: "503895931360"},
+    "ap-northeast-2": {RULES_ECR_REPO_NAME: "578805364391"},
+    "eu-west-2": {RULES_ECR_REPO_NAME: "250201462417"},
+    "ap-northeast-1": {RULES_ECR_REPO_NAME: "430734990657"},
+    "us-west-2": {RULES_ECR_REPO_NAME: "895741380848"},
+    "us-west-1": {RULES_ECR_REPO_NAME: "685455198987"},
+    "ap-southeast-1": {RULES_ECR_REPO_NAME: "972752614525"},
+    "ap-southeast-2": {RULES_ECR_REPO_NAME: "184798709955"},
+    "ca-central-1": {RULES_ECR_REPO_NAME: "519511493484"},
+}
+
+
+def get_rule_container_image_uri(region):
+    """
+    Returns the rule image uri for the given AWS region and rule type
+
+    Args:
+        region: AWS Region
+
+    Returns:
+        str: Formatted image uri for the given region and the rule container type
+    """
+    registry_id = SAGEMAKER_RULE_CONTAINERS_ACCOUNTS_MAP.get(region).get(RULES_ECR_REPO_NAME)
+    return "{}.dkr.ecr.{}.amazonaws.com/{}:latest".format(registry_id, region, RULES_ECR_REPO_NAME)
+
+
 class Rule(object):
     """Rules analyze tensors emitted during the training of a model. They
     monitor conditions that are critical for the success of a training job.
@@ -40,7 +78,7 @@ def __init__(
         name,
         image_uri,
         instance_type,
-        container_local_path,
+        container_local_output_path,
         s3_output_path,
         volume_size_in_gb,
         rule_parameters,
@@ -58,7 +96,7 @@ def __init__(
             image_uri (str): The URI of the image to be used by the debugger rule.
             instance_type (str): Type of EC2 instance to use, for example,
                 'ml.c4.xlarge'.
-            container_local_path (str): The path in the container.
+            container_local_output_path (str): The local path to store the Rule output.
             s3_output_path (str): The location in S3 to store the output.
             volume_size_in_gb (int): Size in GB of the EBS volume
                 to use for storing data.
@@ -68,7 +106,7 @@ def __init__(
         """
         self.name = name
         self.instance_type = instance_type
-        self.container_local_path = container_local_path
+        self.container_local_output_path = container_local_output_path
         self.s3_output_path = s3_output_path
         self.volume_size_in_gb = volume_size_in_gb
         self.rule_parameters = rule_parameters
@@ -80,10 +118,8 @@ def sagemaker(
         cls,
         base_config,
         name=None,
-        instance_type=None,
-        container_local_path=None,
+        container_local_output_path=None,
         s3_output_path=None,
-        volume_size_in_gb=None,
         other_trials_s3_input_paths=None,
         rule_parameters=None,
         collections_to_save=None,
@@ -98,13 +134,8 @@ def sagemaker(
                 built-in list of rules. For example, 'rule_configs.dead_relu()'.
             name (str): The name of the debugger rule. If one is not provided,
                 the name of the base_config will be used.
-            instance_type (str): Type of EC2 instance to use, for example,
-                'ml.c4.xlarge'. If one is not provided, the instance type from
-                the base_config will be used.
-            container_local_path (str): The path in the container.
+            container_local_output_path (str): The path in the container.
             s3_output_path (str): The location in S3 to store the output.
-            volume_size_in_gb (int): Size in GB of the EBS volume
-                to use for storing data.
             other_trials_s3_input_paths ([str]): S3 input paths for other trials.
             rule_parameters (dict): A dictionary of parameters for the rule.
             collections_to_save ([sagemaker.debugger.CollectionConfig]): A list
@@ -113,9 +144,22 @@ def sagemaker(
         Returns:
             sagemaker.debugger.Rule: The instance of the built-in Rule.
         """
-        other_trials_params = {}
+        merged_rule_params = {}
+
+        if rule_parameters is not None and rule_parameters.get("rule_to_invoke") is not None:
+            raise RuntimeError(
+                """You cannot provide a 'rule_to_invoke' for SageMaker rules.
+                Please either remove the rule_to_invoke or use a custom rule.
+                """
+            )
+
         if other_trials_s3_input_paths is not None:
-            other_trials_params["other_trials_s3_input_paths"] = other_trials_s3_input_paths
+            for index, s3_input_path in enumerate(other_trials_s3_input_paths):
+                merged_rule_params["other_trial_{}".format(str(index))] = s3_input_path
+
+        default_rule_params = base_config["DebugRuleConfiguration"].get("RuleParameters", {})
+        merged_rule_params.update(default_rule_params)
+        merged_rule_params.update(rule_parameters or {})
 
         base_config_collections = []
         for config in base_config.get("CollectionConfigurations", []):
@@ -133,16 +177,11 @@ def sagemaker(
         return cls(
             name=name or base_config["DebugRuleConfiguration"].get("RuleConfigurationName"),
             image_uri="DEFAULT_RULE_EVALUATOR_IMAGE",
-            instance_type=instance_type or "t3.medium",
-            # TODO-reinvent-2019 [akarpur]: Remove t3.medium from line above,
-            #  uncomment line below when 1P package updated
-            # or base_config["DebugRuleConfiguration"].get("InstanceType"),
-            container_local_path=container_local_path,
+            instance_type=None,
+            container_local_output_path=container_local_output_path,
             s3_output_path=s3_output_path,
-            volume_size_in_gb=volume_size_in_gb,
-            rule_parameters=other_trials_params.update(
-                rule_parameters or base_config["DebugRuleConfiguration"].get("RuleParameters", {})
-            ),
+            volume_size_in_gb=None,
+            rule_parameters=merged_rule_params,
             collections_to_save=collections_to_save or base_config_collections,
         )
 
@@ -154,7 +193,7 @@ def custom(
         instance_type,
         source=None,
         rule_to_invoke=None,
-        container_local_path=None,
+        container_local_output_path=None,
         s3_output_path=None,
         volume_size_in_gb=None,
         other_trials_s3_input_paths=None,
@@ -175,7 +214,7 @@ def custom(
                 you must also provide rule_to_invoke.
             rule_to_invoke (str): The name of the rule to invoke within the source.
                 If provided, you must also provide source.
-            container_local_path (str): The path in the container.
+            container_local_output_path (str): The path in the container.
             s3_output_path (str): The location in S3 to store the output.
             volume_size_in_gb (int): Size in GB of the EBS volume
                 to use for storing data.
@@ -192,25 +231,28 @@ def custom(
                 "If you provide a source, you must also provide a rule to invoke (and vice versa)."
             )
 
-        source_params = {}
+        merged_rule_params = {}
+
         if source is not None and rule_to_invoke is not None:
-            source_params["source_s3_uri"] = source
-            source_params["rule_to_invoke"] = rule_to_invoke
+            merged_rule_params["source_s3_uri"] = source
+            merged_rule_params["rule_to_invoke"] = rule_to_invoke
 
         other_trials_params = {}
         if other_trials_s3_input_paths is not None:
-            other_trials_params["other_trials_s3_input_paths"] = other_trials_s3_input_paths
+            for index, s3_input_path in enumerate(other_trials_s3_input_paths):
+                other_trials_params["other_trial_{}".format(str(index))] = s3_input_path
 
-        combined_rule_params = source_params.update(other_trials_params) or {}
+        merged_rule_params.update(other_trials_params)
+        merged_rule_params.update(rule_parameters or {})
 
         return cls(
             name=name,
             image_uri=image_uri,
             instance_type=instance_type,
-            container_local_path=container_local_path,
+            container_local_output_path=container_local_output_path,
             s3_output_path=s3_output_path,
             volume_size_in_gb=volume_size_in_gb,
-            rule_parameters=combined_rule_params.update(rule_parameters or {}),
+            rule_parameters=merged_rule_params,
             collections_to_save=collections_to_save or [],
         )
 
@@ -221,21 +263,19 @@ def to_debugger_rule_config_dict(self):
         Returns:
             dict: An portion of an API request as a dictionary.
         """
-        if self.instance_type is None or self.volume_size_in_gb is None:
-            raise RuntimeError(
-                """Cannot create a dictionary if the instance type and volume size are not provided.
-            Please set the instance type and volume size for this Rule object."""
-            )
-
         debugger_rule_config_request = {
             "RuleConfigurationName": self.name,
             "RuleEvaluatorImage": self.image_uri,
-            "InstanceType": self.instance_type,
-            "VolumeSizeInGB": self.volume_size_in_gb,
         }
 
-        if self.container_local_path is not None:
-            debugger_rule_config_request["LocalPath"] = self.container_local_path
+        if self.instance_type is not None:
+            debugger_rule_config_request["InstanceType"] = self.instance_type
+
+        if self.volume_size_in_gb is not None:
+            debugger_rule_config_request["VolumeSizeInGB"] = self.volume_size_in_gb
+
+        if self.container_local_output_path is not None:
+            debugger_rule_config_request["LocalPath"] = self.container_local_output_path
 
         if self.s3_output_path is not None:
             debugger_rule_config_request["S3OutputPath"] = self.s3_output_path
@@ -254,7 +294,7 @@ class DebuggerHookConfig(object):
     def __init__(
         self,
         s3_output_path,
-        container_local_path=None,
+        container_local_output_path=None,
         hook_parameters=None,
         collection_configs=None,
     ):
@@ -264,13 +304,13 @@ def __init__(
 
         Args:
             s3_output_path (str): The location in S3 to store the output.
-            container_local_path (str): The path in the container.
+            container_local_output_path (str): The path in the container.
             hook_parameters (dict): A dictionary of parameters.
             collection_configs ([sagemaker.debugger.CollectionConfig]): A list
                 of CollectionConfig objects to be provided to the API.
         """
         self.s3_output_path = s3_output_path
-        self.container_local_path = container_local_path
+        self.container_local_output_path = container_local_output_path
         self.hook_parameters = hook_parameters
         self.collection_configs = collection_configs
 
@@ -283,8 +323,8 @@ def to_request_dict(self):
         """
         debugger_hook_config_request = {"S3OutputPath": self.s3_output_path}
 
-        if self.container_local_path is not None:
-            debugger_hook_config_request["LocalPath"] = self.container_local_path
+        if self.container_local_output_path is not None:
+            debugger_hook_config_request["LocalPath"] = self.container_local_output_path
 
         if self.hook_parameters is not None:
             debugger_hook_config_request["HookParameters"] = self.hook_parameters
@@ -301,17 +341,17 @@ class TensorBoardOutputConfig(object):
     """TensorBoardOutputConfig provides options to customize
     debugging visualization using TensorBoard."""
 
-    def __init__(self, s3_output_path, container_local_path=None):
+    def __init__(self, s3_output_path, container_local_output_path=None):
         """Initialize an instance of TensorBoardOutputConfig.
         TensorBoardOutputConfig provides options to customize
         debugging visualization using TensorBoard.
 
         Args:
             s3_output_path (str): The location in S3 to store the output.
-            container_local_path (str): The path in the container.
+            container_local_output_path (str): The path in the container.
         """
         self.s3_output_path = s3_output_path
-        self.container_local_path = container_local_path
+        self.container_local_output_path = container_local_output_path
 
     def to_request_dict(self):
         """Generates a request dictionary using the parameters provided
@@ -322,16 +362,16 @@ def to_request_dict(self):
         """
         tensorboard_output_config_request = {"S3OutputPath": self.s3_output_path}
 
-        if self.container_local_path is not None:
-            tensorboard_output_config_request["LocalPath"] = self.container_local_path
+        if self.container_local_output_path is not None:
+            tensorboard_output_config_request["LocalPath"] = self.container_local_output_path
 
         return tensorboard_output_config_request
 
 
 class CollectionConfig(object):
     """CollectionConfig object for SageMaker Debugger."""
 
-    def __init__(self, name, parameters):
+    def __init__(self, name, parameters=None):
         """Initialize a ``CollectionConfig`` object.
 
         Args:
@@ -359,7 +399,7 @@ def __ne__(self, other):
         return self.name != other.name or self.parameters != other.parameters
 
     def __hash__(self):
-        return hash((self.name, tuple(sorted(self.parameters.items()))))
+        return hash((self.name, tuple(sorted((self.parameters or {}).items()))))
 
     def to_request_dict(self):
         """Generates a request dictionary using the parameters provided
@@ -368,9 +408,9 @@ def to_request_dict(self):
         Returns:
             dict: An portion of an API request as a dictionary.
         """
-        collection_config_request = {
-            "CollectionName": self.name,
-            "CollectionParameters": self.parameters,
-        }
+        collection_config_request = {"CollectionName": self.name}
+
+        if self.parameters is not None:
+            collection_config_request["CollectionParameters"] = self.parameters
 
         return collection_config_request

src/sagemaker/estimator.py

@@ -27,7 +27,11 @@
 import sagemaker
 from sagemaker import git_utils
 from sagemaker.analytics import TrainingJobAnalytics
-from sagemaker.debugger import DebuggerHookConfig, TensorBoardOutputConfig
+from sagemaker.debugger import (
+    DebuggerHookConfig,
+    TensorBoardOutputConfig,
+    get_rule_container_image_uri,
+)
 from sagemaker.s3 import S3Uploader
 
 from sagemaker.fw_utils import (
@@ -319,9 +323,7 @@ def _prepare_for_training(self, job_name=None):
 
         # Prepare rules and debugger configs for training.
         if self.rules and not self.debugger_hook_config:
-            self.debugger_hook_config = DebuggerHookConfig(
-                s3_output_path=os.path.join(self.output_path, self._current_job_name, "tensors")
-            )
+            self.debugger_hook_config = DebuggerHookConfig(s3_output_path=self.output_path)
         self._prepare_rules()
         self._prepare_collection_configs()
 
@@ -332,14 +334,16 @@ def _prepare_rules(self):
             # Iterate through each of the provided rules.
             for rule in self.rules:
                 # Set the instance type and volume size using the Estimator's defaults.
-                rule.instance_type = self.train_instance_type
-                rule.volume_size_in_gb = self.train_volume_size
                 # Set the image URI using the default rule evaluator image and the region.
                 if rule.image_uri == "DEFAULT_RULE_EVALUATOR_IMAGE":
-                    # TODO-reinvent-2019 [akarpur]: Replace this image URI with the latest version
-                    rule.image_uri = "453379255795.dkr.ecr.{}.amazonaws.com/script-rule-executor:latest".format(
+                    rule.image_uri = get_rule_container_image_uri(
                         self.sagemaker_session.boto_region_name
                     )
+                    rule.instance_type = None
+                    rule.volume_size_in_gb = None
+                else:
+                    rule.instance_type = self.train_instance_type
+                    rule.volume_size_in_gb = self.train_volume_size
                 # If source was provided as a rule parameter, upload to S3 and save the S3 uri.
                 if "source_s3_uri" in (rule.rule_parameters or {}):
                     parse_result = urlparse(rule.rule_parameters["source_s3_uri"])
@@ -1590,12 +1594,10 @@ def _prepare_for_training(self, job_name=None):
 
         # Set defaults for debugging.
         if self.debugger_hook_config is None:
-            self.debugger_hook_config = DebuggerHookConfig(
-                s3_output_path=os.path.join(self.output_path, self._current_job_name, "tensors")
-            )
+            self.debugger_hook_config = DebuggerHookConfig(s3_output_path=self.output_path)
         if self.tensorboard_output_config is None:
             self.tensorboard_output_config = TensorBoardOutputConfig(
-                s3_output_path=os.path.join(self.output_path, self._current_job_name, "tensorboard")
+                s3_output_path=self.output_path
             )
 
     def _stage_user_code_in_s3(self):