Example: Integrate an Amazon SageMaker model

The below documentation provides an example configuration and model adapter for a custom connection to a model hosted in Amazon SageMaker. Review the benefits of external model integration to make sure this is the right fit for your use case.

For a step-by-step guide, refer to the documentation on how to create a model adapter and how to create a connection to an externally hosted model.

Example Amazon SageMaker tabular model adapter

First, publish and tag a model adapter using the model adapter library in the Code Repositories application. The below model adapter configures a connection to a model hosted in Amazon SageMaker using the AWS SDK for Python (Boto3) ↗ and framework. The below code was tested with versions Python 3.8.17, boto3 1.28.1 and pandas 1.5.3.

Note that this model adapter makes the following assumptions:

• This model adapter assumes that data being provided to this model is tabular.
• This model adapter will serialize the input data to JSON and send this data to the hosted Amazon SageMaker model.
• This model adapter assumes that the response is deserializable from JSON to a pandas dataframe
• This model adapter takes four inputs to construct a Boto3 client.
  • region_name - Provided as connection configuration
  • endpoint_name - Provided as connection configuration
  • access_key_id - Provided as credentials
  • secret_access_key - Provided as credentials

Copied!
1import palantir_models as pm
2import models_api.models_api_executable as executable_api
3
4import boto3
5import json
6import pandas as pd
7import logging
8from typing import Optional
9from botocore.exceptions import ClientError
10
11logger = logging.getLogger(__name__)
12
13
14class SageMakerTabularAdapter(pm.ExternalModelAdapter):
15    """
16    :display-name: SageMaker Tabular Model Adapter
17    :description: Default model adapter for SageMaker models that expect tabular input and output tabular data.
18    """
19
20    def __init__(self, region_name, endpoint_name, access_key_id, secret_access_key):
21        self.endpoint_name = endpoint_name
22        self.runtime = boto3.client(
23            'runtime.sagemaker',
24            aws_access_key_id=access_key_id,
25            aws_secret_access_key=secret_access_key,
26            region_name=region_name
27        )
28
29    @classmethod
30    def init_external(cls, external_context) -> "pm.ExternalModelAdapter":
31        region_name = external_context.connection_config["region_name"]
32        endpoint_name = external_context.connection_config["endpoint_name"]
33        access_key_id = external_context.resolved_credentials["access_key_id"]
34        secret_access_key = external_context.resolved_credentials["secret_access_key"]
35        return cls(
36            region_name,
37            endpoint_name,
38            access_key_id,
39            secret_access_key
40        )
41
42    @classmethod
43    def api(cls):
44        inputs = {"df_in": pm.Pandas()}
45        outputs = {"df_out": pm.Pandas()}
46        return inputs, outputs
47
48    def predict(self, df_in):
49        payload = {
50            "instances": df_in.apply(lambda row: {"features": row.tolist()}, axis=1).tolist()
51        }
52        try:
53            response = self.runtime.invoke_endpoint(
54                EndpointName=self.endpoint_name,
55                ContentType="application/json",
56                Body=json.dumps(payload)
57            )
58        except ClientError as error:
59            logger.error("SageMaker inference call failed. This can indicate an error with this Model's egress "
60                         "policy. Double check your configured egress policy and ensure the remote endpoint is still "
61                         "available.")
62            raise error
63        try:
64            # Output from model is assumed to be json serializable
65            # if result is too large for executor this deserialization may cause an OOM
66            result = json.loads(response['Body'].read().decode())
67        except ValueError as error:
68            logger.error("This SageMakerTabularAdapter expects results to be json serializable.")
69            raise error
70        return pd.json_normalize(result)

Amazon SageMaker tabular model configuration

Next, configure an externally hosted model to use this model adapter and provide the required configuration and credentials as expected by the model adapter. In this example, the model is assumed to be hosted in us-east-1, but this is configurable.

Note that the URL is not required by the above SageMakerTabularAdapter and so is left blank; however, the configuration and credentials maps are completed using the same keys as defined in the Model Adapter.

Select an egress policy

The below uses an egress policy that has been configured for runtime.sagemaker.us-east-1.amazonaws.com (Port 443).

Configure model adapter

Choose the published model adapter in the Connect an externally hosted model dialog.

Configure connection configuration

Define connection configurations as required by the example Amazon SageMaker tabular model adapter.

This adapter requires connection configuration of:

• region_name - The AWS region name where the model is hosted.
• endpoint_name - The unique identifier for the externally hosted model.

Configure credential configuration

Define credential configurations as required by the example Amazon SageMaker tabular model adapter.

This adapter requires credential configuration of:

• access_key_id - This is the unique identifier for the user whose credentials will call the SageMaker model.
• secret_access_key - This is the secret key for the user whose credentials will call the SageMaker model.

Amazon SageMaker tabular model usage

Now that the Amazon SageMaker model has been configured, this model can be hosted in a live deployment or Python transform.

The below image shows an example query made to the Amazon SageMaker model in a live deployment.