support for gated and training unsupported

Author: bencrabtree

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

@@ -22,6 +22,7 @@
     S3ObjectLocation,
 )
 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
 
 
@@ -65,6 +66,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,7 @@
 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:
@@ -34,7 +34,11 @@ def __init__(
         studio_specs: Dict[str, Dict[str, Any]],
     ):
         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
 
@@ -52,6 +56,8 @@ def inference_artifact_s3_reference(self):
     @property
     def training_artifact_s3_reference(self):
         """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)
         )
@@ -66,13 +72,17 @@ def inference_script_s3_reference(self):
     @property
     def training_script_s3_reference(self):
         """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):
         """Retrieves s3 reference for s3 directory containing model training datasets"""
+        if not self.model_specs.training_supported:
+            return None
         return S3ObjectLocation(self._get_bucket_name(), self.__get_training_dataset_prefix())
 
     @property
@@ -95,22 +105,28 @@ def _get_bucket_name(self) -> str:
 
     def __get_training_dataset_prefix(self) -> 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/curated_hub.py

@@ -394,7 +394,6 @@ def _sync_public_model_to_hub(self, model: JumpStartModelInfo, thread_num: int):
         self._sagemaker_session.import_hub_content(
             document_schema_version=HubContentDocument_v2.SCHEMA_VERSION,
             hub_content_name=model.model_id,
-            hub_content_version=model.version,
             hub_name=self.hub_name,
             hub_content_document=hub_content_document,
             hub_content_type=HubContentType.MODEL,

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -133,8 +133,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(
@@ -164,3 +167,8 @@ def create_hub_bucket_if_it_does_not_exist(
     )
 
     return bucket_name
+
+
+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

src/sagemaker/jumpstart/utils.py

@@ -870,6 +870,7 @@ def generate_studio_spec_file_prefix(model_id: str, model_version: str) -> str:
     """Returns the Studio Spec file prefix given a model ID and version."""
     return f"studio_models/{model_id}/studio_specs_v{model_version}.json"
 
+
 def extract_info_from_hub_content_arn(
     arn: str,
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:

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

@@ -127,3 +127,44 @@ 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

@@ -238,3 +238,13 @@ def test_create_hub_bucket_if_it_does_not_exist():
 
     mock_sagemaker_session.boto_session.resource("s3").create_bucketassert_called_once()
     assert created_hub_bucket_name == bucket_name
+
+
+def test_is_gated_bucket():
+    assert utils.is_gated_bucket("jumpstart-private-cache-prod-us-west-2") is True
+
+    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

tests/unit/sagemaker/jumpstart/test_accessors.py

@@ -98,6 +98,7 @@ def test_jumpstart_models_cache_get_model_specs(mock_cache):
         )
     )
 
+
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._cache")
 def test_jumpstart_proprietary_models_cache_get(mock_cache):