Add curated_hub utils unit tests

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -21,6 +21,7 @@
 import botocore
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from sagemaker.utilities.cache import LRUCache
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
@@ -36,6 +37,7 @@
     JUMPSTART_DEFAULT_S3_CACHE_EXPIRATION_HORIZON,
     JUMPSTART_DEFAULT_SEMANTIC_VERSION_CACHE_EXPIRATION_HORIZON,
 )
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.types import (
     JumpStartCachedContentKey,
     JumpStartCachedContentValue,
@@ -45,14 +47,12 @@
     JumpStartVersionedModelId,
 )
 from sagemaker.jumpstart.curated_hub.types import (
-    HubContentType,
     DescribeHubResponse,
     DescribeHubContentsResponse,
+    HubContentType,
 )
-from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.curated_hub import utils as hub_utils
 from sagemaker.jumpstart.curated_hub.curated_hub import CuratedHub
-from sagemaker.utilities.cache import LRUCache
 
 
 class JumpStartModelsCache:

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -17,11 +17,11 @@
 from typing import Any, Dict, Optional
 from sagemaker.session import Session
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.curated_hub import utils as hub_utils
+from sagemaker.jumpstart.curated_hub.utils import create_hub_bucket_if_it_does_not_exist
 from sagemaker.jumpstart.curated_hub.types import (
     DescribeHubResponse,
-    HubContentType,
     DescribeHubContentsResponse,
+    HubContentType,
 )
 
 
@@ -51,9 +51,7 @@ def create(
     ) -> Dict[str, str]:
         """Creates a hub with the given description"""
 
-        bucket_name = hub_utils.create_hub_bucket_if_it_does_not_exist(
-            bucket_name, self._sagemaker_session
-        )
+        bucket_name = create_hub_bucket_if_it_does_not_exist(bucket_name, self._sagemaker_session)
 
         return self._sagemaker_session.create_hub(
             hub_name=self.hub_name,

src/sagemaker/jumpstart/curated_hub/types.py

@@ -15,6 +15,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
+
 from sagemaker.jumpstart.types import JumpStartDataHolderType
 
 

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -15,8 +15,8 @@
 import re
 from typing import Optional
 from sagemaker.session import Session
-from sagemaker.jumpstart import constants
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart import constants
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentType,
     HubArnExtractedInfo,

src/sagemaker/jumpstart/session_utils.py

@@ -15,11 +15,11 @@
 from __future__ import absolute_import
 
 from typing import Optional, Tuple
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
-from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 from sagemaker.session import Session
 from sagemaker.utils import aws_partition
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.utils import get_jumpstart_model_id_version_from_resource_arn
 
 
 def get_model_id_version_from_endpoint(

src/sagemaker/jumpstart/types.py

@@ -15,16 +15,14 @@
 from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, Union
+from sagemaker.session import Session
 from sagemaker.utils import get_instance_type_family, format_tags, Tags
+from sagemaker.enums import EndpointType
 from sagemaker.model_metrics import ModelMetrics
 from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.drift_check_baselines import DriftCheckBaselines
-
-from sagemaker.session import Session
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
-from sagemaker.enums import EndpointType
-from sagemaker.jumpstart.curated_hub import types as hub_types
 
 
 class JumpStartDataHolderType:
@@ -119,14 +117,17 @@ def to_json(self) -> Dict[str, Any]:
         return json_obj
 
 
+from sagemaker.jumpstart.curated_hub.types import HubContentType  # noqa: E402
+
+
 class JumpStartS3FileType(str, Enum):
     """Type of files published in JumpStart S3 distribution buckets."""
 
     MANIFEST = "manifest"
     SPECS = "specs"
 
 
-JumpStartContentDataType = Union[JumpStartS3FileType, hub_types.HubContentType]
+JumpStartContentDataType = Union[JumpStartS3FileType, HubContentType]
 
 
 class JumpStartLaunchedRegionInfo(JumpStartDataHolderType):

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

@@ -11,10 +11,11 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
-from unittest.mock import Mock
-from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
+from unittest.mock import Mock
+from botocore.exceptions import ClientError
 from sagemaker.jumpstart.curated_hub import utils
+from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 from sagemaker.jumpstart.curated_hub.types import HubArnExtractedInfo
 
 
@@ -151,3 +152,35 @@ def test_generate_hub_arn_for_estimator_init_kwargs():
         utils.generate_hub_arn_for_estimator_init_kwargs(hub_arn, None, mock_custom_session)
         == hub_arn
     )
+
+
+def test_generate_default_hub_bucket_name():
+    mock_sagemaker_session = Mock()
+    mock_sagemaker_session.account_id.return_value = "123456789123"
+    mock_sagemaker_session.boto_region_name = "us-east-1"
+
+    assert (
+        utils.generate_default_hub_bucket_name(sagemaker_session=mock_sagemaker_session)
+        == "sagemaker-hubs-us-east-1-123456789123"
+    )
+
+
+def test_create_hub_bucket_if_it_does_not_exist():
+    mock_sagemaker_session = Mock()
+    mock_sagemaker_session.account_id.return_value = "123456789123"
+    mock_sagemaker_session.client("sts").get_caller_identity.return_value = {
+        "Account": "123456789123"
+    }
+    mock_sagemaker_session.boto_session.resource("s3").Bucket().creation_date = None
+    mock_sagemaker_session.boto_region_name = "us-east-1"
+    error = ClientError(
+        error_response={"Error": {"Code": "404", "Message": "Not Found"}},
+        operation_name="foo",
+    )
+    bucket_name = "sagemaker-hubs-us-east-1-123456789123"
+    created_hub_bucket_name = utils.create_hub_bucket_if_it_does_not_exist(
+        sagemaker_session=mock_sagemaker_session
+    )
+
+    mock_sagemaker_session.boto_session.resource("s3").create_bucketassert_called_once()
+    assert created_hub_bucket_name == bucket_name

tests/unit/sagemaker/jumpstart/utils.py

@@ -22,13 +22,13 @@
     JUMPSTART_REGION_NAME_SET,
 )
 from sagemaker.jumpstart.types import (
-    HubContentType,
     JumpStartCachedContentKey,
     JumpStartCachedContentValue,
     JumpStartModelSpecs,
     JumpStartS3FileType,
     JumpStartModelHeader,
 )
+
 from sagemaker.jumpstart.utils import get_formatted_manifest
 from tests.unit.sagemaker.jumpstart.constants import (
     PROTOTYPICAL_MODEL_SPECS_DICT,
@@ -37,6 +37,7 @@
     BASE_HEADER,
     SPECIAL_MODEL_SPECS_DICT,
 )
+from sagemaker.jumpstart.curated_hub.types import HubContentType
 
 
 def get_header_from_base_header(