You've painstakingly engineered a low-latency feature store, perhaps even a zero-copy solution like we explored previously. But as your real-time ML applications scale, you notice a new bottleneck: model inference. Native Python frameworks, while flexible, often introduce unacceptable latency and resource overhead in production, turning your sub-millisecond feature fetches into multi-millisecond predictions. This post will show you how to break through that barrier by converting a PyTorch (or sklearn) model to ONNX, optimizing it with ONNX Runtime, and deploying it on NVIDIA Triton Inference Server to achieve blazing-fast, scalable, and framework-agnostic real-time predictions. My core judgment is that for true sub-millisecond inference at scale, moving beyond native Python serving is not an option, but a necessity.
Key Takeaways
- Native Python model serving can introduce significant latency and resource overhead, making it a bottleneck for real-time applications even after optimizing feature retrieval.
- ONNX provides a universal, optimized intermediate representation for machine learning models, enabling cross-framework deployment and performance gains through specialized runtimes.
- NVIDIA Triton Inference Server is a powerful, production-ready solution for deploying and serving ML models at scale, offering features like dynamic batching, concurrent model execution, and multi-framework support for sub-millisecond inference.
- Preprocessing steps (like text vectorization) often remain in Python, while only the numerical core of the ML model is offloaded to ONNX and Triton, requiring careful orchestration between components.
- Achieving sub-millisecond inference involves not just model conversion, but also strategic deployment on an optimized serving platform like Triton, coupled with efficient client-side interaction.
The Problem
Our journey towards real-time machine learning often starts with data. We build sophisticated pipelines to ingest, transform, and serve features with minimal delay. In a previous post, Beyond Feast: Architecting a Low-Latency, Zero-Copy Feature Store from Scratch in Python, I detailed how to achieve sub-millisecond feature retrieval. But what happens after you've fetched those fresh features? Too often, the next step—model inference—becomes the new bottleneck. You feed those carefully prepared features into a model written in PyTorch or Scikit-learn, and suddenly, your blazing-fast data path grinds to a halt. The overhead of Python interpreters, GIL contention, and the lack of hardware acceleration for generic tensor operations can easily add tens or even hundreds of milliseconds to your prediction latency. For applications demanding true real-time responses, like fraud detection, personalized recommendations, or algorithmic trading, this simply isn't acceptable. We need a way to serve our models that is as performant as our feature store.
Data and Sources
For this post, we'll be classifying the titles of recent blog posts from the GitHub Engineering blog. This provides a real-world, dynamic text dataset to demonstrate our real-time classification pipeline.
- GitHub Engineering RSS Feed: https://github.blog/engineering/feed/
feedparserlibrary: https://pypi.org/project/feedparser/scikit-learndocumentation: https://scikit-learn.org/stable/skl2onnxdocumentation: https://onnx.ai/sklearn-onnx/ONNX Runtimedocumentation: https://onnxruntime.ai/NVIDIA Triton Inference Serverdocumentation: https://developer.nvidia.com/nvidia-triton-inference-serverPython Triton Clientdocumentation: https://github.com/triton-inference-server/client/tree/main/src/python
Data accessed on 2024-05-15.
Step 1 — The Bottleneck Beyond Features: Why Native Frameworks Fall Short for Serving
You've built your model in PyTorch, TensorFlow, or Scikit-learn, trained it, and it performs beautifully offline. But deploying it to production for real-time inference exposes a different set of challenges. Python, while fantastic for development and data manipulation, isn't inherently designed for low-latency, high-throughput serving at scale. Each inference request often means spinning up a Python process, loading the model, and executing the graph, which incurs significant overhead. Even with multi-threading or async approaches, the Global Interpreter Lock (GIL) can limit true parallelism for CPU-bound tasks. Furthermore, these frameworks are often resource-intensive, consuming more memory and CPU cycles than strictly necessary for inference alone. They also lack a standardized, framework-agnostic way to deploy models, leading to bespoke serving solutions that are hard to maintain and scale. This is where we need to introduce a dedicated inference serving layer, optimized for performance and interoperability.
Step 2 — From Scikit-learn to ONNX: The Universal Intermediate Representation
The first step to escaping the native framework bottleneck is to decouple our model from its training ecosystem. This is where ONNX (Open Neural Network Exchange) comes in. ONNX provides a standardized, open format for representing machine learning models. It acts as a universal intermediate representation, allowing you to train a model in one framework (like Scikit-learn or PyTorch), convert it to ONNX, and then run it in another framework or optimized runtime. This is crucial for performance and flexibility.
Here, we'll convert a simple Scikit-learn text classification pipeline into the ONNX format. Note that for text models, the TfidfVectorizer (or similar tokenizer/vectorizer) is usually a preprocessing step that we keep in Python, as its logic is not directly part of the neural network graph that ONNX is designed for. We convert only the numerical classification part of the pipeline.
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import StringTensorType
import joblib
import os
def create_and_convert_model(train_texts, train_labels, model_dir="triton_model_repo/text_classifier/1"):
# Create a simple Scikit-learn pipeline
vectorizer = TfidfVectorizer(max_features=1000)
classifier = LogisticRegression(solver='liblinear', random_state=42)
# Train the vectorizer and classifier separately for ONNX export
X_train_vec = vectorizer.fit_transform(train_texts)
classifier.fit(X_train_vec, train_labels)
# Save the vectorizer separately for Python-side preprocessing
os.makedirs(model_dir, exist_ok=True)
joblib.dump(vectorizer, os.path.join(model_dir, "vectorizer.joblib"))
# Define the initial types for ONNX conversion.
# For a LogisticRegression, the input will be a float tensor.
# The number of features (1000 in this case) must match the TfidfVectorizer output.
initial_type = [('float_input', StringTensorType([None, 1]))] # This is a placeholder, actual will be float
# Convert only the LogisticRegression model to ONNX
# The input to the ONNX model will be the vectorized features, not raw strings.
# skl2onnx expects a numeric input for the classifier.
# We need to create a dummy input to infer the correct ONNX input shape.
dummy_input_for_onnx = X_train_vec[0:1].astype('float32') # Use float32 for ONNX compatibility
# Convert the classifier to ONNX
onnx_model = convert_sklearn(
classifier,
initial_types=[('float_input', StringTensorType([None, dummy_input_for_onnx.shape[1]]))], # Incorrect, should be FloatTensorType
target_opset=13, # Use a recent opset
)
# Corrected initial_types for a numerical input classifier
onnx_model = convert_sklearn(
classifier,
initial_types=[('float_input', FloatTensorType([None, dummy_input_for_onnx.shape[1]]))],
target_opset=13,
)
onnx_path = os.path.join(model_dir, "model.onnx")
with open(onnx_path, "wb") as f:
f.write(onnx_model.SerializeToString())
print(f"ONNX model saved to {onnx_path}")
return onnx_path
The code first trains a TfidfVectorizer and a LogisticRegression model. Crucially, we save the TfidfVectorizer separately using joblib. This is because the text preprocessing (tokenization, TF-IDF transformation) will happen on the client side (in Python) before sending numerical features to Triton. The skl2onnx.convert_sklearn function then takes our trained LogisticRegression model and converts it into an ONNX graph. We specify the initial_types to inform skl2onnx about the expected input shape and type for our numerical features (a float tensor), ensuring the ONNX model expects the output from our TfidfVectorizer. The target_opset ensures compatibility with ONNX Runtime versions.
Step 3 — Accelerating Inference with ONNX Runtime: Beyond Just Conversion
Converting to ONNX is just the first step. To truly unlock performance, we need an optimized runtime. ONNX Runtime (ORT) is a high-performance inference engine for ONNX models. It's designed to maximize performance by leveraging hardware accelerators (CPUs, GPUs, FPGAs) and applying graph optimizations. When we deploy our ONNX model on Triton, Triton will internally use ONNX Runtime (or other optimized backends) to execute the model. However, you can also use ONNX Runtime directly in Python for local testing or if Triton isn't part of your architecture for some reason.
import onnxruntime as ort
import numpy as np
def test_onnx_model_locally(onnx_path, vectorizer_path, sample_text):
# Load the ONNX model
session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
# Load the vectorizer
vectorizer = joblib.load(vectorizer_path)
# Preprocess the sample text
text_vec = vectorizer.transform([sample_text]).astype(np.float32)
# Get input/output names
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
# Run inference
# ONNX Runtime expects input as a dictionary
outputs = session.run([output_name], {input_name: text_vec})
# The output from LogisticRegression in ONNX is usually class probabilities or raw scores
# We need to map it back to labels
predicted_class_idx = np.argmax(outputs[0], axis=1)[0]
print(f"Local ONNX Runtime prediction for '{sample_text}': Class Index {predicted_class_idx}")
return predicted_class_idx
This snippet demonstrates how to load the ONNX model using onnxruntime.InferenceSession. We specify CPUExecutionProvider for CPU inference, but you could use CUDAExecutionProvider for GPU. Before calling session.run(), we preprocess the input text using our saved TfidfVectorizer, converting it into the numerical format the ONNX model expects. This local test confirms our ONNX model and preprocessing steps are working correctly before we introduce Triton.
Step 4 — Architecting for Scale: Introducing NVIDIA Triton Inference Server
While ONNX Runtime offers local acceleration, managing multiple models, dynamic batching, concurrent requests, and hardware acceleration across different frameworks in a production environment is complex. This is where NVIDIA Triton Inference Server shines. Triton is a purpose-built inference serving solution that simplifies the deployment of AI models at scale. It offers a unified interface for models from various frameworks (TensorFlow, PyTorch, ONNX, OpenVINO, etc.), supports multiple execution environments (CPUs, GPUs), and provides advanced features like dynamic batching, concurrent model execution, and model versioning to maximize throughput and minimize latency.
To use Triton, you need to organize your ONNX model in a specific model repository structure:
triton_model_repo/
├── text_classifier/
│ ├── config.pbtxt
│ └── 1/
│ └── model.onnx
│ └── vectorizer.joblib # (Optional, for client-side preprocessing awareness)
The config.pbtxt file is crucial; it tells Triton about your model, its expected inputs/outputs, and how to run it. For our ONNX model, it's quite straightforward:
name: "text_classifier"
platform: "onnxruntime_onnx"
max_batch_size: 64
input [
{
name: "float_input"
data_type: TYPE_FP32
dims: [ -1, 1000 ] # -1 for dynamic batch size, 1000 features from TfidfVectorizer
}
]
output [
{
name: "output_probabilities" # Replace with actual output name from ONNX graph
data_type: TYPE_FP32
dims: [ -1, 2 ] # -1 for dynamic batch size, 2 classes (binary classification example)
}
]
You would adjust dims for the output based on your number of classes and the actual output name from your ONNX model (which you can inspect with tools like Netron). For a multi-class Logistic Regression, the output would be [ -1, num_classes ]. The next step is to run Triton. This is typically done via Docker:
docker run --rm -d --name triton -p 8000:8000 -p 8001:8001 -p 8002:8002 -v $(pwd)/triton_model_repo:/models nvcr.io/nvidia/tritonserver:23.10-py3 tritonserver --model-repository=/models
This command starts a Triton server instance, mapping the local triton_model_repo directory into the container's /models directory. It also exposes the necessary ports for HTTP (8000), gRPC (8001), and metrics (8002) communication. Once this command runs, Triton will load your text_classifier model, making it ready for inference requests.
Step 5 — Deploying and Querying Your ONNX Model on Triton
With Triton running, we can now send inference requests from our Python application. NVIDIA provides a Python client library for this purpose, simplifying interaction with the server. This client allows us to construct inference requests, send them over HTTP or gRPC, and parse the responses efficiently.
import tritonclient.grpc as grpcclient
import numpy as np
import time
def query_triton_server(model_name, input_data, classes, server_url="localhost:8001"):
try:
triton_client = grpcclient.InferenceServerClient(url=server_url)
except Exception as e:
print(f"ERROR: Could not connect to Triton server at {server_url}. Please ensure Triton is running. {e}")
return None, None
inputs = []
outputs = []
# Create input object for Triton
input_name = "float_input" # Must match name in config.pbtxt
output_name = "output_probabilities" # Must match name in config.pbtxt
inputs.append(grpcclient.InferInput(input_name, input_data.shape, "FP32"))
inputs[0].set_content_from_numpy(input_data)
outputs.append(grpcclient.InferRequestedOutput(output_name, binary_data=True))
start_time = time.perf_counter_ns()
response = triton_client.infer(
model_name=model_name,
inputs=inputs,
outputs=outputs,
client_timeout=10.0 # Timeout in seconds
)
end_time = time.perf_counter_ns()
inference_latency_us = (end_time - start_time) / 1000.0
# Get the output probabilities
output_array = response.as_numpy(output_name)
# Post-process: get predicted class index
predicted_class_idx = np.argmax(output_array, axis=1)[0]
predicted_label = classes[predicted_class_idx]
return predicted_label, inference_latency_us
This function establishes a gRPC connection to Triton. It constructs an InferInput object, filling it with our preprocessed numerical features. The set_content_from_numpy method efficiently transfers the data. We then request specific outputs and send the inference request. The time.perf_counter_ns() calls allow us to precisely measure the round-trip latency, including network and server processing time. The response contains the raw model output (probabilities), which we then convert back into a human-readable class label using our defined classes list.
Step 6 — Real-time Categorization Pipeline: Connecting the Pieces
Now, let's tie everything together into a complete real-time text categorization pipeline. This involves fetching new text data, preprocessing it with our saved TfidfVectorizer, sending the numerical features to Triton for inference, and then interpreting the results. This end-to-end flow demonstrates how you can achieve sub-millisecond inference for new, incoming data streams.
import feedparser
import joblib
import numpy as np
import os
import time
import tritonclient.grpc as grpcclient
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
# --- Helper functions (from previous steps) ---
def fetch_github_engineering_posts(feed_url):
try:
feed = feedparser.parse(feed_url)
if feed.status != 200 and feed.status != 301: # 301 is also a success for feedparser
raise Exception(f"Failed to fetch RSS feed, status: {feed.status}")
entries = []
for entry in feed.entries:
entries.append({'title': entry.title, 'link': entry.link})
return entries
except Exception as e:
print(f"Error fetching RSS feed: {e}")
return []
def generate_labels_and_train_model(posts):
titles = [p['title'] for p in posts]
# Heuristic-based labeling for demonstration
labels = []
classes = ["Infrastructure", "Performance", "Security", "Feature/Product", "General"]
keyword_map = {
"Infrastructure": ["Git", "infrastructure", "Azure", "database", "agent-scale", "architecture", "scale"],
"Performance": ["performance", "site speed", "CSS", "optimization", "latency", "improving"],
"Security": ["secure access", "repo hygiene", "collaborator", "security"],
"Feature/Product": ["star history", "new API", "features", "rebuilding", "building"],
}
for title in titles:
assigned_label = "General"
for label, keywords in keyword_map.items():
if any(kw.lower() in title.lower() for kw in keywords):
assigned_label = label
break
labels.append(assigned_label)
# Convert labels to numerical format for LogisticRegression
label_to_int = {label: i for i, label in enumerate(classes)}
int_labels = [label_to_int[label] for label in labels]
vectorizer = TfidfVectorizer(max_features=1000)
classifier = LogisticRegression(solver='liblinear', random_state=42)
X_train_vec = vectorizer.fit_transform(titles)
classifier.fit(X_train_vec, int_labels)
return vectorizer, classifier, classes
def convert_and_save_onnx_model(vectorizer, classifier, model_dir="triton_model_repo/text_classifier/1"):
os.makedirs(model_dir, exist_ok=True)
joblib.dump(vectorizer, os.path.join(model_dir, "vectorizer.joblib"))
dummy_input_for_onnx = vectorizer.transform(["dummy text"]).astype(np.float32)
onnx_model = convert_sklearn(
classifier,
initial_types=[('float_input', FloatTensorType([None, dummy_input_for_onnx.shape[1]]))],
target_opset=13,
)
onnx_path = os.path.join(model_dir, "model.onnx")
with open(onnx_path, "wb") as f:
f.write(onnx_model.SerializeToString())
print(f"ONNX model saved to {onnx_path}")
return onnx_path
def create_triton_config(num_features, num_classes, model_dir="triton_model_repo/text_classifier"):
config_path = os.path.join(model_dir, "config.pbtxt")
config_content = f"""
name: "text_classifier"
platform: "onnxruntime_onnx"
max_batch_size: 64
input [
{{
name: "float_input"
data_type: TYPE_FP32
dims: [ -1, {num_features} ]
}}
]
output [
{{
name: "output_probabilities"
data_type: TYPE_FP32
dims: [ -1, {num_classes} ]
}}
]
"""
os.makedirs(model_dir, exist_ok=True)
with open(config_path, "w") as f:
f.write(config_content)
print(f"Triton config saved to {config_path}")
def perform_triton_inference(model_name, input_data, classes, server_url="localhost:8001"):
try:
triton_client = grpcclient.InferenceServerClient(url=server_url)
except Exception as e:
print(f"ERROR: Could not connect to Triton server at {server_url}. Please ensure Triton is running as per instructions.")
print(f"Details: {e}")
return None, None
inputs = []
outputs = []
input_name = "float_input"
output_name = "output_probabilities"
inputs.append(grpcclient.InferInput(input_name, input_data.shape, "FP32"))
inputs[0].set_content_from_numpy(input_data)
outputs.append(grpcclient.InferRequestedOutput(output_name, binary_data=True))
try:
start_time = time.perf_counter_ns()
response = triton_client.infer(
model_name=model_name,
inputs=inputs,
outputs=outputs,
client_timeout=5.0 # Shorter timeout for real-time
)
end_time = time.perf_counter_ns()
inference_latency_us = (end_time - start_time) / 1000.0
output_array = response.as_numpy(output_name)
predicted_class_idx = np.argmax(output_array, axis=1)[0]
predicted_label = classes[predicted_class_idx]
return predicted_label, inference_latency_us
except Exception as e:
print(f"ERROR: Inference request failed: {e}")
return None, None
# --- Main script logic ---
def main():
RSS_FEED_URL = 'https://github.blog/engineering/feed/'
TRITON_MODEL_DIR = "triton_model_repo"
MODEL_NAME = "text_classifier"
TRITON_SERVER_URL = "localhost:8001"
print("Fetching GitHub Engineering blog posts to train model...")
training_posts = fetch_github_engineering_posts(RSS_FEED_URL)
if not training_posts:
print("No posts fetched for training. Exiting.")
return
# Use a subset for training if many posts, or all for this demo
training_data_limit = 20
training_posts_subset = training_posts[:training_data_limit]
train_titles = [p['title'] for p in training_posts_subset]
print(f"Training model with {len(train_titles)} posts...")
vectorizer, classifier, classes = generate_labels_and_train_model(training_posts_subset)
num_features = vectorizer.max_features
num_classes = len(classes)
print("Converting Scikit-learn model to ONNX and saving artifacts...")
onnx_path = convert_and_save_onnx_model(vectorizer, classifier, os.path.join(TRITON_MODEL_DIR, MODEL_NAME, "1"))
create_triton_config(num_features, num_classes, os.path.join(TRITON_MODEL_DIR, MODEL_NAME))
print("\n--- Model deployment preparation complete ---")
print("Please ensure NVIDIA Triton Inference Server is running with the following command:")
print(f"docker run --rm -d --name triton -p 8000:8000 -p 8001:8001 -p 8002:8002 -v {os.path.abspath(TRITON_MODEL_DIR)}:/models nvcr.io/nvidia/tritonserver:23.10-py3 tritonserver --model-repository=/models")
print("\nWaiting for 10 seconds for Triton to start and load the model...")
time.sleep(10) # Give Triton time to spin up and load
print("\n--- Performing real-time inference on new posts ---")
# Fetch some fresh posts for inference (can be the next set of posts)
inference_posts_raw = fetch_github_engineering_posts(RSS_FEED_URL)
if not inference_posts_raw:
print("No posts fetched for inference. Exiting.")
return
# Take the next 5 posts not used for training, or if less than 20, take the first 5
inference_posts = inference_posts_raw[training_data_limit:training_data_limit+5] if len(inference_posts_raw) > training_data_limit else inference_posts_raw[:5]
if not inference_posts:
print("No posts available for inference. Exiting.")
return
total_latency = 0
inference_count = 0
for i, post in enumerate(inference_posts):
title = post['title']
link = post['link']
# Preprocess the text using the *saved* vectorizer
preprocessed_input = vectorizer.transform([title]).astype(np.float32)
# Perform inference via Triton
predicted_label, latency = perform_triton_inference(MODEL_NAME, preprocessed_input, classes, TRITON_SERVER_URL)
if predicted_label:
print(f"\nPost {i+1}:")
print(f" Title: {title}")
print(f" Link: {link}")
print(f" Predicted Category: {predicted_label}")
print(f" Inference Latency (us): {latency:.2f}")
total_latency += latency
inference_count += 1
else:
print(f"\nPost {i+1}: Failed to get prediction for '{title}'.")
if inference_count > 0:
avg_latency = total_latency / inference_count
print(f"\n--- Summary ---")
print(f"Total inferences: {inference_count}")
print(f"Average Inference Latency (us): {avg_latency:.2f}")
print(f"Average Inference Latency (ms): {avg_latency / 1000.0:.2f}")
if __name__ == "__main__":
main()
Complete Script
The full runnable script combining all steps:
#!/usr/bin/env python3
import feedparser
import joblib
import numpy as np
import os
import time
import tritonclient.grpc as grpcclient
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
# --- Helper functions ---
def fetch_github_engineering_posts(feed_url):
"""Fetches recent GitHub Engineering blog posts from RSS feed."""
try:
feed = feedparser.parse(feed_url)
if feed.status != 200 and feed.status != 301:
raise Exception(f"Failed to fetch RSS feed, status: {feed.status}")
entries = []
for entry in feed.entries:
entries.append({'title': entry.title, 'link': entry.link})
return entries
except Exception as e:
print(f"Error fetching RSS feed: {e}")
return []
def generate_labels_and_train_model(posts):
"""Generates heuristic labels and trains a TfidfVectorizer + LogisticRegression model."""
titles = [p['title'] for p in posts]
# Heuristic-based labeling for demonstration purposes
labels = []
classes = ["Infrastructure", "Performance", "Security", "Feature/Product", "General"]
keyword_map = {
"Infrastructure": ["Git", "infrastructure", "Azure", "database", "agent-scale", "architecture", "scale"],
"Performance": ["performance", "site speed", "CSS", "optimization", "latency", "improving"],
"Security": ["secure access", "repo hygiene", "collaborator", "security"],
"Feature/Product": ["star history", "new API", "features", "rebuilding", "building"],
}
for title in titles:
assigned_label = "General"
for label, keywords in keyword_map.items():
if any(kw.lower() in title.lower() for kw in keywords):
assigned_label = label
break
labels.append(assigned_label)
# Convert labels to numerical format for LogisticRegression
label_to_int = {label: i for i, label in enumerate(classes)}
int_labels = [label_to_int[label] for label in labels]
vectorizer = TfidfVectorizer(max_features=1000)
classifier = LogisticRegression(solver='liblinear', random_state=42)
X_train_vec = vectorizer.fit_transform(titles)
classifier.fit(X_train_vec, int_labels)
return vectorizer, classifier, classes
def convert_and_save_onnx_model(vectorizer, classifier, model_version_dir):
"""Converts the Scikit-learn classifier to ONNX and saves it along with the vectorizer."""
os.makedirs(model_version_dir, exist_ok=True)
joblib.dump(vectorizer, os.path.join(model_version_dir, "vectorizer.joblib"))
# Create a dummy input to infer the correct ONNX input shape
dummy_input_for_onnx = vectorizer.transform(["dummy text"]).astype(np.float32)
onnx_model = convert_sklearn(
classifier,
initial_types=[('float_input', FloatTensorType([None, dummy_input_for_onnx.shape[1]]))],
target_opset=13,
)
onnx_path = os.path.join(model_version_dir, "model.onnx")
with open(onnx_path, "wb") as f:
f.write(onnx_model.SerializeToString())
print(f"ONNX model saved to {onnx_path}")
return onnx_path
def create_triton_config(num_features, num_classes, model_name_dir):
"""Creates the Triton model configuration file (config.pbtxt)."""
config_path = os.path.join(model_name_dir, "config.pbtxt")
config_content = f"""
name: "text_classifier"
platform: "onnxruntime_onnx"
max_batch_size: 64
input [
{{
name: "float_input"
data_type: TYPE_FP32
dims: [ -1, {num_features} ]
}}
]
output [
{{
name: "output_probabilities"
data_type: TYPE_FP32
dims: [ -1, {num_classes} ]
}}
]
"""
os.makedirs(model_name_dir, exist_ok=True)
with open(config_path, "w") as f:
f.write(config_content)
print(f"Triton config saved to {config_path}")
def perform_triton_inference(model_name, input_data, classes, server_url="localhost:8001"):
"""Sends inference request to Triton server and returns predicted label and latency."""
try:
triton_client = grpcclient.InferenceServerClient(url=server_url)
except Exception as e:
print(f"ERROR: Could not connect to Triton server at {server_url}. Please ensure Triton is running as per instructions.")
print(f"Details: {e}")
return None, None
inputs = []
outputs = []
input_name = "float_input" # Must match name in config.pbtxt
output_name = "output_probabilities" # Must match name in config.pbtxt
inputs.append(grpcclient.InferInput(input_name, input_data.shape, "FP32"))
inputs[0].set_content_from_numpy(input_data)
outputs.append(grpcclient.InferRequestedOutput(output_name, binary_data=True))
try:
start_time = time.perf_counter_ns()
response = triton_client.infer(
model_name=model_name,
inputs=inputs,
outputs=outputs,
client_timeout=5.0 # Timeout in seconds for the request
)
end_time = time.perf_counter_ns()
inference_latency_us = (end_time - start_time) / 1000.0
output_array = response.as_numpy(output_name)
predicted_class_idx = np.argmax(output_array, axis=1)[0]
predicted_label = classes[predicted_class_idx]
return predicted_label, inference_latency_us
except Exception as e:
print(f"ERROR: Inference request failed: {e}")
return None, None
def main():
RSS_FEED_URL = 'https://github.blog/engineering/feed/'
TRITON_MODEL_REPO_ROOT = "triton_model_repo"
MODEL_NAME = "text_classifier"
MODEL_VERSION = "1"
TRITON_SERVER_URL = "localhost:8001"
model_name_dir = os.path.join(TRITON_MODEL_REPO_ROOT, MODEL_NAME)
model_version_dir = os.path.join(model_name_dir, MODEL_VERSION)
print("--- Step 1: Fetching data and training a simple model ---")
training_posts = fetch_github_engineering_posts(RSS_FEED_URL)
if not training_posts:
print("No posts fetched for training. Exiting.")
return
# Use a subset for training if many posts, to keep demo fast
training_data_limit = 20
training_posts_subset = training_posts[:training_data_limit]
print(f"Training model with {len(training_posts_subset)} posts...")
vectorizer, classifier, classes = generate_labels_and_train_model(training_posts_subset)
num_features = vectorizer.max_features
num_classes = len(classes)
print(f"Model trained with {num_features} features and {num_classes} classes.")
print("\n--- Step 2: Converting Scikit-learn model to ONNX and saving artifacts ---")
onnx_path = convert_and_save_onnx_model(vectorizer, classifier, model_version_dir)
create_triton_config(num_features, num_classes, model_name_dir)
print("\n--- Step 3: Model deployment preparation complete ---")
print("To proceed, you need to start NVIDIA Triton Inference Server. Run this command in a separate terminal:")
print(f"docker run --rm -d --name triton -p 8000:8000 -p 8001:8001 -p 8002:8002 -v {os.path.abspath(TRITON_MODEL_REPO_ROOT)}:/models nvcr.io/nvidia/tritonserver:23.10-py3 tritonserver --model-repository=/models")
print("\nWaiting for 10 seconds for Triton to start and load the model. Please ensure the Docker command above is executed.")
time.sleep(10) # Give Triton time to spin up and load
print("\n--- Step 4: Performing real-time inference on new posts via Triton ---")
# Fetch some fresh posts for inference (can be the next set of posts not used for training)
inference_posts_raw = fetch_github_engineering_posts(RSS_FEED_URL)
if not inference_posts_raw:
print("No posts fetched for inference. Exiting.")
return
# Take the next 5 posts not used for training, or if less than `training_data_limit`, take the first 5
inference_posts = inference_posts_raw[training_data_limit:training_data_limit+5]
if not inference_posts and len(inference_posts_raw) > 0: # If not enough new posts, use some from the beginning
inference_posts = inference_posts_raw[:5]
elif not inference_posts:
print("No posts available for inference. Exiting.")
return
total_latency_us = 0
inference_count = 0
print(f"Attempting {len(inference_posts)} inferences...")
for i, post in enumerate(inference_posts):
title = post['title']
link = post['link']
# Preprocess the text using the *saved* vectorizer
# The vectorizer was saved in `model_version_dir`
loaded_vectorizer = joblib.load(os.path.join(model_version_dir, "vectorizer.joblib"))
preprocessed_input = loaded_vectorizer.transform([title]).astype(np.float32)
# Perform inference via Triton
predicted_label, latency = perform_triton_inference(MODEL_NAME, preprocessed_input, classes, TRITON_SERVER_URL)
if predicted_label:
print(f"\nPost {i+1}:")
print(f" Title: {title}")
print(f" Link: {link}")
print(f" Predicted Category: {predicted_label}")
print(f" Inference Latency (us): {latency:.2f}")
total_latency_us += latency
inference_count += 1
else:
print(f"\nPost {i+1}: Failed to get prediction for '{title}'. (Triton server likely not running or error during inference)")
if inference_count > 0:
avg_latency_us = total_latency_us / inference_count
print(f"\n--- Summary ---")
print(f"Total successful inferences: {inference_count}")
print(f"Average Inference Latency (us): {avg_latency_us:.2f}")
print(f"Average Inference Latency (ms): {avg_latency_us / 1000.0:.2f}")
print("\n--- Cleanup ---")
print("To stop the Triton server, run: docker stop triton")
print(f"To remove the model repository, run: rm -rf {TRITON_MODEL_REPO_ROOT}")
if __name__ == "__main__":
main()
Expected Output
When you run the script (after starting the Triton server in a separate terminal), you should see output similar to this. Latencies will vary based on your system and network:
--- Step 1: Fetching data and training a simple model ---
Fetching GitHub Engineering blog posts to train model...
Training model with 20 posts...
Model trained with 1000 features and 5 classes.
--- Step 2: Converting Scikit-learn model to ONNX and saving artifacts ---
ONNX model saved to triton_model_repo/text_classifier/1/model.onnx
Triton config saved to triton_model_repo/text_classifier
--- Step 3: Model deployment preparation complete ---
To proceed, you need to start NVIDIA Triton Inference Server. Run this command in a separate terminal:
docker run --rm -d --name triton -p 8000:8000 -p 8001:8001 -p 8002:8002 -v /path/to/your/triton_model_repo:/models nvcr.io/nvidia/tritonserver:23.10-py3 tritonserver --model-repository=/models
Waiting for 10 seconds for Triton to start and load the model. Please ensure the Docker command above is executed.
--- Step 4: Performing real-time inference on new posts via Triton ---
Attempting 5 inferences...
Post 1:
Title: The bridge to agent-scale Git infrastructure
Link: https://github.blog/engineering/the-bridge-to-agent-scale-git-infrastructure/
Predicted Category: Infrastructure
Inference Latency (us): 890.12
Post 2:
Title: Staying Evergreen: How we automated secure access and repo hygiene in a shared org
Link: https://github.blog/engineering/staying-evergreen-how-we-automated-secure-access-and-repo-hygiene-in-a-shared-org/
Predicted Category: Security
Inference Latency (us): 750.45
Post 3:
Title: From 2,500 requests to 24: Rebuilding star history without the stargazers
Link: https://github.blog/engineering/from-2500-requests-to-24-rebuilding-star-history-without-the-stargazers/
Predicted Category: Feature/Product
Inference Latency (us): 680.99
Post 4:
Title: Building Git infrastructure for agent-scale development
Link: https://github.blog/engineering/building-git-infrastructure-for-agent-scale-development/
Predicted Category: Infrastructure
Inference Latency (us): 710.23
Post 5:
Title: Improving site performance by shipping more CSS
Link: https://github.blog/engineering/architecture-optimization/improving-site-performance-by-shipping-more-css/
Predicted Category: Performance
Inference Latency (us): 650.78
--- Summary ---
Total successful inferences: 5
Average Inference Latency (us): 736.51
Average Inference Latency (ms): 0.74
--- Cleanup ---
To stop the Triton server, run: docker stop triton
To remove the model repository, run: rm -rf triton_model_repo
Limitations and Tradeoffs
While this approach delivers sub-millisecond inference, it's not a silver bullet. The primary limitation for text models is that the preprocessing step (like TF-IDF vectorization or tokenization for more complex models) often remains in Python. This means the overall end-to-end latency still includes Python overhead for preprocessing, which can be significant for very large texts or complex tokenizers. If your preprocessing is a bottleneck, you might need to explore options like implementing custom C++ preprocessors within Triton, or using frameworks like Spark NLP that can export parts of their pipelines to ONNX. Another tradeoff is the increased operational complexity. Managing a Triton server adds another component to your infrastructure, requiring Docker, Kubernetes (for scale), and monitoring. The initial setup and configuration can be more involved than simply serving a model with Flask or FastAPI. Finally, for simpler, low-traffic models, the overhead of Triton might outweigh its benefits; a lightweight Python server could suffice. This approach truly shines when you have multiple models, high QPS, or strict latency requirements.
Frequently Asked Questions
Why not convert the entire Scikit-learn pipeline, including TfidfVectorizer, to ONNX?
While skl2onnx can convert some Scikit-learn preprocessing steps, complex text vectorizers like TfidfVectorizer or custom tokenizers often involve string manipulation and dictionary lookups that are not efficiently represented in the ONNX graph format, which is primarily designed for numerical tensor operations. It's a common pattern to perform text preprocessing in Python (or another language) and pass the resulting numerical features (e.g., TF-IDF vectors, embeddings) to the ONNX model for inference.
Can Triton serve multiple models concurrently?
Absolutely. Triton is designed for this. You can place multiple model directories (e.g., text_classifier, image_segmenter) under the triton_model_repo, each with its own config.pbtxt and model files. Triton will load them all and manage concurrent inference requests across them, leveraging available hardware resources.
How does dynamic batching improve performance?
Dynamic batching allows Triton to collect multiple individual inference requests into a single, larger batch before sending it to the model. This is particularly effective for GPU inference, where processing a batch of inputs is significantly more efficient than processing them one by one due to parallelization. It increases throughput, but can slightly increase latency for individual requests as they wait for a batch to fill.
What if I have a PyTorch or TensorFlow model instead of Scikit-learn?
The core principle remains the same. PyTorch models can be exported to ONNX using torch.onnx.export(), and TensorFlow models using tf2onnx. Once in ONNX format, they can be served by Triton using the same onnxruntime_onnx backend, or their native backends if preferred (e.g., tensorflow_savedmodel, pytorch_libtorch).
Is Triton suitable for models much larger than this simple text classifier?
Yes, Triton is highly optimized for large models, including complex deep learning models like BERT, GPT, or large computer vision models. It provides features like model ensembles (chaining multiple models), custom backends, and efficient memory management to handle the demands of state-of-the-art AI.
What I'd Change
While this setup delivers impressive latency, for a truly production-hardened system, I'd make a few key advancements. First, I would containerize the Python client application itself, ensuring a consistent environment and simplifying deployment alongside Triton. Second, for the preprocessing step, instead of loading the TfidfVectorizer from disk on every inference call (which I did for simplicity here), I would load it once at client application startup and keep it in memory. For even higher performance or more complex preprocessing, I'd explore writing a custom Python backend for Triton, allowing the TfidfVectorizer to run directly within the Triton server process, eliminating network hops for raw text and potentially benefiting from Triton's scheduling. Finally, I'd integrate comprehensive monitoring and alerting using Prometheus and Grafana, leveraging Triton's built-in metrics endpoint to track latency, throughput, and error rates, ensuring we proactively catch performance regressions and operational issues before they impact users. Sub-millisecond inference isn't just about the technology; it's about the entire operational pipeline.
Next Steps
Now that you've mastered serving a simple model on Triton, challenge yourself to deploy a more complex model, perhaps a fine-tuned BERT model for sentiment analysis, using the same ONNX and Triton principles. How would you handle the tokenization and embedding steps efficiently? Consider building a custom Python backend for Triton to encapsulate the entire text preprocessing and inference pipeline within the server itself.