feature: Add disable_profiler field in config and propagate changes (#3523)

Co-authored-by: Marius Moisescu <mariumof@dev-dsk-mariumof-1e-9b137084.us-east-1.amazon.com>

Author: mariumof

src/sagemaker/debugger/profiler_config.py

@@ -32,6 +32,7 @@ def __init__(
         s3_output_path: Optional[Union[str, PipelineVariable]] = None,
         system_monitor_interval_millis: Optional[Union[int, PipelineVariable]] = None,
         framework_profile_params: Optional[FrameworkProfile] = None,
+        disable_profiler: Optional[Union[str, PipelineVariable]] = False,
     ):
         """Initialize a ``ProfilerConfig`` instance.
 
@@ -78,6 +79,7 @@ class and SageMaker Framework estimators.
         self.s3_output_path = s3_output_path
         self.system_monitor_interval_millis = system_monitor_interval_millis
         self.framework_profile_params = framework_profile_params
+        self.disable_profiler = disable_profiler
 
     def _to_request_dict(self):
         """Generate a request dictionary using the parameters provided when initializing the object.
@@ -91,6 +93,8 @@ def _to_request_dict(self):
         if self.s3_output_path is not None:
             profiler_config_request["S3OutputPath"] = self.s3_output_path
 
+        profiler_config_request["DisableProfiler"] = self.disable_profiler
+
         if self.system_monitor_interval_millis is not None:
             profiler_config_request[
                 "ProfilingIntervalInMilliseconds"

src/sagemaker/estimator.py

@@ -938,26 +938,29 @@ def _prepare_collection_configs(self):
     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 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(...)]
+        No default profiler rule will be used. The user needs to specify rules explicitly
         """
         if self.disable_profiler:
-            if self.profiler_config:
-                raise RuntimeError("profiler_config cannot be set when disable_profiler is True.")
+            if self.profiler_config and not self.profiler_config.disable_profiler:
+                raise RuntimeError(
+                    "profiler_config.disable_profiler cannot be False"
+                    + " when disable_profiler is True."
+                )
             if self.profiler_rules:
                 raise RuntimeError("ProfilerRule cannot be set when disable_profiler is True.")
         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):
-                self.profiler_rules = [get_default_profiler_rule()]
+                self.profiler_rules = []
 
         if self.profiler_config and not self.profiler_config.s3_output_path:
             self.profiler_config.s3_output_path = self.output_path
 
         self.profiler_rule_configs = self._prepare_profiler_rules()
+        # if profiler_config is still None, it means the job has profiler disabled
+        if self.profiler_config is None:
+            self.profiler_config = ProfilerConfig(disable_profiler=True)
 
     def _prepare_profiler_rules(self):
         """Set any necessary values in profiler rules, if they are provided."""
@@ -1048,7 +1051,7 @@ def latest_job_profiler_artifacts_path(self):
             error_message="""Cannot get the profiling output artifacts path.
         The Estimator is not associated with a training job."""
         )
-        if self.profiler_config is not None:
+        if self.profiler_config is not None and not self.profiler_config.disable_profiler:
             return os.path.join(
                 self.profiler_config.s3_output_path,
                 self.latest_training_job.name,
@@ -1895,8 +1898,8 @@ def enable_default_profiling(self):
         else:
             self.profiler_config = ProfilerConfig(s3_output_path=self.output_path)
 
-        self.profiler_rules = [get_default_profiler_rule()]
-        self.profiler_rule_configs = self._prepare_profiler_rules()
+        self.profiler_rules = []
+        self.profiler_rule_configs = []
 
         _TrainingJob.update(
             self, self.profiler_rule_configs, self.profiler_config._to_request_dict()

tests/integ/sagemaker/workflow/test_workflow.py

@@ -1269,8 +1269,6 @@ def test_caching_behavior(
         # create pipeline
         pipeline.create(role)
         definition = json.loads(pipeline.definition())
-        # delete profiler config for assertions as it will contain a timestamp
-        del definition["Steps"][1]["Arguments"]["ProfilerRuleConfigurations"]
 
         # verify input path
         expected_abalone_input_path = f"{pipeline_name}/{step_process.name}" f"/input/abalone_data"
@@ -1295,7 +1293,6 @@ def test_caching_behavior(
 
         # verify no changes
         definition2 = json.loads(pipeline.definition())
-        del definition2["Steps"][1]["Arguments"]["ProfilerRuleConfigurations"]
         assert definition == definition2
 
         # add dummy file to source_dir
@@ -1306,7 +1303,6 @@ def test_caching_behavior(
 
         # verify changes
         definition3 = json.loads(pipeline.definition())
-        del definition3["Steps"][1]["Arguments"]["ProfilerRuleConfigurations"]
         assert definition != definition3
 
     finally:

tests/integ/test_profiler.py

@@ -13,7 +13,6 @@
 from __future__ import absolute_import
 
 import os
-import re
 import time
 import uuid
 
@@ -22,7 +21,6 @@
 from sagemaker.debugger import (
     DebuggerHookConfig,
     FrameworkProfile,
-    get_rule_container_image_uri,
     ProfilerConfig,
     ProfilerRule,
     Rule,
@@ -93,8 +91,6 @@ def test_mxnet_with_default_profiler_config_and_profiler_rule(
         )
 
         job_description = mx.latest_training_job.describe()
-        if "DisableProfiler" in job_description["ProfilerConfig"]:
-            job_description["ProfilerConfig"].pop("DisableProfiler")
         assert (
             job_description["ProfilerConfig"]
             == ProfilerConfig(
@@ -103,13 +99,6 @@ def test_mxnet_with_default_profiler_config_and_profiler_rule(
         )
         assert job_description.get("ProfilingStatus") == "Enabled"
 
-        profiler_rule_configuration = job_description.get("ProfilerRuleConfigurations")[0]
-        assert re.match(r"ProfilerReport-\d*", profiler_rule_configuration["RuleConfigurationName"])
-        assert profiler_rule_configuration["RuleEvaluatorImage"] == get_rule_container_image_uri(
-            mx.sagemaker_session.boto_region_name
-        )
-        assert profiler_rule_configuration["RuleParameters"] == {"rule_to_invoke": "ProfilerReport"}
-
         with pytest.raises(ValueError) as error:
             mx.enable_default_profiling()
         assert "Debugger monitoring is already enabled." in str(error)
@@ -155,18 +144,9 @@ def test_mxnet_with_custom_profiler_config_then_update_rule_and_config(
         )
 
         job_description = mx.latest_training_job.describe()
-        if "DisableProfiler" in job_description["ProfilerConfig"]:
-            job_description["ProfilerConfig"].pop("DisableProfiler")
         assert job_description.get("ProfilerConfig") == profiler_config._to_request_dict()
         assert job_description.get("ProfilingStatus") == "Enabled"
 
-        profiler_rule_configuration = job_description.get("ProfilerRuleConfigurations")[0]
-        assert re.match(r"ProfilerReport-\d*", profiler_rule_configuration["RuleConfigurationName"])
-        assert profiler_rule_configuration["RuleEvaluatorImage"] == get_rule_container_image_uri(
-            mx.sagemaker_session.boto_region_name
-        )
-        assert profiler_rule_configuration["RuleParameters"] == {"rule_to_invoke": "ProfilerReport"}
-
         _wait_until_training_can_be_updated(sagemaker_session.sagemaker_client, training_job_name)
 
         mx.update_profiler(
@@ -178,13 +158,6 @@ def test_mxnet_with_custom_profiler_config_then_update_rule_and_config(
         assert job_description["ProfilerConfig"]["S3OutputPath"] == profiler_config.s3_output_path
         assert job_description["ProfilerConfig"]["ProfilingIntervalInMilliseconds"] == 500
 
-        profiler_report_rule_config = job_description.get("ProfilerRuleConfigurations")[0]
-        assert re.match(r"ProfilerReport-\d*", profiler_report_rule_config["RuleConfigurationName"])
-        assert profiler_report_rule_config["RuleEvaluatorImage"] == get_rule_container_image_uri(
-            mx.sagemaker_session.boto_region_name
-        )
-        assert profiler_report_rule_config["RuleParameters"] == {"rule_to_invoke": "ProfilerReport"}
-
 
 def test_mxnet_with_built_in_profiler_rule_with_custom_parameters(
     sagemaker_session,
@@ -225,8 +198,6 @@ def test_mxnet_with_built_in_profiler_rule_with_custom_parameters(
         )
 
         job_description = mx.latest_training_job.describe()
-        if "DisableProfiler" in job_description["ProfilerConfig"]:
-            job_description["ProfilerConfig"].pop("DisableProfiler")
         assert job_description.get("ProfilingStatus") == "Enabled"
         assert (
             job_description.get("ProfilerConfig")
@@ -298,8 +269,6 @@ def test_mxnet_with_profiler_and_debugger_then_disable_framework_metrics(
         )
 
         job_description = mx.latest_training_job.describe()
-        if "DisableProfiler" in job_description["ProfilerConfig"]:
-            job_description["ProfilerConfig"].pop("DisableProfiler")
         assert job_description["ProfilerConfig"] == profiler_config._to_request_dict()
         assert job_description["DebugHookConfig"] == debugger_hook_config._to_request_dict()
         assert job_description.get("ProfilingStatus") == "Enabled"
@@ -387,13 +356,6 @@ def test_mxnet_with_enable_framework_metrics_then_update_framework_metrics(
             == updated_framework_profile.profiling_parameters
         )
 
-        profiler_rule_configuration = job_description.get("ProfilerRuleConfigurations")[0]
-        assert re.match(r"ProfilerReport-\d*", profiler_rule_configuration["RuleConfigurationName"])
-        assert profiler_rule_configuration["RuleEvaluatorImage"] == get_rule_container_image_uri(
-            mx.sagemaker_session.boto_region_name
-        )
-        assert profiler_rule_configuration["RuleParameters"] == {"rule_to_invoke": "ProfilerReport"}
-
 
 def test_mxnet_with_disable_profiler_then_enable_default_profiling(
     sagemaker_session,
@@ -431,12 +393,10 @@ def test_mxnet_with_disable_profiler_then_enable_default_profiling(
         )
 
         job_description = mx.latest_training_job.describe()
-        assert job_description.get("ProfilerConfig") is None
         assert job_description.get("ProfilerRuleConfigurations") is None
         assert job_description.get("ProfilingStatus") == "Disabled"
 
         _wait_until_training_can_be_updated(sagemaker_session.sagemaker_client, training_job_name)
-
         mx.enable_default_profiling()
 
         job_description = mx.latest_training_job.describe()

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -143,14 +143,8 @@ def _create_train_job(version, base_framework_version):
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
-        "profiler_rule_configs": [
-            {
-                "RuleConfigurationName": "ProfilerReport-1510006209",
-                "RuleEvaluatorImage": "503895931360.dkr.ecr.us-east-1.amazonaws.com/sagemaker-debugger-rules:latest",
-                "RuleParameters": {"rule_to_invoke": "ProfilerReport"},
-            }
-        ],
         "profiler_config": {
+            "DisableProfiler": False,
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
     }

tests/unit/sagemaker/tensorflow/test_estimator.py

@@ -136,14 +136,8 @@ def _create_train_job(tf_version, horovod=False, ps=False, py_version="py2", smd
         "metric_definitions": None,
         "environment": None,
         "experiment_config": None,
-        "profiler_rule_configs": [
-            {
-                "RuleConfigurationName": "ProfilerReport-1510006209",
-                "RuleEvaluatorImage": "895741380848.dkr.ecr.us-west-2.amazonaws.com/sagemaker-debugger-rules:latest",
-                "RuleParameters": {"rule_to_invoke": "ProfilerReport"},
-            }
-        ],
         "profiler_config": {
+            "DisableProfiler": False,
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
     }

tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py

@@ -145,14 +145,8 @@ def _create_train_job(
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
-        "profiler_rule_configs": [
-            {
-                "RuleConfigurationName": "ProfilerReport-1510006209",
-                "RuleEvaluatorImage": "503895931360.dkr.ecr.us-east-1.amazonaws.com/sagemaker-debugger-rules:latest",
-                "RuleParameters": {"rule_to_invoke": "ProfilerReport"},
-            }
-        ],
         "profiler_config": {
+            "DisableProfiler": False,
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
     }

tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py

@@ -143,14 +143,8 @@ def _create_train_job(
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
-        "profiler_rule_configs": [
-            {
-                "RuleConfigurationName": "ProfilerReport-1510006209",
-                "RuleEvaluatorImage": "503895931360.dkr.ecr.us-east-1.amazonaws.com/sagemaker-debugger-rules:latest",
-                "RuleParameters": {"rule_to_invoke": "ProfilerReport"},
-            }
-        ],
         "profiler_config": {
+            "DisableProfiler": False,
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
     }

tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py

@@ -137,14 +137,10 @@ def _create_train_job(version, instance_type, training_compiler_config, instance
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
-        "profiler_rule_configs": [
-            {
-                "RuleConfigurationName": "ProfilerReport-1510006209",
-                "RuleEvaluatorImage": "503895931360.dkr.ecr.us-east-1.amazonaws.com/sagemaker-debugger-rules:latest",
-                "RuleParameters": {"rule_to_invoke": "ProfilerReport"},
-            }
-        ],
-        "profiler_config": {"S3OutputPath": "s3://{}/".format(BUCKET_NAME)},
+        "profiler_config": {
+            "DisableProfiler": False,
+            "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
+        },
     }
 
 

tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

@@ -145,14 +145,8 @@ def _create_train_job(framework_version, instance_type, training_compiler_config
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
-        "profiler_rule_configs": [
-            {
-                "RuleConfigurationName": "ProfilerReport-1510006209",
-                "RuleEvaluatorImage": "503895931360.dkr.ecr.us-east-1.amazonaws.com/sagemaker-debugger-rules:latest",
-                "RuleParameters": {"rule_to_invoke": "ProfilerReport"},
-            }
-        ],
         "profiler_config": {
+            "DisableProfiler": False,
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),
         },
     }