fix: Add tests

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -51,17 +51,14 @@
     tag_hub_content,
     get_jumpstart_model_and_version,
     find_jumpstart_tags_for_hub_content,
-    get_latest_version_for_model,
     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
@@ -215,13 +212,23 @@ def _is_invalid_model_list_input(self, model_list: List[Dict[str, str]]) -> bool
                 return True
         return False
 
+    def _populate_latest_model_version(self, model: Dict[str, str]) -> Dict[str, str]:
+        """Populates the lastest version of a model from specs no matter what is passed.
+
+        Returns model ({ model_id: str, version: str })
+        """
+        model_specs = utils.verify_model_region_and_return_specs(
+            model["model_id"], "*", JumpStartScriptScope.INFERENCE, self.region
+        )
+        return {"model_id": model["model_id"], "version": model_specs.version}
+
     def _get_jumpstart_models_in_hub(self) -> List[HubContentSummary]:
         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[CuratedHubModelInfo]
-    ) -> List[CuratedHubModelInfo]:
+        self, model_list: List[JumpStartModelInfo], models_in_hub: Dict[str, HubContentSummary]
+    ) -> List[JumpStartModelInfo]:
         """Determines which models from `sync` params to sync into the CuratedHub.
 
         Algorithm:
@@ -232,22 +239,11 @@ 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.
         """
-        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="*"
-                ))
+        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):
+                models_to_sync.append(model)
 
         return models_to_sync
 
@@ -265,30 +261,24 @@ def sync(self, model_list: List[Dict[str, str]]):
             )
 
         # Retrieve latest version of unspecified JumpStart model versions
-        model_version_list: List[JumpStartModelInfo] = []
+        model_version_list = []
         for model in model_list:
-            model_id = model.get("model_id")
             version = model.get("version", "*")
             if version == "*":
-                version = get_latest_version_for_model(model_id=model_id, region=self.region)
+                model = self._populate_latest_model_version(model)
                 JUMPSTART_LOGGER.warning(
                     "No version specified for model %s. Using version %s",
-                    model_id,
-                    version,
+                    model["model_id"],
+                    model["version"],
                 )
-            model_version_list.append(JumpStartModelInfo(model_id, version))
-
-        jumpstart_models_in_hub = self._get_jumpstart_models_in_hub()
-        curated_models = [
-            CuratedHubModelInfo(
-                jumpstart_model_info=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)
+            model_version_list.append(JumpStartModelInfo(model["model_id"], model["version"]))
+
+        js_models_in_hub = self._get_jumpstart_models_in_hub()
+        mapped_models_in_hub = {model.hub_content_name: model for model in js_models_in_hub}
+
+        models_to_sync = self._determine_models_to_sync(model_version_list, mapped_models_in_hub)
         JUMPSTART_LOGGER.warning(
-            "Syncing the following models into Hub %s: %s", self.hub_name, [model.jumpstart_model_info for model in models_to_sync]
+            "Syncing the following models into Hub %s: %s", self.hub_name, models_to_sync
         )
 
         # Delete old models?
@@ -322,11 +312,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: CuratedHubModelInfo, thread_num: int):
+    def _sync_public_model_to_hub(self, model: JumpStartModelInfo, 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.jumpstart_model_info.model_id,
-            version=model.jumpstart_model_info.version,
+            model_id=model.model_id,
+            version=model.version,
             region=self.region,
             scope=JumpStartScriptScope.INFERENCE,
             sagemaker_session=self._sagemaker_session,
@@ -335,7 +325,7 @@ def _sync_public_model_to_hub(self, model: CuratedHubModelInfo, thread_num: int)
 
         dest_location = S3ObjectLocation(
             bucket=self.hub_storage_location.bucket,
-            key=f"{self.hub_storage_location.key}/curated_models/{model.jumpstart_model_info.model_id}/{model.jumpstart_model_info.version}",
+            key=f"{self.hub_storage_location.key}/curated_models/{model.model_id}/{model.version}",
         )
         src_files = file_generator.generate_file_infos_from_model_specs(
             model_specs, studio_specs, self.region, self._s3_client
@@ -357,16 +347,16 @@ def _sync_public_model_to_hub(self, model: CuratedHubModelInfo, thread_num: int)
                 label=dest_location.key,
             ).execute()
         else:
-            JUMPSTART_LOGGER.warning("Nothing to copy for %s v%s", model.jumpstart_model_info.model_id, model.jumpstart_model_info.version)
+            JUMPSTART_LOGGER.warning("Nothing to copy for %s v%s", model.model_id, model.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.jumpstart_model_info.model_id}",
-            f"{JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX}:{model.jumpstart_model_info.version}",
+            f"{JUMPSTART_HUB_MODEL_ID_TAG_PREFIX}:{model.model_id}",
+            f"{JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX}:{model.version}",
             f"{FRAMEWORK_TAG_PREFIX}:{model_specs.get_framework()}",
             f"{TASK_TAG_PREFIX}:TODO: pull from specs",
         ]
@@ -375,8 +365,8 @@ def _sync_public_model_to_hub(self, model: CuratedHubModelInfo, thread_num: int)
 
         self._sagemaker_session.import_hub_content(
             document_schema_version=HubContentDocument_v2.SCHEMA_VERSION,
-            hub_content_name=model.hub_content_model_id,
-            hub_content_version=model.hub_content_version,
+            hub_content_name=model.model_id,
+            hub_content_version=model.version,
             hub_name=self.hub_name,
             hub_content_document=hub_content_document,
             hub_content_type=HubContentType.MODEL,

src/sagemaker/jumpstart/curated_hub/types.py

@@ -72,13 +72,6 @@ class JumpStartModelInfo:
     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

@@ -274,14 +274,6 @@ def get_jumpstart_model_and_version(hub_content_summary: HubContentSummary) -> O
     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):]
     return JumpStartModelInfo(model_id=jumpstart_model_id, version=jumpstart_model_version)
-    
-def get_latest_version_for_model(model_id: str, region: str) -> str:
-    """Returns the latest version of a model from specs."""
-    model_specs = utils.verify_model_region_and_return_specs(
-        model_id, "*", JumpStartScriptScope.INFERENCE, region
-    )
-    return model_specs.version
-
 
 def summary_from_list_api_response(hub_content_summary: Dict[str, Any]) -> HubContentSummary:
     return HubContentSummary(