在 Flask 应用程序内调用 Sagemaker 中读取 CustomAttributes

Question

我对我的 Sagemaker 端点进行了以下 lambda 函数调用：

import os
import boto3
from CustomModules.Logger import setlogging

global logger
logger = setlogging()


def lambda_handler(event, context):
    '''
    We use this lambda to call the sagemaker endpoint
    which lives on a different account.
    To do so, we need to assume a different role with boto3.
    '''
    # grab environment variables
    ENDPOINT_NAME = os.environ['ENDPOINT']
    someaccount = 'testaccount'
    runtime = boto3.client('runtime.sagemaker')
    sts_connection = boto3.client('sts')
    acct_b = sts_connection.assume_role(
        RoleArn=someaccount,
        RoleSessionName="somerole"
    )
    ACCESS_KEY = mydict['Credentials']['AccessKeyId']
    SECRET_KEY = mydict['Credentials']['SecretAccessKey']
    SESSION_TOKEN = mydict['Credentials']['SessionToken']

    # once we have all the info, we open the client connection with the new creds
    runtime = boto3.client(
        'runtime.sagemaker',
        aws_access_key_id=ACCESS_KEY,
        aws_secret_access_key=SECRET_KEY,
        aws_session_token=SESSION_TOKEN,
    )
    input_data = event['body']
    client_name = 'myclient'
    res = runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME,
                                  ContentType='application/json',
                                  Body=input_data,
                                  CustomAttributes=client_name,
                                  Accept='Accept'
                                  )
    response = {
        "statusCode": res['ResponseMetadata']['HTTPStatusCode'],
        "headers": res['ResponseMetadata']['HTTPHeaders'],
        "body": res['Body'].read().decode('utf-8')}
    return response


if __name__ == '__main__':
    lambda_handler({''}, {''})

我想访问我在 sagemaker 端点中为

CustomAttributes

设置的值。 我的 Sagemaker 端点使用 Flask 来处理调用：

import flask

app = flask.Flask(__name__)

@app.route('/invocations', methods=['POST'])
def transformation():
    # Get input JSON data and convert it to a DF
    input_json = flask.request.get_json()
    ### How can I get the CustomAttributes value here?

如何在 Flask 应用程序中获取

CustomAttributes

的值？我在这里缺少什么？

Answer 1

SageMaker 将 customAttributes 作为上下文传递。请参阅此线程以了解如何阅读它。 使用 CustomAttributes 调用 sagemaker 端点

Answer 2

我能够访问 post 请求标头中的 customAttributes。 尝试以下代码：

from flask import Flask, jsonify, request

@app.route('/invocations', methods=['POST'])
def invocations():
  custom_attributes = request.headers.get('X-Amzn-Sagemaker-Custom-Attributes')
  print(custom_attributes)
  response = {}
  response['custom-attributes'] = custom_attributes
  return jsonify(response)