Restrict JS Gated nodels only in SM Endpoint mode

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -31,6 +31,7 @@
     LocalModelInvocationException,
     LocalModelLoadException,
     SkipTuningComboException,
+    JumpStartGatedModelNotSupported,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -443,6 +444,11 @@ def _build_for_jumpstart(self):
 
         logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
 
+        if self._is_gated_model() and self.mode != Mode.SAGEMAKER_ENDPOINT:
+            raise JumpStartGatedModelNotSupported(
+                "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode"
+            )
+
         if "djl-inference" in image_uri:
             logger.info("Building for DJL JumpStart Model ID...")
             self.model_server = ModelServer.DJL_SERVING
@@ -469,3 +475,12 @@ def _build_for_jumpstart(self):
             )
 
         return self.pysdk_model
+
+    def _is_gated_model(self) -> bool:
+        """Determine if ``this`` Model is Gated"""
+
+        s3_uri = self.pysdk_model.model_data
+        if isinstance(s3_uri, dict):
+            s3_uri = s3_uri.get("S3DataSource").get("S3Uri")
+
+        return "private" in s3_uri

src/sagemaker/serve/utils/exceptions.py

@@ -69,3 +69,12 @@ class TaskNotFoundException(ModelBuilderException):
 
     def __init__(self, message):
         super().__init__(message=message)
+
+
+class JumpStartGatedModelNotSupported(ModelBuilderException):
+    """Raise when deploying JumpStart gated model locally"""
+
+    fmt = "Error Message: {message}"
+
+    def __init__(self, message):
+        super().__init__(message=message)