Run Multiple CV ResNet Pytorch Models on GPU with Amazon SageMaker Multi-Model Endpoints (MME)

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 endpoints(MME) provide a scalable and cost-effective way to deploy large number of deep learning models. Previously, customers had limited options to deploy 100s of deep learning models that need accelerated compute with GPUs. Now customers can deploy 1000s of deep learning models behind one SageMaker endpoint. Now, MME will run multiple models on a GPU core, share GPU instances behind an endpoint across multiple models and dynamically load/unload models based on the incoming traffic. With this, customers can significantly save cost and achieve best price performance.

Amazon SageMaker Multi-Model endpoints with GPU Support

Amazon SageMaker multi-model 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. Dynamically loads models from Amazon Simple Storage Service(S3) to the instance’s storage volume. If the invoked models is not available on instance storage volume, the model is downloaded onto instance storage volume. If the instance storage volume reaches capacity, SageMaker deletes any unused models from the storage volume.
3. 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 subsequents requests are served faster as SageMaker does not need to download and load it again.
4. SageMaker takes care of traffic shaping to the MME endpoint, SageMaker continues to routes traffics 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.
5. SageMaker MME 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 MME endpoints.

In this notebook, we will show you how to use the new features Amazon SageMaker MME 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 MME, leverage different model frameworks such as PyTorch.
• Show how to get insights into instance and invocation metrics using Amazon CloudWatch.

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 MME 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 a RGB image with 3 color channels of dimension 224X224. The script for exporting this model can be found here

PyTorch Model Respository

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.

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 multi-model endpoint.

Deploy Models with MME

We will now deploy ResNet-50 model with PyTorch framework to SageMaker MME. You can reproduce all the steps using step by step notebook on GitHub.

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

Define the serving container

In the container definition, define the ModelDataUrl to specify the S3 directory that contains all the models that SageMaker multi-model endpoint will use to load and serve predictions. Set Mode to MultiModel to indicates SageMaker would create the endpoint with MME container specifications. We set the container with an image that supports deploying multi-model endpoints with GPU, see MME container images for more details.

Create a multi 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 multi model endpoint

Create a multi-model endpoint configurations using create_endpoint_config boto3 API. Specify an accelerated GPU computing instance in InstanceType, in this post we will use g4dn.4xlarge instance. We recommend configuring your endpoints with at least two instances. This allows SageMaker to provide a highly available set of predictions across multiple Availability Zones for the models.

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

Setup Autoscaling policies for GPU Multi Model Endpoint

Amazon SageMaker multi-model 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 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 Multi Model Endpoint

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

PyTorch Model prediction

We can also use binary+json as the payload format to get better performance for the inference call. The specification of this format is provided here.

Note: With the binary+json format, we have to specify the length of the request metadata in the header to allow Triton to correctly parse the binary payload. This is done using a custom Content-Type header application/vnd.sagemaker-triton.binary+json;json-header-size={}.

Please not, this is different from using Inference-Header-Content-Length header on a stand-alone Triton server since custom headers are not allowed in SageMaker.

Cloudwatch metrics for GPU Multi Model Endpoints

Amazon SageMaker multi-model endpoints provides instance level metrics to monitor, for more details refer Monitor Amazon SageMaker with Amazon CloudWatch

• Number of models loaded in the containers (LoadedModelCount),
• Precentage of GPU units that are used by the containers (GPUUtilization),
• Precentage of GPU memory used by the containers (GPUMemoryUtilization),
• Precentage of disk space used by the containers (DiskUtilization) etc.

SageMaker MME also provides Model loading metrics such as-

• Time interval for model to be downloaded or loaded (ModelLoadingWaitTime),
• Time interval to unload model from container (ModelUnloadingTime),
• Time to download the model from S3 (ModelDownloadingTime),
• Number of invocations to model that are already loaded onto the container(ModelCacheHit) etc to get model invocation level insights.

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.