chore: emit warning when instance type is chosen with no gated training artifacts

Author: evakravi

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,12 +123,32 @@ 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():
-                raise ValueError(
-                    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)}."
-                )
+
+                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(

src/sagemaker/jumpstart/utils.py

@@ -467,7 +467,7 @@ def update_inference_tags_with_jumpstart_training_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."""
+    """Returns EULA message to display if one is available, else empty string."""
     if model_specs.hosting_eula_key is None:
         return ""
     return (

tests/unit/sagemaker/environment_variables/jumpstart/test_default.py

@@ -179,26 +179,50 @@ def test_jumpstart_sdk_environment_variables(patched_get_model_specs):
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_sdk_environment_variables_no_gated_env_var_available(patched_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 that unsupported instance types raise an exception
-    with pytest.raises(ValueError) as e:
-        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",
-        )
-    assert (
-        str(e.value) == "'gemma-model' does not support ml.p3.2xlarge instance type for "
+    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."
     )