Address edge case

Author: ndodda-amazon

tests/integ/test_debugger.py

@@ -746,11 +746,10 @@ def test_mxnet_with_debugger_hook_config_disabled(
 
         mx.fit({"train": train_input, "test": test_input})
 
-        assert mx.environment.get(DEBUGGER_FLAG) == "0"
-
         job_description = mx.latest_training_job.describe()
 
         assert job_description.get("DebugHookConfig") is None
+        assert job_description.get("Environment", {}).get("USE_SMDEBUG") == "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"},
     }