Merge remote-tracking branch 'origin' into fix/sagemaker-session-region-not-being-used

Author: evakravi

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -12,9 +12,10 @@
 # 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_LOGGER,
     SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
@@ -111,7 +112,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,
@@ -121,6 +124,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,
@@ -601,6 +602,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
@@ -204,6 +205,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):