Using SageMaker models in Baseten

Since Baseten launched, most of our users have been deploying machine learning models through our model deployment service because of its ease-of-use, extensibility, and scalability. However we’ve heard from some people that they are already locked into SageMaker for model deployment, yet still want to use Baseten to create user-facing applications for their SageMaker models. Today we’re making that possible by launching a native SageMaker integration in Baseten.

As a reminder, Baseten makes it easy to build applications with machine learning models. Our mission is to bring models to business value quickly, wherever they are, and it is only natural to support integrations with other model serving platforms. The SageMaker integration is the first step in that direction.

Imported SageMaker models behave just like Baseten models. You’ll be able to use them as a node in Worklets, and through Worklets make them accessible to your UI views. The integration increases the pool of models you can work with in Baseten, and as we continue to build to support other model serving platforms, this pool will only grow.

The SageMaker integration is available to all users in our private beta.

Configuring the SageMaker connection

To import your SageMaker models into Baseten you first need to create a SageMaker connection:

• Click on your User icon at the bottom of the left sidebar, then click on Data connections, and then click on the New data connection button
• Choose SageMaker as the Data Connection type
• Enter your Access Key Id and Secret access key
• Enter the Region where the SageMaker model is deployed
• Test the connection to see if the connection was successful

Importing a model from SageMaker

Once you’ve created the connection, you can bring your SageMaker models into Baseten using the Import model flow. In this example, we’ve deployed a model from SageMaker's jumpstart.

Baseten lists all the models that are available to import with the specified connection. You can choose the model you want to bring in, the endpoint, and the desired input format (with today’s release, only the models connected to endpoints are available for invocation). Baseten supports the two most common input formats: CSV and JSON. If your model doesn’t support CSV or JSON then you can still use them through Python nodes.

Once imported, the model will show up on the Models page and you’ll be able to use it in Baseten. As with other Baseten models, you can add pre-processing and post-processing nodes before and after your model.

Using SageMaker models in Baseten

You can use your imported SageMaker models in Baseten worklets using the Invoke Model node like you would any other model. You can test the node by giving it input as JSON; JSON is automatically converted to CSV if your model needs that.

SageMaker access through a code node

We find the Invoke Model node to be much more user-friendly, but if you need lower level access to the SageMaker API, you can use the SageMaker client directly from the Python code in code nodes. Baseten provides a Boto3 SageMakerRuntime client, through the context argument, to your code node. You can use this client to perform many operations, including invoking endpoints.

Here’s an example of how to invoke an endpoint using the client in your code node, using the blazingtext model from this tutorial.

<script>
import json
import codecs

def do(node_input, env, context):
   sagemaker_client = context.client('My SageMaker Connection')
   model_input = {
       'instances': [
           'Convair was an american aircraft manufacturing company.',
           'Berwick secondary college is situated in a suburb of berwick.',
       ]
   }
   endpoint_response_raw = sagemaker_client.invoke_endpoint(
       EndpointName='blazingtext-2021-08-12-17-57-05-399',
       Body=json.dumps(model_input),
       ContentType='application/json',
       )
   response = json.load(codecs.getreader(‘utf-8’)(endpoint_response_raw['Body']))
   return response
</script>

What’s next

We’re continuously adding integrations to Baseten to accelerate the impact of machine learning with real user-centric use cases. Being able to use Baseten's application-building framework with SageMaker models is a step in that direction. We’ll soon be adding native integrations with other model serving platforms. We’d love to hear your thoughts on the SageMaker integration and what you’d like to see next.