chore: remove support for ecr spec fallbacks for jumpstart models

Author: evakravi

src/sagemaker/jumpstart/artifacts/image_uris.py

@@ -14,14 +14,12 @@
 from __future__ import absolute_import
 
 from typing import Optional
-from sagemaker import image_uris
 from sagemaker.jumpstart.constants import (
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.enums import (
     JumpStartModelType,
     JumpStartScriptScope,
-    ModelFramework,
 )
 from sagemaker.jumpstart.utils import (
     get_region_fallback,
@@ -142,13 +140,11 @@ def _retrieve_image_uri(
             ecr_uri = model_specs.hosting_ecr_uri
             return ecr_uri
 
-        ecr_specs = model_specs.hosting_ecr_specs
-        if ecr_specs is None:
-            raise ValueError(
-                f"No inference ECR configuration found for JumpStart model ID '{model_id}' "
-                f"with {instance_type} instance type in {region}. "
-                "Please try another instance type or region."
-            )
+        raise ValueError(
+            f"No inference ECR configuration found for JumpStart model ID '{model_id}' "
+            f"with {instance_type} instance type in {region}. "
+            "Please try another instance type or region."
+        )
     elif image_scope == JumpStartScriptScope.TRAINING:
         training_instance_type_variants = model_specs.training_instance_type_variants
         if training_instance_type_variants:
@@ -161,65 +157,8 @@ def _retrieve_image_uri(
             ecr_uri = model_specs.training_ecr_uri
             return ecr_uri
 
-        ecr_specs = model_specs.training_ecr_specs
-        if ecr_specs is None:
-            raise ValueError(
-                f"No training ECR configuration found for JumpStart model ID '{model_id}' "
-                f"with {instance_type} instance type in {region}. "
-                "Please try another instance type or region."
-            )
-    if framework is not None and framework != ecr_specs.framework:
-        raise ValueError(
-            f"Incorrect container framework '{framework}' for JumpStart model ID '{model_id}' "
-            f"and version '{model_version}'."
-        )
-
-    if version is not None and version != ecr_specs.framework_version:
-        raise ValueError(
-            f"Incorrect container framework version '{version}' for JumpStart model ID "
-            f"'{model_id}' and version '{model_version}'."
-        )
-
-    if py_version is not None and py_version != ecr_specs.py_version:
         raise ValueError(
-            f"Incorrect python version '{py_version}' for JumpStart model ID '{model_id}' "
-            f"and version '{model_version}'."
+            f"No training ECR configuration found for JumpStart model ID '{model_id}' "
+            f"with {instance_type} instance type in {region}. "
+            "Please try another instance type or region."
         )
-
-    base_framework_version_override: Optional[str] = None
-    version_override: Optional[str] = None
-    if ecr_specs.framework == ModelFramework.HUGGINGFACE:
-        base_framework_version_override = ecr_specs.framework_version
-        version_override = ecr_specs.huggingface_transformers_version
-
-    if image_scope == JumpStartScriptScope.TRAINING:
-        return image_uris.get_training_image_uri(
-            region=region,
-            framework=ecr_specs.framework,
-            framework_version=version_override or ecr_specs.framework_version,
-            py_version=ecr_specs.py_version,
-            image_uri=None,
-            distribution=None,
-            compiler_config=None,
-            tensorflow_version=None,
-            pytorch_version=base_framework_version_override or base_framework_version,
-            instance_type=instance_type,
-        )
-    if base_framework_version_override is not None:
-        base_framework_version_override = f"pytorch{base_framework_version_override}"
-
-    return image_uris.retrieve(
-        framework=ecr_specs.framework,
-        region=region,
-        version=version_override or ecr_specs.framework_version,
-        py_version=ecr_specs.py_version,
-        instance_type=instance_type,
-        hub_arn=hub_arn,
-        accelerator_type=accelerator_type,
-        image_scope=image_scope,
-        container_version=container_version,
-        distribution=distribution,
-        base_framework_version=base_framework_version_override or base_framework_version,
-        training_compiler_config=training_compiler_config,
-        config_name=config_name,
-    )

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

@@ -46,7 +46,13 @@ def test_jumpstart_default_hyperparameters(
         model_version="*",
         sagemaker_session=mock_session,
     )
-    assert params == {"adam-learning-rate": "0.05", "batch-size": "4", "epochs": "3"}
+    assert params == {
+        "train_only_top_layer": "True",
+        "epochs": "5",
+        "learning_rate": "0.001",
+        "batch_size": "4",
+        "reinitialize_top_layer": "Auto",
+    }
 
     patched_get_model_specs.assert_called_once_with(
         region=region,
@@ -66,7 +72,13 @@ def test_jumpstart_default_hyperparameters(
         model_version="1.*",
         sagemaker_session=mock_session,
     )
-    assert params == {"adam-learning-rate": "0.05", "batch-size": "4", "epochs": "3"}
+    assert params == {
+        "train_only_top_layer": "True",
+        "epochs": "5",
+        "learning_rate": "0.001",
+        "batch_size": "4",
+        "reinitialize_top_layer": "Auto",
+    }
 
     patched_get_model_specs.assert_called_once_with(
         region=region,
@@ -88,12 +100,14 @@ def test_jumpstart_default_hyperparameters(
         sagemaker_session=mock_session,
     )
     assert params == {
-        "adam-learning-rate": "0.05",
-        "batch-size": "4",
-        "epochs": "3",
-        "sagemaker_container_log_level": "20",
-        "sagemaker_program": "transfer_learning.py",
+        "train_only_top_layer": "True",
+        "epochs": "5",
+        "learning_rate": "0.001",
+        "batch_size": "4",
+        "reinitialize_top_layer": "Auto",
         "sagemaker_submit_directory": "/opt/ml/input/data/code/sourcedir.tar.gz",
+        "sagemaker_program": "transfer_learning.py",
+        "sagemaker_container_log_level": "20",
     }
 
     patched_get_model_specs.assert_called_once_with(

tests/unit/sagemaker/hyperparameters/jumpstart/test_validate.py

@@ -21,7 +21,7 @@
 from sagemaker.jumpstart.exceptions import JumpStartHyperparametersError
 from sagemaker.jumpstart.types import JumpStartHyperparameter
 
-from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec, get_spec_from_base_spec
 
 region = "us-west-2"
 mock_client = boto3.client("s3")
@@ -34,7 +34,7 @@ def test_jumpstart_validate_provided_hyperparameters(
     patched_get_model_specs, patched_validate_model_id_and_get_type
 ):
     def add_options_to_hyperparameter(*largs, **kwargs):
-        spec = get_spec_from_base_spec(*largs, **kwargs)
+        spec = get_prototype_model_spec(*largs, **kwargs)
         spec.hyperparameters.extend(
             [
                 JumpStartHyperparameter(
@@ -115,7 +115,7 @@ def add_options_to_hyperparameter(*largs, **kwargs):
     patched_get_model_specs.side_effect = add_options_to_hyperparameter
     patched_validate_model_id_and_get_type.return_value = JumpStartModelType.OPEN_WEIGHTS
 
-    model_id, model_version = "pytorch-eqa-bert-base-cased", "*"
+    model_id, model_version = "mxnet-semseg-fcn-resnet50-ade", "*"
     region = "us-west-2"
 
     hyperparameter_to_test = {
@@ -412,7 +412,7 @@ def test_jumpstart_validate_algorithm_hyperparameters(
     patched_get_model_specs, patched_validate_model_id_and_get_type
 ):
     def add_options_to_hyperparameter(*largs, **kwargs):
-        spec = get_spec_from_base_spec(*largs, **kwargs)
+        spec = get_prototype_model_spec(*largs, **kwargs)
         spec.hyperparameters.append(
             JumpStartHyperparameter(
                 {
@@ -429,10 +429,11 @@ def add_options_to_hyperparameter(*largs, **kwargs):
     patched_get_model_specs.side_effect = add_options_to_hyperparameter
     patched_validate_model_id_and_get_type.return_value = JumpStartModelType.OPEN_WEIGHTS
 
-    model_id, model_version = "pytorch-eqa-bert-base-cased", "*"
+    model_id, model_version = "mxnet-semseg-fcn-resnet50-ade", "*"
     region = "us-west-2"
 
     hyperparameter_to_test = {
+        "train-only-top-layer": "True",
         "adam-learning-rate": "0.05",
         "batch-size": "4",
         "epochs": "3",
@@ -488,13 +489,14 @@ def test_jumpstart_validate_all_hyperparameters(
     patched_get_model_specs, patched_validate_model_id_and_get_type
 ):
 
-    patched_get_model_specs.side_effect = get_spec_from_base_spec
+    patched_get_model_specs.side_effect = get_prototype_model_spec
     patched_validate_model_id_and_get_type.return_value = JumpStartModelType.OPEN_WEIGHTS
 
-    model_id, model_version = "pytorch-eqa-bert-base-cased", "*"
+    model_id, model_version = "mxnet-semseg-fcn-resnet50-ade", "*"
     region = "us-west-2"
 
     hyperparameter_to_test = {
+        "train-only-top-layer": "True",
         "adam-learning-rate": "0.05",
         "batch-size": "4",
         "epochs": "3",

tests/unit/sagemaker/image_uris/jumpstart/test_catboost.py

@@ -30,7 +30,7 @@ def test_jumpstart_catboost_image_uri(patched_get_model_specs, session):
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
     model_id, model_version = "catboost-classification-model", "*"
-    instance_type = "ml.p2.xlarge"
+    instance_type = "ml.m5.xlarge"
     region = "us-west-2"
 
     model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
@@ -55,7 +55,7 @@ def test_jumpstart_catboost_image_uri(patched_get_model_specs, session):
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
-    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
+    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:2.0.1-cpu-py310"
 
     # training
     uri = image_uris.retrieve(
@@ -78,4 +78,4 @@ def test_jumpstart_catboost_image_uri(patched_get_model_specs, session):
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri
-    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.9.0-gpu-py38"
+    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.9.0-cpu-py38"

tests/unit/sagemaker/image_uris/jumpstart/test_common.py

@@ -47,7 +47,7 @@ def test_jumpstart_common_image_uri(
         image_scope="training",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="*",
-        instance_type="ml.p2.xlarge",
+        instance_type="ml.m5.xlarge",
         sagemaker_session=mock_session,
     )
     patched_get_model_specs.assert_called_once_with(
@@ -70,7 +70,7 @@ def test_jumpstart_common_image_uri(
         image_scope="inference",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="1.*",
-        instance_type="ml.p2.xlarge",
+        instance_type="ml.m5.xlarge",
         sagemaker_session=mock_session,
     )
     patched_get_model_specs.assert_called_once_with(
@@ -93,7 +93,7 @@ def test_jumpstart_common_image_uri(
         image_scope="training",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="*",
-        instance_type="ml.p2.xlarge",
+        instance_type="ml.m5.xlarge",
         sagemaker_session=mock_session,
     )
     patched_get_model_specs.assert_called_once_with(
@@ -116,7 +116,7 @@ def test_jumpstart_common_image_uri(
         image_scope="inference",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="1.*",
-        instance_type="ml.p2.xlarge",
+        instance_type="ml.m5.xlarge",
         sagemaker_session=mock_session,
     )
     patched_get_model_specs.assert_called_once_with(
@@ -137,7 +137,7 @@ def test_jumpstart_common_image_uri(
             image_scope="BAD_SCOPE",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
-            instance_type="ml.p2.xlarge",
+            instance_type="ml.m5.xlarge",
         )
 
     with pytest.raises(KeyError):
@@ -147,7 +147,7 @@ def test_jumpstart_common_image_uri(
             image_scope="training",
             model_id="blah",
             model_version="*",
-            instance_type="ml.p2.xlarge",
+            instance_type="ml.m5.xlarge",
         )
 
     with pytest.raises(ValueError):
@@ -157,7 +157,7 @@ def test_jumpstart_common_image_uri(
             image_scope="training",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
-            instance_type="ml.p2.xlarge",
+            instance_type="ml.m5.xlarge",
         )
 
     with pytest.raises(ValueError):
@@ -166,7 +166,7 @@ def test_jumpstart_common_image_uri(
             region="us-west-2",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
-            instance_type="ml.p2.xlarge",
+            instance_type="ml.m5.xlarge",
         )
 
     with pytest.raises(ValueError):
@@ -175,7 +175,7 @@ def test_jumpstart_common_image_uri(
             region="us-west-2",
             image_scope="training",
             model_version="*",
-            instance_type="ml.p2.xlarge",
+            instance_type="ml.m5.xlarge",
         )
 
     with pytest.raises(ValueError):
@@ -184,5 +184,5 @@ def test_jumpstart_common_image_uri(
             framework=None,
             image_scope="training",
             model_id="pytorch-ic-mobilenet-v2",
-            instance_type="ml.p2.xlarge",
+            instance_type="ml.m5.xlarge",
         )

tests/unit/sagemaker/image_uris/jumpstart/test_huggingface.py

@@ -28,7 +28,8 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs, session):
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
     model_id, model_version = "huggingface-spc-bert-base-cased", "*"
-    instance_type = "ml.p2.xlarge"
+    instance_type = "ml.m5.xlarge"
+    training_instance_type = "ml.p3.2xlarge"
     region = "us-west-2"
 
     model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
@@ -55,7 +56,7 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs, session):
 
     assert (
         uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:"
-        "1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04"
+        "1.7.1-transformers4.6.1-cpu-py36-ubuntu18.04"
     )
 
     # training
@@ -65,7 +66,7 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs, session):
         image_scope="training",
         model_id=model_id,
         model_version=model_version,
-        instance_type=instance_type,
+        instance_type=training_instance_type,
     )
 
     framework_class_uri = HuggingFace(
@@ -75,7 +76,7 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs, session):
         entry_point="some_entry_point",
         transformers_version=model_specs.training_ecr_specs.huggingface_transformers_version,
         pytorch_version=model_specs.training_ecr_specs.framework_version,
-        instance_type=instance_type,
+        instance_type=training_instance_type,
         instance_count=1,
         sagemaker_session=session,
     ).training_image_uri(region=region)

tests/unit/sagemaker/image_uris/jumpstart/test_lightgbm.py

@@ -28,7 +28,7 @@ def test_jumpstart_lightgbm_image_uri(patched_get_model_specs, session):
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
     model_id, model_version = "lightgbm-classification-model", "*"
-    instance_type = "ml.p2.xlarge"
+    instance_type = "ml.m5.xlarge"
     region = "us-west-2"
 
     model_specs = accessors.JumpStartModelsAccessor.get_model_specs(region, model_id, model_version)
@@ -53,7 +53,7 @@ def test_jumpstart_lightgbm_image_uri(patched_get_model_specs, session):
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
-    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
+    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:2.0.1-cpu-py310"
 
     # training
     uri = image_uris.retrieve(
@@ -76,4 +76,4 @@ def test_jumpstart_lightgbm_image_uri(patched_get_model_specs, session):
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri
-    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.9.0-gpu-py38"
+    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.9.0-cpu-py38"