feature: JumpStart CuratedHub class creation and function definitions (aws#4448)

Author: jinyoung-lim

src/sagemaker/jumpstart/cache.py

@@ -34,6 +34,7 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     MODEL_TYPE_TO_MANIFEST_MAP,
     MODEL_TYPE_TO_SPECS_MAP,
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.exceptions import (
     get_wildcard_model_version_msg,

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

@@ -147,6 +147,38 @@ def test_generate_hub_arn_for_init_kwargs():
         utils.generate_hub_arn_for_init_kwargs(hub_arn, "us-east-1", mock_custom_session) == hub_arn
     )
 
+    assert (
+        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"
+    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
     assert utils.generate_hub_arn_for_init_kwargs(hub_arn, None, mock_custom_session) == hub_arn
 
 

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -28,6 +28,7 @@
     JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY,
     JumpStartModelsCache,
 )
+from sagemaker.session_settings import SessionSettings
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
@@ -1133,7 +1134,7 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
 
     mocked_is_dir.assert_any_call("/some/directory/metadata/manifest/root")
     assert mocked_is_dir.call_count == 2
-    mocked_open.assert_not_called()
+    assert mocked_open.call_count == 2
     mocked_get_json_file_and_etag_from_s3.assert_has_calls(
         calls=[
             call("models_manifest.json"),

tests/unit/sagemaker/jumpstart/utils.py

@@ -22,7 +22,6 @@
     JUMPSTART_REGION_NAME_SET,
 )
 from sagemaker.jumpstart.types import (
-    HubContentType,
     JumpStartCachedContentKey,
     JumpStartCachedContentValue,
     JumpStartModelSpecs,
@@ -32,6 +31,7 @@
     HubContentType,
 )
 from sagemaker.jumpstart.enums import JumpStartModelType
+
 from sagemaker.jumpstart.utils import get_formatted_manifest
 from tests.unit.sagemaker.jumpstart.constants import (
     PROTOTYPICAL_MODEL_SPECS_DICT,
@@ -254,7 +254,7 @@ def patched_retrieval_function(
             )
         )
     # TODO: Implement
-    if datatype == HubContentType.HUB:
+    if datatype == HubType.HUB:
         return None
 
     raise ValueError(f"Bad value for datatype: {datatype}")