Run deep learning models on GPUs with Amazon SageMaker endpoints (SME) with ONNX backend on Triton Server

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 provides a scalable and cost-effective way to deploy large number of deep learning models.

Amazon SageMaker Model endpoints with GPU Support

Amazon SageMaker endpoints with GPU works using NVIDIA Triton Inference Server. NVIDIA Triton Inference Server is open-source inference serving software that simplifies the inference serving process and provides high inference performance. Triton supports all major training and inference frameworks, such as TensorFlow, NVIDIA TensorRT, PyTorch, MXNet, Python, ONNX, XGBoost, scikit-learn, RandomForest, OpenVINO, custom C++, and more. It offers dynamic batching, concurrent execution, post-training quantization, optimal model configuration to achieve high performance inference. When SageMaker receives an invocation request for a particular model, it does the following:

How it works?

1. SageMaker routes traffic to the right instance behind the endpoint where the target model is loaded. SageMaker takes care of model management behind the endpoint, loads model to the container's memory and unloads the model based on the endpoint's traffic pattern.
2. SageMaker loads the model to NVIDIA Triton container’s memory on GPU accelerated instance and serve the inference request. If the model is already loaded in the container memory, the subsequent requests are served faster as SageMaker does not need to download and load it again.
3. SageMaker takes care of traffic shaping to the SME endpoint, SageMaker continues to route traffic to the instance where the model is loaded. If the instance resources reach capacity due to high utilization, SageMaker unloads least used models from the container to free up resource to load more frequently used models.
4. SageMaker SME can horizontally scale using auto-scaling policy, provision additional GPU compute instances based on metrics such as GPU utilization, memory utilization etc. to serve spiky traffic to SME endpoints.

In this notebook, we will show you how to use the new features Amazon SageMaker SME with GPU with a computer vision use case. For demonstration purpose, we will use a ResNet-50 convolutional neural network pre-trained model that can classify images into 1000 categories. We will -

• Show how to use NVIDIA Triton inference container on SageMaker SME, leverage different model frameworks such and ONNXRuntime, PyTorch and TensorRT backends.
• Walk you through steps to convert ResNet-50 models to optimized ONNXRuntime backend engine format along with TensorRT engine format and deploy it with SageMaker SME.

Installs

Installs the dependencies required to package the model and run inferences using Triton server. Update SageMaker, boto3, awscli etc

Imports and variables

Creating Model Artifacts

This section presents overview of steps to prepare ResNet-50 pre-trained model to be deployed on SageMaker SME using Triton Inference server model configurations.

Prepare PyTorch Model

generate_model_pytorch.sh file in the workspace directory contains scripts to generate a PyTorch model. First, we load a pre-trained ResNet50 model using torchvision models package. We save the model as model.pt file in TorchScript optimized and serialized format. TorchScript needs an example inputs to do a model forward pass, so we pass one instance of an RGB image with 3 color channels of dimension 224X224. The script for exporting this model can be found here

PyTorch Model Repository

The model repository contains model to serve, in our case it will be the model.pt and configuration file with input/output specifications and metadata.

PyTorch Model configuration

Model configuration file config.pbtxt must specify name of the model(resnet), the platform and backend properties (pytorch_libtorch), max_batch_size(128) and the input and output tensors along with the data type(TYPE_FP32) information. Additionally, you can specify instance_group and dynamic_batching properties to achieve high performance inference.

Prepare TensorRT Model

1. We export the pre-trained ResNet50 model into an ONNX file, which runs the model once to trace its execution and then export the traced model to the specified file. It is one of the better options in terms model conversion and deployment when converting using ONNX.

2. We use trtexec to automatically convert ONNX model to TensorRT plan. As ONNX is framework agnostic it works with models in TF, PyTorch and more. You will export the weights of your model from the framework and load them into your TensorRT network.

In this step, we load pre-trained ResNet50 model from torch and convert to onnx representation using torch onnx exporter. Once onnx model is created, we use TensorRT trtexec command to create the model plan to be hosted with Triton. The script for exporting this model can be found here. This is run as part of the generate_model_trt.sh script from the below cell.

TensorRT Model Repository

The model repository contains model to serve, for TensorRT model it will be the model.plan(created in above steps) and configuration file with input/output specifications and metadata.

TensorRT Model configuration

For the TensorRT model, we specify tensorrt_plan as platform, input tensor specification of the image of dimension 224X224 which has 3 color channels. Output tensor with 1000 dimensions of type TYPE_FP32 corresponding the different object categories.

ONNX Model configuration

For the ONNX model, we specify tensorrt_plan as platform, input tensor specification of the image of dimension 224X224 which has 3 color channels. Output tensor with 1000 dimensions of type TYPE_FP32 corresponding the different object categories.

from https://github.com/triton-inference-server/server/blob/main/docs/examples/model_repository/densenet_onnx/config.pbtxt

Note that the ONNX model is already generated by the script "generate_model_trt.sh" as a substep.

3. Export model artifacts to S3

SageMaker expects the model artifacts in below format, it should also satisfy Triton container requirements such as model name, version, config.pbtxt files etc. tar the folder containing the model file as model.tar.gz and upload it to s3

Now that we have uploaded the model artifacts to S3, we can create a SageMaker endpoint.

Deploy Models with SME

We will now deploy ResNet-50 model with ONNX framework backends.

We will use AWS SDK for Python (Boto) APIs create_model, create_endpoint_config and create_endpoint to create a mulit-model endpoint.

Create a model object

Using the SageMaker boto3 client, create the model using create_model API. We will pass the container definition to the create model API along with ModelName and ExecutionRoleArn.

Define configuration for the endpoint

Create Single 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 InService once the deployment is successful.

Setup Autoscaling policies for GPU endpoint

Amazon SageMaker inference endpoints supports automatic scaling (auto scaling) for your hosted models. Auto scaling dynamically adjusts the number of instances provisioned for a model in response to changes in your workload. When the workload increases, auto scaling brings more instances online. When the workload decreases, auto scaling removes unnecessary instances so that you don't pay for provisioned instances that you aren't using.

In the below scaling policy, use a custom metric GPUUtilization in TargetTrackingScalingPolicyConfiguration configuration and set a TargetValue of 60.0 for the target value of that metric. This autoscaling policy will provision additional instances upto MaxCapacity when GPU Utilization is more than 60%.

Prepare Input Payload for PyTorch and TensorRT model

The following method transforms a sample image we will be using for inference into the payload that can be sent for inference to the Triton server.

The tritonclient package provides utility methods to generate the payload without having to know the details of the specification. We'll use the following methods to convert our inference request into a binary format which provides lower latencies for inference.

Invoke target model on Sagemaker Endpoint

Once the endpoint is successfully created, we can send inference request to endpoint using invoke_endpoint API. We specify the TargetModel in the invocation call and pass in the payload for each model type. Sample invocation for PyTorch model and TensorRT model is shown below

TensorRT model prediction

ONNX model prediction

Terminate endpoint and clean up artifacts

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.