API Gateway as a Proxy for SageMaker Endpoints

Question

Pulumi · Accepted Answer

To create an API Gateway as a proxy for SageMaker endpoints, you'd typically create an AWS API Gateway REST API and configure it with an integration to the SageMaker endpoint. I'll guide you through setting up this infrastructure using Pulumi in Python.

Here's a high-level overview of what we'll do:

Define a REST API with AWS API Gateway - This will be the front door for our clients to send requests.
Define an Integration resource - This tells the API Gateway how to send the request to the SageMaker endpoint.
Define a Method for that Integration - This configures the actual HTTP method (GET, POST, etc.) that clients will use.
Define a Deployment and a Stage for the REST API - This makes the API accessible over the internet.
Attach the necessary permissions - Ensure that API Gateway has permissions to invoke the SageMaker endpoint.

Below is the Pulumi program that provisions the above resources:

import pulumi
import pulumi_aws as aws

# This is the ARN of your existing SageMaker endpoint. You would replace this with your actual SageMaker endpoint ARN.
sagemaker_endpoint_arn = "arn:aws:sagemaker:<region>:<account>:endpoint/<endpoint-name>"

# Define the AWS API Gateway REST API
api = aws.apigateway.RestApi("apiGatewayProxy",
    description="API Gateway to proxy requests to SageMaker Endpoint",
    # Reference to the documentation for pulumi_aws.apigateway.RestApi:
    # https://www.pulumi.com/registry/packages/aws/api-docs/apigateway/restapi/
)

# Define the integration between the API Gateway and the SageMaker endpoint
integration = aws.apigateway.Integration("apiGatewayIntegration",
    rest_api=api.id,
    resource_id=api.root_resource_id,
    http_method="POST",
    integration_http_method="POST",
    type="AWS_PROXY",
    uri=pulumi.Output.concat("arn:aws:apigateway:", pulumi_aws.get_region().name, ":sagemaker:path//endpoints/",
                             sagemaker_endpoint_arn.split(":")[-1], "/invocations"),
    # Reference to the documentation for pulumi_aws.apigateway.Integration:
    # https://www.pulumi.com/registry/packages/aws/api-docs/apigateway/integration/
)

# Define the method that clients will use to communicate with your endpoint.
# In this case, we're assuming clients will POST to the API.
method = aws.apigateway.Method("apiGatewayMethod",
    rest_api=api.id,
    resource_id=api.root_resource_id,
    http_method="POST",
    authorization="NONE",
    integration=integration.id,
    # Reference to the documentation for pulumi_aws.apigateway.Method:
    # https://www.pulumi.com/registry/packages/aws/api-docs/apigateway/method/
)

# Deploy the REST API to make it publicly accessible
deployment = aws.apigateway.Deployment("apiGatewayDeployment",
    rest_api=api.id,
    stage_name="prod",
    # The following attribute provides a way to redeploy the api when the swagger configuration changes.
    triggers={
        "redeployment": pulumi.Output.concat(
            "Integration ID: ", integration.id
        ),
    },
    # Reference to the documentation for pulumi_aws.apigateway.Deployment:
    # https://www.pulumi.com/registry/packages/aws/api-docs/apigateway/deployment/
)

# Export the endpoint URL for easy access
pulumi.export("endpoint_url", deployment.invoke_url.apply(lambda url: url + "prod"))

Explanation:

AWS API Gateway REST API: This component defines the REST API in AWS API Gateway. We give it a name (apiGatewayProxy) and a brief description. It acts as the gateway for all the requests that want to reach the SageMaker endpoint.
Integration: Here, the API Gateway needs to know how to interact with the SageMaker endpoint. We define an integration (apiGatewayIntegration) where we specify the HTTP method and the type, which, in this case, is AWS_PROXY. The URI is composed dynamically to point to the SageMaker endpoint invocation URL.
Method: Clients use HTTP methods to communicate with the API. Here, we declare a POST method (apiGatewayMethod), assuming clients will be sending a POST request to some data for inferencing. We also link it to the integration defined above.
Deployment: To make the API accessible to clients, we need to deploy it (apiGatewayDeployment). We choose a stage name prod to denote a production environment.
Permissions: In this example, permissions are not set explicitly, but you would need to make sure that API Gateway has the necessary permissions attached to its role to invoke the SageMaker endpoint.

Finally, we export the URL of the deployed endpoint so that you can easily access your new API Gateway from anywhere.

Ensure that the AWS credentials have been set up correctly for Pulumi to deploy the resources to your AWS account. Once you run this Pulumi program, it will provision the necessary resources on AWS, and you'll get the URL of the API Gateway endpoint which acts as a proxy for your SageMaker endpoints.