fix: jumpstart estimator for gated uncompressed training

Author: evakravi

src/sagemaker/estimator.py

@@ -1310,6 +1310,7 @@ def fit(
         logs: str = "All",
         job_name: Optional[str] = None,
         experiment_config: Optional[Dict[str, str]] = None,
+        accept_eula: Optional[bool] = None,
     ):
         """Train a model using the input training dataset.
 
@@ -1363,14 +1364,21 @@ def fit(
                 * Both `ExperimentName` and `TrialName` will be ignored if the Estimator instance
                 is built with :class:`~sagemaker.workflow.pipeline_context.PipelineSession`.
                 However, the value of `TrialComponentDisplayName` is honored for display in Studio.
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                    False to indicate whether model terms of use have been accepted.
+                    The `accept_eula` value must be explicitly defined as `True` in order to
+                    accept the end-user license agreement (EULA) that some
+                    models require. (Default: None).
         Returns:
             None or pipeline step arguments in case the Estimator instance is built with
             :class:`~sagemaker.workflow.pipeline_context.PipelineSession`
         """
         self._prepare_for_training(job_name=job_name)
 
         experiment_config = check_and_get_run_experiment_config(experiment_config)
-        self.latest_training_job = _TrainingJob.start_new(self, inputs, experiment_config)
+        self.latest_training_job = _TrainingJob.start_new(
+            self, inputs, experiment_config, accept_eula
+        )
         self.jobs.append(self.latest_training_job)
         forward_to_mlflow_tracking_server = False
         if os.environ.get("MLFLOW_TRACKING_URI") and self.enable_network_isolation():
@@ -2484,7 +2492,7 @@ class _TrainingJob(_Job):
     """Placeholder docstring"""
 
     @classmethod
-    def start_new(cls, estimator, inputs, experiment_config):
+    def start_new(cls, estimator, inputs, experiment_config, accept_eula):
         """Create a new Amazon SageMaker training job from the estimator.
 
         Args:
@@ -2504,19 +2512,24 @@ def start_new(cls, estimator, inputs, experiment_config):
                 will be unassociated.
                 * `TrialComponentDisplayName` is used for display in Studio.
                 * `RunName` is used to record an experiment run.
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                False to indicate whether model terms of use have been accepted.
+                The `accept_eula` value must be explicitly defined as `True` in order to
+                accept the end-user license agreement (EULA) that some
+                models require. (Default: None).
         Returns:
             sagemaker.estimator._TrainingJob: Constructed object that captures
             all information about the started training job.
         """
-        train_args = cls._get_train_args(estimator, inputs, experiment_config)
+        train_args = cls._get_train_args(estimator, inputs, experiment_config, accept_eula)
 
         logger.debug("Train args after processing defaults: %s", train_args)
         estimator.sagemaker_session.train(**train_args)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
 
     @classmethod
-    def _get_train_args(cls, estimator, inputs, experiment_config):
+    def _get_train_args(cls, estimator, inputs, experiment_config, accept_eula):
         """Constructs a dict of arguments for an Amazon SageMaker training job from the estimator.
 
         Args:
@@ -2536,6 +2549,11 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
                 will be unassociated.
                 * `TrialComponentDisplayName` is used for display in Studio.
                 * `RunName` is used to record an experiment run.
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                False to indicate whether model terms of use have been accepted.
+                The `accept_eula` value must be explicitly defined as `True` in order to
+                accept the end-user license agreement (EULA) that some
+                models require. (Default: None).
 
         Returns:
             Dict: dict for `sagemaker.session.Session.train` method
@@ -2652,6 +2670,9 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         if estimator.get_session_chaining_config() is not None:
             train_args["session_chaining_config"] = estimator.get_session_chaining_config()
 
+        if accept_eula is not None:
+            cls._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
         return train_args
 
     @classmethod
@@ -2674,6 +2695,42 @@ def _add_spot_checkpoint_args(cls, local_mode, estimator, train_args):
                 raise ValueError("Setting checkpoint_local_path is not supported in local mode.")
             train_args["checkpoint_local_path"] = estimator.checkpoint_local_path
 
+    @classmethod
+    def _set_accept_eula_for_input_data_config(cls, train_args, accept_eula):
+        """Set the AcceptEula flag for all input data configurations.
+
+        This method sets the AcceptEula flag in the ModelAccessConfig for all S3DataSources
+        in the InputDataConfig array. It handles cases where keys might not exist in the
+        nested dictionary structure.
+
+        Args:
+            train_args (dict): The training job arguments dictionary
+            accept_eula (bool): The value to set for AcceptEula flag
+        """
+        if "InputDataConfig" not in train_args:
+            return
+
+        eula_count = 0
+        s3_uris = []
+
+        for idx in range(len(train_args["InputDataConfig"])):
+            if "DataSource" in train_args["InputDataConfig"][idx]:
+                data_source = train_args["InputDataConfig"][idx]["DataSource"]
+                if "S3DataSource" in data_source:
+                    s3_data_source = data_source["S3DataSource"]
+                    if "ModelAccessConfig" not in s3_data_source:
+                        s3_data_source["ModelAccessConfig"] = {}
+                    s3_data_source["ModelAccessConfig"]["AcceptEula"] = accept_eula
+                    eula_count += 1
+
+                    # Collect S3 URI if available
+                    if "S3Uri" in s3_data_source:
+                        s3_uris.append(s3_data_source["S3Uri"])
+
+        # Log info if more than one EULA needs to be accepted
+        if eula_count > 1:
+            logger.info("Accepting EULA for %d S3 data sources: %s", eula_count, ", ".join(s3_uris))
+
     @classmethod
     def _is_local_channel(cls, input_uri):
         """Placeholder docstring"""

src/sagemaker/jumpstart/estimator.py

@@ -713,6 +713,7 @@ def fit(
             sagemaker_session=self.sagemaker_session,
             config_name=self.config_name,
             hub_access_config=self.hub_access_config,
+            accept_eula=accept_eula,
         )
         remove_env_var_from_estimator_kwargs_if_model_access_config_present(
             self.init_kwargs, self.model_access_config

src/sagemaker/jumpstart/factory/estimator.py

@@ -266,6 +266,7 @@ def get_fit_kwargs(
     sagemaker_session: Optional[Session] = None,
     config_name: Optional[str] = None,
     hub_access_config: Optional[Dict] = None,
+    accept_eula: Optional[bool] = None,
 ) -> JumpStartEstimatorFitKwargs:
     """Returns kwargs required call `fit` on `sagemaker.estimator.Estimator` object."""
 
@@ -283,6 +284,7 @@ def get_fit_kwargs(
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        accept_eula=accept_eula,
     )
 
     estimator_fit_kwargs, _ = _set_temp_sagemaker_session_if_not_set(kwargs=estimator_fit_kwargs)

src/sagemaker/jumpstart/types.py

@@ -1940,9 +1940,6 @@ def use_inference_script_uri(self) -> bool:
 
     def use_training_model_artifact(self) -> bool:
         """Returns True if the model should use a model uri when kicking off training job."""
-        # gated model never use training model artifact
-        if self.gated_bucket:
-            return False
 
         # otherwise, return true is a training model package is not set
         return len(self.training_model_package_artifact_uris or {}) == 0
@@ -2595,6 +2592,7 @@ class JumpStartEstimatorFitKwargs(JumpStartKwargs):
         "sagemaker_session",
         "config_name",
         "specs",
+        "accept_eula",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2625,6 +2623,7 @@ def __init__(
         tolerate_vulnerable_model: Optional[bool] = None,
         sagemaker_session: Optional[Session] = None,
         config_name: Optional[str] = None,
+        accept_eula: Optional[bool] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -2642,6 +2641,7 @@ def __init__(
         self.tolerate_vulnerable_model = tolerate_vulnerable_model
         self.sagemaker_session = sagemaker_session
         self.config_name = config_name
+        self.accept_eula = accept_eula
 
 
 class JumpStartEstimatorDeployKwargs(JumpStartKwargs):