support CustomAttributes in local mode; doc fixes

Author: Jonathan Esterhazy

CHANGELOG.rst

@@ -5,6 +5,7 @@ CHANGELOG
 1.14.2-dev
 ==========
 
+* bug-fix: support ``CustomAttributes`` argument in local mode ``invoke_endpoint`` requests
 * enhancement: add ``content_type`` parameter to ``sagemaker.tensorflow.serving.Predictor``
 * doc-fix: add TensorFlow Serving Container docs
 * doc-fix: fix rendering error in README.rst

src/sagemaker/local/local_session.py

@@ -164,14 +164,19 @@ def __init__(self, config=None):
         self.config = config
         self.serving_port = get_config_value('local.serving_port', config) or 8080
 
-    def invoke_endpoint(self, Body, EndpointName, ContentType, Accept):  # pylint: disable=unused-argument
+    def invoke_endpoint(self, Body, EndpointName,  # pylint: disable=unused-argument
+                        ContentType, Accept, CustomAttributes):
         url = "http://localhost:%s/invocations" % self.serving_port
         headers = {
             'Content-type': ContentType
         }
+
         if Accept is not None:
             headers['Accept'] = Accept
 
+        if CustomAttributes is not None:
+            headers['X-Amzn-SageMaker-Custom-Attributes'] = CustomAttributes
+
         r = self.http.request('POST', url, body=Body, preload_content=False,
                               headers=headers)
 

src/sagemaker/tensorflow/README.rst

@@ -617,7 +617,7 @@ Note that TensorBoard is not supported when passing wait=False to ``fit``.
 Deploying TensorFlow Serving models
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-After a TensorFlow estimator has been fit, it saves a TensorFlow ``SavedModel`` in
+After a TensorFlow estimator has been fit, it saves a TensorFlow SavedModel in
 the S3 location defined by ``output_path``. You can call ``deploy`` on a TensorFlow
 estimator to create a SageMaker Endpoint.
 

src/sagemaker/tensorflow/deploying_python.rst

@@ -94,8 +94,9 @@ The following code adds a prediction request to the previous code example:
 
 The ``predictor.predict`` method call takes one parameter, the input ``data`` for which you want the SageMaker Endpoint
 to provide inference. ``predict`` will serialize the input data, and send it in as request to the SageMaker Endpoint by
-an ``InvokeEndpoint`` SageMaker operation. ``InvokeEndpoint`` operation requests can be made by ``predictor.predict``, by
-boto3 ``sageMaker.runtime`` client or by AWS CLI.
+an ``InvokeEndpoint`` SageMaker operation. ``InvokeEndpoint`` operation requests can be made by ``predictor.predict``,
+by boto3 `SageMakerRuntime <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker-runtime.html>`_
+client or by AWS CLI.
 
 The SageMaker Endpoint web server will process the request, make an inference using the deployed model, and return a response.
 The ``result`` returned by ``predict`` is

src/sagemaker/tensorflow/deploying_tensorflow_serving.rst

@@ -267,7 +267,7 @@ TensorFlow Serving Endpoints allow you to deploy multiple models to the same End
 To use this feature, you will need to:
 
 #. create a multi-model archive file
-#. create a SageMaker and deploy it to an Endpoint
+#. create a SageMaker Model and deploy it to an Endpoint
 #. create Predictor instances that direct requests to a specific model
 
 Creating a multi-model archive file
@@ -386,6 +386,13 @@ additional ``Predictor`` instances. Here's how:
   # get a predictor for 'model2'
   model2_predictor = Predictor(endpoint, model_name='model2')
 
+  # note: that will for actual SageMaker endpoints, but if you are using
+  # local mode you need to create the new Predictor this way:
+  #
+  # model2_predictor = Predictor(endpoint, model_name='model2'
+  #                              sagemaker_session=predictor.sagemaker_session)
+
+
   # result is prediction from 'model2'
   result = model2_predictor.predict(...)