Invoking sagemaker endpoint with CustomAttributes - amazon-web-services

Question

I am trying to invoke my SageMaker endpoint and pass CustomAttributes argument specified here.
what I want to know is how to retrieve the CustomAttributes in the model endpoint?
I created an inference.py file for my endpoint that has the following structure:
imports
def get_device():
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
return device
def model_fn(model_dir):
return model
def transform_fn(model, request_body, content_type, accept)
return json.dumps(predictions)
just as request_body, content_type, accept are passed to transform_fn, i want to pass the CustomAttributes. Is this possible and if so how can i do it ?
Thanks in advance!

Answer 1

It is convenient to structure the inference code by having full control of all the underlying steps.
As the documentation "Adapting Your Own Inference Container" suggests, you can arrange 4 functions: model_fn, input_fn, predict_fn and output_fn.
Beyond these, you can create your own handler to handle all the attributes you pass when you invoke an endpoint ("How to implement the pre- and/or post-processing handler(s)").
You can have the input_handler / output_handler pair of functions or a single handler function.
Below is an example of code for a generic inference script.
Inside your block of endpoint's invocation:
import boto3
import json
runtime = boto3.Session().client('sagemaker-runtime')
runtime_client.invoke_endpoint(
EndpointName = your_endpoint_name,
Body = your_data,
CustomAttributes = json.dumps(your_attributes_dict),
ContentType = your_content_type
)
and inside your inference.py:
import json
def handler(data, context):
processed_input = _process_input(data, context)
custom_attrs = json.loads(context.custom_attributes)
# here place your function to parse and use your custom_attrs json
response = requests.post(context.rest_uri, data=processed_input)
return _process_output(response, context)
def _process_input(data, context):
# your _process_input to decode the request_content_type
if context.request_content_type == YOUR_CONTEXT_TYPE:
return your_process_func(data)
raise ValueError('{{"error": "unsupported content type {}"}}'.format(
context.request_content_type or "unknown"))
def _process_output(data, context):
if data.status_code != 200:
raise ValueError(data.content.decode('utf-8'))
response_content_type = context.accept_header
prediction = data.content
return prediction, response_content_type
Pay attention to this note:
Note that if handler function is implemented, input_handler and
output_handler are ignored.
It means that if you use a framework like TensorFlow or PyTorch, you will have to see how to override these methods starting from their default handlers.