feat: SDK defaults add disable profiler to createTrainingJob (#3936)

* feat: add disable profiler to createTrainingJob

added profilerConfig to the traningjob schema
for sagemaker_config

disableprofiler attribute is injected via config
in session,estimator

* fix: Update comments based on PR

* fix: spaces, linting, based on black-check, flake8

* fix: comment on PR

add support for update_training_job

doc change

added unitTests

---------

Co-authored-by: Ashwin Krishna <ashwikri@amazon.com>

Author: akrishna1995

src/sagemaker/config/__init__.py

@@ -24,6 +24,8 @@
     TRAINING_JOB_VPC_CONFIG_PATH,
     TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH,
     TRAINING_JOB_RESOURCE_CONFIG_PATH,
+    TRAINING_JOB_PROFILE_CONFIG_PATH,
+    TRAINING_JOB_DISABLE_PROFILER_PATH,
     PROCESSING_JOB_INPUTS_PATH,
     PROCESSING_JOB,
     PROCESSING_JOB_ENVIRONMENT_PATH,
@@ -141,6 +143,8 @@
     CLUSTER_ROLE_ARN,
     PROCESSING_OUTPUT_CONFIG,
     PROCESSING_RESOURCES,
+    PROFILER_CONFIG,
+    DISABLE_PROFILER,
     RESOURCE_CONFIG,
     EXECUTION_ROLE_ARN,
     ASYNC_INFERENCE_CONFIG,

src/sagemaker/config/config_schema.py

@@ -98,6 +98,8 @@
 CONTAINERS = "Containers"
 PRIMARY_CONTAINER = "PrimaryContainer"
 INFERENCE_SPECIFICATION = "InferenceSpecification"
+PROFILER_CONFIG = "ProfilerConfig"
+DISABLE_PROFILER = "DisableProfiler"
 
 
 def _simple_path(*args: str):
@@ -128,6 +130,10 @@ def _simple_path(*args: str):
     TRAINING_JOB_VPC_CONFIG_PATH, SECURITY_GROUP_IDS
 )
 TRAINING_JOB_SUBNETS_PATH = _simple_path(TRAINING_JOB_VPC_CONFIG_PATH, SUBNETS)
+TRAINING_JOB_PROFILE_CONFIG_PATH = _simple_path(SAGEMAKER, TRAINING_JOB, PROFILER_CONFIG)
+TRAINING_JOB_DISABLE_PROFILER_PATH = _simple_path(
+    TRAINING_JOB_PROFILE_CONFIG_PATH, DISABLE_PROFILER
+)
 EDGE_PACKAGING_KMS_KEY_ID_PATH = _simple_path(
     SAGEMAKER, EDGE_PACKAGING_JOB, OUTPUT_CONFIG, KMS_KEY_ID
 )
@@ -1001,6 +1007,11 @@ def _simple_path(*args: str):
                             ADDITIONAL_PROPERTIES: False,
                             PROPERTIES: {VOLUME_KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"}},
                         },
+                        PROFILER_CONFIG: {
+                            TYPE: OBJECT,
+                            ADDITIONAL_PROPERTIES: False,
+                            PROPERTIES: {DISABLE_PROFILER: {TYPE: "boolean"}},
+                        },
                         ROLE_ARN: {"$ref": "#/definitions/roleArn"},
                         VPC_CONFIG: {"$ref": "#/definitions/vpcConfig"},
                         TAGS: {"$ref": "#/definitions/tags"},

src/sagemaker/estimator.py

@@ -37,6 +37,7 @@
     TRAINING_JOB_ROLE_ARN_PATH,
     TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
     TRAINING_JOB_ENVIRONMENT_PATH,
+    TRAINING_JOB_DISABLE_PROFILER_PATH,
     TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
 )
 from sagemaker.debugger import (  # noqa: F401 # pylint: disable=unused-import
@@ -157,7 +158,7 @@ def __init__(
         enable_sagemaker_metrics: Optional[Union[bool, PipelineVariable]] = None,
         enable_network_isolation: Union[bool, PipelineVariable] = None,
         profiler_config: Optional[ProfilerConfig] = None,
-        disable_profiler: bool = False,
+        disable_profiler: bool = None,
         environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         max_retry_attempts: Optional[Union[int, PipelineVariable]] = None,
         source_dir: Optional[Union[str, PipelineVariable]] = None,
@@ -687,7 +688,12 @@ def __init__(
         )
 
         self.profiler_config = profiler_config
-        self.disable_profiler = disable_profiler
+        self.disable_profiler = resolve_value_from_config(
+            direct_input=disable_profiler,
+            config_path=TRAINING_JOB_DISABLE_PROFILER_PATH,
+            default_value=False,
+            sagemaker_session=self.sagemaker_session,
+        )
 
         self.environment = resolve_value_from_config(
             direct_input=environment,

src/sagemaker/session.py

@@ -45,6 +45,7 @@
     TRAINING_JOB_VPC_CONFIG_PATH,
     TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH,
     TRAINING_JOB_RESOURCE_CONFIG_PATH,
+    TRAINING_JOB_PROFILE_CONFIG_PATH,
     PROCESSING_JOB_INPUTS_PATH,
     PROCESSING_JOB,
     PROCESSING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
@@ -833,6 +834,9 @@ def train(  # noqa: C901
         inferred_resource_config = update_nested_dictionary_with_values_from_config(
             resource_config, TRAINING_JOB_RESOURCE_CONFIG_PATH, sagemaker_session=self
         )
+        inferred_profiler_config = update_nested_dictionary_with_values_from_config(
+            profiler_config, TRAINING_JOB_PROFILE_CONFIG_PATH, sagemaker_session=self
+        )
         if (
             not customer_supplied_kms_key
             and "InstanceType" in inferred_resource_config
@@ -875,7 +879,7 @@ def train(  # noqa: C901
             tensorboard_output_config=tensorboard_output_config,
             enable_sagemaker_metrics=enable_sagemaker_metrics,
             profiler_rule_configs=profiler_rule_configs,
-            profiler_config=profiler_config,
+            profiler_config=inferred_profiler_config,
             environment=environment,
             retry_strategy=retry_strategy,
         )
@@ -1152,11 +1156,13 @@ def update_training_job(
         # No injections from sagemaker_config because the UpdateTrainingJob API's resource_config
         # object accepts fewer parameters than the CreateTrainingJob API, and none that the
         # sagemaker_config currently supports
-
+        inferred_profiler_config = update_nested_dictionary_with_values_from_config(
+            profiler_config, TRAINING_JOB_PROFILE_CONFIG_PATH, sagemaker_session=self
+        )
         update_training_job_request = self._get_update_training_job_request(
             job_name=job_name,
             profiler_rule_configs=profiler_rule_configs,
-            profiler_config=profiler_config,
+            profiler_config=inferred_profiler_config,
             resource_config=resource_config,
         )
         LOGGER.info("Updating training job with name %s", job_name)

tests/data/config/config.yaml

@@ -167,6 +167,8 @@ SageMaker:
     ResourceConfig:
       VolumeKmsKeyId: 'volumekmskeyid1'
     RoleArn: 'arn:aws:iam::555555555555:role/IMRole'
+    ProfilerConfig:
+      DisableProfiler: false
     VpcConfig:
       SecurityGroupIds:
         - 'sg123'

tests/unit/__init__.py

@@ -68,6 +68,8 @@
     PROCESSING_OUTPUT_CONFIG,
     PROCESSING_RESOURCES,
     TRAINING_JOB,
+    PROFILER_CONFIG,
+    DISABLE_PROFILER,
     RESOURCE_CONFIG,
     TRANSFORM_JOB,
     EXECUTION_ROLE_ARN,
@@ -327,6 +329,7 @@
             ENVIRONMENT: {"configEnvVar1": "value1", "configEnvVar2": "value2"},
             OUTPUT_DATA_CONFIG: {KMS_KEY_ID: "TestKms"},
             RESOURCE_CONFIG: {VOLUME_KMS_KEY_ID: "volumekey"},
+            PROFILER_CONFIG: {DISABLE_PROFILER: False},
             ROLE_ARN: "arn:aws:iam::111111111111:role/ConfigRole",
             VPC_CONFIG: {SUBNETS: ["subnets-123"], SECURITY_GROUP_IDS: ["sg-123"]},
             TAGS: [{KEY: "some-tag", VALUE: "value-for-tag"}],

tests/unit/sagemaker/config/conftest.py

@@ -142,6 +142,7 @@ def valid_training_job_config(valid_iam_role_arn, valid_vpc_config, valid_enviro
         "Environment": valid_environment_config,
         "OutputDataConfig": {"KmsKeyId": "kmskeyid1"},
         "ResourceConfig": {"VolumeKmsKeyId": "volumekmskeyid1"},
+        "ProfilerConfig": {"DisableProfiler": False},
         "RoleArn": valid_iam_role_arn,
         "VpcConfig": valid_vpc_config,
     }

tests/unit/test_estimator.py

@@ -407,6 +407,9 @@ def test_framework_initialization_with_sagemaker_config_injection(sagemaker_sess
         "TrainingJob"
     ]["EnableInterContainerTrafficEncryption"]
     expected_environment = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"]["TrainingJob"]["Environment"]
+    expected_disable_profiler_attribute = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"]["TrainingJob"][
+        "ProfilerConfig"
+    ]["DisableProfiler"]
     assert framework.role == expected_role_arn
     assert framework.enable_network_isolation() == expected_enable_network_isolation
     assert (
@@ -418,6 +421,7 @@ def test_framework_initialization_with_sagemaker_config_injection(sagemaker_sess
     assert framework.security_group_ids == expected_security_groups
     assert framework.subnets == expected_subnets
     assert framework.environment == expected_environment
+    assert framework.disable_profiler == expected_disable_profiler_attribute
 
 
 def test_estimator_initialization_with_sagemaker_config_injection(sagemaker_session):
@@ -453,6 +457,9 @@ def test_estimator_initialization_with_sagemaker_config_injection(sagemaker_sess
         "TrainingJob"
     ]["EnableInterContainerTrafficEncryption"]
     expected_environment = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"]["TrainingJob"]["Environment"]
+    expected_disable_profiler_attribute = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"]["TrainingJob"][
+        "ProfilerConfig"
+    ]["DisableProfiler"]
     assert estimator.role == expected_role_arn
     assert estimator.enable_network_isolation() == expected_enable_network_isolation
     assert (
@@ -464,6 +471,7 @@ def test_estimator_initialization_with_sagemaker_config_injection(sagemaker_sess
     assert estimator.security_group_ids == expected_security_groups
     assert estimator.subnets == expected_subnets
     assert estimator.environment == expected_environment
+    assert estimator.disable_profiler == expected_disable_profiler_attribute
 
 
 def test_estimator_initialization_with_sagemaker_config_injection_no_kms_supported(
@@ -498,6 +506,9 @@ def test_estimator_initialization_with_sagemaker_config_injection_no_kms_support
     expected_enable_inter_container_traffic_encryption = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"][
         "TrainingJob"
     ]["EnableInterContainerTrafficEncryption"]
+    expected_disable_profiler_attribute = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"]["TrainingJob"][
+        "ProfilerConfig"
+    ]["DisableProfiler"]
     assert estimator.role == expected_role_arn
     assert estimator.enable_network_isolation() == expected_enable_network_isolation
     assert (
@@ -508,6 +519,7 @@ def test_estimator_initialization_with_sagemaker_config_injection_no_kms_support
     assert estimator.volume_kms_key is None
     assert estimator.security_group_ids == expected_security_groups
     assert estimator.subnets == expected_subnets
+    assert estimator.disable_profiler == expected_disable_profiler_attribute
 
 
 def test_estimator_initialization_with_sagemaker_config_injection_partial_kms_support(
@@ -545,6 +557,9 @@ def test_estimator_initialization_with_sagemaker_config_injection_partial_kms_su
     expected_enable_inter_container_traffic_encryption = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"][
         "TrainingJob"
     ]["EnableInterContainerTrafficEncryption"]
+    expected_disable_profiler_attribute = SAGEMAKER_CONFIG_TRAINING_JOB["SageMaker"]["TrainingJob"][
+        "ProfilerConfig"
+    ]["DisableProfiler"]
     assert estimator.role == expected_role_arn
     assert estimator.enable_network_isolation() == expected_enable_network_isolation
     assert (
@@ -555,6 +570,7 @@ def test_estimator_initialization_with_sagemaker_config_injection_partial_kms_su
     assert estimator.volume_kms_key == expected_volume_kms_key_id
     assert estimator.security_group_ids == expected_security_groups
     assert estimator.subnets == expected_subnets
+    assert estimator.disable_profiler == expected_disable_profiler_attribute
 
 
 def test_framework_with_heterogeneous_cluster(sagemaker_session):