使用 CustomAttributes 调用 sagemaker 端点

我正在尝试调用我的 SageMaker 端点并传递指定的 CustomAttributes 参数here https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker-runtime.html#SageMakerRuntime.Client.invoke_endpoint.

我想知道的是如何检索模型端点中的 CustomAttributes？

我为端点创建了一个 inference.py 文件，其结构如下：

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)

正如 request_body、content_type、accept 被传递给transform_fn 一样，我想传递 CustomAttributes。这可能吗？如果可以我该怎么做？

提前致谢！

通过完全控制所有底层步骤，可以方便地构建推理代码。

正如文档“调整您自己的推理容器 https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-inference-container.html”建议，您可以安排 4 个功能：model_fn, input_fn, predict_fn and output_fn.

除此之外，您还可以创建自己的处理程序来处理调用端点时传递的所有属性（“如何实现预处理和/或后处理处理程序 https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/deploying_tensorflow_serving.html#how-to-implement-the-pre-and-or-post-processing-handler-s").

您可以拥有input_handler / output_handler pair的函数或single handler功能。

以下是通用推理脚本的代码示例。

在端点调用块内：

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
)

在你的 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

请注意此注释：

请注意，如果实现了处理函数，则 input_handler 和 output_handler 被忽略。

这意味着，如果您使用 TensorFlow 或 PyTorch 等框架，您将必须了解如何从默认处理程序开始重写这些方法。