Deploy a TensorFlow Model using NVIDIA Triton on SageMaker

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 example will showcase how to deploy a pre-trained TensorFlow model using NVIDIA Triton on SageMaker.

The model used here was pre-trained on the MNIST dataset. See this Deploy a Trained TensorFlow V2 Model example for the training of the model.

Contents

1. Introduction to NVIDIA Triton Server
2. Set up the environment
3. Transform TensorFlow Model structure
4. Inspect the model using the saved_model_cli
5. Create the config.pbtxt
6. Create the tar ball in the required Triton structure
7. Deploy model to SageMaker Endpoint
8. 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 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

Download the pre-trained TensorFlow model from a public S3 bucket. Also define the IAM role that will give SageMaker access to the model artifacts and the NVIDIA Triton ECR image.

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 the saved_model_cli

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

Create the config.pbtxt

Triton requires a Model Configuration file known as a config.pbtxt. We create one below in the correct directory.

The name in the config.pbtxt must match the name of our model directory. In this case we will use MNIST.

Create the tar ball in the required Triton structure

Upload the new tar ball containing the Triton model structure to s3

Deploy model to SageMaker Endpoint

We start off by creating a sagemaker model from the model files we uploaded to s3 in the previous step.

In this step we also provide an additional Environment Variable i.e. SAGEMAKER_TRITON_DEFAULT_MODEL_NAME which specifies the name of the model to be loaded by Triton. The value of this key should match the folder name in the model package uploaded to s3. This variable is optional in case of a single model. In case of ensemble models, this key has to be specified for Triton to startup in SageMaker.

Additionally, customers can set SAGEMAKER_TRITON_BUFFER_MANAGER_THREAD_COUNT and SAGEMAKER_TRITON_THREAD_COUNT for optimizing the thread counts.

Clean up

We strongly recommend to delete 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.