change: cleanup code

Author: evakravi

src/sagemaker/image_uris.py

@@ -82,9 +82,9 @@ def retrieve(
         model_version (str): Version of the JumpStart model for which to retrieve the
             image URI (default: None).
         tolerate_vulnerable_model (bool): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+            not raised). False if these models should raise an exception. (Default: None).
         tolerate_deprecated_model (bool): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+            not raised). False if these models should raise an exception. (Default: None).
 
     Returns:
         str: the ECR URI for the corresponding SageMaker Docker image.

src/sagemaker/jumpstart/artifacts.py

@@ -16,8 +16,7 @@
 from sagemaker import image_uris
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
-    INFERENCE,
-    TRAINING,
+    JumpStartScriptScope,
     ModelFramework,
     VariableScope,
 )
@@ -77,9 +76,9 @@ def _retrieve_image_uri(
         training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
             A configuration class for the SageMaker Training Compiler.
         tolerate_vulnerable_model (bool): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception.
+            not raised). False if these models should raise an exception.
         tolerate_deprecated_model (bool): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception.
+            not raised). False if these models should raise an exception.
 
     Returns:
         str: the ECR URI for the corresponding SageMaker Docker image.
@@ -103,9 +102,9 @@ def _retrieve_image_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
-    if image_scope == INFERENCE:
+    if image_scope == JumpStartScriptScope.INFERENCE.value:
         ecr_specs = model_specs.hosting_ecr_specs
-    elif image_scope == TRAINING:
+    elif image_scope == JumpStartScriptScope.TRAINING.value:
         assert model_specs.training_ecr_specs is not None
         ecr_specs = model_specs.training_ecr_specs
 
@@ -133,7 +132,7 @@ def _retrieve_image_uri(
         base_framework_version_override = ecr_specs.framework_version
         version_override = ecr_specs.huggingface_transformers_version
 
-    if image_scope == TRAINING:
+    if image_scope == JumpStartScriptScope.TRAINING.value:
         return image_uris.get_training_image_uri(
             region=region,
             framework=ecr_specs.framework,
@@ -183,9 +182,9 @@ def _retrieve_model_uri(
             Valid values: "training" and "inference".
         region (str): Region for which to retrieve model S3 URI.
         tolerate_vulnerable_model (bool): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception.
+            not raised). False if these models should raise an exception.
         tolerate_deprecated_model (bool): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception.
+            not raised). False if these models should raise an exception.
     Returns:
         str: the model artifact S3 URI for the corresponding model.
 
@@ -208,9 +207,9 @@ def _retrieve_model_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
-    if model_scope == INFERENCE:
+    if model_scope == JumpStartScriptScope.INFERENCE.value:
         model_artifact_key = model_specs.hosting_artifact_key
-    elif model_scope == TRAINING:
+    elif model_scope == JumpStartScriptScope.TRAINING.value:
         assert model_specs.training_artifact_key is not None
         model_artifact_key = model_specs.training_artifact_key
 
@@ -240,9 +239,9 @@ def _retrieve_script_uri(
             Valid values: "training" and "inference".
         region (str): Region for which to retrieve model script S3 URI.
         tolerate_vulnerable_model (bool): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception.
+            not raised). False if these models should raise an exception.
         tolerate_deprecated_model (bool): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception.
+            not raised). False if these models should raise an exception.
     Returns:
         str: the model script URI for the corresponding model.
 
@@ -265,9 +264,9 @@ def _retrieve_script_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
-    if script_scope == INFERENCE:
+    if script_scope == JumpStartScriptScope.INFERENCE.value:
         model_script_key = model_specs.hosting_script_key
-    elif script_scope == TRAINING:
+    elif script_scope == JumpStartScriptScope.TRAINING.value:
         assert model_specs.training_script_key is not None
         model_script_key = model_specs.training_script_key
 

src/sagemaker/jumpstart/constants.py

@@ -116,14 +116,21 @@
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 
-INFERENCE = "inference"
-TRAINING = "training"
-SUPPORTED_JUMPSTART_SCOPES = set([INFERENCE, TRAINING])
 
 INFERENCE_ENTRYPOINT_SCRIPT_NAME = "inference.py"
 TRAINING_ENTRYPOINT_SCRIPT_NAME = "transfer_learning.py"
 
 
+class JumpStartScriptScope(str, Enum):
+    """Enum class for JumpStart script scopes."""
+
+    INFERENCE = "inference"
+    TRAINING = "training"
+
+
+SUPPORTED_JUMPSTART_SCOPES = set(scope.value for scope in JumpStartScriptScope)
+
+
 class ModelFramework(str, Enum):
     """Enum class for JumpStart model framework.
 

src/sagemaker/jumpstart/exceptions.py

@@ -12,46 +12,79 @@
 # language governing permissions and limitations under the License.
 """This module stores exceptions related to SageMaker JumpStart."""
 
+from __future__ import absolute_import
 from typing import List, Optional
 
+from sagemaker.jumpstart.constants import JumpStartScriptScope
+
 
 class VulnerableJumpStartModelError(Exception):
-    """Exception raised for errors with vulnerable JumpStart models."""
+    """Exception raised when trying to access a JumpStart model specs flagged as vulnerable.
+
+    Raise this exception only if the scope of attributes accessed in the specifications have
+    vulnerabilities. For example, a model training script may have vulnerabilities, but not
+    the hosting scripts. In such a case, raise a ``VulnerableJumpStartModelError`` only when
+    accessing the training specifications.
+    """
 
     def __init__(
         self,
         model_id: Optional[str] = None,
         version: Optional[str] = None,
         vulnerabilities: Optional[List[str]] = None,
-        inference: Optional[bool] = None,
+        scope: Optional[JumpStartScriptScope] = None,
         message: Optional[str] = None,
     ):
+        """Instantiates VulnerableJumpStartModelError exception.
+
+        Args:
+            model_id (Optional[str]): model id of vulnerable JumpStart model.
+                (Default: None).
+            version (Optional[str]): version of vulnerable JumpStart model.
+                (Default: None).
+            vulnerabilities (Optional[List[str]]): vulnerabilities associated with
+                model. (Default: None).
+
+        """
         if message:
             self.message = message
         else:
-            if None in [model_id, version, vulnerabilities, inference]:
+            if None in [model_id, version, vulnerabilities, scope]:
                 raise ValueError(
-                    "Must specify `model_id`, `version`, `vulnerabilities`, "
-                    "and inference arguments."
+                    "Must specify `model_id`, `version`, `vulnerabilities`, " "and scope arguments."
                 )
-            if inference is True:
+            if scope == JumpStartScriptScope.INFERENCE:
                 self.message = (
-                    f"JumpStart model '{model_id}' and version '{version}' has at least 1 "
-                    "vulnerable dependency in the inference scripts. "
-                    f"List of vulnerabilities: {', '.join(vulnerabilities)}"
+                    f"Version '{version}' of JumpStart model '{model_id}' "  # type: ignore
+                    "has at least 1 vulnerable dependency in the inference script. "
+                    "Please try targetting a higher version of the model. "
+                    f"List of vulnerabilities: {', '.join(vulnerabilities)}"  # type: ignore
                 )
-            else:
+            elif scope == JumpStartScriptScope.TRAINING:
                 self.message = (
-                    f"JumpStart model '{model_id}' and version '{version}' has at least 1 "
-                    "vulnerable dependency in the training scripts. "
-                    f"List of vulnerabilities: {', '.join(vulnerabilities)}"
+                    f"Version '{version}' of JumpStart model '{model_id}' "  # type: ignore
+                    "has at least 1 vulnerable dependency in the training script. "
+                    "Please try targetting a higher version of the model. "
+                    f"List of vulnerabilities: {', '.join(vulnerabilities)}"  # type: ignore
+                )
+            else:
+                raise NotImplementedError(
+                    "Unsupported scope for VulnerableJumpStartModelError: "  # type: ignore
+                    f"'{scope.value}'"
                 )
 
         super().__init__(self.message)
 
 
 class DeprecatedJumpStartModelError(Exception):
-    """Exception raised for errors with deprecated JumpStart models."""
+    """Exception raised when trying to access a JumpStart model deprecated specifications.
+
+    A deprecated specification for a JumpStart model does not mean the whole model is
+    deprecated. There may be more recent specifications available for this model. For
+    example, all specification before version ``2.0.0`` may be deprecated, in such a
+    case, the SDK would raise this exception only when specifications ``1.*`` are
+    accessed.
+    """
 
     def __init__(
         self,
@@ -64,6 +97,9 @@ def __init__(
         else:
             if None in [model_id, version]:
                 raise ValueError("Must specify `model_id` and `version` arguments.")
-            self.message = f"JumpStart model '{model_id}' and version '{version}' is deprecated."
+            self.message = (
+                f"Version '{version}' of JumpStart model '{model_id}' is deprecated. "
+                "Please try targetting a higher version of the model."
+            )
 
         super().__init__(self.message)

src/sagemaker/jumpstart/utils.py

@@ -21,7 +21,11 @@
     DeprecatedJumpStartModelError,
     VulnerableJumpStartModelError,
 )
-from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartVersionedModelId
+from sagemaker.jumpstart.types import (
+    JumpStartModelHeader,
+    JumpStartModelSpecs,
+    JumpStartVersionedModelId,
+)
 
 
 def get_jumpstart_launched_regions_message() -> str:
@@ -149,12 +153,9 @@ def verify_model_region_and_return_specs(
     region: str,
     tolerate_vulnerable_model: Optional[bool] = None,
     tolerate_deprecated_model: Optional[bool] = None,
-):
+) -> JumpStartModelSpecs:
     """Verifies that an acceptable model_id, version, scope, and region combination is provided.
 
-    If the scope is not supported, the model id/region/version has no spec, or the model is vulnerable
-    or deprecated, an exception will be raised.
-
     Args:
         model_id (Optional[str]): model id of the JumpStart model to verify and
             obtains specs.
@@ -163,10 +164,19 @@ def verify_model_region_and_return_specs(
         scope (Optional[str]): scope of the JumpStart model to verify.
         region (Optional[str]): region of the JumpStart model to verify and
             obtains specs.
-        tolerate_vulnerable_model (Optional[bool]): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
-        tolerate_deprecated_model (Optional[bool]): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+        tolerate_vulnerable_model (Optional[bool]): True if vulnerable models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: None).
+        tolerate_deprecated_model (Optional[bool]): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: None).
+
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+        NotImplementedError: If the scope is not supported.
+        VulnerableJumpStartModelError: If the model is vulnerable.
+        DeprecatedJumpStartModelError: If the model is deprecated.
     """
 
     if tolerate_vulnerable_model is None:
@@ -182,15 +192,22 @@ def verify_model_region_and_return_specs(
         )
 
     if scope not in constants.SUPPORTED_JUMPSTART_SCOPES:
-        raise ValueError(
-            f"JumpStart models only support scopes: {', '.join(constants.SUPPORTED_JUMPSTART_SCOPES)}."
+        raise NotImplementedError(
+            "JumpStart models only support scopes: "
+            f"{', '.join(constants.SUPPORTED_JUMPSTART_SCOPES)}."
         )
 
+    assert model_id is not None
+    assert version is not None
+
     model_specs = accessors.JumpStartModelsAccessor.get_model_specs(
         region=region, model_id=model_id, version=version
     )
 
-    if scope == constants.TRAINING and not model_specs.training_supported:
+    if (
+        scope == constants.JumpStartScriptScope.TRAINING.value
+        and not model_specs.training_supported
+    ):
         raise ValueError(
             f"JumpStart model ID '{model_id}' and version '{version}' " "does not support training."
         )
@@ -199,27 +216,27 @@ def verify_model_region_and_return_specs(
         raise DeprecatedJumpStartModelError(model_id=model_id, version=version)
 
     if (
-        scope == constants.INFERENCE
+        scope == constants.JumpStartScriptScope.INFERENCE.value
         and model_specs.inference_vulnerable
         and not tolerate_vulnerable_model
     ):
         raise VulnerableJumpStartModelError(
             model_id=model_id,
             version=version,
             vulnerabilities=model_specs.inference_vulnerabilities,
-            inference=True,
+            scope=constants.JumpStartScriptScope.INFERENCE,
         )
 
     if (
-        scope == constants.TRAINING
+        scope == constants.JumpStartScriptScope.TRAINING.value
         and model_specs.training_vulnerable
         and not tolerate_vulnerable_model
     ):
         raise VulnerableJumpStartModelError(
             model_id=model_id,
             version=version,
             vulnerabilities=model_specs.training_vulnerabilities,
-            inference=False,
+            scope=constants.JumpStartScriptScope.TRAINING,
         )
 
     return model_specs

src/sagemaker/model_uris.py

@@ -42,9 +42,9 @@ def retrieve(
         model_scope (str): The model type, i.e. what it is used for.
             Valid values: "training" and "inference".
         tolerate_vulnerable_model (bool): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+            not raised). False if these models should raise an exception. (Default: None).
         tolerate_deprecated_model (bool): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+            not raised). False if these models should raise an exception. (Default: None).
     Returns:
         str: the model artifact S3 URI for the corresponding model.
 

src/sagemaker/script_uris.py

@@ -42,9 +42,9 @@ def retrieve(
         script_scope (str): The script type, i.e. what it is used for.
             Valid values: "training" and "inference".
         tolerate_vulnerable_model (bool): True if vulnerable models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+            not raised). False if these models should raise an exception. (Default: None).
         tolerate_deprecated_model (bool): True if deprecated models should be tolerated (exception
-            not thrown). False if these models should throw an exception. (Default: None).
+            not raised). False if these models should raise an exception. (Default: None).
     Returns:
         str: the model script URI for the corresponding model.
 

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

@@ -96,7 +96,7 @@ def test_jumpstart_common_image_uri(
     )
     patched_verify_model_region_and_return_specs.assert_called_once()
 
-    with pytest.raises(ValueError):
+    with pytest.raises(NotImplementedError):
         image_uris.retrieve(
             framework=None,
             region="us-west-2",