chore: raise client exception if accept_eula flag is not set for gated models

Author: evakravi

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -120,7 +120,7 @@ def _retrieve_default_environment_variables(
                 instance_type=instance_type,
             )
 
-            if gated_model_env_var is None and model_specs.gated_bucket:
+            if gated_model_env_var is None and model_specs.is_gated_model():
                 raise ValueError(
                     f"'{model_id}' does not support {instance_type} instance type for training. "
                     "Please use one of the following instance types: "

src/sagemaker/jumpstart/factory/estimator.py

@@ -62,6 +62,7 @@
 )
 from sagemaker.jumpstart.utils import (
     add_jumpstart_model_id_version_tags,
+    get_eula_message,
     update_dict_if_key_not_present,
     resolve_estimator_sagemaker_config_field,
     verify_model_region_and_return_specs,
@@ -595,6 +596,25 @@ def _add_env_to_kwargs(
             value,
         )
 
+    environment = getattr(kwargs, "environment", {}) or {}
+    if (
+        environment.get(SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY)
+        and str(environment.get("accept_eula", "")).lower() != "true"
+    ):
+        model_specs = verify_model_region_and_return_specs(
+            model_id=kwargs.model_id,
+            version=kwargs.model_version,
+            region=kwargs.region,
+            scope=JumpStartScriptScope.TRAINING,
+            tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+            tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+            sagemaker_session=kwargs.sagemaker_session,
+        )
+        if model_specs.is_gated_model():
+            raise ValueError(
+                f"Need to define ‘accept_eula'='true' within Environment. {get_eula_message(model_specs, kwargs.region)}"
+            )
+
     return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -956,6 +956,10 @@ def use_training_model_artifact(self) -> bool:
         # otherwise, return true is a training model package is not set
         return len(self.training_model_package_artifact_uris or {}) == 0
 
+    def is_gated_model(self) -> bool:
+        """Returns True if the model has a EULA key or the model bucket is gated."""
+        return self.gated_bucket or self.hosting_eula_key is not None
+
     def supports_incremental_training(self) -> bool:
         """Returns True if the model supports incremental training."""
         return self.incremental_training_supported

src/sagemaker/jumpstart/utils.py

@@ -466,21 +466,25 @@ def update_inference_tags_with_jumpstart_training_tags(
     return inference_tags
 
 
+def get_eula_message(model_specs: JumpStartModelSpecs, region: str) -> str:
+    """Returns EULA message to display to customers if one is available, else empty string."""
+    if model_specs.hosting_eula_key is None:
+        return ""
+    return (
+        f"Model '{model_specs.model_id}' requires accepting end-user license agreement (EULA). "
+        f"See https://{get_jumpstart_content_bucket(region=region)}.s3.{region}."
+        f"amazonaws.com{'.cn' if region.startswith('cn-') else ''}"
+        f"/{model_specs.hosting_eula_key} for terms of use."
+    )
+
+
 def emit_logs_based_on_model_specs(
     model_specs: JumpStartModelSpecs, region: str, s3_client: boto3.client
 ) -> None:
     """Emits logs based on model specs and region."""
 
     if model_specs.hosting_eula_key:
-        constants.JUMPSTART_LOGGER.info(
-            "Model '%s' requires accepting end-user license agreement (EULA). "
-            "See https://%s.s3.%s.amazonaws.com%s/%s for terms of use.",
-            model_specs.model_id,
-            get_jumpstart_content_bucket(region=region),
-            region,
-            ".cn" if region.startswith("cn-") else "",
-            model_specs.hosting_eula_key,
-        )
+        constants.JUMPSTART_LOGGER.info(get_eula_message(model_specs, region))
 
     full_version: str = model_specs.version
 

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

@@ -315,39 +315,21 @@ def test_gated_model_s3_uri(
         mock_session_estimator.return_value = sagemaker_session
         mock_session_model.return_value = sagemaker_session
 
-        JumpStartEstimator(
-            model_id=model_id,
-            environment={
-                "accept_eula": "false",
-                "what am i": "doing",
-                "SageMakerGatedModelS3Uri": "none of your business",
-            },
-        )
-
-        mock_estimator_init.assert_called_once_with(
-            instance_type="ml.p3.2xlarge",
-            instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.3-cu117",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
-            "meta/transfer_learning/textgeneration/v1.0.0/sourcedir.tar.gz",
-            entry_point="transfer_learning.py",
-            role=execution_role,
-            sagemaker_session=sagemaker_session,
-            max_run=360000,
-            enable_network_isolation=True,
-            encrypt_inter_container_traffic=True,
-            environment={
-                "accept_eula": "false",
-                "what am i": "doing",
-                "SageMakerGatedModelS3Uri": "none of your business",
-            },
-            tags=[
-                {
-                    "Key": "sagemaker-sdk:jumpstart-model-id",
-                    "Value": "js-gated-artifact-trainable-model",
+        with pytest.raises(ValueError) as e:
+            JumpStartEstimator(
+                model_id=model_id,
+                environment={
+                    "accept_eula": "false",
+                    "what am i": "doing",
+                    "SageMakerGatedModelS3Uri": "none of your business",
                 },
-                {"Key": "sagemaker-sdk:jumpstart-model-version", "Value": "2.0.0"},
-            ],
+            )
+        assert str(e.value) == (
+            "Need to define ‘accept_eula'='true' within Environment. "
+            "Model 'meta-textgeneration-llama-2-7b-f' requires accepting end-user "
+            "license agreement (EULA). See "
+            "https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com/fmhMetadata/eula/llamaEula.txt"
+            " for terms of use."
         )
 
         mock_estimator_init.reset_mock()

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -905,13 +905,8 @@ def make_accept_eula_inference_spec(*largs, **kwargs):
             make_accept_eula_inference_spec(), "us-east-1", MOCK_CLIENT
         )
         mocked_info_log.assert_any_call(
-            "Model '%s' requires accepting end-user license agreement (EULA). "
-            "See https://%s.s3.%s.amazonaws.com%s/%s for terms of use.",
-            "pytorch-eqa-bert-base-cased",
-            "jumpstart-cache-prod-us-east-1",
-            "us-east-1",
-            "",
-            "read/the/fine/print.txt",
+            "Model 'pytorch-eqa-bert-base-cased' requires accepting end-user license agreement (EULA). "
+            "See https://jumpstart-cache-prod-us-east-1.s3.us-east-1.amazonaws.com/read/the/fine/print.txt for terms of use.",
         )