Serving with NVIDIA's Triton Server
NVIDIDA's Triton Server is a serving framework that simplifies deploying AI models at scale in production environments. It provides a consistent and optimized interface for deploying models trained with various frameworks like PyTorch and ONNX. Triton Server offers features like model versioning, multi-model deployment, and support for various deployment environments, making it a versatile solution for AI inference deployment.
This page shows you how to try serving MAX Engine with Triton by running it in a Docker container, and then send inference requests with a Python program.
Deploy an example
Here’s quick-start guide to run an inference with MAX Engine and Triton Server on a local system.
We’ve created a Docker container that includes both Triton and MAX
Engine, which you can either download and run, or build yourself from a
Dockerfile. After the Docker container is running, you can send inference
requests from an HTTP/gRPC client via tritonclient
, as you'll see in our
example.
To make it as simple as possible to run, we’ve created an example script that downloads and runs the Docker container, and then sends an inference request with a Python program. All you need to do is run a bash script.
First, clone the code examples (if you haven’t already):
git clone git@github.com:modularml/max.git
Then, run one of the deploy.sh
scripts. For example,
here’s how to run an inference with a BERT model:
-
Install the Python requirements:
cd max/examples/inference/bert-python-torchscript
python3 -m pip install -r requirements.txt
-
Run the Docker container and send an inference request:
bash deploy.sh
It might take a few minutes for the model to compile. But once it's done, execution is nearly instant, and subsequent loads will be faster.
You should see output like this:
input text: Paris is the [MASK] of France.
filled mask: Paris is the capital of France.
Deploy your own Docker image
If you look in the
deploy.sh
script, you’ll see that it builds a Docker image that's hosted in the following
container registry, which is always the latest version of MAX:
public.ecr.aws/modular/max-serving
However, you might want to customize this image yourself. In that case, you can instead build the Docker image using the Dockerfile that’s included in the MAX SDK, which will always correspond to the version of MAX that you have installed.
The following sections describe how to do that, run it, and send requests.
Get the model
First, create a directory for your model that you can mount inside the container.
For example, let’s use the BERT model from our GitHub examples.
If you didn’t already clone the max
repo, do that now:
git clone https://github.com/modularml/max.git
Then, with your terminal in that same location, download the model from HuggingFace with our Python script, which also converts the model to TorchScript format, and save it in the path that you’ll mount in the container:
MODEL_REPOSITORY=~/model-repository \
&& BERT_DIR=$MODEL_REPOSITORY/bert-mlm \
&& mkdir -p $BERT_DIR/1
python3 max/examples/inference/common/bert-torchscript/download-model.py \
--output-path $BERT_DIR/1/bert-mlm.torchscript --mlm
Add the Triton config
Next, you need to add a model configuration file.
For example, here's the configuration file for the BERT model:
input {
name: "input_ids"
data_type: TYPE_INT64
dims: [1,128]
}
input {
name: "attention_mask"
data_type: TYPE_INT64
dims: [1,128]
}
input {
name: "token_type_ids"
data_type: TYPE_INT64
dims: [1,128]
}
output {
name: "result0"
data_type: TYPE_FP32
dims: [1, 128, 768]
}
instance_group {
kind: KIND_CPU
}
backend: "max"
default_model_filename: "bert-mlm.torchscript"
You can copy this from our repo into your project with this command:
cp max/examples/inference/bert-python-torchscript/config.pbtxt $BERT_DIR
Now make sure all the files are in the model repository:
tree $MODEL_REPOSITORY
/home/ubuntu/model-repository
└── bert-mlm
├── 1
│ └── bert-mlm.torchscript
└── config.pbtxt
2 directories, 2 files
Build the Docker image
When you installed the MAX SDK, it saved the Dockerfile for MAX Engine with Triton on your system. Here's how you can build it:
MAX_INSTALL_DIR=$(modular config max.path)
docker buildx build \
--file ${MAX_INSTALL_DIR}/Dockerfile \
--tag max_serving_local \
--load \
${MAX_INSTALL_DIR}
Run the Docker image
Now start the container and Triton with this command:
docker run -it --rm --net=host \
-v $MODEL_REPOSITORY:/models \
max_serving_local \
tritonserver --model-repository=/models \
--model-control-mode=explicit \
--load-model=bert-mlm
As the Docker container starts up, MAX Engine will compile the model. Once that’s done and the server is running, you’ll see the usual Triton logs, including the endpoints that are running:
I0215 23:16:18.906473 1 grpc_server.cc:2451] Started GRPCInferenceService at 0.0.0.0:8001
I0215 23:16:18.906671 1 http_server.cc:3558] Started HTTPService at 0.0.0.0:8000
I0215 23:16:18.948484 1 http_server.cc:187] Started Metrics Service at 0.0.0.0:8002
Leave this terminal running.
Send inference requests
Open a second terminal on the same machine and send an inference request with this example client:
python3 max/examples/inference/bert-python-torchscript/triton-inference.py \
--text "Paris is the [MASK] of France."
You should see the results like this:
input text: Paris is the [MASK] of France.
filled mask: Paris is the capital of France.
That's it! You're now running a Triton Server with MAX Engine as the backend.
You can also fetch some metadata about the service:
-
Check if the container is started and ready:
curl -v localhost:8000/v2/health/ready | python3 -m json.tool
-
Get the model metadata (input/output parameters):
curl localhost:8000/v2/models/bert-mlm | python3 -m json.tool
-
Get the loaded model configuration:
curl localhost:8000/v2/models/bert-mlm/config | python3 -m json.tool
This page has been a brief walkthrough of serving MAX Engine using Nvidia's Triton Server. This implementation is not commercially-licensed for production deployment. However, you can freely tinker with it and evaluate it for your AI usecases.
For a more turn-key solution, we'll also make MAX available on AWS Marketplace, so you can quickly deploy a managed container on AWS.
All of this and more is coming soon, when we release MAX for production workloads. Sign up for updates.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!
If you'd like to share more information, please report an issue on GitHub
😔 What went wrong?