feat: Enhance jumpstart logging for wildcard version identifiers and old model versions #1273 (#1301)

Author: evakravi

src/sagemaker/jumpstart/cache.py

@@ -26,8 +26,10 @@
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
     JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
     JUMPSTART_DEFAULT_REGION_NAME,
+    JUMPSTART_LOGGER,
     MODEL_ID_LIST_WEB_URL,
 )
+from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
 from sagemaker.jumpstart.parameters import (
     JUMPSTART_DEFAULT_MAX_S3_CACHE_ITEMS,
     JUMPSTART_DEFAULT_MAX_SEMANTIC_VERSION_CACHE_ITEMS,
@@ -172,7 +174,7 @@ def _get_manifest_key_from_model_id_semantic_version(
 
         manifest = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_content
+        )[0].formatted_content
 
         sm_version = utils.get_sagemaker_version()
 
@@ -330,7 +332,7 @@ def _retrieval_function(
         if file_type == JumpStartS3FileType.SPECS:
             formatted_body, _ = self._get_json_file(s3_key, file_type)
             model_specs = JumpStartModelSpecs(formatted_body)
-            utils.emit_logs_based_on_model_specs(model_specs, self.get_region())
+            utils.emit_logs_based_on_model_specs(model_specs, self.get_region(), self._s3_client)
             return JumpStartCachedS3ContentValue(
                 formatted_content=model_specs
             )
@@ -343,7 +345,7 @@ def get_manifest(self) -> List[JumpStartModelHeader]:
 
         manifest_dict = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_content
+        )[0].formatted_content
         manifest = list(manifest_dict.values())  # type: ignore
         return manifest
 
@@ -403,10 +405,11 @@ def _get_header_impl(
 
         versioned_model_id = self._model_id_semantic_version_manifest_key_cache.get(
             JumpStartVersionedModelId(model_id, semantic_version_str)
-        )
+        )[0]
+
         manifest = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.MANIFEST, self._manifest_file_s3_key)
-        ).formatted_content
+        )[0].formatted_content
         try:
             header = manifest[versioned_model_id]  # type: ignore
             return header
@@ -427,10 +430,18 @@ def get_specs(self, model_id: str, semantic_version_str: str) -> JumpStartModelS
 
         header = self.get_header(model_id, semantic_version_str)
         spec_key = header.spec_key
-        specs = self._s3_cache.get(
+        specs, cache_hit = self._s3_cache.get(
             JumpStartCachedS3ContentKey(JumpStartS3FileType.SPECS, spec_key)
-        ).formatted_content
-        return specs  # type: ignore
+        )
+        if not cache_hit and "*" in semantic_version_str:
+            JUMPSTART_LOGGER.warning(
+                get_wildcard_model_version_msg(
+                    header.model_id,
+                    semantic_version_str,
+                    header.version
+                )
+            )
+        return specs.formatted_content
 
     def clear(self) -> None:
         """Clears the model ID/version and s3 cache."""

src/sagemaker/jumpstart/exceptions.py

@@ -30,6 +30,35 @@
 )
 
 
+_MAJOR_VERSION_WARNING_MSG = (
+    "Note that models may have different input/output signatures after a major version upgrade."
+)
+
+
+def get_wildcard_model_version_msg(
+    model_id: str, wildcard_model_version: str, full_model_version: str
+) -> str:
+    """Returns customer-facing message for using a model version with a wildcard character."""
+
+    return (
+        f"Using model '{model_id}' with wildcard version identifier '{wildcard_model_version}'. "
+        f"Please consider pinning to version '{full_model_version}' to "
+        f"ensure stable results. {_MAJOR_VERSION_WARNING_MSG}"
+    )
+
+
+def get_old_model_version_msg(
+    model_id: str, current_model_version: str, latest_model_version: str
+) -> str:
+    """Returns customer-facing message associated with using an old model version."""
+
+    return (
+        f"Using model '{model_id}' with old version '{current_model_version}'. "
+        f"Please consider upgrading to version '{latest_model_version}'"
+        f". {_MAJOR_VERSION_WARNING_MSG}"
+    )
+
+
 class JumpStartHyperparametersError(ValueError):
     """Exception raised for bad hyperparameters of a JumpStart model."""
 

src/sagemaker/jumpstart/utils.py

@@ -16,6 +16,7 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
+import boto3
 from packaging.version import Version
 import sagemaker
 from sagemaker.config.config_schema import (
@@ -31,6 +32,7 @@
 from sagemaker.jumpstart.exceptions import (
     DeprecatedJumpStartModelError,
     VulnerableJumpStartModelError,
+    get_old_model_version_msg,
 )
 from sagemaker.jumpstart.types import (
     JumpStartModelHeader,
@@ -462,7 +464,9 @@ def update_inference_tags_with_jumpstart_training_tags(
     return inference_tags
 
 
-def emit_logs_based_on_model_specs(model_specs: JumpStartModelSpecs, region: str) -> None:
+def emit_logs_based_on_model_specs(
+    model_specs: JumpStartModelSpecs, region: str, s3_client: boto3.client
+) -> None:
     """Emits logs based on model specs and region."""
 
     if model_specs.hosting_eula_key:
@@ -476,6 +480,24 @@ def emit_logs_based_on_model_specs(model_specs: JumpStartModelSpecs, region: str
             model_specs.hosting_eula_key,
         )
 
+    full_version: str = model_specs.version
+
+    models_manifest_list = accessors.JumpStartModelsAccessor._get_manifest(
+        region=region, s3_client=s3_client
+    )
+    max_version_for_model_id: Optional[str] = None
+    for header in models_manifest_list:
+        if header.model_id == model_specs.model_id:
+            if max_version_for_model_id is None or Version(header.version) > Version(
+                max_version_for_model_id
+            ):
+                max_version_for_model_id = header.version
+
+    if full_version != max_version_for_model_id:
+        constants.JUMPSTART_LOGGER.info(
+            get_old_model_version_msg(model_specs.model_id, full_version, max_version_for_model_id)
+        )
+
     if model_specs.deprecated:
         deprecated_message = model_specs.deprecated_message or (
             "Using deprecated JumpStart model "

src/sagemaker/utilities/cache.py

@@ -15,7 +15,7 @@
 
 import datetime
 import collections
-from typing import TypeVar, Generic, Callable, Optional
+from typing import Tuple, TypeVar, Generic, Callable, Optional
 
 KeyType = TypeVar("KeyType")
 ValType = TypeVar("ValType")
@@ -86,23 +86,26 @@ def clear(self) -> None:
         """Deletes all elements from the cache."""
         self._lru_cache.clear()
 
-    def get(self, key: KeyType, data_source_fallback: Optional[bool] = True) -> ValType:
-        """Returns value corresponding to key in cache.
+    def get(
+        self, key: KeyType, data_source_fallback: Optional[bool] = True
+    ) -> Tuple[ValType, bool]:
+        """Returns value corresponding to key in cache and boolean indicating cache hit.
 
         Args:
             key (KeyType): Key in cache to retrieve.
             data_source_fallback (Optional[bool]): True if data should be retrieved if
                 it's stale or not in cache. Default: True.
-            Raises:
-                KeyError: If key is not found in cache or is outdated and
-                ``data_source_fallback`` is False.
+
+        Raises:
+            KeyError: If key is not found in cache or is outdated and
+            ``data_source_fallback`` is False.
         """
         if data_source_fallback:
             if key in self._lru_cache:
-                return self._get_item(key, False)
+                return self._get_item(key, False), True
             self.put(key)
-            return self._get_item(key, False)
-        return self._get_item(key, True)
+            return self._get_item(key, False), False
+        return self._get_item(key, True), True
 
     def put(self, key: KeyType, value: Optional[ValType] = None) -> None:
         """Adds key to cache using ``retrieval_function``.

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -42,6 +42,7 @@
     BASE_MANIFEST,
     BASE_SPEC,
 )
+from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 
 
 @patch.object(JumpStartModelsCache, "_retrieval_function", patched_retrieval_function)
@@ -525,11 +526,14 @@ def test_jumpstart_cache_evaluates_md5_hash(mock_boto3_client):
         )
 
 
+@patch("sagemaker.jumpstart.cache.utils.emit_logs_based_on_model_specs")
 @patch("boto3.client")
-def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
+def test_jumpstart_cache_makes_correct_s3_calls(
+    mock_boto3_client, mock_emit_logs_based_on_model_specs
+):
 
     # test get_header
-    mock_json = json.dumps(
+    mock_manifest_json = json.dumps(
         [
             {
                 "model_id": "pytorch-ic-imagenet-inception-v3-classification-4",
@@ -542,17 +546,17 @@ def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
     )
     mock_boto3_client.return_value.get_object.return_value = {
         "Body": botocore.response.StreamingBody(
-            io.BytesIO(bytes(mock_json, "utf-8")), content_length=len(mock_json)
+            io.BytesIO(bytes(mock_manifest_json, "utf-8")), content_length=len(mock_manifest_json)
         ),
         "ETag": "etag",
     }
 
     mock_boto3_client.return_value.head_object.return_value = {"ETag": "some-hash"}
 
-    bucket_name = "bucket_name"
+    bucket_name = get_jumpstart_content_bucket("us-west-2")
     client_config = botocore.config.Config(signature_version="my_signature_version")
     cache = JumpStartModelsCache(
-        s3_bucket_name=bucket_name, s3_client_config=client_config, region="my_region"
+        s3_bucket_name=bucket_name, s3_client_config=client_config, region="us-west-2"
     )
     cache.get_header(
         model_id="pytorch-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
@@ -563,7 +567,7 @@ def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
     )
     mock_boto3_client.return_value.head_object.assert_not_called()
 
-    mock_boto3_client.assert_called_with("s3", region_name="my_region", config=client_config)
+    mock_boto3_client.assert_called_with("s3", region_name="us-west-2", config=client_config)
 
     # test get_specs. manifest already in cache, so only s3 call will be to get specs.
     mock_json = json.dumps(BASE_SPEC)
@@ -576,9 +580,22 @@ def test_jumpstart_cache_makes_correct_s3_calls(mock_boto3_client):
         ),
         "ETag": "etag",
     }
-    cache.get_specs(
-        model_id="pytorch-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
-    )
+
+    with patch("logging.Logger.warning") as mocked_warning_log:
+        cache.get_specs(
+            model_id="pytorch-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
+        )
+        mocked_warning_log.assert_called_once_with(
+            "Using model 'pytorch-ic-imagenet-inception-v3-classification-4' with wildcard "
+            "version identifier '*'. Please consider pinning to version '2.0.0' to "
+            "ensure stable results. Note that models may have different input/output "
+            "signatures after a major version upgrade."
+        )
+        mocked_warning_log.reset_mock()
+        cache.get_specs(
+            model_id="pytorch-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
+        )
+        mocked_warning_log.assert_not_called()
 
     mock_boto3_client.return_value.get_object.assert_called_with(
         Bucket=bucket_name,
@@ -595,10 +612,18 @@ def test_jumpstart_cache_handles_bad_semantic_version_manifest_key_cache():
     cache.clear = MagicMock()
     cache._model_id_semantic_version_manifest_key_cache = MagicMock()
     cache._model_id_semantic_version_manifest_key_cache.get.side_effect = [
-        JumpStartVersionedModelId(
-            "tensorflow-ic-imagenet-inception-v3-classification-4", "999.0.0"
+        (
+            JumpStartVersionedModelId(
+                "tensorflow-ic-imagenet-inception-v3-classification-4", "999.0.0"
+            ),
+            True,
+        ),
+        (
+            JumpStartVersionedModelId(
+                "tensorflow-ic-imagenet-inception-v3-classification-4", "1.0.0"
+            ),
+            True,
         ),
-        JumpStartVersionedModelId("tensorflow-ic-imagenet-inception-v3-classification-4", "1.0.0"),
     ]
 
     assert JumpStartModelHeader(
@@ -616,11 +641,17 @@ def test_jumpstart_cache_handles_bad_semantic_version_manifest_key_cache():
     cache.clear.reset_mock()
 
     cache._model_id_semantic_version_manifest_key_cache.get.side_effect = [
-        JumpStartVersionedModelId(
-            "tensorflow-ic-imagenet-inception-v3-classification-4", "999.0.0"
+        (
+            JumpStartVersionedModelId(
+                "tensorflow-ic-imagenet-inception-v3-classification-4", "999.0.0"
+            ),
+            True,
         ),
-        JumpStartVersionedModelId(
-            "tensorflow-ic-imagenet-inception-v3-classification-4", "987.0.0"
+        (
+            JumpStartVersionedModelId(
+                "tensorflow-ic-imagenet-inception-v3-classification-4", "987.0.0"
+            ),
+            True,
         ),
     ]
     with pytest.raises(KeyError):
@@ -735,6 +766,7 @@ def test_jumpstart_local_metadata_override_header(
     mocked_get_json_file_and_etag_from_s3.assert_not_called()
 
 
+@patch("sagemaker.jumpstart.cache.utils.emit_logs_based_on_model_specs")
 @patch.object(JumpStartModelsCache, "_get_json_file_and_etag_from_s3")
 @patch("sagemaker.jumpstart.utils.get_sagemaker_version", lambda: "2.68.3")
 @patch.dict(
@@ -747,7 +779,10 @@ def test_jumpstart_local_metadata_override_header(
 @patch("sagemaker.jumpstart.cache.os.path.isdir")
 @patch("builtins.open")
 def test_jumpstart_local_metadata_override_specs(
-    mocked_open: Mock, mocked_is_dir: Mock, mocked_get_json_file_and_etag_from_s3: Mock
+    mocked_open: Mock,
+    mocked_is_dir: Mock,
+    mocked_get_json_file_and_etag_from_s3: Mock,
+    mock_emit_logs_based_on_model_specs,
 ):
 
     mocked_open.side_effect = [
@@ -776,6 +811,7 @@ def test_jumpstart_local_metadata_override_specs(
     mocked_get_json_file_and_etag_from_s3.assert_not_called()
 
 
+@patch("sagemaker.jumpstart.cache.utils.emit_logs_based_on_model_specs")
 @patch.object(JumpStartModelsCache, "_get_json_file_and_etag_from_s3")
 @patch("sagemaker.jumpstart.utils.get_sagemaker_version", lambda: "2.68.3")
 @patch.dict(
@@ -791,6 +827,7 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
     mocked_open: Mock,
     mocked_is_dir: Mock,
     mocked_get_json_file_and_etag_from_s3: Mock,
+    mocked_emit_logs_based_on_model_specs,
 ):
     model_id, version = "tensorflow-ic-imagenet-inception-v3-classification-4", "2.0.0"