chore: address PR comments

Author: evakravi

src/sagemaker/base_predictor.py

@@ -14,7 +14,7 @@
 from __future__ import print_function, absolute_import
 
 import abc
-from typing import Any, Tuple
+from typing import Any, Optional, Tuple, Union
 
 from sagemaker.deprecations import (
     deprecated_class,
@@ -205,27 +205,29 @@ def _create_request_args(
     ):
         """Placeholder docstring"""
 
-        js_accept = None
+        jumpstart_serialized_data: Optional[Union[str, bytes]] = None
+        jumpstart_accept: Optional[str] = None
+        jumpstart_content_type: Optional[str] = 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(
+            jumpstart_serialized_data = PayloadSerializer(
                 bucket=bucket, region=region, s3_client=s3_client
             ).serialize(data)
-            js_content_type = data.content_type
-            js_accept = data.accept
+            jumpstart_content_type = data.content_type
+            jumpstart_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:
-            if isinstance(data, JumpStartSerializablePayload):
-                args["ContentType"] = js_content_type
+            if isinstance(data, JumpStartSerializablePayload) and jumpstart_content_type:
+                args["ContentType"] = jumpstart_content_type
             else:
                 args["ContentType"] = (
                     self.content_type
@@ -234,8 +236,8 @@ def _create_request_args(
                 )
 
         if "Accept" not in args:
-            if isinstance(data, JumpStartSerializablePayload) and js_accept:
-                args["Accept"] = js_accept
+            if isinstance(data, JumpStartSerializablePayload) and jumpstart_accept:
+                args["Accept"] = jumpstart_accept
             else:
                 args["Accept"] = (
                     self.accept if isinstance(self.accept, str) else ", ".join(self.accept)
@@ -254,9 +256,9 @@ def _create_request_args(
             args["CustomAttributes"] = custom_attributes
 
         data = (
-            self.serializer.serialize(data)
-            if not isinstance(data, JumpStartSerializablePayload)
-            else js_serialized_data
+            jumpstart_serialized_data
+            if isinstance(data, JumpStartSerializablePayload) and jumpstart_serialized_data
+            else self.serializer.serialize(data)
         )
 
         args["Body"] = data

src/sagemaker/jumpstart/accessors.py

@@ -39,12 +39,19 @@ def get_sagemaker_version() -> str:
 
 
 class JumpStartS3Accessor(object):
-    """Static class for storing and retrieving auxilliary s3 artifacts."""
+    """Static class for storing and retrieving auxilliary S3 artifacts."""
+
+    @staticmethod
+    def clear_cache() -> None:
+        """Clears LRU caches associated with S3 client and retrieved objects."""
+
+        JumpStartS3Accessor._get_default_s3_client.cache_clear()
+        JumpStartS3Accessor.get_object_cached.cache_clear()
 
     @staticmethod
     @functools.lru_cache()
     def _get_default_s3_client(region: str = JUMPSTART_DEFAULT_REGION_NAME) -> boto3.client:
-        """Returns default s3 client associated with the region.
+        """Returns default S3 client associated with the region.
 
         Result is cached so multiple clients in memory are not created.
         """
@@ -58,9 +65,9 @@ def get_object_cached(
         region: str = JUMPSTART_DEFAULT_REGION_NAME,
         s3_client: Optional[boto3.client] = None,
     ) -> bytes:
-        """Returns s3 object located at the bucket and key.
+        """Returns S3 object located at the bucket and key.
 
-        Requests are cached so that the same s3 request is never made more
+        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(
@@ -74,7 +81,7 @@ def get_object(
         region: str = JUMPSTART_DEFAULT_REGION_NAME,
         s3_client: Optional[boto3.client] = None,
     ) -> bytes:
-        """Returns s3 object located at the bucket and key."""
+        """Returns S3 object located at the bucket and key."""
         if s3_client is None:
             s3_client = JumpStartS3Accessor._get_default_s3_client(region)
 

src/sagemaker/jumpstart/artifacts/payloads.py

@@ -10,7 +10,7 @@
 # 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."""
+"""This module contains functions for obtaining example payloads for JumpStart models."""
 from __future__ import absolute_import
 from copy import deepcopy
 from typing import Dict, Optional

src/sagemaker/jumpstart/model.py

@@ -315,11 +315,11 @@ 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.
+        """Returns the 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(
+        payload_options: Optional[List[JumpStartSerializablePayload]] = payloads.retrieve_options(
             model_id=self.model_id,
             model_version=self.model_version,
             region=self.region,
@@ -328,12 +328,10 @@ def retrieve_default_payload(self) -> JumpStartSerializablePayload:
             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}'."
-            )
+        if payload_options is None or len(payload_options) == 0:
+            return None
 
-        return sample_payloads[0]
+        return payload_options[0]
 
     def _create_sagemaker_model(
         self,

src/sagemaker/jumpstart/payload_utils.py

@@ -10,11 +10,11 @@
 # 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 stores payload utilities for SageMaker JumpStart."""
+"""This module stores stores inference payload utilities for JumpStart models."""
 from __future__ import absolute_import
 import base64
 import json
-from typing import Any, Optional, Union
+from typing import Optional, Union
 import re
 import boto3
 
@@ -50,14 +50,14 @@ def get_bytes_payload_with_s3_references(
         self,
         payload_str: str,
     ) -> bytes:
-        """Returns bytes object corresponding to referenced s3 object.
+        """Returns bytes object corresponding to referenced S3 object.
 
         Raises:
             ValueError: If the raw bytes payload is not formatted correctly.
         """
         s3_keys = re.compile(S3_BYTES_REGEX).findall(payload_str)
         if len(s3_keys) != 1:
-            raise ValueError(f"Invalid bytes payload: {payload_str}")
+            raise ValueError("Invalid bytes payload.")
 
         s3_key = s3_keys[0]
         serialized_s3_object = JumpStartS3Accessor.get_object_cached(
@@ -70,7 +70,10 @@ def embed_s3_references_in_str_payload(
         self,
         payload: str,
     ) -> str:
-        """Embeds s3 references in string payloads."""
+        """Inserts serialized S3 content into string payload.
+
+        If no S3 content is embedded in payload, original string is returned.
+        """
         return self._embed_s3_b64_references_in_str_payload(payload_body=payload)
 
     def _embed_s3_b64_references_in_str_payload(
@@ -98,10 +101,12 @@ def _embed_s3_b64_references_in_str_payload(
     def embed_s3_references_in_json_payload(
         self, payload_body: Union[list, dict, str, int, float]
     ) -> Union[list, dict, str, int, float]:
-        """Finds all s3 references in payload and embeds serialized s3 data.
+        """Finds all S3 references in payload and embeds serialized S3 data.
 
-        S3 bucket is assumed to be the default JumpStart content bucket. If no s3 references
-        are found, the payload is returned un-modified.
+        If no S3 references are found, the payload is returned un-modified.
+
+        Raises:
+            ValueError: If the payload has an unrecognized type.
         """
         if isinstance(payload_body, str):
             return self.embed_s3_references_in_str_payload(payload_body)
@@ -116,8 +121,12 @@ def embed_s3_references_in_json_payload(
             }
         raise ValueError(f"Payload has unrecognized type: {type(payload_body)}")
 
-    def serialize(self, payload: JumpStartSerializablePayload) -> Any:
-        """Returns payload bytes that can be inputted to inference endpoint."""
+    def serialize(self, payload: JumpStartSerializablePayload) -> Union[str, bytes]:
+        """Returns payload string or bytes that can be inputted to inference endpoint.
+
+        Raises:
+            ValueError: If the payload has an unrecognized type.
+        """
         content_type = MIMEType.from_suffixed_type(payload.content_type)
         body = payload.body
 

src/sagemaker/jumpstart/types.py

@@ -348,6 +348,9 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
         Args:
             json_obj (Dict[str, Any]): Dictionary representation of serializable
                 payload specs.
+
+        Raises:
+            KeyError: If the dictionary is missing keys.
         """
 
         if json_obj is None:

src/sagemaker/payloads.py

@@ -27,7 +27,7 @@
 logger = logging.getLogger(__name__)
 
 
-def retrieve_samples(
+def retrieve_options(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
@@ -78,7 +78,7 @@ def retrieve_samples(
         Dict[str, JumpStartSerializablePayload]
     ] = artifacts._retrieve_default_payloads(
         model_id,
-        model_version,  # type: ignore
+        model_version,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -86,31 +86,6 @@ def test_prepacked_jumpstart_model(setup):
     assert response is not None
 
 
-def test_default_payload_jumpstart_model(setup):
-
-    # DO NOT COMMIT THIS LINE
-    os.environ.update({"AWS_JUMPSTART_CONTENT_BUCKET_OVERRIDE": "jumpstart-cache-alpha-us-west-2"})
-
-    model_id = "model-depth2img-stable-diffusion-v1-5-controlnet-v1-1-fp16"
-
-    model = JumpStartModel(
-        model_id=model_id,
-        role=get_sm_session().get_caller_identity_arn(),
-        sagemaker_session=get_sm_session(),
-    )
-
-    default_payload = model.retrieve_default_payload()
-
-    # uses ml.g5.8xlarge instance
-    predictor = model.deploy(
-        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
-    )
-
-    response = predictor.predict(default_payload)
-
-    assert response is not None
-
-
 @pytest.mark.skipif(
     tests.integ.test_region() not in GATED_INFERENCE_MODEL_SUPPORTED_REGIONS,
     reason=f"JumpStart gated inference models unavailable in {tests.integ.test_region()}.",