fix: create profiler specific unsupported regions (#2101)

Co-authored-by: Song Jiang <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: 3 people

src/sagemaker/estimator.py

@@ -49,6 +49,7 @@
     UploadedCode,
     validate_source_dir,
     _region_supports_debugger,
+    _region_supports_profiler,
     get_mp_parameters,
 )
 from sagemaker.inputs import TrainingInput
@@ -494,7 +495,7 @@ def _prepare_profiler_for_training(self):
         """Set necessary values and do basic validations in profiler config and profiler rules.
 
         When user explicitly set rules to an empty list, default profiler rule won't be enabled.
-        Default profiler rule will be enabled when either:
+        Default profiler rule will be enabled in supported regions when either:
         1. user doesn't specify any rules, i.e., rules=None; or
         2. user only specify debugger rules, i.e., rules=[Rule.sagemaker(...)]
         """
@@ -503,7 +504,7 @@ def _prepare_profiler_for_training(self):
                 raise RuntimeError("profiler_config cannot be set when disable_profiler is True.")
             if self.profiler_rules:
                 raise RuntimeError("ProfilerRule cannot be set when disable_profiler is True.")
-        elif _region_supports_debugger(self.sagemaker_session.boto_region_name):
+        elif _region_supports_profiler(self.sagemaker_session.boto_region_name):
             if self.profiler_config is None:
                 self.profiler_config = ProfilerConfig(s3_output_path=self.output_path)
             if self.rules is None or (self.rules and not self.profiler_rules):

src/sagemaker/fw_utils.py

@@ -49,6 +49,8 @@
 )
 
 DEBUGGER_UNSUPPORTED_REGIONS = ("us-iso-east-1",)
+PROFILER_UNSUPPORTED_REGIONS = ("us-iso-east-1", "cn-north-1", "cn-northwest-1")
+
 SINGLE_GPU_INSTANCE_TYPES = ("ml.p2.xlarge", "ml.p3.2xlarge")
 SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES = (
     "ml.p3.16xlarge",
@@ -550,6 +552,19 @@ def _region_supports_debugger(region_name):
     return region_name.lower() not in DEBUGGER_UNSUPPORTED_REGIONS
 
 
+def _region_supports_profiler(region_name):
+    """Returns bool indicating whether region supports Amazon SageMaker Debugger profiling feature.
+
+    Args:
+        region_name (str): Name of the region to check against.
+
+    Returns:
+        bool: Whether or not the region supports Amazon SageMaker Debugger profiling feature.
+
+    """
+    return region_name.lower() not in PROFILER_UNSUPPORTED_REGIONS
+
+
 def validate_version_or_image_args(framework_version, py_version, image_uri):
     """Checks if version or image arguments are specified.
 

tests/unit/test_estimator.py

@@ -35,6 +35,7 @@
     Rule,
 )
 from sagemaker.estimator import Estimator, EstimatorBase, Framework, _TrainingJob
+from sagemaker.fw_utils import PROFILER_UNSUPPORTED_REGIONS
 from sagemaker.inputs import ShuffleConfig
 from sagemaker.model import FrameworkModel
 from sagemaker.predictor import Predictor
@@ -632,6 +633,32 @@ def test_framework_with_profiler_config_without_s3_output_path(time, sagemaker_s
     ]
 
 
+@pytest.mark.parametrize("region", PROFILER_UNSUPPORTED_REGIONS)
+def test_framework_with_no_default_profiler_in_unsupported_region(region):
+    boto_mock = Mock(name="boto_session", region_name=region)
+    sms = MagicMock(
+        name="sagemaker_session",
+        boto_session=boto_mock,
+        boto_region_name=region,
+        config=None,
+        local_mode=False,
+        s3_client=None,
+        s3_resource=None,
+    )
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sms,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+    )
+    f.fit("s3://mydata")
+    sms.train.assert_called_once()
+    _, args = sms.train.call_args
+    assert args.get("profiler_config") is None
+    assert args.get("profiler_rule_configs") is None
+
+
 def test_framework_with_profiler_config_and_profiler_disabled(sagemaker_session):
     with pytest.raises(RuntimeError) as error:
         f = DummyFramework(