feature: Support profiler config in the pipeline training job step (#2183)

Author: navaj0

src/sagemaker/workflow/steps.py

@@ -160,9 +160,6 @@ def arguments(self) -> RequestType:
         NOTE: The CreateTrainingJob request is not quite the args list that workflow needs.
         The TrainingJobName and ExperimentConfig attributes cannot be included.
         """
-        self.estimator.disable_profiler = True
-        self.estimator.profiler_config = None
-        self.estimator.profiler_rules = None
 
         self.estimator._prepare_for_training()
         train_args = _TrainingJob._get_train_args(

tests/integ/test_workflow.py

@@ -450,7 +450,7 @@ def test_conditional_pytorch_training_model_registration(
             pass
 
 
-def test_training_job_with_debugger(
+def test_training_job_with_debugger_and_profiler(
     sagemaker_session,
     pipeline_name,
     role,
@@ -535,6 +535,9 @@ def test_training_job_with_debugger(
                 config["RuleParameters"]["rule_to_invoke"] == rule.rule_parameters["rule_to_invoke"]
             )
         assert job_description["DebugHookConfig"] == debugger_hook_config._to_request_dict()
+
+        assert job_description["ProfilingStatus"] == "Enabled"
+        assert job_description["ProfilerConfig"]["ProfilingIntervalInMilliseconds"] == 500
     finally:
         try:
             pipeline.delete()

tests/unit/sagemaker/workflow/test_steps.py

@@ -21,6 +21,7 @@
     PropertyMock,
 )
 
+from sagemaker.debugger import ProfilerConfig
 from sagemaker.estimator import Estimator
 from sagemaker.inputs import TrainingInput, TransformInput, CreateModelInput
 from sagemaker.model import Model
@@ -112,6 +113,8 @@ def test_training_step(sagemaker_session):
         role=ROLE,
         instance_count=1,
         instance_type="c4.4xlarge",
+        profiler_config=ProfilerConfig(system_monitor_interval_millis=500),
+        rules=[],
         sagemaker_session=sagemaker_session,
     )
     inputs = TrainingInput(f"s3://{BUCKET}/train_manifest")
@@ -144,6 +147,10 @@ def test_training_step(sagemaker_session):
             },
             "RoleArn": ROLE,
             "StoppingCondition": {"MaxRuntimeInSeconds": 86400},
+            "ProfilerConfig": {
+                "ProfilingIntervalInMilliseconds": 500,
+                "S3OutputPath": f"s3://{BUCKET}/",
+            },
         },
         "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }