Triton TensorRT Sentence Transformer

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.

Prerequisites

Install the necessary Python modules to use and interact with NVIDIA Triton Inference Server.

Part 1 - Setup

Part 2 - Generate TensorRT Model

In the following cells, we are using HuggingFace Auto Classes to load a pre-trained model from the HuggingFace Model Hub. We then convert the model to the ONNX format, and compile it using NVIDIA TensorRT - namely its command-line wrapper tool, trtexec -, using the scripts provided in the official AWS Sample for SageMaker Triton.

NVIDIA TensorRT is an SDK that facilitates high-performance machine learning inference. You can use it to create engines from models that have already been trained, optimizing for a selected GPU architecture. Triton natively supports the TensorRT runtime, which enables you to easily deploy a TensorRT engine and pair it with the rich features that Triton provides.

Parameters:

• model_name: Model identifier from the Hugging Face model hub library

Option 1 - TensorRT Model with Amazon SageMaker Studio

WARNING: The next cell will only work if you have first created a custom Studio image, described in Step 2 of this repository's README. Change the RUNNING_IN_STUDIO to True if this is the case.

Option 2 - TensorRT Model with SageMaker Notebook Instances

To make sure we use TensorRT version and dependencies that are compatible with the ones in our Triton container, we compile the model using the corresponding version of NVIDIA's PyTorch container image.

If you take a look at the python files within the workspace folder, you will see that we are first convert the model into ONNX format, specifying dynamic axis indexes so that inputs with a different batch size and sequence length can be passed to the model. TensorRT will treat other input dimensions as fixed, and optimize for those.

Explore the output logs of the compilation process; at the very end, we get a section headlined "=== Performance summary ===" which gives us a series of metrics on the obtained engine's performance (latency, throughput, etc...).

Part 3 - Run Local Triton Inference Server

When using the Triton Python backend, you can include your own environment and dependencies. The recommended way to do this is to use conda pack to generate a conda environment archive in tar.gz format, include it in your model repository, and point to it in the config.pbtxt file of python models that should use it, adding the snippet:

parameters: {
  key: "EXECUTION_ENV_PATH",
  value: {string_value: "$$TRITON_MODEL_DIRECTORY/your_env.tar.gz"}
}

Let's create this file and save it to the pre and postprocessing model repositories.

WARNING: The cells under part 3 will only work if run within a SageMaker Notebook Instance!

The following cells run the Triton Inference Server container in the background and load all the models within the folder /ensemble_hf. The docker won't fail if one or more of the model fails because of --exit-on-error=false, which is useful for iterative code and model repository building. Remove -d to see the logs.

Uncomment the next cell and run it to view the container logs and understand Triton model loading.

Test TensorRT model by invoking the local Triton Server

Part 4 - Deploy Triton to SageMaker Real-Time Endpoint

Deploy with SageMaker Triton container

First we get the URI for the Sagemaker Triton container image that matches the one we used for TensorRT model compilation.

Register the model with Sagemaker.

Create an endpoint configuration.

Deploy the endpoint.

Wait for the endpoint to be up and running.

Test the SageMaker Triton Endpoint

Part 5 - 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.