Merge branch 'master-jumpstart-curated-hub' into curated_hub_tagris_copy

Author: chrstfu

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

@@ -24,6 +24,7 @@
 from sagemaker.jumpstart.curated_hub.accessors.public_model_data import (
     PublicModelDataAccessor,
 )
+from sagemaker.jumpstart.curated_hub.utils import is_gated_bucket
 from sagemaker.jumpstart.types import JumpStartModelSpecs
 
 
@@ -67,6 +68,10 @@ def generate_file_infos_from_model_specs(
     files = []
     for dependency in HubContentDependencyType:
         location: S3ObjectLocation = public_model_data_accessor.get_s3_reference(dependency)
+        # Training dependencies will return as None if training is unsupported
+        if not location or is_gated_bucket(location.bucket):
+            continue
+
         location_type = "prefix" if location.key.endswith("/") else "object"
 
         if location_type == "prefix":

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

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module accessors for the SageMaker JumpStart Public Hub."""
 from __future__ import absolute_import
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 from sagemaker import model_uris, script_uris
 from sagemaker.jumpstart.curated_hub.types import (
     HubContentDependencyType,
@@ -21,7 +21,10 @@
 from sagemaker.jumpstart.curated_hub.utils import create_s3_object_reference_from_uri
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.types import JumpStartModelSpecs
-from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.jumpstart.utils import (
+    get_jumpstart_content_bucket,
+    get_jumpstart_gated_content_bucket,
+)
 
 
 class PublicModelDataAccessor:
@@ -35,7 +38,11 @@ def __init__(
     ):
         """Creates a PublicModelDataAccessor."""
         self._region = region
-        self._bucket = get_jumpstart_content_bucket(region)
+        self._bucket = (
+            get_jumpstart_gated_content_bucket(region)
+            if model_specs.gated_bucket
+            else get_jumpstart_content_bucket(region)
+        )
         self.model_specs = model_specs
         self.studio_specs = studio_specs  # Necessary for SDK - Studio metadata drift
 
@@ -44,47 +51,53 @@ def get_s3_reference(self, dependency_type: HubContentDependencyType):
         return getattr(self, dependency_type.value)
 
     @property
-    def inference_artifact_s3_reference(self):
+    def inference_artifact_s3_reference(self) -> Optional[S3ObjectLocation]:
         """Retrieves s3 reference for model inference artifact"""
         return create_s3_object_reference_from_uri(
             self._jumpstart_artifact_s3_uri(JumpStartScriptScope.INFERENCE)
         )
 
     @property
-    def training_artifact_s3_reference(self):
+    def training_artifact_s3_reference(self) -> Optional[S3ObjectLocation]:
         """Retrieves s3 reference for model training artifact"""
+        if not self.model_specs.training_supported:
+            return None
         return create_s3_object_reference_from_uri(
             self._jumpstart_artifact_s3_uri(JumpStartScriptScope.TRAINING)
         )
 
     @property
-    def inference_script_s3_reference(self):
+    def inference_script_s3_reference(self) -> Optional[S3ObjectLocation]:
         """Retrieves s3 reference for model inference script"""
         return create_s3_object_reference_from_uri(
             self._jumpstart_script_s3_uri(JumpStartScriptScope.INFERENCE)
         )
 
     @property
-    def training_script_s3_reference(self):
+    def training_script_s3_reference(self) -> Optional[S3ObjectLocation]:
         """Retrieves s3 reference for model training script"""
+        if not self.model_specs.training_supported:
+            return None
         return create_s3_object_reference_from_uri(
             self._jumpstart_script_s3_uri(JumpStartScriptScope.TRAINING)
         )
 
     @property
-    def default_training_dataset_s3_reference(self):
+    def 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())
+        if not self.model_specs.training_supported:
+            return None
+        return S3ObjectLocation(self._get_bucket_name(), self._get_training_dataset_prefix())
 
     @property
-    def demo_notebook_s3_reference(self):
+    def demo_notebook_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model demo jupyter notebook"""
         framework = self.model_specs.get_framework()
         key = f"{framework}-notebooks/{self.model_specs.model_id}-inference.ipynb"
         return S3ObjectLocation(self._get_bucket_name(), key)
 
     @property
-    def markdown_s3_reference(self):
+    def markdown_s3_reference(self) -> S3ObjectLocation:
         """Retrieves s3 reference for model markdown"""
         framework = self.model_specs.get_framework()
         key = f"{framework}-metadata/{self.model_specs.model_id}.md"
@@ -94,24 +107,30 @@ def _get_bucket_name(self) -> str:
         """Retrieves s3 bucket"""
         return self._bucket
 
-    def __get_training_dataset_prefix(self) -> str:
+    def _get_training_dataset_prefix(self) -> Optional[str]:
         """Retrieves training dataset location"""
-        return self.studio_specs["defaultDataKey"]
+        return self.studio_specs.get("defaultDataKey")
 
-    def _jumpstart_script_s3_uri(self, model_scope: str) -> str:
+    def _jumpstart_script_s3_uri(self, model_scope: str) -> Optional[str]:
         """Retrieves JumpStart script s3 location"""
-        return script_uris.retrieve(
-            region=self._region,
-            model_id=self.model_specs.model_id,
-            model_version=self.model_specs.version,
-            script_scope=model_scope,
-        )
+        try:
+            return script_uris.retrieve(
+                region=self._region,
+                model_id=self.model_specs.model_id,
+                model_version=self.model_specs.version,
+                script_scope=model_scope,
+            )
+        except ValueError:
+            return None
 
-    def _jumpstart_artifact_s3_uri(self, model_scope: str) -> str:
+    def _jumpstart_artifact_s3_uri(self, model_scope: str) -> Optional[str]:
         """Retrieves JumpStart artifact s3 location"""
-        return model_uris.retrieve(
-            region=self._region,
-            model_id=self.model_specs.model_id,
-            model_version=self.model_specs.version,
-            model_scope=model_scope,
-        )
+        try:
+            return model_uris.retrieve(
+                region=self._region,
+                model_id=self.model_specs.model_id,
+                model_version=self.model_specs.version,
+                model_scope=model_scope,
+            )
+        except ValueError:
+            return None

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -143,8 +143,11 @@ def generate_default_hub_bucket_name(
     return f"sagemaker-hubs-{region}-{account_id}"
 
 
-def create_s3_object_reference_from_uri(s3_uri: str) -> S3ObjectLocation:
+def create_s3_object_reference_from_uri(s3_uri: Optional[str]) -> Optional[S3ObjectLocation]:
     """Utiity to help generate an S3 object reference"""
+    if not s3_uri:
+        return None
+
     bucket, key = parse_s3_url(s3_uri)
 
     return S3ObjectLocation(
@@ -303,3 +306,8 @@ def get_jumpstart_model_and_version(
         len(JUMPSTART_HUB_MODEL_VERSION_TAG_PREFIX) :
     ]
     return JumpStartModelInfo(model_id=jumpstart_model_id, version=jumpstart_model_version)
+
+  
+def is_gated_bucket(bucket_name: str) -> bool:
+    """Returns true if the bucket name is the JumpStart gated bucket."""
+    return bucket_name in constants.JUMPSTART_GATED_BUCKET_NAME_SET

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

@@ -127,3 +127,46 @@ def test_s3_path_file_generator_with_no_objects(s3_client):
 
     s3_client.list_objects_v2.assert_called_once()
     assert response == []
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_specs_file_generator_training_unsupported(patched_get_model_specs, s3_client):
+    specs = Mock()
+    specs.model_id = "mock_model_123"
+    specs.training_supported = False
+    specs.gated_bucket = False
+    specs.hosting_prepacked_artifact_key = "/my/inference/tarball.tgz"
+    specs.hosting_script_key = "/my/inference/script.py"
+
+    response = generate_file_infos_from_model_specs(specs, {}, "us-west-2", s3_client)
+
+    assert response == [
+        FileInfo(
+            "jumpstart-cache-prod-us-west-2",
+            "/my/inference/tarball.tgz",
+            123456789,
+            "08-14-1997 00:00:00",
+        ),
+        FileInfo(
+            "jumpstart-cache-prod-us-west-2",
+            "/my/inference/script.py",
+            123456789,
+            "08-14-1997 00:00:00",
+        ),
+    ]
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_specs_file_generator_gated_model(patched_get_model_specs, s3_client):
+    specs = Mock()
+    specs.model_id = "mock_model_123"
+    specs.gated_bucket = True
+    specs.training_supported = True
+    specs.hosting_prepacked_artifact_key = "/my/inference/tarball.tgz"
+    specs.hosting_script_key = "/my/inference/script.py"
+    specs.training_prepacked_artifact_key = "/my/training/tarball.tgz"
+    specs.training_script_key = "/my/training/script.py"
+
+    response = generate_file_infos_from_model_specs(specs, {}, "us-west-2", s3_client)
+
+    assert response == []

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

@@ -175,15 +175,14 @@ def test_create_hub_bucket_if_it_does_not_exist_hub_arn():
     assert utils.generate_hub_arn_for_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"
+def test_is_gated_bucket():
+    assert utils.is_gated_bucket("jumpstart-private-cache-prod-us-west-2") is True
 
-    assert (
-        utils.generate_default_hub_bucket_name(sagemaker_session=mock_sagemaker_session)
-        == "sagemaker-hubs-us-east-1-123456789123"
-    )
+    assert utils.is_gated_bucket("jumpstart-private-cache-prod-us-east-1") is True
+
+    assert utils.is_gated_bucket("jumpstart-cache-prod-us-west-2") is False
+    
+    assert utils.is_gated_bucket("") is False
 
 
 def test_create_hub_bucket_if_it_does_not_exist():

tests/unit/sagemaker/jumpstart/test_accessors.py

@@ -98,9 +98,6 @@ def test_jumpstart_models_cache_get_model_specs(mock_cache):
         )
     )
 
-    # necessary because accessors is a static module
-    reload(accessors)
-
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._cache")
 def test_jumpstart_proprietary_models_cache_get(mock_cache):

tests/unit/sagemaker/jumpstart/utils.py

@@ -254,16 +254,6 @@ def patched_retrieval_function(
             )
         )
 
-    if datatype == HubContentType.MODEL:
-        _, _, _, model_name, model_version = id_info.split("/")
-        return JumpStartCachedContentValue(
-            formatted_content=get_spec_from_base_spec(model_id=model_name, version=model_version)
-        )
-
-    # TODO: Implement
-    if datatype == HubType.HUB:
-        return None
-
     raise ValueError(f"Bad value for datatype: {datatype}")