feat: jumpstart default payloads

Author: evakravi

src/sagemaker/base_predictor.py

@@ -32,6 +32,9 @@
     StreamDeserializer,
     StringDeserializer,
 )
+from sagemaker.jumpstart.payload_utils import PayloadSerializer
+from sagemaker.jumpstart.types import JumpStartSerializablePayload
+from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 from sagemaker.model_monitor import (
     DataCaptureConfig,
     DefaultModelMonitor,
@@ -201,20 +204,42 @@ def _create_request_args(
         custom_attributes=None,
     ):
         """Placeholder docstring"""
+
+        js_accept = None
+
+        if isinstance(data, JumpStartSerializablePayload):
+            s3_client = self.sagemaker_session.s3_client
+            region = self.sagemaker_session._region_name
+            bucket = get_jumpstart_content_bucket(region)
+
+            js_serialized_data = PayloadSerializer(
+                bucket=bucket, region=region, s3_client=s3_client
+            ).serialize(data)
+            js_content_type = data.content_type
+            js_accept = data.accept
+
         args = dict(initial_args) if initial_args else {}
 
         if "EndpointName" not in args:
             args["EndpointName"] = self.endpoint_name
 
         if "ContentType" not in args:
-            args["ContentType"] = (
-                self.content_type
-                if isinstance(self.content_type, str)
-                else ", ".join(self.content_type)
-            )
+            if isinstance(data, JumpStartSerializablePayload):
+                args["ContentType"] = js_content_type
+            else:
+                args["ContentType"] = (
+                    self.content_type
+                    if isinstance(self.content_type, str)
+                    else ", ".join(self.content_type)
+                )
 
         if "Accept" not in args:
-            args["Accept"] = self.accept if isinstance(self.accept, str) else ", ".join(self.accept)
+            if isinstance(data, JumpStartSerializablePayload) and js_accept:
+                args["Accept"] = js_accept
+            else:
+                args["Accept"] = (
+                    self.accept if isinstance(self.accept, str) else ", ".join(self.accept)
+                )
 
         if target_model:
             args["TargetModel"] = target_model
@@ -228,7 +253,11 @@ def _create_request_args(
         if custom_attributes:
             args["CustomAttributes"] = custom_attributes
 
-        data = self.serializer.serialize(data)
+        data = (
+            self.serializer.serialize(data)
+            if not isinstance(data, JumpStartSerializablePayload)
+            else js_serialized_data
+        )
 
         args["Body"] = data
         return args

src/sagemaker/jumpstart/accessors.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains accessors related to SageMaker JumpStart."""
 from __future__ import absolute_import
+import functools
 from typing import Any, Dict, List, Optional
 import boto3
 
@@ -37,6 +38,49 @@ def get_sagemaker_version() -> str:
         return SageMakerSettings._parsed_sagemaker_version
 
 
+class JumpStartS3Accessor(object):
+    """Static class for storing and retrieving auxilliary s3 artifacts."""
+
+    @functools.cache
+    @staticmethod
+    def _get_default_s3_client(region: str = JUMPSTART_DEFAULT_REGION_NAME) -> boto3.client:
+        """Returns default s3 client associated with the region.
+
+        Result is cached so multiple clients in memory are not created.
+        """
+        return boto3.client("s3", region_name=region)
+
+    @functools.lru_cache
+    @staticmethod
+    def get_object_cached(
+        bucket: str,
+        key: str,
+        region: str = JUMPSTART_DEFAULT_REGION_NAME,
+        s3_client: Optional[boto3.client] = None,
+    ) -> bytes:
+        """Returns s3 object located at the bucket and key.
+
+        Requests are cached so that the same s3 request is never made more
+        than once, unless a different region or client is used.
+        """
+        return JumpStartS3Accessor.get_object(
+            bucket=bucket, key=key, region=region, s3_client=s3_client
+        )
+
+    @staticmethod
+    def get_object(
+        bucket: str,
+        key: str,
+        region: str = JUMPSTART_DEFAULT_REGION_NAME,
+        s3_client: Optional[boto3.client] = None,
+    ) -> bytes:
+        """Returns s3 object located at the bucket and key."""
+        if s3_client is None:
+            s3_client = JumpStartS3Accessor._get_default_s3_client(region)
+
+        return s3_client.get_object(Bucket=bucket, Key=key)["Body"].read()
+
+
 class JumpStartModelsAccessor(object):
     """Static class for storing the JumpStart models cache."""
 

src/sagemaker/jumpstart/artifacts/__init__.py

@@ -61,3 +61,6 @@
     _retrieve_model_package_arn,
     _retrieve_model_package_model_artifact_s3_uri,
 )
+from sagemaker.jumpstart.artifacts.payloads import (  # noqa: F401
+    _retrieve_default_payloads,
+)

src/sagemaker/jumpstart/artifacts/payloads.py

@@ -0,0 +1,84 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains functions for obtaining JumpStart payloads."""
+from __future__ import absolute_import
+from copy import deepcopy
+from typing import Dict, Optional
+from sagemaker.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
+from sagemaker.jumpstart.enums import (
+    JumpStartScriptScope,
+)
+from sagemaker.jumpstart.types import JumpStartSerializablePayload
+from sagemaker.jumpstart.utils import (
+    verify_model_region_and_return_specs,
+)
+from sagemaker.session import Session
+
+
+def _retrieve_default_payloads(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+) -> Optional[Dict[str, JumpStartSerializablePayload]]:
+    """Returns default payloads.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            get default payloads.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            default resource name.
+        region (Optional[str]): Region for which to retrieve the
+            default resource name.
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+    Returns:
+        str: the default payload.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        scope=JumpStartScriptScope.INFERENCE,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+    )
+
+    default_payloads = model_specs.default_payloads
+
+    if default_payloads:
+        for payload in default_payloads.values():
+            payload.accept = getattr(
+                payload, "accept", model_specs.predictor_specs.default_accept_type
+            )
+
+    return deepcopy(default_payloads) if default_payloads else None

src/sagemaker/jumpstart/model.py

@@ -16,6 +16,7 @@
 import re
 
 from typing import Dict, List, Optional, Union
+from sagemaker import payloads
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.base_deserializers import BaseDeserializer
 from sagemaker.base_serializers import BaseSerializer
@@ -28,6 +29,7 @@
     get_deploy_kwargs,
     get_init_kwargs,
 )
+from sagemaker.jumpstart.types import JumpStartSerializablePayload
 from sagemaker.jumpstart.utils import is_valid_model_id
 from sagemaker.utils import stringify_object
 from sagemaker.model import MODEL_PACKAGE_ARN_PATTERN, Model
@@ -312,6 +314,27 @@ def _is_valid_model_id_hook():
 
         super(JumpStartModel, self).__init__(**model_init_kwargs.to_kwargs_dict())
 
+    def retrieve_default_payload(self) -> JumpStartSerializablePayload:
+        """Returns default payload associated with the model.
+
+        Payload can be directly used with the `sagemaker.predictor.Predictor.predict(...)` function.
+        """
+        sample_payloads: Optional[List[JumpStartSerializablePayload]] = payloads.retrieve_samples(
+            model_id=self.model_id,
+            model_version=self.model_version,
+            region=self.region,
+            tolerate_deprecated_model=self.tolerate_deprecated_model,
+            tolerate_vulnerable_model=self.tolerate_vulnerable_model,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        if sample_payloads is None or len(sample_payloads) == 0:
+            raise NotImplementedError(
+                f"No default payload supported for model ID '{self.model_id}'."
+            )
+
+        return sample_payloads[0]
+
     def _create_sagemaker_model(
         self,
         instance_type=None,