fix: Addressing naming comments

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/constants.py

@@ -13,7 +13,7 @@
 """This module stores constants related to SageMaker JumpStart CuratedHub."""
 from __future__ import absolute_import
 
-JUMPSTART_HUB_MODEL_ID_TAG_PREFIX = "@jumpstart-model-id:"
-JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX = "@jumpstart-model-version:"
-FRAMEWORK_TAG_PREFIX = "@framework:"
-TASK_TAG_PREFIX = "@mltask:"
+JUMPSTART_HUB_MODEL_ID_TAG_PREFIX = "@jumpstart-model-id"
+JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX = "@jumpstart-model-version"
+FRAMEWORK_TAG_PREFIX = "@framework"
+TASK_TAG_PREFIX = "@mltask"

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -50,22 +50,22 @@
     create_s3_object_reference_from_uri,
     tag_hub_content,
     get_jumpstart_model_and_version,
-    find_jumpstart_tags_for_hub_content,
+    find_unsupported_flags_for_hub_content_versions,
     summary_list_from_list_api_response,
 )
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentDocument_v2,
     JumpStartModelInfo,
     S3ObjectLocation,
-    CuratedHubTag,
     HubContentSummary,
 )
-
-LIST_HUB_CACHE = None
+from sagemaker.utils import TagsDict
 
 class CuratedHub:
     """Class for creating and managing a curated JumpStart hub"""
 
+    _list_hubs_cache: Dict[str, Any] = None
+
     def __init__(
         self,
         hub_name: str,
@@ -153,25 +153,20 @@ def describe(self) -> DescribeHubResponse:
 
 
     def list_models(self, clear_cache: bool = True, **kwargs) -> List[Dict[str, Any]]:
-        """Lists the models in this Curated Hub
+        """Lists the models in this Curated Hub.
 
-        **kwargs: Passed to invocation of ``Session:list_hub_contents``.
-        """
-        if clear_cache:
-            LIST_HUB_CACHE = None
-        if LIST_HUB_CACHE:
-          return LIST_HUB_CACHE
-        return self._list_models(**kwargs)
-    
-    def _list_models(self, **kwargs) -> Dict[str, Any]:
-        """Lists the models in this Curated Hub
+        This function caches the models in local memory 
 
         **kwargs: Passed to invocation of ``Session:list_hub_contents``.
         """
-        hub_content_summaries = self._sagemaker_session.list_hub_contents(
+        if clear_cache:
+            self._list_hubs_cache = None
+        if self._list_hubs_cache is None:
+          hub_content_summaries = self._sagemaker_session.list_hub_contents(
             hub_name=self.hub_name, hub_content_type=HubContentType.MODEL, **kwargs
-        )
-        return hub_content_summaries
+          )
+          self._list_hubs_cache = hub_content_summaries
+        return self._list_hubs_cache
 
     def describe_model(
         self, model_name: str, model_version: str = "*"
@@ -242,9 +237,29 @@ def _determine_models_to_sync(
         models_to_sync = []
         for model in model_list:
             matched_model = models_in_hub.get(model.model_id)
-            if not matched_model or Version(matched_model.hub_content_version) < Version(model.version):
+
+            # Model does not exist in Hub, sync
+            if not matched_model:
                 models_to_sync.append(model)
 
+            if matched_model:
+                model_version = Version(model.version)
+                hub_model_version = Version(matched_model.hub_content_version)
+
+                # 1. Model version exists in Hub, pass
+                if hub_model_version == model_version:
+                    pass
+
+                # 2. Invalid model version exists in Hub, pass
+                # This will only happen if something goes wrong in our metadata
+                if hub_model_version > model_version:
+                    pass
+
+                # 3. Old model version exists in Hub, update
+                if hub_model_version < model_version:
+                    # Check minSDKVersion against current SDK version, emit log
+                    models_to_sync.append(model)
+
         return models_to_sync
 
     def sync(self, model_list: List[Dict[str, str]]):
@@ -416,9 +431,8 @@ def scan_and_tag_models(self, model_list: List[Dict[str, str]] = None) -> None:
 
         models_to_scan = model_list if model_list else self.list_models()
         js_models_in_hub = [model for model in models_to_scan if get_jumpstart_model_and_version(model) is not None]
-        tags_added: Dict[str, List[CuratedHubTag]] = {}
         for model in js_models_in_hub:
-            tags_to_add: List[CuratedHubTag] = find_jumpstart_tags_for_hub_content(
+            tags_to_add: List[TagsDict] = find_unsupported_flags_for_hub_content_versions(
                 hub_name=self.hub_name,
                 hub_content_name=model.hub_content_name,
                 region=self.region,
@@ -429,9 +443,4 @@ def scan_and_tag_models(self, model_list: List[Dict[str, str]] = None) -> None:
                 tags=tags_to_add,
                 session=self._sagemaker_session
             )
-            tags_added.update({model.hub_content_arn: tags_to_add})
-
-        output_string = "No tags were added!"
-        if len(tags_added) > 0:
-          output_string = f"Added the following tags: {tags_added}"
-        JUMPSTART_LOGGER.info(output_string)
+        JUMPSTART_LOGGER.info("Tagging complete!")

src/sagemaker/jumpstart/curated_hub/types.py

@@ -19,7 +19,7 @@
 
 from sagemaker.jumpstart.types import JumpStartDataHolderType, JumpStartModelSpecs, HubContentType
 
-class CuratedHubTagName(str, Enum):
+class CuratedHubUnsupportedFlag(str, Enum):
     """Enum class for Curated Hub tag names."""
     DEPRECATED_VERSIONS = "deprecated_versions"
     TRAINING_VULNERABLE_VERSIONS = "training_vulnerable_versions"
@@ -38,14 +38,6 @@ class HubContentSummary:
   hub_content_display_name: str = None
   hub_content_description: str = None
   hub_content_search_keywords: List[str] = None
-  
-
-@dataclass
-class CuratedHubTag:
-  """Dataclass to store Curated Hub-specific tags."""
-  key: CuratedHubTagName
-  value: str
-
 
 @dataclass
 class S3ObjectLocation:

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -15,6 +15,7 @@
 import re
 from typing import Optional
 from sagemaker.jumpstart.curated_hub.types import S3ObjectLocation
+from sagemaker.jumpstart.constants import JUMPSTART_LOGGER
 from sagemaker.s3_utils import parse_s3_url
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
@@ -25,8 +26,7 @@
     HubArnExtractedInfo
 )
 from sagemaker.jumpstart.curated_hub.types import (
-    CuratedHubTag,
-    CuratedHubTagName,
+    CuratedHubUnsupportedFlag,
     HubContentSummary,
     JumpStartModelInfo
 )
@@ -40,7 +40,10 @@
     TASK_TAG_PREFIX,
     FRAMEWORK_TAG_PREFIX,
 )
-from uuid import uuid4
+from sagemaker.utils import (
+  format_tags,
+  TagsDict
+)
 
 
 def get_info_from_hub_resource_arn(
@@ -183,13 +186,14 @@ def create_hub_bucket_if_it_does_not_exist(
 
     return bucket_name
 
-def tag_hub_content(hub_content_arn: str, tags: List[CuratedHubTag], session: Session) -> None:
+def tag_hub_content(hub_content_arn: str, tags: List[TagsDict], session: Session) -> None:
     session.add_tags(
       ResourceArn=hub_content_arn,
-      Tags=[tag_to_add_tags_api_call(tag) for tag in tags]
+      Tags=str(tags)
     )
+    JUMPSTART_LOGGER.info(f"Added tags to HubContentArn %s: %s", hub_content_arn, TagsDict)
 
-def find_jumpstart_tags_for_hub_content(hub_name: str, hub_content_name: str, region: str, session: Session) -> List[CuratedHubTag]:
+def find_unsupported_flags_for_hub_content_versions(hub_name: str, hub_content_name: str, region: str, session: Session) -> List[TagsDict]:
     """Finds the JumpStart public hub model for a HubContent and calculates relevant tags.
     
     Since tags are the same for all versions of a HubContent, these tags will map from the key to a list of versions impacted.
@@ -203,33 +207,33 @@ def find_jumpstart_tags_for_hub_content(hub_name: str, hub_content_name: str, re
     )
     hub_content_versions: List[HubContentSummary] = summary_list_from_list_api_response(list_versions_response)
 
-    tag_name_to_versions_map: Dict[CuratedHubTagName, List[str]] = {}
+    unsupported_hub_content_versions_map: Dict[str, List[str]] = {}
     for hub_content_version_summary in hub_content_versions:
         jumpstart_model = get_jumpstart_model_and_version(hub_content_version_summary)
         if jumpstart_model is None:
             continue
-        tag_names_to_add: List[CuratedHubTagName] = find_jumpstart_tags_for_model_version(
+        tag_names_to_add: List[CuratedHubUnsupportedFlag] = find_unsupported_flags_for_model_version(
             model_id=jumpstart_model.model_id,
             version=jumpstart_model.version,
             region=region,
             session=session
         )
 
         for tag_name in tag_names_to_add:
-          if tag_name not in tag_name_to_versions_map:
-              tag_name_to_versions_map[tag_name] = []
-          tag_name_to_versions_map[tag_name].append(hub_content_version_summary.hub_content_version)
+          if tag_name not in unsupported_hub_content_versions_map:
+              unsupported_hub_content_versions_map[tag_name.value] = []
+          unsupported_hub_content_versions_map[tag_name.value].append(hub_content_version_summary.hub_content_version)
 
-    return [CuratedHubTag(tag_name, str(versions)) for (tag_name, versions) in tag_name_to_versions_map.items()]
+    return format_tags(unsupported_hub_content_versions_map)
 
 
-def find_jumpstart_tags_for_model_version(model_id: str, version: str, region: str, session: Session) -> List[CuratedHubTagName]:
+def find_unsupported_flags_for_model_version(model_id: str, version: str, region: str, session: Session) -> List[CuratedHubUnsupportedFlag]:
     """Finds relevant CuratedHubTags for a version of a JumpStart public hub model.
     
     For example, if the public hub model is deprecated, this utility will return a `deprecated` tag.
     Since tags are the same for all versions of a HubContent, these tags will map from the key to a list of versions impacted.
     """
-    tags_to_add: List[CuratedHubTagName] = []
+    flags_to_add: List[CuratedHubUnsupportedFlag] = []
     jumpstart_model_specs = utils.verify_model_region_and_return_specs(
         model_id=model_id, 
         version=version,
@@ -241,13 +245,13 @@ def find_jumpstart_tags_for_model_version(model_id: str, version: str, region: s
     )
 
     if (jumpstart_model_specs.deprecated):
-        tags_to_add.append(CuratedHubTagName.DEPRECATED_VERSIONS)
+        flags_to_add.append(CuratedHubUnsupportedFlag.DEPRECATED_VERSIONS)
     if (jumpstart_model_specs.inference_vulnerable):
-        tags_to_add.append(CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS)
+        flags_to_add.append(CuratedHubUnsupportedFlag.INFERENCE_VULNERABLE_VERSIONS)
     if (jumpstart_model_specs.training_vulnerable):
-        tags_to_add.append(CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS)
+        flags_to_add.append(CuratedHubUnsupportedFlag.TRAINING_VULNERABLE_VERSIONS)
 
-    return tags_to_add
+    return flags_to_add
 
 
 
@@ -272,7 +276,7 @@ def get_jumpstart_model_and_version(hub_content_summary: HubContentSummary) -> O
     if jumpstart_model_id_tag is None or jumpstart_model_version_tag is None:
         return None
     jumpstart_model_id = jumpstart_model_id_tag[len(JUMPSTART_HUB_MODEL_ID_TAG_PREFIX):] # Need to remove the tag_prefix and ":"
-    jumpstart_model_version = jumpstart_model_version_tag[len(JUMPSTART_HUB_MODEL_ID_TAG_PREFIX):]
+    jumpstart_model_version = jumpstart_model_version_tag[len(JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX):]
     return JumpStartModelInfo(model_id=jumpstart_model_id, version=jumpstart_model_version)
 
 def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubContentSummary:
@@ -292,11 +296,3 @@ def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubCo
 def summary_list_from_list_api_response(list_hub_contents_response: Dict[str, Any]) -> List[HubContentSummary]:
     return list(map(summary_from_list_api_response, list_hub_contents_response["HubContentSummaries"]))
 
-def tag_to_add_tags_api_call(tag: CuratedHubTag) -> Dict[str, str]:
-    return {
-          'Key': tag.key,
-          'Value': tag.value
-    }
-
-def generate_unique_hub_content_model_name(model_id: str) -> str:
-    return f"{model_id}-{uuid4()}"

tests/unit/sagemaker/jumpstart/curated_hub/test_curated_hub.py

@@ -177,7 +177,6 @@ def test_sync_kicks_off_parallel_syncs(
 
     hub.sync([model_one, model_two])
 
-    # mock_get_model_specs.assert_called_once()
     mock_sync_public_models.assert_has_calls(
         [
             mock.call(JumpStartModelInfo("mock-model-one-huggingface", "*"), 0),

tests/unit/sagemaker/jumpstart/curated_hub/test_utils.py

@@ -19,8 +19,7 @@
 from sagemaker.jumpstart.curated_hub import utils
 from unittest.mock import patch
 from sagemaker.jumpstart.curated_hub.types import (
-    CuratedHubTag,
-    CuratedHubTagName,
+    CuratedHubUnsupportedFlag,
     HubContentSummary
 )
 from sagemaker.jumpstart.types import JumpStartDataHolderType, JumpStartModelSpecs, HubContentType
@@ -187,7 +186,7 @@ def test_find_tags_for_jumpstart_model_version(mock_spec_util):
     mock_specs.training_vulnerable = True
     mock_spec_util.return_value = mock_specs
 
-    tags = utils.find_jumpstart_tags_for_model_version(
+    tags = utils.find_unsupported_flags_for_model_version(
         model_id="test",
         version="test",
         region="test",
@@ -204,7 +203,7 @@ def test_find_tags_for_jumpstart_model_version(mock_spec_util):
         sagemaker_session=mock_sagemaker_session,
     )
 
-    assert tags == [CuratedHubTagName.DEPRECATED_VERSIONS, CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS, CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS]
+    assert tags == [CuratedHubUnsupportedFlag.DEPRECATED_VERSIONS, CuratedHubUnsupportedFlag.INFERENCE_VULNERABLE_VERSIONS, CuratedHubUnsupportedFlag.TRAINING_VULNERABLE_VERSIONS]
 
 @patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")
 def test_find_tags_for_jumpstart_model_version_some_false(mock_spec_util):
@@ -215,7 +214,7 @@ def test_find_tags_for_jumpstart_model_version_some_false(mock_spec_util):
     mock_specs.training_vulnerable = False
     mock_spec_util.return_value = mock_specs
 
-    tags = utils.find_jumpstart_tags_for_model_version(
+    tags = utils.find_unsupported_flags_for_model_version(
         model_id="test",
         version="test",
         region="test",
@@ -232,7 +231,7 @@ def test_find_tags_for_jumpstart_model_version_some_false(mock_spec_util):
         sagemaker_session=mock_sagemaker_session,
     )
 
-    assert tags == [CuratedHubTagName.DEPRECATED_VERSIONS]
+    assert tags == [CuratedHubUnsupportedFlag.DEPRECATED_VERSIONS]
 
 @patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")
 def test_find_tags_for_jumpstart_model_version_all_false(mock_spec_util):
@@ -243,7 +242,7 @@ def test_find_tags_for_jumpstart_model_version_all_false(mock_spec_util):
     mock_specs.training_vulnerable = False
     mock_spec_util.return_value = mock_specs
 
-    tags = utils.find_jumpstart_tags_for_model_version(
+    tags = utils.find_unsupported_flags_for_model_version(
         model_id="test",
         version="test",
         region="test",
@@ -294,7 +293,7 @@ def test_find_all_tags_for_jumpstart_model_filters_non_jumpstart_models(mock_spe
     mock_specs.training_vulnerable = True
     mock_spec_util.return_value = mock_specs
 
-    tags = utils.find_jumpstart_tags_for_hub_content(
+    tags = utils.find_unsupported_flags_for_hub_content_versions(
         hub_name="test",
         hub_content_name="test",
         region="test",
@@ -308,9 +307,9 @@ def test_find_all_tags_for_jumpstart_model_filters_non_jumpstart_models(mock_spe
     )
 
     assert tags == [
-        CuratedHubTag(key=CuratedHubTagName.DEPRECATED_VERSIONS, value=str(["1.0.0", "2.0.0"])),
-        CuratedHubTag(key=CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS, value=str(["1.0.0", "2.0.0"])),
-        CuratedHubTag(key=CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS, value=str(["1.0.0", "2.0.0"]))
+        {"Key":CuratedHubUnsupportedFlag.DEPRECATED_VERSIONS.value, "Value":str(["1.0.0", "2.0.0"])},
+        {"Key":CuratedHubUnsupportedFlag.INFERENCE_VULNERABLE_VERSIONS.value, "Value":str(["1.0.0", "2.0.0"])},
+        {"Key":CuratedHubUnsupportedFlag.TRAINING_VULNERABLE_VERSIONS.value, "Value":str(["1.0.0", "2.0.0"])}
     ]
 
 @patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")