feature: add HealthCheckConfig support (#4105)

Co-authored-by: Xin Wang <[email protected]>

Author: emeraldbay

src/sagemaker/estimator.py

@@ -173,6 +173,7 @@ def __init__(
         instance_groups: Optional[List[InstanceGroup]] = None,
         training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
@@ -536,6 +537,8 @@ def __init__(
                 a training job.
             disable_output_compression (bool): Optional. When set to true, Model is uploaded
                 to Amazon S3 without compression after training finishes.
+            enable_infra_check (bool or PipelineVariable): Optional.
+                Specifies whether it is running Sagemaker built-in infra check jobs.
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -665,6 +668,7 @@ def __init__(
             training_repository_credentials_provider_arn
         )
 
+        self.enable_infra_check = enable_infra_check
         # container entry point / arguments configs
         self.container_entry_point = container_entry_point
         self.container_arguments = container_arguments
@@ -1904,6 +1908,11 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
                 "EnableInterContainerTrafficEncryption"
             ]
 
+        if "InfraCheckConfig" in job_details:
+            init_params["enable_infra_check"] = job_details["InfraCheckConfig"].get(
+                "EnableInfraCheck"
+            )
+
         subnets, security_group_ids = vpc_utils.from_dict(job_details.get(vpc_utils.VPC_CONFIG_KEY))
         if subnets:
             init_params["subnets"] = subnets
@@ -2446,6 +2455,10 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
                 ] = estimator.training_repository_credentials_provider_arn
             train_args["training_image_config"] = training_image_config
 
+        if estimator.enable_infra_check is not None:
+            infra_check_config = {"EnableInfraCheck": estimator.enable_infra_check}
+            train_args["infra_check_config"] = infra_check_config
+
         if estimator.container_entry_point is not None:
             train_args["container_entry_point"] = estimator.container_entry_point
 
@@ -2661,6 +2674,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -3020,6 +3034,8 @@ def __init__(
                 a training job.
             disable_output_compression (bool): Optional. When set to true, Model is uploaded
                 to Amazon S3 without compression after training finishes.
+            enable_infra_check (bool or PipelineVariable): Optional.
+                Specifies whether it is running Sagemaker built-in infra check jobs.
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}

src/sagemaker/session.py

@@ -674,6 +674,7 @@ def train(  # noqa: C901
         enable_network_isolation=None,
         image_uri=None,
         training_image_config=None,
+        infra_check_config=None,
         container_entry_point=None,
         container_arguments=None,
         algorithm_arn=None,
@@ -803,6 +804,15 @@ def train(  # noqa: C901
             retry_strategy(dict): Defines RetryStrategy for InternalServerFailures.
                 * max_retry_attsmpts (int): Number of times a job should be retried.
                 The key in RetryStrategy is 'MaxRetryAttempts'.
+            infra_check_config(dict): Infra check configuration.
+                Optionally, the dict can contain 'EnableInfraCheck'(bool).
+                For example,
+
+                .. code:: python
+
+                    infra_check_config = {
+                        "EnableInfraCheck": True,
+                    }
         Returns:
             str: ARN of the training job, if it is created.
         """
@@ -866,6 +876,7 @@ def train(  # noqa: C901
             enable_network_isolation=enable_network_isolation,
             image_uri=image_uri,
             training_image_config=training_image_config,
+            infra_check_config=infra_check_config,
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
             algorithm_arn=algorithm_arn,
@@ -907,6 +918,7 @@ def _get_train_request(  # noqa: C901
         enable_network_isolation=False,
         image_uri=None,
         training_image_config=None,
+        infra_check_config=None,
         container_entry_point=None,
         container_arguments=None,
         algorithm_arn=None,
@@ -1063,6 +1075,9 @@ def _get_train_request(  # noqa: C901
         if training_image_config is not None:
             train_request["AlgorithmSpecification"]["TrainingImageConfig"] = training_image_config
 
+        if infra_check_config is not None:
+            train_request["InfraCheckConfig"] = infra_check_config
+
         if container_entry_point is not None:
             train_request["AlgorithmSpecification"]["ContainerEntrypoint"] = container_entry_point
 

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -677,7 +677,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
         Please add the new argument to the skip set below,
         and reach out to JumpStart team."""
 
-        init_args_to_skip: Set[str] = set(["kwargs"])
+        init_args_to_skip: Set[str] = set(["kwargs", "enable_infra_check"])
         fit_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])
 

tests/unit/test_estimator.py

@@ -103,6 +103,7 @@
 REPO_DIR = "/tmp/repo_dir"
 ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
 TRAINING_REPOSITORY_ACCESS_MODE = "VPC"
+ENABLE_INFRA_CHECK = True
 TRAINING_REPOSITORY_CREDENTIALS_PROVIDER_ARN = "arn:aws:lambda:us-west-2:1234567890:function:test"
 CONTAINER_ENTRY_POINT = ["entry_point1", "entry_point2"]
 CONTAINER_ARGUMENTS = ["container_arg1", "container_arg2"]
@@ -769,6 +770,39 @@ def test_framework_without_training_repository_config(sagemaker_session):
     assert args.get("training_image_config") is None
 
 
+def test_framework_without_infra_check_config(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args.get("health_check_config") is None
+
+
+def test_framework_with_infra_check_config(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+        enable_infra_check=ENABLE_INFRA_CHECK,
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args["infra_check_config"]["EnableInfraCheck"] == ENABLE_INFRA_CHECK
+
+
 def test_framework_with_container_entry_point(sagemaker_session):
     f = DummyFramework(
         entry_point=SCRIPT_PATH,

tests/unit/test_session.py

@@ -1899,6 +1899,7 @@ def test_train_with_sagemaker_config_injection(sagemaker_session):
             "TrainingRepositoryCredentialsProviderArn": "arn:aws:lambda:us-west-2:1234567897:function:test"
         },
     }
+    INFRA_CHECK_CONFIG = {"EnableInfraCheck": True}
     CONTAINER_ENTRY_POINT = ["bin/bash", "test.sh"]
     CONTAINER_ARGUMENTS = ["--arg1", "value1", "--arg2", "value2"]
 
@@ -1920,6 +1921,7 @@ def test_train_with_sagemaker_config_injection(sagemaker_session):
         training_image_config=TRAINING_IMAGE_CONFIG,
         container_entry_point=CONTAINER_ENTRY_POINT,
         container_arguments=CONTAINER_ARGUMENTS,
+        infra_check_config=INFRA_CHECK_CONFIG,
     )
 
     _, _, actual_train_args = sagemaker_session.sagemaker_client.method_calls[0]
@@ -1966,6 +1968,7 @@ def test_train_with_sagemaker_config_injection(sagemaker_session):
         actual_train_args["AlgorithmSpecification"]["ContainerEntrypoint"] == CONTAINER_ENTRY_POINT
     )
     assert actual_train_args["AlgorithmSpecification"]["ContainerArguments"] == CONTAINER_ARGUMENTS
+    assert actual_train_args["InfraCheckConfig"] == INFRA_CHECK_CONFIG
     assert actual_train_args["RoleArn"] == expected_role_arn
     assert actual_train_args["ResourceConfig"] == {
         "InstanceCount": INSTANCE_COUNT,