fix: Missing JumpStart estimator args (#3978)

Author: evakravi

src/sagemaker/jumpstart/estimator.py

@@ -104,6 +104,9 @@ 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,
+        container_entry_point: Optional[List[str]] = None,
+        container_arguments: Optional[List[str]] = None,
+        disable_output_compression: Optional[bool] = None,
     ):
         """Initializes a ``JumpStartEstimator``.
 
@@ -484,6 +487,13 @@ 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).
+            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.
+            container_arguments (Optional[List[str]]): The arguments for a container used to run
+                a training job.
+            disable_output_compression (Optional[bool]): When set to true, Model is uploaded
+                to Amazon S3 without compression after training finishes.
 
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
@@ -553,6 +563,9 @@ def _is_valid_model_id_hook():
                 training_repository_credentials_provider_arn
             ),
             image_uri=image_uri,
+            container_entry_point=container_entry_point,
+            container_arguments=container_arguments,
+            disable_output_compression=disable_output_compression,
         )
 
         self.model_id = estimator_init_kwargs.model_id

src/sagemaker/jumpstart/factory/estimator.py

@@ -116,6 +116,9 @@ def get_init_kwargs(
     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,
+    container_entry_point: Optional[List[str]] = None,
+    container_arguments: Optional[List[str]] = None,
+    disable_output_compression: Optional[bool] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -168,6 +171,9 @@ def get_init_kwargs(
         training_repository_access_mode=training_repository_access_mode,
         training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,
         image_uri=image_uri,
+        container_entry_point=container_entry_point,
+        container_arguments=container_arguments,
+        disable_output_compression=disable_output_compression,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)

src/sagemaker/jumpstart/types.py

@@ -777,6 +777,9 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "training_repository_credentials_provider_arn",
         "tolerate_deprecated_model",
         "tolerate_vulnerable_model",
+        "container_entry_point",
+        "container_arguments",
+        "disable_output_compression",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -837,6 +840,9 @@ def __init__(
         training_repository_credentials_provider_arn: Optional[Union[str, Any]] = None,
         tolerate_vulnerable_model: Optional[bool] = None,
         tolerate_deprecated_model: Optional[bool] = None,
+        container_entry_point: Optional[List[str]] = None,
+        container_arguments: Optional[List[str]] = None,
+        disable_output_compression: Optional[bool] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -890,6 +896,9 @@ def __init__(
         )
         self.tolerate_vulnerable_model = tolerate_vulnerable_model
         self.tolerate_deprecated_model = tolerate_deprecated_model
+        self.container_entry_point = container_entry_point
+        self.container_arguments = container_arguments
+        self.disable_output_compression = disable_output_compression
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

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

@@ -520,11 +520,9 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
 
         """If you add arguments to <Estimator constructor>, this test will fail.
         Please add the new argument to the skip set below,
-        and cut a ticket sev-3 to JumpStart team: AWS > SageMaker > JumpStart"""
+        and reach out to JumpStart team."""
 
-        init_args_to_skip: Set[str] = set(
-            ["container_entry_point", "container_arguments", "disable_output_compression", "kwargs"]
-        )
+        init_args_to_skip: Set[str] = set(["kwargs"])
         fit_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])