feature: support remote debug for sagemaker training job

Author: xixinyu-aws

src/sagemaker/estimator.py

@@ -178,6 +178,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
+        enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -540,6 +541,8 @@ def __init__(
                 to Amazon S3 without compression after training finishes.
             enable_infra_check (bool or PipelineVariable): Optional.
                 Specifies whether it is running Sagemaker built-in infra check jobs.
+            enable_remote_debug (bool or PipelineVariable): Optional.
+                Specifies whether RemoteDebug is enabled for the training job
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -766,6 +769,8 @@ def __init__(
 
         self.tensorboard_app = TensorBoardApp(region=self.sagemaker_session.boto_region_name)
 
+        self.enable_remote_debug = enable_remote_debug
+
     @abstractmethod
     def training_image_uri(self):
         """Return the Docker image to use for training.
@@ -1947,6 +1952,11 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
             max_wait = job_details.get("StoppingCondition", {}).get("MaxWaitTimeInSeconds")
             if max_wait:
                 init_params["max_wait"] = max_wait
+
+        if "RemoteDebugConfig" in job_details:
+            init_params["enable_remote_debug"] = job_details["RemoteDebugConfig"].get(
+                "EnableRemoteDebug"
+            )
         return init_params
 
     def _get_instance_type(self):
@@ -2281,6 +2291,22 @@ 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.
+
+        This method updates the ``enable_remote_debug`` parameter
+        and enables or disables remote debug for a training job
+
+        Args:
+            enable_remote_debug (bool):
+            Specifies whether RemoteDebug is to be enabled for the 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}
+        )
+
     def get_app_url(
         self,
         app_type,
@@ -2509,6 +2535,9 @@ 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}
+
         return train_args
 
     @classmethod
@@ -2538,7 +2567,12 @@ def _is_local_channel(cls, input_uri):
 
     @classmethod
     def update(
-        cls, estimator, profiler_rule_configs=None, profiler_config=None, resource_config=None
+        cls,
+        estimator,
+        profiler_rule_configs=None,
+        profiler_config=None,
+        resource_config=None,
+        remote_debug_config=None,
     ):
         """Update a running Amazon SageMaker training job.
 
@@ -2551,20 +2585,31 @@ def update(
             resource_config (dict): Configuration of the resources for the training job. You can
                 update the keep-alive period if the warm pool status is `Available`. No other fields
                 can be updated. (default: None).
+            remote_debug_config (dict): Configuration for RemoteDebug. (default: ``None``)
+                The dict can contain 'EnableRemoteDebug'(bool).
+                For example,
+
+                .. code:: python
+
+                    remote_debug_config = {
+                        "EnableRemoteDebug": True,
+                    } (default: None).
 
         Returns:
             sagemaker.estimator._TrainingJob: Constructed object that captures
             all information about the updated training job.
         """
         update_args = cls._get_update_args(
-            estimator, profiler_rule_configs, profiler_config, resource_config
+            estimator, profiler_rule_configs, profiler_config, resource_config, remote_debug_config
         )
         estimator.sagemaker_session.update_training_job(**update_args)
 
         return estimator.latest_training_job
 
     @classmethod
-    def _get_update_args(cls, estimator, profiler_rule_configs, profiler_config, resource_config):
+    def _get_update_args(
+        cls, estimator, profiler_rule_configs, profiler_config, resource_config, remote_debug_config
+    ):
         """Constructs a dict of arguments for updating an Amazon SageMaker training job.
 
         Args:
@@ -2585,6 +2630,7 @@ def _get_update_args(cls, estimator, profiler_rule_configs, profiler_config, res
         update_args.update(build_dict("profiler_rule_configs", profiler_rule_configs))
         update_args.update(build_dict("profiler_config", profiler_config))
         update_args.update(build_dict("resource_config", resource_config))
+        update_args.update(build_dict("remote_debug_config", remote_debug_config))
 
         return update_args
 
@@ -2683,6 +2729,7 @@ def __init__(
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
+        enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -3044,6 +3091,8 @@ def __init__(
                 to Amazon S3 without compression after training finishes.
             enable_infra_check (bool or PipelineVariable): Optional.
                 Specifies whether it is running Sagemaker built-in infra check jobs.
+            enable_remote_debug (bool or PipelineVariable): Optional.
+                Specifies whether RemoteDebug is enabled for the training job
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -3095,6 +3144,7 @@ def __init__(
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
+            enable_remote_debug=enable_remote_debug,
             **kwargs,
         )
 

src/sagemaker/session.py

@@ -741,6 +741,7 @@ def train(  # noqa: C901
         profiler_config=None,
         environment: Optional[Dict[str, str]] = None,
         retry_strategy=None,
+        remote_debug_config=None,
     ):
         """Create an Amazon SageMaker training job.
 
@@ -851,6 +852,15 @@ def train(  # noqa: C901
                 configurations.src/sagemaker/lineage/artifact.py:285
             profiler_config (dict): Configuration for how profiling information is emitted
                 with SageMaker Profiler. (default: ``None``).
+            remote_debug_config(dict): Configuration for RemoteDebug. (default: ``None``)
+                The dict can contain 'EnableRemoteDebug'(bool).
+                For example,
+
+                .. code:: python
+
+                    remote_debug_config = {
+                        "EnableRemoteDebug": True,
+                    }
             environment (dict[str, str]) : Environment variables to be set for
                 use during training job (default: ``None``)
             retry_strategy(dict): Defines RetryStrategy for InternalServerFailures.
@@ -943,6 +953,7 @@ def train(  # noqa: C901
             enable_sagemaker_metrics=enable_sagemaker_metrics,
             profiler_rule_configs=profiler_rule_configs,
             profiler_config=inferred_profiler_config,
+            remote_debug_config=remote_debug_config,
             environment=environment,
             retry_strategy=retry_strategy,
         )
@@ -985,6 +996,7 @@ def _get_train_request(  # noqa: C901
         enable_sagemaker_metrics=None,
         profiler_rule_configs=None,
         profiler_config=None,
+        remote_debug_config=None,
         environment=None,
         retry_strategy=None,
     ):
@@ -1096,6 +1108,15 @@ def _get_train_request(  # noqa: C901
             profiler_rule_configs (list[dict]): A list of profiler rule configurations.
             profiler_config(dict): Configuration for how profiling information is emitted with
                 SageMaker Profiler. (default: ``None``).
+            remote_debug_config(dict): Configuration for RemoteDebug. (default: ``None``)
+                The dict can contain 'EnableRemoteDebug'(bool).
+                For example,
+
+                .. code:: python
+
+                    remote_debug_config = {
+                        "EnableRemoteDebug": True,
+                    }
             environment (dict[str, str]) : Environment variables to be set for
                 use during training job (default: ``None``)
             retry_strategy(dict): Defines RetryStrategy for InternalServerFailures.
@@ -1199,6 +1220,9 @@ def _get_train_request(  # noqa: C901
         if profiler_config is not None:
             train_request["ProfilerConfig"] = profiler_config
 
+        if remote_debug_config is not None:
+            train_request["RemoteDebugConfig"] = remote_debug_config
+
         if retry_strategy is not None:
             train_request["RetryStrategy"] = retry_strategy
 
@@ -1210,6 +1234,7 @@ def update_training_job(
         profiler_rule_configs=None,
         profiler_config=None,
         resource_config=None,
+        remote_debug_config=None,
     ):
         """Calls the UpdateTrainingJob API for the given job name and returns the response.
 
@@ -1221,6 +1246,15 @@ def update_training_job(
             resource_config (dict): Configuration of the resources for the training job. You can
                 update the keep-alive period if the warm pool status is `Available`. No other fields
                 can be updated. (default: ``None``).
+            remote_debug_config(dict): Configuration for RemoteDebug. (default: ``None``)
+                The dict can contain 'EnableRemoteDebug'(bool).
+                For example,
+
+                .. code:: python
+
+                    remote_debug_config = {
+                        "EnableRemoteDebug": True,
+                    }
         """
         # No injections from sagemaker_config because the UpdateTrainingJob API's resource_config
         # object accepts fewer parameters than the CreateTrainingJob API, and none that the
@@ -1233,6 +1267,7 @@ def update_training_job(
             profiler_rule_configs=profiler_rule_configs,
             profiler_config=inferred_profiler_config,
             resource_config=resource_config,
+            remote_debug_config=remote_debug_config,
         )
         LOGGER.info("Updating training job with name %s", job_name)
         LOGGER.debug("Update request: %s", json.dumps(update_training_job_request, indent=4))
@@ -1244,6 +1279,7 @@ def _get_update_training_job_request(
         profiler_rule_configs=None,
         profiler_config=None,
         resource_config=None,
+        remote_debug_config=None,
     ):
         """Constructs a request compatible for updating an Amazon SageMaker training job.
 
@@ -1255,6 +1291,15 @@ def _get_update_training_job_request(
             resource_config (dict): Configuration of the resources for the training job. You can
                 update the keep-alive period if the warm pool status is `Available`. No other fields
                 can be updated. (default: ``None``).
+            remote_debug_config(dict): Configuration for RemoteDebug. (default: ``None``)
+                The dict can contain 'EnableRemoteDebug'(bool).
+                For example,
+
+                .. code:: python
+
+                    remote_debug_config = {
+                        "EnableRemoteDebug": True,
+                    }
 
         Returns:
             Dict: an update training request dict
@@ -1272,6 +1317,9 @@ def _get_update_training_job_request(
         if resource_config is not None:
             update_training_job_request["ResourceConfig"] = resource_config
 
+        if remote_debug_config is not None:
+            update_training_job_request["RemoteDebugConfig"] = remote_debug_config
+
         return update_training_job_request
 
     def process(

tests/unit/test_estimator.py

@@ -1994,6 +1994,43 @@ def test_sagemaker_model_custom_channel_name(sagemaker_session):
     ]
 
 
+def test_framework_with_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),
+        ],
+        enable_remote_debug=True,
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args["remote_debug_config"]["EnableRemoteDebug"]
+
+
+def test_framework_update_remote_debug(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        enable_remote_debug=True,
+    )
+    f.fit("s3://mydata")
+    f.update_remote_debug(False)
+
+    sagemaker_session.update_training_job.assert_called_once()
+    _, args = sagemaker_session.update_training_job.call_args
+    assert args["remote_debug_config"] == {
+        "EnableRemoteDebug": False,
+    }
+    assert len(args) == 2
+
+
 @patch("time.strftime", return_value=TIMESTAMP)
 def test_custom_code_bucket(time, sagemaker_session):
     code_bucket = "codebucket"

tests/unit/test_session.py

@@ -1876,6 +1876,15 @@ def test_update_training_job_with_sagemaker_config_injection(sagemaker_session):
     )
 
 
+def test_update_training_job_with_remote_debug_config(sagemaker_session):
+    sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_TRAINING_JOB
+    sagemaker_session.update_training_job(
+        job_name="MyTestJob", remote_debug_config={"EnableRemoteDebug": False}
+    )
+    _, _, actual_train_args = sagemaker_session.sagemaker_client.method_calls[0]
+    assert not actual_train_args["RemoteDebugConfig"]["EnableRemoteDebug"]
+
+
 def test_train_with_sagemaker_config_injection(sagemaker_session):
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_TRAINING_JOB
 
@@ -2128,6 +2137,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
     }
     CONTAINER_ENTRY_POINT = ["bin/bash", "test.sh"]
     CONTAINER_ARGUMENTS = ["--arg1", "value1", "--arg2", "value2"]
+    remote_debug_config = {"EnableRemoteDebug": True}
 
     sagemaker_session.train(
         image_uri=IMAGE,
@@ -2152,6 +2162,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
         training_image_config=TRAINING_IMAGE_CONFIG,
         container_entry_point=CONTAINER_ENTRY_POINT,
         container_arguments=CONTAINER_ARGUMENTS,
+        remote_debug_config=remote_debug_config,
     )
 
     _, _, actual_train_args = sagemaker_session.sagemaker_client.method_calls[0]
@@ -2174,6 +2185,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
         actual_train_args["AlgorithmSpecification"]["ContainerEntrypoint"] == CONTAINER_ENTRY_POINT
     )
     assert actual_train_args["AlgorithmSpecification"]["ContainerArguments"] == CONTAINER_ARGUMENTS
+    assert actual_train_args["RemoteDebugConfig"]["EnableRemoteDebug"]
 
 
 def test_create_transform_job_with_sagemaker_config_injection(sagemaker_session):