Implement a SageMaker Multi-Model Endpoint for TensorFlow Vision models on a Triton Server from NVIDIA

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.

Amazon SageMaker Multi-Model Endpoint (MME) is a cost-effective way of running multiple models behind a single endpoint. SageMaker manages the process of loading the target model into memory when needed, which leads to better utilization of the container resources and reduces cost.

Multi-model endpoints are ideal when you have infrequently used models that can handle minor delays introduced by an occasional cold start.

NVIDIA Triton Inference Server is an open source software that provides high performance inference on a wide variety of CPU and GPU hardware and supports all the major ML frameworks. It has many built-in features to improve inference throughput and achieves better utilization of the resources.

Now the NVIDIA Triton Inference Server can be deployed on GPU based SageMaker ML instances. It supports the SageMaker MME API to for dynamic loading and unloading of models for implementing SageMaker multi-model endpoints.

This notebook shows how to deploy multiple TensorFlow models trained on the MNIST dataset to a SageMaker MME using the NVIDIA Triton Server.

Here we use two different instances of an existing model artifact. The model used here was pre-trained on the MNIST dataset. If you want to learn how to train the model, please See TensorFlow script mode training and serving.

Contents

1. Introduction to NVIDIA Triton Server
2. Set up the environment
3. Transform TensorFlow Model structure
   1. Inspect the model using a CLI command
   2. Create the model configuration file
   3. Create the tar ball in the required Triton structure
   4. Upload model artifact to S3
   5. Create additional instances of the model in S3 for testing MME
4. Deploy model to SageMaker Triton Server MME
5. Test the SageMaker Triton Server MME
6. Clean up

Introduction to NVIDIA Triton Server

NVIDIA Triton Inference Server was developed specifically to enable scalable, cost-effective, and easy deployment of models in production. NVIDIA Triton Inference Server is open-source inference serving software that simplifies the inference serving process and provides high inference performance.

Some key features of Triton are:

• Support for Multiple frameworks: Triton can be used to deploy models from all major frameworks. Triton supports TensorFlow, ONNX, PyTorch, and many other model formats.
• Model pipelines: Triton model ensemble represents a pipeline of one or more models or pre- / post-processing logic and the connection of input and output tensors between them. A single inference request to an ensemble will trigger the execution of the entire pipeline.
• Concurrent model execution: Multiple models (or multiple instances of the same model) can run simultaneously on the same GPU or on multiple GPUs for different model management needs.
• Dynamic batching: For models that support batching, Triton has multiple built-in scheduling and batching algorithms that combine individual inference requests together to improve inference throughput. These scheduling and batching decisions are transparent to the client requesting inference.
• Diverse CPUs and GPUs: The models can be executed on CPUs or GPUs for maximum flexibility and to support heterogeneous computing requirements.

Set up the environment

This notebook uses the Python 3 (Data Science) kernel.

Install TensorFlow. This notebook is tested with version 2.11.

For this exercise we download a TensorFlow model pre-trained on the MNIST data set from an Amazon S3 bucket. The model artifact is saved locally.

You should have already configured the default IAM role for running this notebook with access to the model artifacts and the NVIDIA Triton Server image in Amazon Elastic Container Registry (ECR).

Download the Triton Server image from Amazon ECR.

Transform TensorFlow Model structure

The model that we want to deploy currently has the following structure:

00000000
        ├── saved_model.pb
        ├── assets/
        └── variables/
            ├── variables.data-00000-of-00001
            └── variables.index

For Triton, the model needs to have the following structure:

<model-name>
├── config.pbtxt
└── 1
    └── model.savedmodel
        ├── saved_model.pb
        ├── assets/
        └── variables/
            ├── variables.data-00000-of-00001
            └── variables.index
            

Inspect the model using a CLI command

In order to create the config.pbtxt we need to confirm the model inputs and outputs (Signature). We use a CLI command to inspect the model and take note of the input and output shape.

Create the config.pbtxt file

Triton requires a Model Configuration file called config.pbtxt.

We create one below in the local folder for uploading with the model artifact.

Create a tar ball of the model in the required folder structure for the Triton Server

Upload model artifact to S3.

Create additional instances of the model in S3 for Inference using SageMaker MME

For testing the MME we create two copies of the compressed model artifact under different names in the S3 folder. In the cell below, you can specify how many model instances to create.

Deploy TensorFlow models to a Multi-Model Endpoint for Triton Server

Define the serving container

In the container definition below, we need to pass in the following parameters.

• Image: Triton server image URI that supports deploying multi-model endpoints with GPUs.
• URI to S3 folder that contains all the models that SageMaker multi-model endpoint will use to load and serve predictions.
• Mode: Set to MultiModel

Create a model object using the container defined above

Create the model object using the Boto3 create_model API. We pass the container definition to the create model API along with the model name and execution role.

Deploy and test the Multi-Model endpoint

Create a multi-model endpoint configurations using the create_endpoint_config Boto3 API. We specify an accelerated GPU computing instance as the instance type. For testing we specify a single instance. In real scenarios we recommend the value of initial instance count to be two or higher for high availability.

Create endpoint configuration

Create Multi-Model endpoint

Using the above endpoint configuration we create a new SageMaker endpoint and wait for the deployment to finish. The status will change to In Service once the deployment is successful.

Invoke target models behind the Multi-Model endpoint

Once the endpoint is successfully created, we can send inference requests to the multi-model endpoint using invoke_endpoint API. We specify the target model in the invocation call and pass in the payload for each model type.

Let's download some test data

Clean up

We strongly recommend deleting the real-time endpoint created to stop incurring cost when finished with the example.

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.