Triton on SageMaker - this notebook shows how you can take a RoBERTA 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: RoBERTA Model

• 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

Start RoBERTA Base for Triton

Run for Triton server

Note: Amazon SageMaker expects the model tarball file to have a top level directory with the same name as the model defined in the config.pbtxt. Below is the sample model directory structure

roberta-large
├── 1
│   └── model.pt
└── config.pbtxt

Have to use the same Tokenizer to generate the input to test as BERT uncased

Create the RoBERTA Model in Torch Script mode -- .pt model

use the ore trained and use torchscript flag here

Run a simple test for RoBERTA base

* We run multiple tests
    * First we token ize and then de tokenize to make sure the vaues match
    * Then we use the model and run predictions to get values
    * Then we run on the traced Model and run predictions to get values 
    * Check to make sure they match

Prepare some dummy inputs for tracing

Test encoders various methods

Test the HuggingFace and then the scripted model locally

Now test the Scripted model -- Scripted model gives us tensors back

Upload the Model.tar after it has been created correctly by the above scripted and the config.pbtxt files

Upload the model.tar.gz to S3 location

Start Single Model Triton for starting

Triton Image download and sagemaker variables

Model creation

Endpoint config

Endpoint

Now Invoke The endpoint

Run the JSON invocation

Invoke using the Binary Format

Stress Test it

Clean up

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.