fix: Set flag when debugger is disabled (#2520)

* Set flag when debugger is disabled

* Address edge case

* Address edge case

* Address edge case

* Fix formatting errors

* retrigger CI

Co-authored-by: Shreya Pandit <[email protected]>

Author: ndodda-amazon

src/sagemaker/debugger/__init__.py

@@ -15,6 +15,7 @@
 
 from sagemaker.debugger.debugger import (  # noqa: F401
     CollectionConfig,
+    DEBUGGER_FLAG,
     DebuggerHookConfig,
     framework_name,
     get_default_profiler_rule,

src/sagemaker/debugger/debugger.py

@@ -32,6 +32,7 @@
 from sagemaker.utils import build_dict
 
 framework_name = "debugger"
+DEBUGGER_FLAG = "USE_SMDEBUG"
 
 
 def get_rule_container_image_uri(region):

src/sagemaker/estimator.py

@@ -29,6 +29,7 @@
 from sagemaker.analytics import TrainingJobAnalytics
 from sagemaker.debugger import TensorBoardOutputConfig  # noqa: F401 # pylint: disable=unused-import
 from sagemaker.debugger import (
+    DEBUGGER_FLAG,
     DebuggerHookConfig,
     FrameworkProfile,
     get_default_profiler_rule,
@@ -2269,6 +2270,11 @@ def _validate_and_set_debugger_configs(self):
                     )
                     self.debugger_hook_config = False
 
+        if self.debugger_hook_config is False:
+            if self.environment is None:
+                self.environment = {}
+            self.environment[DEBUGGER_FLAG] = "0"
+
     def _stage_user_code_in_s3(self):
         """Upload the user training script to s3 and return the location.
 

tests/integ/test_debugger.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker.debugger.debugger import (
+    DEBUGGER_FLAG,
     DebuggerHookConfig,
     Rule,
     rule_configs,
@@ -748,6 +749,7 @@ def test_mxnet_with_debugger_hook_config_disabled(
         job_description = mx.latest_training_job.describe()
 
         assert job_description.get("DebugHookConfig") is None
+        assert job_description.get("Environment", {}).get(DEBUGGER_FLAG) == "0"
 
 
 def _get_rule_evaluation_statuses(job_description):

tests/unit/sagemaker/workflow/test_steps.py

@@ -23,7 +23,7 @@
     patch,
 )
 
-from sagemaker.debugger import ProfilerConfig
+from sagemaker.debugger import DEBUGGER_FLAG, ProfilerConfig
 from sagemaker.estimator import Estimator
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.inputs import TrainingInput, TransformInput, CreateModelInput
@@ -275,6 +275,7 @@ def test_training_step_tensorflow(sagemaker_session):
                 "sagemaker_distributed_dataparallel_custom_mpi_options": '""',
             },
             "ProfilerConfig": {"S3OutputPath": "s3://my-bucket/"},
+            "Environment": {DEBUGGER_FLAG: "0"},
         },
         "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }