feat: jumpstart estimator enable infra check flag

Author: evakravi

src/sagemaker/estimator.py

@@ -3082,6 +3082,7 @@ def __init__(
             hyperparameters=hyperparameters,
             instance_groups=instance_groups,
             training_repository_access_mode=training_repository_access_mode,
+            enable_infra_check=enable_infra_check,
             training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,  # noqa: E501 # pylint: disable=line-too-long
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,

src/sagemaker/jumpstart/estimator.py

@@ -102,6 +102,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: Optional[bool] = None,
@@ -485,6 +486,8 @@ def __init__(
                 private Docker registry where your training image is hosted (Default: None).
                 When it's set to None, SageMaker will not do authentication before pulling the image
                 in the private Docker registry. (Default: None).
+            enable_infra_check (Optional[Union[bool, PipelineVariable]]):
+                Specifies whether it is running Sagemaker built-in infra check jobs.
             container_entry_point (Optional[List[str]]): The entrypoint script for a Docker
                 container used to run a training job. This script takes precedence over
                 the default train processing instructions.
@@ -565,6 +568,7 @@ def _is_valid_model_id_hook():
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
+            enable_infra_check=enable_infra_check,
         )
 
         self.model_id = estimator_init_kwargs.model_id

src/sagemaker/jumpstart/factory/estimator.py

@@ -123,6 +123,7 @@ def get_init_kwargs(
     container_entry_point: Optional[List[str]] = None,
     container_arguments: Optional[List[str]] = None,
     disable_output_compression: Optional[bool] = None,
+    enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -178,6 +179,7 @@ def get_init_kwargs(
         container_entry_point=container_entry_point,
         container_arguments=container_arguments,
         disable_output_compression=disable_output_compression,
+        enable_infra_check=enable_infra_check,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)

src/sagemaker/jumpstart/types.py

@@ -18,6 +18,7 @@
 from sagemaker.utils import get_instance_type_family
 
 from sagemaker.session import Session
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class JumpStartDataHolderType:
@@ -946,6 +947,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "container_entry_point",
         "container_arguments",
         "disable_output_compression",
+        "enable_infra_check",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1009,6 +1011,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1065,6 +1068,7 @@ def __init__(
         self.container_entry_point = container_entry_point
         self.container_arguments = container_arguments
         self.disable_output_compression = disable_output_compression
+        self.enable_infra_check = enable_infra_check
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

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", "enable_infra_check"])
+        init_args_to_skip: Set[str] = set(["kwargs"])
         fit_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])