feature: add inference_id to predict

Author: akionakamura

src/sagemaker/local/local_session.py

@@ -426,17 +426,27 @@ def invoke_endpoint(
         CustomAttributes=None,
         TargetModel=None,
         TargetVariant=None,
+        InferenceId=None,
     ):
         """Invoke the endpoint.
 
         Args:
-          Body:
-          EndpointName:
-          Accept:  (Default value = None)
-          CustomAttributes:  (Default value = None)
+            Body: Input data for which you want the model to provide inference.
+            EndpointName: The name of the endpoint that you specified when you
+                created the endpoint using the CreateEndpoint API.
+            ContentType: The MIME type of the input data in the request body (Default value = None)
+            Accept: The desired MIME type of the inference in the response (Default value = None)
+            CustomAttributes: Provides additional information about a request for an inference
+                submitted to a model hosted at an Amazon SageMaker endpoint (Default value = None)
+            TargetModel: The model to request for inference when invoking a multi-model endpoint
+                (Default value = None)
+            TargetVariant: Specify the production variant to send the inference request to when
+                invoking an endpoint that is running two or more variants (Default value = None)
+            InferenceId: If you provide a value, it is added to the captured data when you enable
+               data capture on the endpoint (Default value = None)
 
         Returns:
-
+            object: Inference for the given input.
         """
         url = "http://localhost:%s/invocations" % self.serving_port
         headers = {}
@@ -456,6 +466,9 @@ def invoke_endpoint(
         if TargetVariant is not None:
             headers["X-Amzn-SageMaker-Target-Variant"] = TargetVariant
 
+        if InferenceId is not None:
+            headers["X-Amzn-SageMaker-Inference-Id"] = InferenceId
+
         r = self.http.request("POST", url, body=Body, preload_content=False, headers=headers)
 
         return {"Body": r, "ContentType": Accept}

src/sagemaker/predictor.py

@@ -95,7 +95,9 @@ def __init__(
         self._model_names = self._get_model_names()
         self._context = None
 
-    def predict(self, data, initial_args=None, target_model=None, target_variant=None):
+    def predict(
+        self, data, initial_args=None, target_model=None, target_variant=None, inference_id=None
+    ):
         """Return the inference from the specified endpoint.
 
         Args:
@@ -111,8 +113,10 @@ def predict(self, data, initial_args=None, target_model=None, target_variant=Non
                 in case of a multi model endpoint. Does not apply to endpoints hosting
                 single model (Default: None)
             target_variant (str): The name of the production variant to run an inference
-            request on (Default: None). Note that the ProductionVariant identifies the model
-            you want to host and the resources you want to deploy for hosting it.
+                request on (Default: None). Note that the ProductionVariant identifies the
+                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).
 
         Returns:
             object: Inference for the given input. If a deserializer was specified when creating
@@ -121,7 +125,9 @@ def predict(self, data, initial_args=None, target_model=None, target_variant=Non
                 as is.
         """
 
-        request_args = self._create_request_args(data, initial_args, target_model, target_variant)
+        request_args = self._create_request_args(
+            data, initial_args, target_model, target_variant, inference_id
+        )
         response = self.sagemaker_session.sagemaker_runtime_client.invoke_endpoint(**request_args)
         return self._handle_response(response)
 
@@ -131,7 +137,9 @@ def _handle_response(self, response):
         content_type = response.get("ContentType", "application/octet-stream")
         return self.deserializer.deserialize(response_body, content_type)
 
-    def _create_request_args(self, data, initial_args=None, target_model=None, target_variant=None):
+    def _create_request_args(
+        self, data, initial_args=None, target_model=None, target_variant=None, inference_id=None
+    ):
         """Placeholder docstring"""
         args = dict(initial_args) if initial_args else {}
 
@@ -150,6 +158,9 @@ def _create_request_args(self, data, initial_args=None, target_model=None, targe
         if target_variant:
             args["TargetVariant"] = target_variant
 
+        if inference_id:
+            args["InferenceId"] = inference_id
+
         data = self.serializer.serialize(data)
 
         args["Body"] = data

tests/unit/test_predictor.py

@@ -31,6 +31,7 @@
 RETURN_VALUE = 0
 CSV_RETURN_VALUE = "1,2,3\r\n"
 PRODUCTION_VARIANT_1 = "PRODUCTION_VARIANT_1"
+INFERENCE_ID = "inference-id"
 
 ENDPOINT_DESC = {"EndpointArn": "foo", "EndpointConfigName": ENDPOINT}
 
@@ -98,6 +99,29 @@ def test_predict_call_with_target_variant():
     assert result == RETURN_VALUE
 
 
+def test_predict_call_with_inference_id():
+    sagemaker_session = empty_sagemaker_session()
+    predictor = Predictor(ENDPOINT, sagemaker_session)
+
+    data = "untouched"
+    result = predictor.predict(data, inference_id=INFERENCE_ID)
+
+    assert sagemaker_session.sagemaker_runtime_client.invoke_endpoint.called
+
+    expected_request_args = {
+        "Accept": DEFAULT_ACCEPT,
+        "Body": data,
+        "ContentType": DEFAULT_CONTENT_TYPE,
+        "EndpointName": ENDPOINT,
+        "InferenceId": INFERENCE_ID,
+    }
+
+    call_args, kwargs = sagemaker_session.sagemaker_runtime_client.invoke_endpoint.call_args
+    assert kwargs == expected_request_args
+
+    assert result == RETURN_VALUE
+
+
 def test_multi_model_predict_call():
     sagemaker_session = empty_sagemaker_session()
     predictor = Predictor(ENDPOINT, sagemaker_session)