Merge branch 'dev' into fix-cond-step

staubhp · web-flow · commit a2456410f08d · 2022-02-21T11:15:46.000-08:00
diff --git a/doc/overview.rst b/doc/overview.rst
@@ -746,6 +746,7 @@ see `Model <https://sagemaker.readthedocs.io/en/stable/api/inference/model.html
 .. code:: python
 
    from sagemaker.model import Model
+   from sagemaker.predictor import Predictor
    from sagemaker.session import Session
 
    # Create the SageMaker model instance
@@ -755,6 +756,7 @@ see `Model <https://sagemaker.readthedocs.io/en/stable/api/inference/model.html
        source_dir=script_uri,
        entry_point="inference.py",
        role=Session().get_caller_identity_arn(),
+       predictor_cls=Predictor,
    )
 
 Save the output from deploying the model to a variable named
@@ -766,12 +768,9 @@ Deployment may take about 5 minutes.
 
 .. code:: python
 
-   from sagemaker.predictor import Predictor
-
    predictor = model.deploy(
        initial_instance_count=instance_count,
        instance_type=instance_type,
-       predictor_cls=Predictor
    )
 
 Because ``catboost`` and ``lightgbm`` rely on the PyTorch Deep Learning Containers
diff --git a/src/sagemaker/jumpstart/cache.py b/src/sagemaker/jumpstart/cache.py
@@ -13,6 +13,7 @@
 """This module defines the JumpStartModelsCache class."""
 from __future__ import absolute_import
 import datetime
+from difflib import get_close_matches
 from typing import List, Optional
 import json
 import boto3
@@ -204,14 +205,34 @@ def _get_manifest_key_from_model_id_semantic_version(
             sm_version_to_use = sm_version_to_use_list[0]
 
             error_msg = (
-                f"Unable to find model manifest for {model_id} with version {version} "
-                f"compatible with your SageMaker version ({sm_version}). "
+                f"Unable to find model manifest for '{model_id}' with version '{version}' "
+                f"compatible with your SageMaker version ('{sm_version}'). "
                 f"Consider upgrading your SageMaker library to at least version "
-                f"{sm_version_to_use} so you can use version "
-                f"{model_version_to_use_incompatible_with_sagemaker} of {model_id}."
+                f"'{sm_version_to_use}' so you can use version "
+                f"'{model_version_to_use_incompatible_with_sagemaker}' of '{model_id}'."
             )
             raise KeyError(error_msg)
-        error_msg = f"Unable to find model manifest for {model_id} with version {version}."
+
+        error_msg = f"Unable to find model manifest for '{model_id}' with version '{version}'. "
+        error_msg += (
+            "Visit https://sagemaker.readthedocs.io/en/stable/doc_utils/jumpstart.html"
+            " for updated list of models. "
+        )
+
+        other_model_id_version = self._select_version(
+            "*", versions_incompatible_with_sagemaker
+        )  # all versions here are incompatible with sagemaker
+        if other_model_id_version is not None:
+            error_msg += (
+                f"Consider using model ID '{model_id}' with version "
+                f"'{other_model_id_version}'."
+            )
+
+        else:
+            possible_model_ids = [header.model_id for header in manifest.values()]
+            closest_model_id = get_close_matches(model_id, possible_model_ids, n=1, cutoff=0)[0]
+            error_msg += f"Did you mean to use model ID '{closest_model_id}'?"
+
         raise KeyError(error_msg)
 
     def _get_file_from_s3(
diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py
@@ -122,3 +122,5 @@
 TRAINING_ENTRY_POINT_SCRIPT_NAME = "transfer_learning.py"
 
 SUPPORTED_JUMPSTART_SCOPES = set(scope.value for scope in JumpStartScriptScope)
+
+ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE = "AWS_JUMPSTART_CONTENT_BUCKET_OVERRIDE"
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -135,12 +135,12 @@ def from_json(self, json_obj: Dict[str, str]) -> None:
 class JumpStartECRSpecs(JumpStartDataHolderType):
     """Data class for JumpStart ECR specs."""
 
-    __slots__ = {
+    __slots__ = [
         "framework",
         "framework_version",
         "py_version",
         "huggingface_transformers_version",
-    }
+    ]
 
     def __init__(self, spec: Dict[str, Any]):
         """Initializes a JumpStartECRSpecs object from its json representation.
@@ -173,7 +173,7 @@ def to_json(self) -> Dict[str, Any]:
 class JumpStartHyperparameter(JumpStartDataHolderType):
     """Data class for JumpStart hyperparameter definition in the training container."""
 
-    __slots__ = {
+    __slots__ = [
         "name",
         "type",
         "options",
@@ -183,7 +183,7 @@ class JumpStartHyperparameter(JumpStartDataHolderType):
         "max",
         "exclusive_min",
         "exclusive_max",
-    }
+    ]
 
     def __init__(self, spec: Dict[str, Any]):
         """Initializes a JumpStartHyperparameter object from its json representation.
@@ -234,12 +234,12 @@ def to_json(self) -> Dict[str, Any]:
 class JumpStartEnvironmentVariable(JumpStartDataHolderType):
     """Data class for JumpStart environment variable definitions in the hosting container."""
 
-    __slots__ = {
+    __slots__ = [
         "name",
         "type",
         "default",
         "scope",
-    }
+    ]
 
     def __init__(self, spec: Dict[str, Any]):
         """Initializes a JumpStartEnvironmentVariable object from its json representation.
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
@@ -13,6 +13,7 @@
 """This module contains utilities related to SageMaker JumpStart."""
 from __future__ import absolute_import
 import logging
+import os
 from typing import Dict, List, Optional
 from urllib.parse import urlparse
 from packaging.version import Version
@@ -60,6 +61,14 @@ def get_jumpstart_content_bucket(region: str) -> str:
     Raises:
         RuntimeError: If JumpStart is not launched in ``region``.
     """
+
+    if (
+        constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE in os.environ
+        and len(os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]) > 0
+    ):
+        bucket_override = os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]
+        LOGGER.info("Using JumpStart bucket override: '%s'", bucket_override)
+        return bucket_override
     try:
         return constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[region].content_bucket
     except KeyError:
diff --git a/src/sagemaker/jumpstart/validators.py b/src/sagemaker/jumpstart/validators.py
@@ -49,7 +49,7 @@ def _validate_hyperparameter(
 
     if len(hyperparameter_spec) > 1:
         raise JumpStartHyperparametersError(
-            f"Unable to perform validation -- found multiple hyperparameter "
+            "Unable to perform validation -- found multiple hyperparameter "
             f"'{hyperparameter_name}' in model specs."
         )
 
@@ -76,35 +76,35 @@ def _validate_hyperparameter(
             if hyperparameter_value not in hyperparameter_spec.options:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must have one of the following "
-                    f"values: {', '.join(hyperparameter_spec.options)}"
+                    f"values: {', '.join(hyperparameter_spec.options)}."
                 )
 
         if hasattr(hyperparameter_spec, "min"):
             if len(hyperparameter_value) < hyperparameter_spec.min:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must have length no less than "
-                    f"{hyperparameter_spec.min}"
+                    f"{hyperparameter_spec.min}."
                 )
 
         if hasattr(hyperparameter_spec, "exclusive_min"):
             if len(hyperparameter_value) <= hyperparameter_spec.exclusive_min:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must have length greater than "
-                    f"{hyperparameter_spec.exclusive_min}"
+                    f"{hyperparameter_spec.exclusive_min}."
                 )
 
         if hasattr(hyperparameter_spec, "max"):
             if len(hyperparameter_value) > hyperparameter_spec.max:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must have length no greater than "
-                    f"{hyperparameter_spec.max}"
+                    f"{hyperparameter_spec.max}."
                 )
 
         if hasattr(hyperparameter_spec, "exclusive_max"):
             if len(hyperparameter_value) >= hyperparameter_spec.exclusive_max:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must have length less than "
-                    f"{hyperparameter_spec.exclusive_max}"
+                    f"{hyperparameter_spec.exclusive_max}."
                 )
 
     # validate numeric types
@@ -125,35 +125,35 @@ def _validate_hyperparameter(
             if not hyperparameter_value_str[start_index:].isdigit():
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must be integer type "
-                    "('{hyperparameter_value}')."
+                    f"('{hyperparameter_value}')."
                 )
 
         if hasattr(hyperparameter_spec, "min"):
             if numeric_hyperparam_value < hyperparameter_spec.min:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' can be no less than "
-                    "{hyperparameter_spec.min}."
+                    f"{hyperparameter_spec.min}."
                 )
 
         if hasattr(hyperparameter_spec, "max"):
             if numeric_hyperparam_value > hyperparameter_spec.max:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' can be no greater than "
-                    "{hyperparameter_spec.max}."
+                    f"{hyperparameter_spec.max}."
                 )
 
         if hasattr(hyperparameter_spec, "exclusive_min"):
             if numeric_hyperparam_value <= hyperparameter_spec.exclusive_min:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must be greater than "
-                    "{hyperparameter_spec.exclusive_min}."
+                    f"{hyperparameter_spec.exclusive_min}."
                 )
 
         if hasattr(hyperparameter_spec, "exclusive_max"):
             if numeric_hyperparam_value >= hyperparameter_spec.exclusive_max:
                 raise JumpStartHyperparametersError(
                     f"Hyperparameter '{hyperparameter_name}' must be less than "
-                    "{hyperparameter_spec.exclusive_max}."
+                    f"{hyperparameter_spec.exclusive_max}."
                 )
 
 
diff --git a/tests/unit/sagemaker/hyperparameters/jumpstart/test_validate.py b/tests/unit/sagemaker/hyperparameters/jumpstart/test_validate.py
diff --git a/tests/unit/sagemaker/jumpstart/test_cache.py b/tests/unit/sagemaker/jumpstart/test_cache.py
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py
