repository·main·Indexed 21 days ago
https://github.com/triton-inference-server/pytritonPyTriton is a Flask/FastAPI-like framework that simplifies the deployment of Python functions (PyTorch, TensorFlow, JAX, etc.) as high-performance HTTP/gRPC APIs using NVIDIA's Triton Inference Server. It allows developers to wrap arbitrary Python functions as Inference Callables, supporting blocking and background modes, asynchronous execution, and streaming partial results in decoupled mode. The library manages the communication layer between Triton and Python, enabling features like dynamic batching and response caching.
The Identity Python Model example is composed of three main scripts that demonstrate the full lifecycle of a PyTriton model:
install.sh: Installs additional dependencies required for this specific example.server.py: Starts the model within the Triton Inference Server environment.client.py: Acts as a client to execute HTTP or gRPC requests against the deployed model to verify its behavior.This example demonstrates how to deploy large JAX-based language models (like HuggingFace OPT) across multiple nodes and GPUs using PyTriton. It leverages jax.distributed and jax.experimental.pjit for model sharding and distribution.
server.py: Runs the Triton server (on rank 0) or a JAX worker (on ranks > 0). It handles input distribution to workers.client.py: A simple client for testing inference.opt_utils.py: Contains sharding strategies, parameter copying logic, and inference execution.modeling_flax_opt.py: A modified Flax OPT model optimized for FP16 storage and operations.You can use various pretrained or random models ranging from facebook/opt-125m up to random/530B parameters.
The Add-Sub Vertex AI example is composed of the following scripts and files:
install.sh: Installs additional required dependencies.server.py: Starts the model using Triton Inference Server configured in Vertex AI mode.client.py: Executes HTTP/gRPC requests to the deployed model for testing.data.json: Contains example input data used for inference.infer_vertext_ai_endpoint.sh: A script to perform inference directly on the Vertex AI endpoint.health_vertex_ai_endpoint.sh: A script to check if the Vertex AI endpoint is ready and healthy.The MNIST online learning example is composed of the following scripts:
| Script | Purpose |
|---|---|
install.sh | Installs additional dependencies required for the example |
server.py | Starts the model deployment on Triton Inference Server |
client_infer.py | Executes HTTP/gRPC requests for inference using the test dataset |
client_train.py | Executes HTTP/gRPC requests for training using the training dataset |
model.py | Defines the model logic for both inference and training |
PyTriton provides a Flask/FastAPI-like framework for NVIDIA Triton Inference Server with the following capabilities:
@batch) to handle pre-processing and batching logic.asyncio) APIs.The HuggingFace BERT JAX example is composed of three primary scripts:
install.sh: Handles the installation of additional dependencies required to download the model from HuggingFace and sets up the JAX library.server.py: The script used to start and manage the model within the Triton Inference Server environment.client.py: A client script used to execute HTTP or gRPC requests against the deployed model to perform inference.PyTriton acts as a Flask/FastAPI-like interface for NVIDIA's Triton Inference Server. It allows you to define Python functions that execute machine learning model predictions and exposes them via HTTP/gRPC APIs.
PyTriton installs Triton Inference Server in your environment and manages the communication layer between Triton and your Python functions. This enables you to use Triton's performance features (like dynamic batching and response caching) for models implemented in Python (PyTorch, TensorFlow, JAX, etc.) without needing to change your model's environment.
When building DALI pipelines for inference within PyTriton, keep the following technical constraints in mind:
While DALI's default prefetch_queue_depth = 2 is ideal for training to overlap data loading with model execution, for inference, it is often better to set prefetch_queue_depth = 1 to minimize latency and process data as quickly as possible.
In DALI, the batch dimension is often implicit. When processing video, the pipeline returns data in NFCHW layout (Batch, Frame, Channel, Height, Width). To form a standard batch for inference, this must be flattened to an (N*F)CHW layout.
There are two primary ways to decode video in DALI:
fn.decoders.video: Receives an encoded buffer via fn.external_source and decodes the entire video at once. This is the recommended method for PyTriton because it does not require the model to be a decoupled model.fn.inputs.video: Acts as a standalone input and decodes specific portions of a video using the sequence_length operator. This is more memory-efficient for very long videos but requires the DALI model to be configured as a decoupled model to generate multiple responses per request. Note: PyTriton does not currently support decoupled models.To expose a Python model through PyTriton, you must define an Inference Callable. This is a Python function that acts as the entry point for predictions.
When the Triton Inference Server receives a request at the v2/models/<model name>/infer endpoint, it routes the data to this callable. The data is provided to your function as numpy arrays. The integration layer binds this callable to the Triton Server and exposes it under the name you specify during configuration.
A decoupled model is a model that is independent of Triton's standard batching logic. It can receive multiple requests in parallel and send multiple responses to a single request (e.g., streaming tokens in an LLM).
Important Constraints:
ModelClient, AsyncioModelClient, FuturesModelClient) cannot be used with decoupled models and will raise an exception.DecoupledModelClient or AsyncioDecoupledModelClient.decoupled=True in the ModelConfig during model binding.An Inference Callable is the entry point for handling inference requests in PyTriton. It receives a list of Request objects and must return a list of response dictionaries. Each response dictionary maps output names to NumPy ndarrays.
The pytriton.proxy.types.Request object provides:
data: A mapping of input names to NumPy ndarrays. You can also use the request mapping protocol (e.g., request["input_name"]).parameters: A mapping containing combined parameters and HTTP/gRPC headers.You can implement an Inference Callable as a simple function or as a class method (useful for managing model initialization state).
import numpy as np
from typing import Dict, List
from pytriton.proxy.types import Request
# Function implementation
def infer_fn(requests: List[Request]) -> List[Dict[str, np.ndarray]]:
# Process requests and return list of dicts
...
# Class implementation (recommended for stateful initialization)
class InferCallable:
def __init__(self, *args, **kwargs):
# model initialization logic here
...
def __call__(self, requests: List[Request]) -> List[Dict[str, np.ndarray]]:
# inference logic here
...PyTriton uses a structured configuration system to define how models are served. The primary entry point for defining a model's behavior is the pytriton.model_config.ModelConfig class. This class allows you to specify input and output tensors, batching strategies, and queueing policies.
Key components used within ModelConfig include:
Tensor: Defines the name, data type, and shape of inputs and outputs.DeviceKind: Specifies the target device (e.g., CPU, GPU).DynamicBatcher: Configures how multiple requests are grouped into a single batch.QueuePolicy: Defines how requests are handled in the model's queue.TimeoutAction: Determines the behavior when a request exceeds its timeout limit.