optimize function and add hub bucket prefix

Author: bencrabtree

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

@@ -35,16 +35,16 @@ def generate_file_infos_from_s3_location(
     """
     parameters = {"Bucket": location.bucket, "Prefix": location.key}
     response = s3_client.list_objects_v2(**parameters)
-    contents = response.get("Contents", None)
+    contents = response.get("Contents")
 
     if not contents:
         return []
 
     files = []
     for s3_obj in contents:
-        key: str = s3_obj.get("Key")
-        size: bytes = s3_obj.get("Size", None)
-        last_modified: str = s3_obj.get("LastModified", None)
+        key = s3_obj.get("Key")
+        size = s3_obj.get("Size")
+        last_modified = s3_obj.get("LastModified")
         files.append(FileInfo(location.bucket, key, size, last_modified))
     return files
 
@@ -71,28 +71,26 @@ def generate_file_infos_from_model_specs(
         if location_type == "prefix":
             parameters = {"Bucket": location.bucket, "Prefix": location.key}
             response = s3_client.list_objects_v2(**parameters)
-            contents = response.get("Contents", None)
+            contents = response.get("Contents")
             for s3_obj in contents:
-                key: str = s3_obj.get("Key")
-                size: bytes = s3_obj.get("Size", None)
-                last_modified: datetime = s3_obj.get("LastModified", None)
-                dependency_type: HubContentDependencyType = dependency
+                key = s3_obj.get("Key")
+                size = s3_obj.get("Size")
+                last_modified = s3_obj.get("LastModified")
                 files.append(
                     FileInfo(
                         location.bucket,
                         key,
                         size,
                         last_modified,
-                        dependency_type,
+                        dependency,
                     )
                 )
         elif location_type == "object":
             parameters = {"Bucket": location.bucket, "Key": location.key}
             response = s3_client.head_object(**parameters)
-            size: bytes = response.get("ContentLength", None)
-            last_updated: datetime = response.get("LastModified", None)
-            dependency_type: HubContentDependencyType = dependency
+            size = response.get("ContentLength")
+            last_updated = response.get("LastModified")
             files.append(
-                FileInfo(location.bucket, location.key, size, last_updated, dependency_type)
+                FileInfo(location.bucket, location.key, size, last_updated, dependency)
             )
     return files

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

@@ -62,6 +62,7 @@ class MultiPartCopyHandler(object):
     """Multi Part Copy Handler class."""
 
     WORKERS = 20
+    # Config values from in S3:Copy
     MULTIPART_CONFIG = 8 * (1024**2)
 
     def __init__(

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -13,6 +13,7 @@
 """This module provides the JumpStart Curated Hub class."""
 from __future__ import absolute_import
 from concurrent import futures
+from datetime import datetime
 import json
 import traceback
 from typing import Optional, Dict, List, Any
@@ -50,6 +51,7 @@
     HubContentDocument_v2,
     JumpStartModelInfo,
     S3ObjectLocation,
+    create_s3_object_reference_from_uri,
 )
 
 
@@ -73,20 +75,21 @@ def __init__(
         self.region = sagemaker_session.boto_region_name
         self._sagemaker_session = sagemaker_session
         self._default_thread_pool_size = 20
-        self.hub_bucket_name = bucket_name or self._fetch_hub_bucket_name()
         self._s3_client = self._get_s3_client()
+        self.hub_storage_location = self._generate_hub_storage_location(bucket_name)
 
     def _get_s3_client(self) -> BaseClient:
-        """Returns an S3 client."""
+        """Returns an S3 client used for creating a HubContentDocument."""
         return boto3.client("s3", region_name=self.region)
 
     def _fetch_hub_bucket_name(self) -> str:
         """Retrieves hub bucket name from Hub config if exists"""
         try:
             hub_response = self._sagemaker_session.describe_hub(hub_name=self.hub_name)
-            hub_bucket_prefix = hub_response["S3StorageConfig"].get("S3OutputPath", None)
-            if hub_bucket_prefix:
-                return hub_bucket_prefix.replace("s3://", "")
+            hub_output_location = hub_response["S3StorageConfig"].get("S3OutputPath")
+            if hub_output_location:
+                location = create_s3_object_reference_from_uri(hub_output_location)
+                return location.bucket
             default_bucket_name = generate_default_hub_bucket_name(self._sagemaker_session)
             JUMPSTART_LOGGER.warning(
                 "There is not a Hub bucket associated with %s. Using %s",
@@ -103,6 +106,12 @@ def _fetch_hub_bucket_name(self) -> str:
             )
             return hub_bucket_name
 
+    def _generate_hub_storage_location(self, bucket_name: Optional[str] = None) -> None:
+        """Generates an ``S3ObjectLocation`` given a Hub name."""
+        hub_bucket_name = bucket_name or self._fetch_hub_bucket_name()
+        curr_timestamp = datetime.now().timestamp()
+        return S3ObjectLocation(bucket=hub_bucket_name, key=f"{self.hub_name}-{curr_timestamp}")
+
     def create(
         self,
         description: str,
@@ -112,16 +121,16 @@ def create(
     ) -> Dict[str, str]:
         """Creates a hub with the given description"""
 
-        bucket_name = create_hub_bucket_if_it_does_not_exist(
-            self.hub_bucket_name, self._sagemaker_session
+        create_hub_bucket_if_it_does_not_exist(
+            self.hub_storage_location.bucket, self._sagemaker_session
         )
 
         return self._sagemaker_session.create_hub(
             hub_name=self.hub_name,
             hub_description=description,
             hub_display_name=display_name,
             hub_search_keywords=search_keywords,
-            s3_storage_config={"S3OutputPath": f"s3://{bucket_name}"},
+            s3_storage_config={"S3OutputPath": self.hub_storage_location.get_uri()},
             tags=tags,
         )
 
@@ -226,7 +235,7 @@ def _get_jumpstart_models_in_hub(self) -> List[Dict[str, Any]]:
         return js_models_in_hub
 
     def _determine_models_to_sync(
-        self, model_list: List[JumpStartModelInfo], models_in_hub
+        self, model_list: List[JumpStartModelInfo], models_in_hub: Dict[str, Any]
     ) -> List[JumpStartModelInfo]:
         """Determines which models from `sync` params to sync into the CuratedHub.
 
@@ -240,14 +249,7 @@ def _determine_models_to_sync(
         """
         models_to_sync = []
         for model in model_list:
-            matched_model = next(
-                (
-                    hub_model
-                    for hub_model in models_in_hub
-                    if hub_model and hub_model["name"] == model.model_id
-                ),
-                None,
-            )
+            matched_model = models_in_hub.get(model.model_id)
 
             # Model does not exist in Hub, sync
             if not matched_model:
@@ -300,8 +302,9 @@ def sync(self, model_list: List[Dict[str, str]]):
             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["name"]: model for model in js_models_in_hub }
 
-        models_to_sync = self._determine_models_to_sync(model_version_list, 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, models_to_sync
         )
@@ -349,7 +352,8 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
         studio_specs = self._fetch_studio_specs(model_specs=model_specs)
 
         dest_location = S3ObjectLocation(
-            bucket=self.hub_bucket_name, key=f"{model.model_id}/{model.version}"
+            bucket=self.hub_storage_location.bucket,
+            key=f"{self.hub_storage_location.key}/{model.model_id}/{model.version}"
         )
         src_files = file_generator.generate_file_infos_from_model_specs(
             model_specs, studio_specs, self.region, self._s3_client

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

@@ -12,22 +12,25 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 from copy import deepcopy
+import datetime
 from unittest import mock
 from unittest.mock import patch
 import pytest
 from mock import Mock
 from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.jumpstart.curated_hub.types import JumpStartModelInfo
+from sagemaker.jumpstart.curated_hub.types import JumpStartModelInfo, S3ObjectLocation
 from sagemaker.jumpstart.types import JumpStartModelSpecs
 from tests.unit.sagemaker.jumpstart.constants import BASE_SPEC
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
 
+
 REGION = "us-east-1"
 ACCOUNT_ID = "123456789123"
 HUB_NAME = "mock-hub-name"
 
 MODULE_PATH = "sagemaker.jumpstart.curated_hub.curated_hub.CuratedHub"
 
+FAKE_TIME = datetime.datetime(1997, 8, 14, 00, 00, 00)
 
 @pytest.fixture()
 def sagemaker_session():
@@ -39,7 +42,7 @@ def sagemaker_session():
         "Boto3/1.9.69 Python/3.6.5 Linux/4.14.77-70.82.amzn1.x86_64 Botocore/1.12.69 Resource"
     )
     sagemaker_session_mock.describe_hub.return_value = {
-        "S3StorageConfig": {"S3OutputPath": "mock-bucket-123"}
+        "S3StorageConfig": {"S3OutputPath": "s3://mock-bucket-123"}
     }
     sagemaker_session_mock.account_id.return_value = ACCOUNT_ID
     return sagemaker_session_mock
@@ -66,7 +69,9 @@ def test_instantiates(sagemaker_session):
         ),
     ],
 )
+@patch("sagemaker.jumpstart.curated_hub.curated_hub.CuratedHub._generate_hub_storage_location")
 def test_create_with_no_bucket_name(
+    mock_generate_hub_storage_location,
     sagemaker_session,
     hub_name,
     hub_description,
@@ -75,18 +80,22 @@ def test_create_with_no_bucket_name(
     hub_search_keywords,
     tags,
 ):
+    storage_location = S3ObjectLocation("sagemaker-hubs-us-east-1-123456789123",f"{hub_name}-{FAKE_TIME.timestamp()}")
+    mock_generate_hub_storage_location.return_value = storage_location
     create_hub = {"HubArn": f"arn:aws:sagemaker:us-east-1:123456789123:hub/{hub_name}"}
     sagemaker_session.create_hub = Mock(return_value=create_hub)
     sagemaker_session.describe_hub.return_value = {
-        "S3StorageConfig": {"S3OutputPath": hub_bucket_name}
+        "S3StorageConfig": {"S3OutputPath": f"s3://{hub_bucket_name}/{storage_location.key}"}
     }
     hub = CuratedHub(hub_name=hub_name, sagemaker_session=sagemaker_session)
     request = {
         "hub_name": hub_name,
         "hub_description": hub_description,
         "hub_display_name": hub_display_name,
         "hub_search_keywords": hub_search_keywords,
-        "s3_storage_config": {"S3OutputPath": "s3://sagemaker-hubs-us-east-1-123456789123"},
+        "s3_storage_config": {
+            "S3OutputPath": f"s3://sagemaker-hubs-us-east-1-123456789123/{storage_location.key}"
+        },
         "tags": tags,
     }
     response = hub.create(
@@ -113,7 +122,9 @@ def test_create_with_no_bucket_name(
         ),
     ],
 )
+@patch("sagemaker.jumpstart.curated_hub.curated_hub.CuratedHub._generate_hub_storage_location")
 def test_create_with_bucket_name(
+    mock_generate_hub_storage_location,
     sagemaker_session,
     hub_name,
     hub_description,
@@ -122,6 +133,8 @@ def test_create_with_bucket_name(
     hub_search_keywords,
     tags,
 ):
+    storage_location = S3ObjectLocation(hub_bucket_name,f"{hub_name}-{FAKE_TIME.timestamp()}")
+    mock_generate_hub_storage_location.return_value = storage_location
     create_hub = {"HubArn": f"arn:aws:sagemaker:us-east-1:123456789123:hub/{hub_name}"}
     sagemaker_session.create_hub = Mock(return_value=create_hub)
     hub = CuratedHub(
@@ -132,7 +145,7 @@ def test_create_with_bucket_name(
         "hub_description": hub_description,
         "hub_display_name": hub_display_name,
         "hub_search_keywords": hub_search_keywords,
-        "s3_storage_config": {"S3OutputPath": "s3://mock-bucket-123"},
+        "s3_storage_config": {"S3OutputPath": f"s3://mock-bucket-123/{storage_location.key}"},
         "tags": tags,
     }
     response = hub.create(
@@ -397,92 +410,92 @@ def test_determine_models_to_sync(sagemaker_session):
     hub_name = "mock_hub_name"
     hub = CuratedHub(hub_name=hub_name, sagemaker_session=sagemaker_session)
 
-    js_models_in_hub = [
-        {
+    js_model_map = {
+        "mock-model-two-pytorch": {
             "name": "mock-model-two-pytorch",
             "version": "1.0.1",
             "search_keywords": [
                 "@jumpstart-model-id:model-two-pytorch",
                 "@jumpstart-model-version:1.0.2",
             ],
         },
-        {
+        "mock-model-four-huggingface": {
             "name": "mock-model-four-huggingface",
             "version": "2.0.2",
             "search_keywords": [
                 "@jumpstart-model-id:model-four-huggingface",
                 "@jumpstart-model-version:2.0.2",
             ],
         },
-    ]
+    }
     model_one = JumpStartModelInfo("mock-model-one-huggingface", "1.2.3")
     model_two = JumpStartModelInfo("mock-model-two-pytorch", "1.0.2")
     # No model_one, older model_two
-    res = hub._determine_models_to_sync([model_one, model_two], js_models_in_hub)
+    res = hub._determine_models_to_sync([model_one, model_two], js_model_map)
     assert res == [model_one, model_two]
 
-    js_models_in_hub = [
-        {
+    js_model_map = {
+        "mock-model-two-pytorch": {
             "name": "mock-model-two-pytorch",
             "version": "1.0.3",
             "search_keywords": [
                 "@jumpstart-model-id:model-two-pytorch",
                 "@jumpstart-model-version:1.0.3",
             ],
         },
-        {
+        "mock-model-four-huggingface": {
             "name": "mock-model-four-huggingface",
             "version": "2.0.2",
             "search_keywords": [
                 "@jumpstart-model-id:model-four-huggingface",
                 "@jumpstart-model-version:2.0.2",
             ],
         },
-    ]
+    }
     # No model_one, newer model_two
-    res = hub._determine_models_to_sync([model_one, model_two], js_models_in_hub)
+    res = hub._determine_models_to_sync([model_one, model_two], js_model_map)
     assert res == [model_one]
 
-    js_models_in_hub = [
-        {
+    js_model_map = {
+        "mock-model-one-huggingface": {
             "name": "mock-model-one-huggingface",
             "version": "1.2.3",
             "search_keywords": [
                 "@jumpstart-model-id:model-one-huggingface",
                 "@jumpstart-model-version:1.2.3",
             ],
         },
-        {
+        "mock-model-two-pytorch": {
             "name": "mock-model-two-pytorch",
             "version": "1.0.2",
             "search_keywords": [
                 "@jumpstart-model-id:model-two-pytorch",
                 "@jumpstart-model-version:1.0.2",
             ],
         },
-    ]
+    }
     # Same model_one, same model_two
-    res = hub._determine_models_to_sync([model_one, model_two], js_models_in_hub)
+    res = hub._determine_models_to_sync([model_one, model_two], js_model_map)
     assert res == []
 
-    js_models_in_hub = [
-        {
+    js_model_map = {
+        "mock-model-one-huggingface": {
             "name": "mock-model-one-huggingface",
             "version": "1.2.1",
             "search_keywords": [
                 "@jumpstart-model-id:model-one-huggingface",
                 "@jumpstart-model-version:1.2.1",
             ],
         },
-        {
+        "mock-model-two-pytorch": {
             "name": "mock-model-two-pytorch",
             "version": "1.0.2",
             "search_keywords": [
                 "@jumpstart-model-id:model-two-pytorch",
                 "@jumpstart-model-version:1.0.2",
             ],
         },
-    ]
+    }
     # Old model_one, same model_two
-    res = hub._determine_models_to_sync([model_one, model_two], js_models_in_hub)
+    res = hub._determine_models_to_sync([model_one, model_two], js_model_map)
     assert res == [model_one]