Amazon SageMaker Multi-Model Endpoints using TorchServe

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Contents

With Amazon SageMaker multi-model endpoints, customers can create an endpoint that seamlessly hosts up to thousands of models. These endpoints are well suited to use cases where any one of many models, which can be served from a common inference container, needs to be called on-demand and where it is acceptable for infrequently invoked models to incur some additional latency. For applications which require consistently low inference latency, a traditional endpoint is still the best choice.

At a high level, Amazon SageMaker manages the loading and unloading of models for a multi-model endpoint, as they are needed. When an invocation request is made for a particular model, Amazon SageMaker routes the request to an instance assigned to that model, downloads the model artifacts from S3 onto that instance, and initiates loading of the model into the memory of the container. As soon as the loading is complete, Amazon SageMaker performs the requested invocation and returns the result. If the model is already loaded in memory on the selected instance, the downloading and loading steps are skipped, and the invocation is performed immediately.

This notebook uses SageMaker notebook instance conda_python3 kernel, demonstrates how to use TorchServe on SageMaker MME. In this example, there are 3 distinct models, each with its own set of dependencies, handler implementation and model configuration.

Create Model Artifacts

This example creates a TorchServe model artifact for each model.

Install torch-model-archiver

Model 1: Segment Anything Model(SAM)

A new AI model from Meta that can segment any object in any image with a single click. No additional training needed. We are downloading one of the checkpoints

Download Segment Anything Model(SAM)

Implement customized handler

This step can be skipped if your model uses TorchServe default handler. Here we follow TorchServe instruction to create a customized handler for this model.

Config model

Custom dependencies

Create and upload sam.tar.gz file

Model 2: Stable Diffusion In Paint (SD)

Import and Save Stable Diffusion Model

Implement customized handler

This step can be skipped if your model uses TorchServe default handler. Here we follow TorchServe instruction to create a customized handler for this model.

Config model

Custom dependencies

Create sd.tar.gz file

Model 3: Large Mask In Painting Model (Lama)

Download Pre-Trained Model

Clone Lama Repo

Implement customized handler

This step can be skipped if your model uses TorchServe default handler. Here we follow TorchServe instruction to create a customized handler for this model.

Config model

Custom dependencies

Create lama.tar.gz file

Create the Multi-Model Endpoint with the SageMaker SDK

Create the Amazon SageMaker MultiDataModel entity

We create the multi-model endpoint using the MultiDataModel class.

You can create a MultiDataModel by directly passing in a sagemaker.model.Model object - in which case, the Endpoint will inherit information about the image to use, as well as any environmental variables, network isolation, etc., once the MultiDataModel is deployed.

In addition, a MultiDataModel can also be created without explicitly passing a sagemaker.model.Model object. Please refer to the documentation for additional details.

Deploy the Multi-Model Endpoint

You need to consider the appropriate instance type and number of instances for the projected prediction workload across all the models you plan to host behind your multi-model endpoint. The number and size of the individual models will also drive memory requirements.

Our endpoint has launched! Let's look at what models are available to the endpoint!

By 'available', what we mean is, what model artifacts are currently stored under the S3 prefix we defined when setting up the MultiDataModel above i.e. model_data_prefix.

Currently, since we only have one artifact (i.e. sam.tar.gz files) stored under our defined S3 prefix.

Dynamically deploying models to the endpoint

The .add_model() method of the MultiDataModel will copy over our model artifacts from where they were initially stored, by training, to where our endpoint will source model artifacts for inference requests.

Note that we can continue using this method, as shown below, to dynamically deploy more models to our live endpoint as required!

model_data_source refers to the location of our model artifact (i.e. where it was deposited on S3 after training completed)

model_data_path is the relative path to the S3 prefix we specified above (i.e. model_data_prefix) where our endpoint will source models for inference requests. Since this is a relative path, we can simply pass the name of what we wish to call the model artifact at inference time.

Our models are ready to invoke!

We can see that the S3 prefix we specified when setting up MultiDataModel now has model artifacts listed. As such, the endpoint can now serve up inference requests for these models.

Get predictions from the endpoint

Recall that mme.deploy() returns a Real Time Predictor that we saved in a variable called predictor.

That predictor can now be used as usual to request inference - but specifying which model to call:

Model Segment Anything Inference Request

Model Stable Diffusion In Paint Inference Request

Large Mask In Painting Model Inference Request

Updating a model

To update a model, you would follow the same approach as above and add it as a new model. For example, ModelA-2.

You should avoid overwriting model artifacts in Amazon S3, because the old version of the model might still be loaded in the endpoint's running container(s) or on the storage volume of instances on the endpoint: This would lead invocations to still use the old version of the model.

Alternatively, you could stop the endpoint and re-deploy a fresh set of models.

Clean up

Endpoints should be deleted when no longer in use, since (per the SageMaker pricing page) they're billed by time deployed. Here we'll also delete the endpoint configuration - to keep things tidy.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.