chore: emit warning when no instance specific gated training env var is available, and raise exception when accept_eula flag is not supplied

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):
 

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart environment variables."""
 from __future__ import absolute_import
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional, Set
 from sagemaker.jumpstart.constants import (
  DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
  JUMPSTART_DEFAULT_REGION_NAME,
+ JUMPSTART_LOGGER,
  SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -110,7 +111,9 @@ def _retrieve_default_environment_variables(
 
  default_environment_variables.update(instance_specific_environment_variables)
 
- gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+ retrieve_gated_env_var_for_instance_type: Callable[
+ [str], Optional[str]
+ ] = lambda instance_type: _retrieve_gated_model_uri_env_var_value(
  model_id=model_id,
  model_version=model_version,
  region=region,
@@ -120,6 +123,33 @@ def _retrieve_default_environment_variables(
  instance_type=instance_type,
  )
 
+ gated_model_env_var: Optional[str] = retrieve_gated_env_var_for_instance_type(
+ instance_type
+ )
+
+ if gated_model_env_var is None and model_specs.is_gated_model():
+
+ possible_env_vars: Set[str] = {
+ retrieve_gated_env_var_for_instance_type(instance_type)
+ for instance_type in model_specs.supported_training_instance_types
+ }
+
+ # If all officially supported instance types have the same underlying artifact,
+ # we can use this artifact with high confidence that it'll succeed with
+ # an arbitrary instance.
+ if len(possible_env_vars) == 1:
+ gated_model_env_var = list(possible_env_vars)[0]
+
+ # If this model does not have 1 artifact for all supported instance types,
+ # we cannot determine which artifact to use for an arbitrary instance.
+ else:
+ log_msg = (
+ f"'{model_id}' does not support {instance_type} instance type"
+ " for training. Please use one of the following instance types: "
+ f"{', '.join(model_specs.supported_training_instance_types)}."
+ )
+ JUMPSTART_LOGGER.warning(log_msg)
+
  if gated_model_env_var is not None:
  default_environment_variables.update(
  {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}

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,
@@ -597,6 +598,26 @@ 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(
+ "Need to define ‘accept_eula'='true' within Environment. "
+ f"{get_eula_message(model_specs, kwargs.region)}"
+ )
+
  return kwargs
 
 

src/sagemaker/jumpstart/types.py

@@ -963,6 +963,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

@@ -476,21 +476,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 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/environment_variables/jumpstart/test_default.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker import environment_variables
+from sagemaker.jumpstart.utils import get_jumpstart_gated_content_bucket
 from sagemaker.jumpstart.enums import JumpStartModelType
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
@@ -203,6 +204,70 @@ def test_jumpstart_sdk_environment_variables(
  )
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_1_artifact_all_variants(patched_get_model_specs):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model-1-artifact"
+ region = "us-west-2"
+
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+
+@patch("sagemaker.jumpstart.artifacts.environment_variables.JUMPSTART_LOGGER")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(
+ patched_get_model_specs, patched_jumpstart_logger
+):
+
+ patched_get_model_specs.side_effect = get_special_model_spec
+
+ model_id = "gemma-model"
+ region = "us-west-2"
+
+ assert {} == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.p3.2xlarge",
+ script="training",
+ )
+
+ patched_jumpstart_logger.warning.assert_called_once_with(
+ "'gemma-model' does not support ml.p3.2xlarge instance type for "
+ "training. Please use one of the following instance types: "
+ "ml.g5.12xlarge, ml.g5.24xlarge, ml.g5.48xlarge, ml.p4d.24xlarge."
+ )
+
+ # assert that supported instance types succeed
+ assert {
+ "SageMakerGatedModelS3Uri": f"s3://{get_jumpstart_gated_content_bucket(region)}/"
+ "huggingface-training/g5/v1.0.0/train-huggingface-llm-gemma-7b-instruct.tar.gz"
+ } == environment_variables.retrieve_default(
+ region=region,
+ model_id=model_id,
+ model_version="*",
+ include_aws_sdk_env_vars=False,
+ sagemaker_session=mock_session,
+ instance_type="ml.g5.24xlarge",
+ script="training",
+ )
+
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_sdk_environment_variables_instance_type_overrides(patched_get_model_specs):