fix: Adding more alterations

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -52,16 +52,19 @@
     get_jumpstart_model_and_version,
     find_jumpstart_tags_for_hub_content,
     get_latest_version_for_model,
-    summary_list_from_list_api_response
+    summary_list_from_list_api_response,
+    generate_unique_hub_content_model_name
 )
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentDocument_v2,
     JumpStartModelInfo,
     S3ObjectLocation,
     CuratedHubTag,
     HubContentSummary,
+    CuratedHubModelInfo,
 )
 
+LIST_HUB_CACHE = None
 
 class CuratedHub:
     """Class for creating and managing a curated JumpStart hub"""
@@ -152,16 +155,17 @@ def describe(self) -> DescribeHubResponse:
         return hub_description
 
 
-    def list_models(self, clear_cache: bool = True, **kwargs) -> Dict[str, Any]:
+    def list_models(self, clear_cache: bool = True, **kwargs) -> List[Dict[str, Any]]:
         """Lists the models in this Curated Hub
 
         **kwargs: Passed to invocation of ``Session:list_hub_contents``.
         """
         if clear_cache:
-            self._list_models.cache_clear()
+            LIST_HUB_CACHE = None
+        if LIST_HUB_CACHE:
+          return LIST_HUB_CACHE
         return self._list_models(**kwargs)
 
-    @lru_cache(maxsize=5)
     def _list_models(self, **kwargs) -> Dict[str, Any]:
         """Lists the models in this Curated Hub
 
@@ -212,13 +216,12 @@ def _is_invalid_model_list_input(self, model_list: List[Dict[str, str]]) -> bool
         return False
 
     def _get_jumpstart_models_in_hub(self) -> List[HubContentSummary]:
-        """Returns list of `HubContent` that have been created from a JumpStart model."""
-        hub_models: List[HubContentSummary] = summary_list_from_list_api_response(self.list_models())
+        hub_models = summary_list_from_list_api_response(self.list_models())
         return [model for model in hub_models if get_jumpstart_model_and_version(model) is not None]
 
     def _determine_models_to_sync(
-        self, model_to_sync: List[JumpStartModelInfo], models_in_hub: List[JumpStartModelInfo]
-    ) -> List[JumpStartModelInfo]:
+        self, model_to_sync: List[JumpStartModelInfo], models_in_hub: List[CuratedHubModelInfo]
+    ) -> List[CuratedHubModelInfo]:
         """Determines which models from `sync` params to sync into the CuratedHub.
 
         Algorithm:
@@ -229,12 +232,22 @@ def _determine_models_to_sync(
             in Hub, don't sync. If newer version in Hub, don't sync. If older version in Hub,
             sync that model.
         """
-        jumpstart_model_id_to_model_info_map = {model.model_id: model for model in models_in_hub}
-        models_to_sync: List[JumpStartModelInfo] = []
-        for model in model_to_sync:
-            matched_model = jumpstart_model_id_to_model_info_map.get(model.model_id)
-            if not matched_model or Version(matched_model.version) < Version(model.version):
-                models_to_sync.append(model)
+        hub_content_jumpstart_model_id_map = {model.jumpstart_model_info.model_id: model for model in models_in_hub}
+        models_to_sync: List[CuratedHubModelInfo] = []
+        for public_hub_model in model_to_sync:
+            matched_model = hub_content_jumpstart_model_id_map.get(public_hub_model.model_id)
+            if not matched_model:
+                models_to_sync.append(CuratedHubModelInfo(
+                    jumpstart_model_info=public_hub_model,
+                    hub_content_model_id=generate_unique_hub_content_model_name(public_hub_model.model_id),
+                    hub_content_version="*"
+                ))
+            elif Version(matched_model.jumpstart_model_info.version) < Version(public_hub_model.version):
+                models_to_sync.append(CuratedHubModelInfo(
+                    jumpstart_model_info=public_hub_model,
+                    hub_content_model_id=matched_model.hub_content_model_id,
+                    hub_content_version="*"
+                ))
 
         return models_to_sync
 
@@ -265,12 +278,17 @@ def sync(self, model_list: List[Dict[str, str]]):
                 )
             model_version_list.append(JumpStartModelInfo(model_id, version))
 
-        js_models_in_hub = self._get_jumpstart_models_in_hub()
-        js_models_in_hub = [get_jumpstart_model_and_version(model) for model in js_models_in_hub]
-
-        models_to_sync = self._determine_models_to_sync(model_version_list, js_models_in_hub)
+        jumpstart_models_in_hub = self._get_jumpstart_models_in_hub()
+        curated_models = [
+            CuratedHubModelInfo(
+                jumpstart_model_id=get_jumpstart_model_and_version(model),
+                hub_content_model_id=model.hub_content_name,
+                hub_content_version=model.hub_content_version
+            ) for model in jumpstart_models_in_hub
+        ]
+        models_to_sync = self._determine_models_to_sync(model_version_list, curated_models)
         JUMPSTART_LOGGER.warning(
-            "Syncing the following models into Hub %s: %s", self.hub_name, models_to_sync
+            "Syncing the following models into Hub %s: %s", self.hub_name, [model.jumpstart_model_info for model in models_to_sync]
         )
 
         # Delete old models?
@@ -304,11 +322,11 @@ def sync(self, model_list: List[Dict[str, str]]):
                 f"Failures when importing models to curated hub in parallel: {failed_imports}"
             )
 
-    def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
+    def _sync_public_model_to_hub(self, model: CuratedHubModelInfo, thread_num: int):
         """Syncs a public JumpStart model version to the Hub. Runs in parallel."""
         model_specs = utils.verify_model_region_and_return_specs(
-            model_id=model.model_id,
-            version=model.version,
+            model_id=model.jumpstart_model_info.model_id,
+            version=model.jumpstart_model_info.version,
             region=self.region,
             scope=JumpStartScriptScope.INFERENCE,
             sagemaker_session=self._sagemaker_session,
@@ -317,7 +335,7 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
 
         dest_location = S3ObjectLocation(
             bucket=self.hub_storage_location.bucket,
-            key=f"{self.hub_storage_location.key}/curated_models/{model.model_id}/{model.version}",
+            key=f"{self.hub_storage_location.key}/curated_models/{model.jumpstart_model_info.model_id}/{model.jumpstart_model_info.version}",
         )
         src_files = file_generator.generate_file_infos_from_model_specs(
             model_specs, studio_specs, self.region, self._s3_client
@@ -339,16 +357,16 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
                 label=dest_location.key,
             ).execute()
         else:
-            JUMPSTART_LOGGER.warning("Nothing to copy for %s v%s", model.model_id, model.version)
+            JUMPSTART_LOGGER.warning("Nothing to copy for %s v%s", model.jumpstart_model_info.model_id, model.jumpstart_model_info.version)
 
         # TODO: Tag model if specs say it is deprecated or training/inference
         # vulnerable. Update tag of HubContent ARN without version.
         # Versioned ARNs are not onboarded to Tagris.
         tags = []
 
         search_keywords = [
-            f"{JUMPSTART_HUB_MODEL_ID_TAG_PREFIX}:{model.model_id}",
-            f"{JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX}:{model.version}",
+            f"{JUMPSTART_HUB_MODEL_ID_TAG_PREFIX}:{model.jumpstart_model_info.model_id}",
+            f"{JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX}:{model.jumpstart_model_info.version}",
             f"{FRAMEWORK_TAG_PREFIX}:{model_specs.get_framework()}",
             f"{TASK_TAG_PREFIX}:TODO: pull from specs",
         ]
@@ -357,8 +375,8 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
 
         self._sagemaker_session.import_hub_content(
             document_schema_version=HubContentDocument_v2.SCHEMA_VERSION,
-            hub_content_name=model.model_id,
-            hub_content_version=model.version,
+            hub_content_name=model.hub_content_model_id,
+            hub_content_version=model.hub_content_version,
             hub_name=self.hub_name,
             hub_content_document=hub_content_document,
             hub_content_type=HubContentType.MODEL,
@@ -398,9 +416,9 @@ def scan_and_tag_models(self) -> None:
         JUMPSTART_LOGGER.info(
             "Tagging models in hub: %s", self.hub_name
         )
-        models_in_hub: List[HubContentSummary] = self._get_jumpstart_models_in_hub()
+        js_models_in_hub = [model for model in self.list_models() if get_jumpstart_model_and_version(model) is not None]
         tags_added: Dict[str, List[CuratedHubTag]] = {}
-        for model in models_in_hub:
+        for model in js_models_in_hub:
             tags_to_add: List[CuratedHubTag] = find_jumpstart_tags_for_hub_content(
                 hub_name=self.hub_name,
                 hub_content_name=model.hub_content_name,

src/sagemaker/jumpstart/curated_hub/types.py

@@ -20,26 +20,29 @@
 from sagemaker.jumpstart.types import JumpStartDataHolderType, JumpStartModelSpecs, HubContentType
 
 class CuratedHubTagName(str, Enum):
-    """Enum class for Curated Hub """
+    """Enum class for Curated Hub tag names."""
     DEPRECATED_VERSIONS = "deprecated_versions"
     TRAINING_VULNERABLE_VERSIONS = "training_vulnerable_versions"
     INFERENCE_VULNERABLE_VERSIONS = "inference_vulnerable_versions"
 
 @dataclass
 class HubContentSummary:
+  """Dataclass to store HubContentSummary from List APIs."""
   hub_content_arn: str
   hub_content_name: str
   hub_content_version: str
   hub_content_type: HubContentType
   document_schema_version: str
   hub_content_status: str
-  hub_content_display_name: str
-  hub_content_description: str
-  hub_content_search_keywords: List[str]
   creation_time: str
+  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
 
@@ -62,14 +65,20 @@ def get_uri(self) -> str:
         """Returns the s3 URI"""
         return f"s3://{self.bucket}/{self.key}"
 
-
 @dataclass
 class JumpStartModelInfo:
     """Helper class for storing JumpStart model info."""
 
     model_id: str
     version: str
 
+@dataclass
+class CuratedHubModelInfo:
+    """Helper class to store Curated Hub model info."""
+    jumpstart_model_info: JumpStartModelInfo
+    hub_content_model_id: str
+    hub_content_version: str
+
 
 class HubContentDependencyType(str, Enum):
     """Enum class for HubContent dependency names"""

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -40,6 +40,7 @@
     TASK_TAG_PREFIX,
     FRAMEWORK_TAG_PREFIX,
 )
+from uuid import uuid4
 
 
 def get_info_from_hub_resource_arn(
@@ -253,16 +254,16 @@ def find_jumpstart_tags_for_model_version(model_id: str, version: str, region: s
 def get_jumpstart_model_and_version(hub_content_summary: HubContentSummary) -> Optional[JumpStartModelInfo]:
     jumpstart_model_id = next(
         (
-            tag
-            for tag in hub_content_summary.search_keywords
+            tag[len(JUMPSTART_HUB_MODEL_ID_TAG_PREFIX)+1:] # Need to remove the tag_prefix and ":"
+            for tag in hub_content_summary.hub_content_search_keywords
             if tag.startswith(JUMPSTART_HUB_MODEL_ID_TAG_PREFIX)
         ),
         None,
     )
     jumpstart_model_version = next(
         (
-            tag
-            for tag in hub_content_summary.search_keywords
+            tag[len(JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX)+1:]
+            for tag in hub_content_summary.hub_content_search_keywords
             if tag.startswith(JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX)
         ),
         None,
@@ -282,7 +283,7 @@ def get_latest_version_for_model(model_id: str, region: str) -> str:
 
 def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubContentSummary:
     return HubContentSummary(
-        hub_content_arn=hub_content_summary.get("HubContentSummary"),
+        hub_content_arn=hub_content_summary.get("HubContentArn"),
         hub_content_name=hub_content_summary.get("HubContentName"),
         hub_content_version=hub_content_summary.get("HubContentVersion"),
         hub_content_type=hub_content_summary.get("HubContentType"),
@@ -301,4 +302,7 @@ 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

@@ -187,48 +187,48 @@ def test_sync_kicks_off_parallel_syncs(
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 @patch(f"{MODULE_PATH}._sync_public_model_to_hub")
-@patch(f"{MODULE_PATH}.list_models")
+@patch(f"{MODULE_PATH}._list_models")
 def test_sync_filters_models_that_exist_in_hub(
     mock_list_models, mock_sync_public_models, mock_get_model_specs, sagemaker_session
 ):
     mock_get_model_specs.side_effect = get_spec_from_base_spec
     mock_list_models.return_value = {
         "HubContentSummaries": [
             {
-                "name": "mock-model-two-pytorch",
-                "version": "1.0.2",
-                "search_keywords": [
-                    "@jumpstart-model-id:model-two-pytorch",
+                "HubContentName": "mock-model-two-pytorch",
+                "HubContentVersion": "1.0.2",
+                "HubContentSearchKeywords": [
+                    "@jumpstart-model-id:mock-pytorch-model-already-exists-in-hub",
                     "@jumpstart-model-version:1.0.2",
                 ],
             },
-            {"name": "mock-model-three-nonsense", "version": "1.0.2", "search_keywords": []},
+            {"HubContentName": "mock-model-three-nonsense", "HubContentVersion": "1.0.2", "HubContentSearchKeywords": []},
             {
-                "name": "mock-model-four-huggingface",
-                "version": "2.0.2",
-                "search_keywords": [
+                "HubContentName": "mock-model-four-huggingface",
+                "HubContentVersion": "2.0.2",
+                "HubContentSearchKeywords": [
                     "@jumpstart-model-id:model-four-huggingface",
                     "@jumpstart-model-version:2.0.2",
                 ],
             },
         ]
     }
     hub_name = "mock_hub_name"
-    model_one = {"model_id": "mock-model-one-huggingface"}
-    model_two = {"model_id": "mock-model-two-pytorch", "version": "1.0.2"}
+    model_one = {"model_id": "mock-pytorch-model-does-not-exist"}
+    model_two = {"model_id": "mock-pytorch-model-already-exists-in-hub", "version": "1.0.2"}
     mock_sync_public_models.return_value = ""
     hub = CuratedHub(hub_name=hub_name, sagemaker_session=sagemaker_session)
 
     hub.sync([model_one, model_two])
 
     mock_sync_public_models.assert_called_once_with(
-        JumpStartModelInfo("mock-model-one-huggingface", "*"), 0
+        JumpStartModelInfo("mock-model-does-not-exist", "*"), 0
     )
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 @patch(f"{MODULE_PATH}._sync_public_model_to_hub")
-@patch(f"{MODULE_PATH}.list_models")
+@patch(f"{MODULE_PATH}._list_models")
 def test_sync_updates_old_models_in_hub(
     mock_list_models, mock_sync_public_models, mock_get_model_specs, sagemaker_session
 ):