Truss

The simplest way to serve AI/ML models in production

Why Truss?

• Write once, run anywhere: Package and test model code, weights, and dependencies with a model server that behaves the same in development and production.
• Fast developer loop: Implement your model with fast feedback from a live reload server, and skip Docker and Kubernetes configuration with a batteries-included model serving environment.
• Support for all Python frameworks: From transformers and diffusors to PyTorch and Tensorflow to XGBoost and sklearn, Truss supports models created with any framework, even entirely custom models.

See Trusses for popular models including:

• 🦅 Falcon 40B
• 🧙 WizardLM
• 🎨 Stable Diffusion
• 🗣 Whisper

and dozens more examples.

Installation

Install Truss with:

pip install --upgrade truss

Quickstart

As a quick example, we'll package a text classification pipeline from the open-source transformers package.

Create a Truss

To get started, create a Truss with the following terminal command:

truss init text-classification

This will create an empty Truss at ./text-classification.

Implement the model

The model serving code goes in ./text-classification/model/model.py in your newly created Truss.

from typing import List
from transformers import pipeline


class Model:
    def __init__(self, **kwargs) -> None:
        self._model = None

    def load(self):
        self._model = pipeline("text-classification")

    def predict(self, model_input: str) -> List:
        return self._model(model_input)

There are two functions to implement:

• load() runs once when the model is spun up and is responsible for initializing self._model
• predict() runs each time the model is invoked and handles the inference. It can use any JSON-serializable type as input and output.

Add model dependencies

The pipeline model relies on Transformers and PyTorch. These dependencies must be specified in the Truss config.

In ./text-classification/config.yaml, find the line requirements. Replace the empty list with:

requirements:
  - torch==2.0.1
  - transformers==4.30.0

No other configuration needs to be changed.

Deployment

You can deploy a Truss to your Baseten account with:

cd ./text-classification
truss push

Truss will support other remotes soon, starting with AWS SageMaker.

Truss contributors

Truss is backed by Baseten and built in collaboration with ML engineers worldwide. Special thanks to Stephan Auerhahn @ stability.ai and Daniel Sarfati @ Salad Technologies for their contributions.

We enthusiastically welcome contributions in accordance with our contributors' guide and code of conduct.