change: Replace update_remote_config with 2 helper methods for enable and disable respectively

Author: xixinyu-aws

src/sagemaker/estimator.py

@@ -769,7 +769,7 @@ def __init__(
 
         self.tensorboard_app = TensorBoardApp(region=self.sagemaker_session.boto_region_name)
 
-        self.enable_remote_debug = enable_remote_debug
+        self._enable_remote_debug = enable_remote_debug
 
     @abstractmethod
     def training_image_uri(self):
@@ -2291,21 +2291,31 @@ def update_profiler(
 
         _TrainingJob.update(self, profiler_rule_configs, profiler_config_request_dict)
 
-    def update_remote_debug(self, enable_remote_debug: bool):
-        """Update training jobs to enable remote debug.
+    def get_remote_debug_config(self):
+        """dict: Return the configuration of RemoteDebug"""
+        return (
+            None
+            if self._enable_remote_debug is None
+            else {"EnableRemoteDebug": self._enable_remote_debug}
+        )
 
-        This method updates the ``enable_remote_debug`` parameter
-        and enables or disables remote debug for a training job
+    def enable_remote_debug(self):
+        """Enable remote debug for a training job."""
+        self._update_remote_debug(True)
 
-        Args:
-            enable_remote_debug (bool):
-            Specifies whether RemoteDebug is to be enabled for the training job
+    def disable_remote_debug(self):
+        """Disable remote debug for a training job."""
+        self._update_remote_debug(False)
+
+    def _update_remote_debug(self, enable_remote_debug: bool):
+        """Update to enable or disable remote debug for a training job.
+
+        This method updates the ``_enable_remote_debug`` parameter
+        and enables or disables remote debug for a training job
         """
         self._ensure_latest_training_job()
-        self.enable_remote_debug = enable_remote_debug
-        _TrainingJob.update(
-            self, remote_debug_config={"EnableRemoteDebug": self.enable_remote_debug}
-        )
+        _TrainingJob.update(self, remote_debug_config={"EnableRemoteDebug": enable_remote_debug})
+        self._enable_remote_debug = enable_remote_debug
 
     def get_app_url(
         self,
@@ -2535,8 +2545,8 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         if estimator.profiler_config:
             train_args["profiler_config"] = estimator.profiler_config._to_request_dict()
 
-        if estimator.enable_remote_debug is not None:
-            train_args["remote_debug_config"] = {"EnableRemoteDebug": estimator.enable_remote_debug}
+        if estimator.get_remote_debug_config() is not None:
+            train_args["remote_debug_config"] = estimator.get_remote_debug_config()
 
         return train_args
 

tests/unit/test_estimator.py

@@ -2009,9 +2009,47 @@ def test_framework_with_remote_debug_config(sagemaker_session):
     sagemaker_session.train.assert_called_once()
     _, args = sagemaker_session.train.call_args
     assert args["remote_debug_config"]["EnableRemoteDebug"]
+    assert f.get_remote_debug_config()["EnableRemoteDebug"]
 
 
-def test_framework_update_remote_debug(sagemaker_session):
+def test_framework_without_remote_debug_config(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args.get("remote_debug_config") is None
+    assert f.get_remote_debug_config() is None
+
+
+def test_framework_enable_remote_debug(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+    )
+    f.fit("s3://mydata")
+    f.enable_remote_debug()
+
+    sagemaker_session.update_training_job.assert_called_once()
+    _, args = sagemaker_session.update_training_job.call_args
+    assert args["remote_debug_config"] == {
+        "EnableRemoteDebug": True,
+    }
+    assert f.get_remote_debug_config()["EnableRemoteDebug"]
+    assert len(args) == 2
+
+
+def test_framework_disable_remote_debug(sagemaker_session):
     f = DummyFramework(
         entry_point=SCRIPT_PATH,
         role=ROLE,
@@ -2021,13 +2059,14 @@ def test_framework_update_remote_debug(sagemaker_session):
         enable_remote_debug=True,
     )
     f.fit("s3://mydata")
-    f.update_remote_debug(False)
+    f.disable_remote_debug()
 
     sagemaker_session.update_training_job.assert_called_once()
     _, args = sagemaker_session.update_training_job.call_args
     assert args["remote_debug_config"] == {
         "EnableRemoteDebug": False,
     }
+    assert not f.get_remote_debug_config()["EnableRemoteDebug"]
     assert len(args) == 2