adding tests and update clases

Author: bencrabtree

src/sagemaker/jumpstart/curated_hub/accessors/filegenerator.py

@@ -35,19 +35,28 @@ def __init__(
 
     @singledispatchmethod
     def format(self, file_input) -> List[FileInfo]:
-        """Implement."""
+        """Dispatch method that takes in an input of either ``S3ObjectLocation`` or
+        ``JumpStartModelSpecs`` and is implemented in below registered functions.
+        """
         # pylint: disable=W0107
         pass
 
     @format.register
     def _(self, file_input: S3ObjectLocation) -> List[FileInfo]:
-        """Something."""
+        """Implements ``.format`` when the input is of type ``S3ObjectLocation``.
+
+        Returns a list of ``FileInfo`` objects from the specified bucket location.
+        """
         files = self.s3_format(file_input)
         return files
 
     @format.register
     def _(self, file_input: JumpStartModelSpecs) -> List[FileInfo]:
-        """Something."""
+        """Implements ``.format`` when the input is of type ``JumpStartModelSpecs``.
+
+        Returns a list of ``FileInfo`` objects from dependencies found in the public
+            model specs.
+        """
         files = self.specs_format(file_input, self.studio_specs)
         return files
 
@@ -72,36 +81,16 @@ def s3_format(self, file_input: S3ObjectLocation) -> List[FileInfo]:
     def specs_format(
         self, file_input: JumpStartModelSpecs, studio_specs: Dict[str, Any]
     ) -> List[FileInfo]:
-        """Collects data locations from JumpStart public model specs and
-        converts into FileInfo.
+        """
+        Collects data locations from JumpStart public model specs and
+            converts into FileInfo.
         """
         public_model_data_accessor = PublicModelDataAccessor(
             region=self.region, model_specs=file_input, studio_specs=studio_specs
         )
-        function_table = {
-            HubContentDependencyType.INFERENCE_ARTIFACT: (
-                public_model_data_accessor.get_inference_artifact_s3_reference
-            ),
-            HubContentDependencyType.TRAINING_ARTIFACT: (
-                public_model_data_accessor.get_training_artifact_s3_reference
-            ),
-            HubContentDependencyType.INFERNECE_SCRIPT: (
-                public_model_data_accessor.get_inference_script_s3_reference
-            ),
-            HubContentDependencyType.TRAINING_SCRIPT: (
-                public_model_data_accessor.get_training_script_s3_reference
-            ),
-            HubContentDependencyType.DEFAULT_TRAINING_DATASET: (
-                public_model_data_accessor.get_default_training_dataset_s3_reference
-            ),
-            HubContentDependencyType.DEMO_NOTEBOOK: (
-                public_model_data_accessor.get_demo_notebook_s3_reference
-            ),
-            HubContentDependencyType.MARKDOWN: public_model_data_accessor.get_markdown_s3_reference,
-        }
         files = []
         for dependency in HubContentDependencyType:
-            location = function_table[dependency]()
+            location = public_model_data_accessor.get_s3_reference(dependency)
             parameters = {"Bucket": location.bucket, "Prefix": location.key}
             response = self.s3_client.head_object(**parameters)
             key: str = location.key

src/sagemaker/jumpstart/curated_hub/accessors/fileinfo.py

@@ -21,13 +21,13 @@
 class HubContentDependencyType(str, Enum):
     """Enum class for HubContent dependency names"""
 
-    INFERENCE_ARTIFACT = "INFERENCE_ARTIFACT"
-    TRAINING_ARTIFACT = "TRAINING_ARTIFACT"
-    INFERNECE_SCRIPT = "INFERENCE_SCRIPT"
-    TRAINING_SCRIPT = "TRAINING_SCRIPT"
-    DEFAULT_TRAINING_DATASET = "DEFAULT_TRAINING_DATASET"
-    DEMO_NOTEBOOK = "DEMO_NOTEBOOK"
-    MARKDOWN = "MARKDOWN"
+    INFERENCE_ARTIFACT = "inference_artifact_s3_reference"
+    TRAINING_ARTIFACT = "training_artifact_s3_reference"
+    INFERENCE_SCRIPT = "inference_script_s3_reference"
+    TRAINING_SCRIPT = "training_script_s3_reference"
+    DEFAULT_TRAINING_DATASET = "default_training_dataset_s3_reference"
+    DEMO_NOTEBOOK = "demo_notebook_s3_reference"
+    MARKDOWN = "markdown_s3_reference"
 
 
 @dataclass

src/sagemaker/jumpstart/curated_hub/accessors/public_model_data.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 from typing import Dict, Any
 from sagemaker import model_uris, script_uris
+from sagemaker.jumpstart.curated_hub.accessors.fileinfo import HubContentDependencyType
 from sagemaker.jumpstart.curated_hub.utils import (
     get_model_framework,
 )
@@ -40,53 +41,92 @@ def __init__(
         self.model_specs = model_specs
         self.studio_specs = studio_specs  # Necessary for SDK - Studio metadata drift
 
-    def get_bucket_name(self) -> str:
+    def get_s3_reference(self, dependency_type: HubContentDependencyType):
+        """Retrieves S3 reference given a HubContentDependencyType."""
+        return getattr(self, dependency_type.value)
+
+    @property
+    def inference_artifact_s3_reference(self):
+        """Retrieves s3 reference for model inference artifact"""
+        return self._get_inference_artifact_s3_reference()
+
+    @property
+    def training_artifact_s3_reference(self):
+        """Retrieves s3 reference for model training artifact"""
+        return self._get_training_artifact_s3_reference()
+
+    @property
+    def inference_script_s3_reference(self):
+        """Retrieves s3 reference for model inference script"""
+        return self._get_inference_script_s3_reference()
+
+    @property
+    def training_script_s3_reference(self):
+        """Retrieves s3 reference for model training script"""
+        return self._get_training_script_s3_reference()
+
+    @property
+    def default_training_dataset_s3_reference(self):
+        """Retrieves s3 reference for s3 directory containing model training datasets"""
+        return self._get_default_training_dataset_s3_reference()
+
+    @property
+    def demo_notebook_s3_reference(self):
+        """Retrieves s3 reference for model demo jupyter notebook"""
+        return self._get_demo_notebook_s3_reference()
+
+    @property
+    def markdown_s3_reference(self):
+        """Retrieves s3 reference for model markdown"""
+        return self._get_markdown_s3_reference()
+
+    def _get_bucket_name(self) -> str:
         """Retrieves s3 bucket"""
         return self._bucket
 
-    def get_inference_artifact_s3_reference(self) -> S3ObjectLocation:
+    def _get_inference_artifact_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model inference artifact"""
         return create_s3_object_reference_from_uri(
             self._jumpstart_artifact_s3_uri(JumpStartScriptScope.INFERENCE)
         )
 
-    def get_training_artifact_s3_reference(self) -> S3ObjectLocation:
+    def _get_training_artifact_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model training artifact"""
         return create_s3_object_reference_from_uri(
             self._jumpstart_artifact_s3_uri(JumpStartScriptScope.TRAINING)
         )
 
-    def get_inference_script_s3_reference(self) -> S3ObjectLocation:
+    def _get_inference_script_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model inference script"""
         return create_s3_object_reference_from_uri(
             self._jumpstart_script_s3_uri(JumpStartScriptScope.INFERENCE)
         )
 
-    def get_training_script_s3_reference(self) -> S3ObjectLocation:
+    def _get_training_script_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model training script"""
         return create_s3_object_reference_from_uri(
             self._jumpstart_script_s3_uri(JumpStartScriptScope.TRAINING)
         )
 
-    def get_default_training_dataset_s3_reference(self) -> S3ObjectLocation:
+    def _get_default_training_dataset_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for s3 directory containing model training datasets"""
-        return S3ObjectLocation(self.get_bucket_name(), self._get_training_dataset_prefix())
+        return S3ObjectLocation(self._get_bucket_name(), self.__get_training_dataset_prefix())
 
-    def _get_training_dataset_prefix(self) -> str:
+    def __get_training_dataset_prefix(self) -> str:
         """Retrieves training dataset location"""
         return self.studio_specs["defaultDataKey"]
 
-    def get_demo_notebook_s3_reference(self) -> S3ObjectLocation:
+    def _get_demo_notebook_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model demo jupyter notebook"""
         framework = get_model_framework(self.model_specs)
         key = f"{framework}-notebooks/{self.model_specs.model_id}-inference.ipynb"
-        return S3ObjectLocation(self.get_bucket_name(), key)
+        return S3ObjectLocation(self._get_bucket_name(), key)
 
-    def get_markdown_s3_reference(self) -> S3ObjectLocation:
+    def _get_markdown_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model markdown"""
         framework = get_model_framework(self.model_specs)
         key = f"{framework}-metadata/{self.model_specs.model_id}.md"
-        return S3ObjectLocation(self.get_bucket_name(), key)
+        return S3ObjectLocation(self._get_bucket_name(), key)
 
     def _jumpstart_script_s3_uri(self, model_scope: str) -> str:
         """Retrieves JumpStart script s3 location"""

src/sagemaker/jumpstart/curated_hub/accessors/sync.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""This module provides a class to help copy HubContent dependencies."""
+"""This module provides a class that perfrms functionalities similar to ``S3:Copy``."""
 from __future__ import absolute_import
 from typing import Generator, List
 

src/sagemaker/jumpstart/curated_hub/accessors/synccomparator.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module provides comparators for syncing s3 files."""
 from __future__ import absolute_import
+from datetime import timedelta
 
 from sagemaker.jumpstart.curated_hub.accessors.fileinfo import FileInfo
 
@@ -36,14 +37,6 @@ def determine_should_sync(self, src_file: FileInfo, dest_file: FileInfo) -> bool
             )
         return should_sync
 
-    def total_seconds(self, td):
-        """
-        timedelta's time_seconds() function for python 2.6 users
-
-        :param td: The difference between two datetime objects.
-        """
-        return (td.microseconds + (td.seconds + td.days * 24 * 3600) * 10**6) / 10**6
-
     def compare_size(self, src_file: FileInfo, dest_file: FileInfo):
         """
         :returns: True if the sizes are the same.
@@ -62,7 +55,7 @@ def compare_time(self, src_file: FileInfo, dest_file: FileInfo):
         dest_time = dest_file.last_updated
         delta = dest_time - src_time
         # pylint: disable=R1703,R1705
-        if self.total_seconds(delta) >= 0:
+        if timedelta.total_seconds(delta) >= 0:
             # Destination is newer than source.
             return True
         else:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -18,6 +18,7 @@
 
 import boto3
 from botocore.client import BaseClient
+from packaging.version import Version
 
 from sagemaker.jumpstart.curated_hub.accessors.filegenerator import FileGenerator
 from sagemaker.jumpstart.curated_hub.accessors.objectlocation import S3ObjectLocation
@@ -176,53 +177,69 @@ def sync(self, model_list: List[Dict[str, str]]):
         # Retrieve latest version of unspecified JumpStart model versions
         model_version_list = []
         for model in model_list:
-            # TODO: Uncomment and implement
-            # if not model["version"] or model["version"] == "*":
-            #     model["version"] = self._find_latest_version(model_name=model.model_id)
+            version = model.get("version", "*")
+            if not version or version == "*":
+                model_specs = verify_model_region_and_return_specs(
+                    model["model_id"], version, JumpStartScriptScope.INFERENCE, self.region
+                )
+                model["version"] = model_specs.version
             model_version_list.append(model)
 
         # Find synced JumpStart model versions in the Hub
         js_models_in_hub = []
-        for model in hub_models:
+        for hub_model in hub_models:
             # TODO: extract both in one pass
             jumpstart_model_id = next(
-                (tag for tag in model.search_keywords if tag.startswith("@jumpstart-model-id")),
+                (
+                    tag
+                    for tag in hub_model["search_keywords"]
+                    if tag.startswith("@jumpstart-model-id")
+                ),
                 None,
             )
             jumpstart_model_version = next(
                 (
                     tag
-                    for tag in model.search_keywords
+                    for tag in hub_model["search_keywords"]
                     if tag.startswith("@jumpstart-model-version")
                 ),
                 None,
             )
 
             if jumpstart_model_id and jumpstart_model_version:
-                js_models_in_hub.append(model)
+                js_models_in_hub.append(hub_model)
 
         # Match inputted list of model versions with synced JumpStart model versions in the Hub
         models_to_sync = []
         for model in model_version_list:
-            model_id, version = model
-            matched_model = next((model for model in js_models_in_hub if model.name == model_id))
+            matched_model = next(
+                (
+                    hub_model
+                    for hub_model in js_models_in_hub
+                    if hub_model and hub_model["name"] == model["model_id"]
+                ),
+                None,
+            )
 
             # 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["version"])
+
                 # 1. Model version exists in Hub, pass
-                if matched_model.version == version:
+                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 matched_model.version > version:
+                if hub_model_version > model_version:
                     pass
 
                 # 3. Old model version exists in Hub, update
-                if matched_model.version < version:
+                if hub_model_version < model_version:
                     # Check minSDKVersion against current SDK version, emit log
                     models_to_sync.append(model)
 
@@ -269,6 +286,7 @@ def _sync_public_model_to_hub(self, model: Dict[str, str]):
             scope=JumpStartScriptScope.INFERENCE,
             sagemaker_session=self._sagemaker_session,
         )
+
         # TODO: Uncomment and implement
         # studio_specs = self.fetch_studio_specs(model_id=model_name, version=model_version)
         studio_specs = {}
@@ -282,12 +300,17 @@ def _sync_public_model_to_hub(self, model: Dict[str, str]):
         src_files = file_generator.format(model_specs)
         dest_files = file_generator.format(dest_location)
 
-        files_to_copy = list(FileSync(src_files, dest_files, dest_location).call())
+        files_to_copy = FileSync(src_files, dest_files, dest_location).call()
 
         if len(files_to_copy) > 0:
-            # Copy files with MPU
+            # TODO: Copy files with MPU
             print("hi")
 
+        # 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 = []
+
         hub_content_document = HubContentDocument_v2(spec=model_specs)
 
         self._sagemaker_session.import_hub_content(
@@ -301,4 +324,5 @@ def _sync_public_model_to_hub(self, model: Dict[str, str]):
             hub_content_description="",
             hub_content_markdown="",
             hub_content_search_keywords=[],
+            tags=tags,
         )