fix: disable DebuggerHook and Rules for TF distributions (#290)

* fix: Disable hook for ps or mpi in TF
Co-Authored-By: Denis Davydenko <[email protected]>

Author: 2 people

src/sagemaker/estimator.py

@@ -1656,7 +1656,12 @@ def _prepare_for_training(self, job_name=None):
         self._hyperparameters[JOB_NAME_PARAM_NAME] = self._current_job_name
         self._hyperparameters[SAGEMAKER_REGION_PARAM_NAME] = self.sagemaker_session.boto_region_name
 
-        # Set defaults for debugging.
+        self._validate_and_set_debugger_configs()
+
+    def _validate_and_set_debugger_configs(self):
+        """
+        Set defaults for debugging
+        """
         if self.debugger_hook_config is None:
             self.debugger_hook_config = DebuggerHookConfig(s3_output_path=self.output_path)
 

src/sagemaker/tensorflow/estimator.py

@@ -22,6 +22,7 @@
 import threading
 import time
 
+from sagemaker.debugger import DebuggerHookConfig
 from sagemaker.estimator import Framework
 import sagemaker.fw_utils as fw
 from sagemaker.tensorflow.defaults import TF_VERSION
@@ -696,6 +697,31 @@ def _script_mode_enabled(self):
         """Placeholder docstring"""
         return self.py_version == "py3" or self.script_mode
 
+    def _validate_and_set_debugger_configs(self):
+        """
+        Disable Debugger Hook Config for PS and Horovod as they are not
+        supported in smdebug 0.4.13, the current latest version of smdebug
+
+        Else, set default HookConfig
+        """
+        ps_enabled = "parameter_server" in self.distributions and self.distributions[
+            "parameter_server"
+        ].get("enabled", False)
+        mpi_enabled = "mpi" in self.distributions and self.distributions["mpi"].get(
+            "enabled", False
+        )
+        if ps_enabled or mpi_enabled:
+            if self.debugger_hook_config is not None or self.debugger_rule_configs is not None:
+                logger.info(
+                    "Amazon SageMaker Debugger does not currently support "
+                    "Parameter Server and MPI distributions"
+                )
+            self.debugger_hook_config = None
+            self.debugger_rule_configs = None
+        elif self.debugger_hook_config is None:
+            # Set defaults for debugging.
+            self.debugger_hook_config = DebuggerHookConfig(s3_output_path=self.output_path)
+
     def train_image(self):
         """Placeholder docstring"""
         if self.image_name:

tests/unit/test_tf_estimator.py

@@ -123,9 +123,14 @@ def _hyperparameters(script_mode=False, horovod=False):
 
 
 def _create_train_job(
-    tf_version, script_mode=False, horovod=False, repo_name=IMAGE_REPO_NAME, py_version="py2"
+    tf_version,
+    script_mode=False,
+    horovod=False,
+    ps=False,
+    repo_name=IMAGE_REPO_NAME,
+    py_version="py2",
 ):
-    return {
+    conf = {
         "image": _get_full_cpu_image_uri(tf_version, repo=repo_name, py_version=py_version),
         "input_mode": "File",
         "input_config": [
@@ -153,11 +158,15 @@ def _create_train_job(
         "vpc_config": None,
         "metric_definitions": None,
         "experiment_config": None,
-        "debugger_hook_config": {
+    }
+
+    if not ps and not horovod:
+        conf["debugger_hook_config"] = {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
-        },
-    }
+        }
+
+    return conf
 
 
 def _build_tf(
@@ -1116,7 +1125,7 @@ def test_tf_script_mode_ps(time, strftime, sagemaker_session):
     assert call_names == ["train", "logs_for_job"]
 
     expected_train_args = _create_train_job(
-        "1.11", script_mode=True, repo_name=SM_IMAGE_REPO_NAME, py_version="py3"
+        "1.11", script_mode=True, ps=True, repo_name=SM_IMAGE_REPO_NAME, py_version="py3"
     )
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
     expected_train_args["hyperparameters"][TensorFlow.LAUNCH_PS_ENV_NAME] = json.dumps(True)