Triton JIT MME Sample

Triton on SageMaker - this notebook shows how you can take a RESNET model and create a traced model and leverage the Pytorch back end for Triton

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 is a fully managed service for data science and machine learning workflows. It helps data scientists and developers to prepare, build, train, and deploy high-quality ML models quickly by bringing together a broad set of capabilities purpose-built for ML.

Now, NVIDIA Triton Inference Server can be used to serve models for inference in Amazon SageMaker. Thanks to the new NVIDIA Triton container image, you can easily serve ML models and benefit from the performance optimizations, dynamic batching, and multi-framework support provided by NVIDIA Triton. Triton helps maximize the utilization of GPU and CPU, further lowering the cost of inference.

This notebook was tested on Studio with ml.g4dn.xlarge which comes with 1 GPU and with ml.m5.large which is a CPU based machine only Contents

Introduction to NVIDIA Triton Server
Set up the environment
Basic: PyTorch resnet50 model for Object detection
PyTorch: JIT Trace the model and create a Scripted model
PyTorch: Testing the JIT Traced model 
PyTorch: Packaging model files and uploading to s3
PyTorch: Create SageMaker Endpoint
PyTorch: Run inference
PyTorch: Leverage the Predictions to view the results for Object detection
PyTorch: Terminate endpoint and clean up artifacts

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 GraphDef, TensorFlow SavedModel, ONNX, PyTorch TorchScript, TensorRT, RAPIDS FIL for tree based models, and OpenVINO 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.

Note: This initial release of NVIDIA Triton on SageMaker will only support a single model. Future releases will have multi-model support. A minimal config.pbtxt configuration file is required in the model artifacts. This release doesn't support inferring the model config automatically. Set up the environment

Installs the dependencies required to package the model and run inferences using Triton server.

Also define the IAM role that will give SageMaker access to the model artifacts and the NVIDIA Triton ECR image.

The purpose of this file is to show the ability to take a pytorch computer vision model and create a scripted model which can then be leveraged by Triton using the pytorch back end.

The other option is to build using a python back end but in that we loose some performance gains by compilation to native format

In this example we will take a open source computer vision model and convert into a JIT Traced format and leverage pytorch back end to serve. The model will return 3 tensors for bounded box,

Install some libraries we need to convert the model.

The format is to create a wrapper on top of the model and trace it with dummy inputs to create a traced model

To create the JIT Trace model you need

Torch version 1.12.1 Torch Vision 0.13.1 and Cuda library 11.3

Secondly you need a GPU instance to run the notebook - this has been tested on a ml.g4dn.xlarge which comes with 1 gpu

Create a JIT Traced model

Few points to note: The traced models is provided in the zip file which can be used as is

1. The Model after trace is now returing outputs like OUTPUT_0 , 1 etc
2. To change them to named outputs can be done for now we will leave them as is since it is just a label
3. To full create a jit traced model we will need to provide a sample inputs and hence for now we have created a scripted model
4. The TORCH and the TORCHSCRIPT libraries would need to match the container and hence we use the specific ones mentioned above

Convert the model into Serving mode

Approaches

Approach # 1 Create a stack of the predictions

Approach # 2 can be to ad a dimension along the batch axis

For now we are stack predictions and return back

Approach for doing JIT SCRIPT where in at run time we pass in the required tensors

This has the advantage of no needing to do cuda translations

Use the Open source Model for object detection

We will JIT Trace the model and save the compiled version and use that for the Triton back end

Generate some local predictions - add a batch dimetnion to the 1 image we have

JIT Trace the model

Script could be the other option to create a model Script could be useful to avoid Device conversion

Show the traced results

Create the required Triton config files

Now create the TAR ball and upload

Create the End point

Compare the Model prediction to the original one

We can see All 3 methods match for Object Detection

For any stress test we can use this method below

Optional Stress test the end point

Leverage multi process to stress test the end point

Use a Thread Pool

Use a Multi Process Pool

Clean up

Helper methods as sample to load images in triton format

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.