feat: jumpstart model id suggestions (#2899)

Co-authored-by: Navin Soni <[email protected]>
Co-authored-by: Mufaddal Rohawala <[email protected]>
Co-authored-by: qidewenwhen <[email protected]>
Co-authored-by: HappyAmazonian <[email protected]>

Author: 5 people

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(

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.

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}."
                 )