chore: add support for custom attributes to predictor class

Author: evakravi

src/sagemaker/base_predictor.py

@@ -133,6 +133,7 @@ def predict(
         target_model=None,
         target_variant=None,
         inference_id=None,
+        custom_attributes=None,
     ):
         """Return the inference from the specified endpoint.
 
@@ -153,6 +154,18 @@ def predict(
                 model you want to host and the resources you want to deploy for hosting it.
             inference_id (str): If you provide a value, it is added to the captured data
                 when you enable data capture on the endpoint (Default: None).
+            custom_attributes (str): Provides additional information about a request for an
+                inference submitted to a model hosted at an Amazon SageMaker endpoint.
+                The information is an opaque value that is forwarded verbatim. You could use this
+                value, for example, to provide an ID that you can use to track a request or to provide
+                other metadata that a service endpoint was programmed to process. The value must
+                consist of no more than 1024 visible US-ASCII characters.
+
+                The code in your model is responsible for setting or updating any custom attributes in
+                the response. If your code does not set this value in the response, an empty value is
+                returned. For example, if a custom attribute represents the trace ID, your model can
+                prepend the custom attribute with Trace ID: in your post-processing function
+                (Default: None).
 
         Returns:
             object: Inference for the given input. If a deserializer was specified when creating
@@ -162,7 +175,12 @@ def predict(
         """
 
         request_args = self._create_request_args(
-            data, initial_args, target_model, target_variant, inference_id
+            data,
+            initial_args,
+            target_model,
+            target_variant,
+            inference_id,
+            custom_attributes,
         )
         response = self.sagemaker_session.sagemaker_runtime_client.invoke_endpoint(**request_args)
         return self._handle_response(response)
@@ -180,6 +198,7 @@ def _create_request_args(
         target_model=None,
         target_variant=None,
         inference_id=None,
+        custom_attributes=None,
     ):
         """Placeholder docstring"""
         args = dict(initial_args) if initial_args else {}
@@ -206,6 +225,9 @@ def _create_request_args(
         if inference_id:
             args["InferenceId"] = inference_id
 
+        if custom_attributes:
+            args["CustomAttributes"] = custom_attributes
+
         data = self.serializer.serialize(data)
 
         args["Body"] = data

src/sagemaker/jumpstart/artifacts/__init__.py

@@ -52,4 +52,4 @@
     _retrieve_supported_accept_types,
     _retrieve_supported_content_types,
 )
-from sagemaker.jumpstart.artifacts.model_packages import _retrieve_model_package_arn
+from sagemaker.jumpstart.artifacts.model_packages import _retrieve_model_package_arn  # noqa: F401

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -12,8 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart model packages."""
 from __future__ import absolute_import
-from copy import deepcopy
-from typing import Dict, List, Optional
+from typing import Optional
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
 )

src/sagemaker/jumpstart/model.py

@@ -31,7 +31,7 @@
 )
 from sagemaker.jumpstart.utils import is_valid_model_id
 from sagemaker.utils import stringify_object
-from sagemaker.model import MODEL_PACKAGE_ARN_PATTERN, Model, ModelPackage
+from sagemaker.model import MODEL_PACKAGE_ARN_PATTERN, Model
 from sagemaker.model_monitor.data_capture_config import DataCaptureConfig
 from sagemaker.predictor import PredictorBase
 from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig