triteia is a Python client for performing inference on the NVIDIA Triton Inference Server. It provides model deployment, configuration, and optimization capabilities for the TensorFlow, ONNX, and Python Triton backends directly from Python. This was developed to address limitations in the PyTriton package, and to provide a simple-to-use interface for Whole-Slide Image inference.
- Quick start
- Command-line interfaces
- Model wrappers
- Model configuration
- Model control
- Developer guide
triteia requires histomcs_stream and large_image packages with the TIFF reader
git clone https://github.com/PathologyDataScience/triteia.git
cd triteia
pip install --editable .
--editableensures that updates to thetriteiapackage (aftergit pull) immediately take effect.
Alternatively, to use the Docker image, run the following from the repository root (same directory as this file is located):
# optional: download test data
python download_test_data.py
# triteia will be the name of the Docker image
docker build -f client.Dockerfile . -t triteia:latest --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3)
docker run \
--security-opt seccomp:unconfined --network=host \
--rm -it --shm-size=1g \
-v ${PWD}/test_data:/data:ro \
--name tritonclient triteia:latestNOTE: The
--network=hostoption lets the Docker container access ports from other containers or the host. Since Docker defaults to 64MB of shared memory, use--shm-size=to increase shared memory for processing large whole-slide images. The--security-opt seccomp:unconfinedoption may be required to enable OpenBLAS threading, especially on larger machines. Add the--rmoption to have the container removed after it stops, but avoid it if you need to persist data.
The client Docker container uses the server's Docker network, so any client-required ports must be exposed by the server container. For example, to use Jupyter notebooks on the client, ensure the Jupyter port (8888) is exposed with
-p <your port>:8888when launching the server container.
- feature_extraction demonstrates whole-slide image feature extraction.
- patch_inferencepatch_inference shows feature extraction with image lists instead of WSIs.
- export_pytorch_model shows how to export PyTorch models for Triton.
These examples require a running TritonServer container on the client machine.
This launches a Jupyter Notebook on port 8888. To access it, look for the URL that will be displayed in the terminal and open it in your web browser:
docker run \
--security-opt seccomp:unconfined --network=host \
-v ${PWD}/test_data:/data:ro \
-v ${PWD}/examples:/examples:rw \
--user $UID --rm -it \
--shm-size=1g \
--name tritonclient triteia:latest \
bash -c "jupyter-lab --notebook-dir /examples/ --no-browser"triteia is tested with Triton version 25.02. Support for Tensorflow is deprecated in later versions, but other model backends (like PyTorch) should still work.
To run the server, we recommend using the ./launch_server.sh script. Use ./launch_server --help to see available options. It will prompt you for a model to start with. Additional models can be loaded (or unloaded) with the client.
Some models require a HuggingFace token to access the HuggingFace model hub for download. See HuggingFace token setup for details.
If you want to run the server manually, use the command below:
# cd: the current directory should be triteia, where the models folder is located
docker run \
--gpus=all \
-d \
--rm \
-p 8000:8000 -p 8001:8001 -p 8002:8002 -p 8003:8003 \
--shm-size=4g \
--ulimit memlock=-1 \
--ipc=host \
-v $PWD/models:/models \
nvcr.io/nvidia/tritonserver:25.02-py3 \
tritonserver --model-repository=/models --model-control-mode=explicit --exit-on-error=falseNote: The
--ipc,--shm-size, and--ulimitmemlock options are recommended when using shared memory for client/server communication. These settings allow Triton to access the host's shared memory, increase the default shared memory limit (64MB by default), and prevent RAM from being swapped to disk. To load and modify models while the server is running, use the--model-control-mode=explicitargument. If you are running the client in a Docker container, using--network=hostallows the client container to communicate with the server container over the host's network.
A command-line interface is provided for inference with single or multiple slides, with control over tiling, masking, data loading, and serialization parameters. Models must be loaded prior to inference.
Perform inference with the EfficientNetV2S model on a single slide, outputting serialized embeddings to your home directory:
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflowOptional parameters allow restricting inference to a tissue mask (-m):
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -m TCGA-AN-A0G0-01Z-00-DX1.mask.pngStore features in float32 precision rather than the default float16 (-f):
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -fModification tile size (-t), add tile overlap (-o), and change magnification (-M):
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -t 256 -o 128 -M 10Adjustment of tile reading parameters, including ICC correction (-i), read chunk size (-c), batch size (-b), prefetch (-p), and multiprocessing workers (-w):
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -i -c 8 -b 128 -p 2 -w 16Provide image source (-n) and target (-r) parameters for Macenko color normalization:
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -n ~/TCGA-AN-A0G0-01Z-00-DX1.stain.npy -r ~/standard_stain.npyChange the Triton inference server's address:
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -a "foo.edu:8001"Increase the precision of serialized features to float (default is half float):
python feature_extraction.py ~/TCGA-AN-A0G0-01Z-00-DX1.svs ~/ EfficientNetV2S.tensorflow -fFor large jobs, use a tab-delimited file containing input images and, optionally, their masks and normalization stain profiles:
more ~/inputs.tsv
TCGA-AN-A0G0-01Z-00-DX1.svs TCGA-AN-A0G0-01Z-00-DX1.mask.py TCGA-AN-A0G0-01Z-00-DX1.stain.npy
TCGA-AN-A0G0-01Z-00-DX2.svs TCGA-AN-A0G0-01Z-00-DX2.mask.py TCGA-AN-A0G0-01Z-00-DX2.stain.npy
TCGA-AN-A0G0-01Z-00-DX3.svs TCGA-AN-A0G0-01Z-00-DX3.mask.py TCGA-AN-A0G0-01Z-00-DX3.stain.npy
TCGA-AN-A0G0-01Z-00-DX4.svs TCGA-AN-A0G0-01Z-00-DX4.mask.py TCGA-AN-A0G0-01Z-00-DX4.stain.npy
python feature_extraction.py ~/inputs.tsv ~/ EfficientNetV2S.tensorflowSkip images where output already exists:
python feature_extraction.py ~/inputs.tsv ~/ EfficientNetV2S.tensorflow -striteia provides wrappers for serving popular digital pathology models, including CONCH, UNI, Prov-GigaPath, hibou-L, Phikon, Virchow, and Virchow2 on the Python backend. All models are served using mixed precision.
| Model | Input | Output | Size |
|---|---|---|---|
| conch | (224, 224, 3) | 512 | 0.802 GB |
| gigapath | (224, 224, 3) | 1536 | 4.54 GB |
| hibou-L | (224, 224, 3) | 1024 | 1.21 GB |
| phikon | (224, 224, 3) | 768 | 0.346 GB |
| resnet50 | (3, 224, 224) | 1000 | 0.1 GB |
| SegFormer | (3,dynamic) | (150,dynamic) | 0.01 GB |
| uni | (224, 224, 3) | 1024 | 1.21 GB |
| uni2 | (224, 224, 3) | 1536 | 2.73 GB |
| virchow | (224, 224, 3) | 2560 | 2.53 GB |
| virchow2 | (224, 224, 3) | 2560 | 2.53 GB |
A dockerfile built on Triton Server container encapsulates all requirements for serving these models
docker build -t model-tritonserver -f server.Dockerfile .Each folder in the /models directory contains a model.py file that implements model loading, inference, and cleanup. The model.py files must be placed into the model repository for serving ("config.pbtxt" file is not required).
models
└── phikon # Phikon model folder
└── 1 # version 1
└── model.py # TritonPythonModel Class adapted for Phikon
A HuggingFace token is required to access the UNI, gigapath, and hibou-L models. This is passed to the server by setting the environment variable HF_TOKEN before launch:
export HF_TOKEN=hf_**********************************
docker run \
-e HF_TOKEN=$HF_TOKEN \
--gpus=all \
-p 8000:8000 -p 8001:8001 -p 8002:8002 -p 8003:8003 \
--shm-size=1g --ulimit memlock=-1 --ipc=host \
-v ${PWD}/models/:/models \
model-tritonserver tritonserver --model-repository=/models --model-control-mode=explicit --exit-on-error=falsetriteia.config includes model configuration classes that implement backend-specific configuration options. These classes enable configuration of batching behavior, specification of model input/output shapes and types, and backend optimizations. Refer to the Triton documentation on model configuration and optimization for further details.
Configuration classes like PythonConfiguration and TensorflowConfiguration take additional data classes as inputs that configure batching and caching behavior, hardware resources, and model input/output signatures.
When Triton is launched with --strict-model-config=false, the server automatically configures basic information such as input/output signatures, and the configuration can omit them.
from triteia.config import TensorflowConfig
from triteia.model import TritonModel
name = "mymodel.tensorflow"
config = TensorflowConfig(name, max_batch_size=64)
model = TritonModel(name, "localhost:8001")
model.load(config=config.json())Alternatively, model inputs and output signatures can be defined using the ModelInput and ModelOutput classes.
from triteia.config import ModelInput
input = [ModelInput(name="input_0", shape=[224, 224, 3], dtype=np.float32, optional=False)]
config = TensorflowConfig(name, max_batch_size=64, input=input)Variable-sized input dimensions can be indicated using a value of -1.
The InstanceGroup class configures the use of CPU or GPU resources and the number of model instances per GPU.
from triteia.config import InstanceGroup
instances = InstanceGroup(count=2, kind="gpu", gpus=[0,1,2,3])
config = TensorflowConfig(name=name, instance_group=instances)TensorflowOptimization can be used with TensorflowMixedPrecision, TensorflowXla, and TensorRt to activate automatic mixed precision, XLA compilation, or TensorRT optimization.
from triteia.config import TensorflowMixedPrecision, TensorflowXla, TensorflowOptimization
amp = TensorflowMixedPrecision()
xla = TensorflowXla(level=2)
optimizer = TensorflowOptimization(amp=amp, xla=xla)
ampxla_config = TensorflowConfig(name=name, max_batch_size=64, optimization=optimization)TensorRT cannot be used concurrently with XLA and mixed precision.
For a non-batching model, set the maximum batch size to zero.
config = TensorflowConfig(name, max_batch_size=0)Configuration classes can also save their configuration to a config.pbtxt for automatic file-based configuration.
config.save("/model_repository/mymodel.tensorflow/")File-based configuration is useful for distributing models. When configuring and loading models directly in Python, a JSON-formatted dictionary is used. Configuration files use the protocol buffer format. Configuration classes handle conversion between these formats.
The TritonModel class can be used to load/unload models, retrieve model configurations or metadata, or check whether a model is idle or loaded. A model is defined by a model name and a server URL.
from triteia.model import TritonModel
model = TritonModel("EfficientNetV2S.tensorflow", "localhost:8001")When unloading a model, triteia will check that the model is idle.
# load model with auto-generated configuration
# block and timeout after 1 second
model.load(block=True, timeout=1.)Model status can also be checked.
# check if model is loaded
assert model.is_loaded()
# check if model is idle
assert model.is_idle()The current hosted model configuration can also be queried.
model.get_config()Inference is performed using feature_extraction.inference. This function consumes data from a simple iterator that emits data/metadata pairs. For single-input models, data is provided as a numpy array. For multi-input models, data is provided as a dict of key-value pairs that map numpy arrays to model input names (visible in Model.get_config()).
inference can apply preprocessing functions to the data after loading and before inference by passing a callable to the pre argument.
Note: For multi-input models, all inputs require uniform batch dimensions. You should duplicate singleton values where necessary to satisfy this requirement.
Testing and code formatting are automated with tox and pytest, and can be run with python -m tox run. Running this will evaluate the tests in the environments defined in tox.ini and will format the source using black. Following testing, a coverage.html file will be located in .tox/coverage.
Testing requires running a Triton server on the local machine. Tests are run using an EfficientNetV2S.tensorflow model that can be downloaded using pooch:
import pooch
pooch.retrieve(
fname="EfficientNetV2S.tensorflow.zip",
url="https://drive.usercontent.google.com/download?id=1Mmm2sRGzdzCEAODjABiiPIiBdg40EPwC&export=download&confirm=t",
known_hash="b115917b7d0e480fe080077d60913d90c4b596c5a1e3c74745c22f21ed14df63",
path=host_model_repository
)To test using the Docker container, launch and build the client as follows:
docker build -f client.Dockerfile . -t triteia:latest --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3)
./launch_test_container.sh
You can now run tests inside the container using pytest tests.
Models usually come as PyTorch, TensorFlow, or ONNX backends. NVIDIA has a format called "TRT" (https://github.com/NVIDIA/TensorRT) that is optimized for inference. In examples/resnet50_trt, we show how to download a PyTorch model, convert it to ONNX, and then convert the ONNX file to a TRT backend. The approach can be modified to work with different types of models.
Be aware that extra performance is not guaranteed. PyTorch offers many optimizations that may not be available in the ONNX or TRT backends.
To reproduce the (TBD) paper
- Start the NVIDIA Triton server with all GPUs and models available:
./launch_server.sh --num-gpus <num_gpu> --http-port 7984 --grpc-port 7985 --metrics-port 7986 OUTPUT_DIR="./results" ./benchmarking/paper_benchmarks.sh $OUTPUT_DIR- Inspect results:
tensorboard --logdir=.... To view results as figures:
# convert TensorBoard to CSV
./benchmarking/tensorboard_to_csv.py
# CSV to plot files
./benchmarking/tensorboard_csv_to_plot.pyTo do the benchmarks using Docker, use the benchmark_client.Dockerfile in the benchmarking directory.
It is identical to the client.Dockerfile in this directory, except it has access to CUDA, allowing it to export GPU metrics and automatically start and stop Triton with GPUs.
To build it from the git root directory: docker build -f benchmarking/benchmark_client.Dockerfile . -t triteia:benchmark