Adding Triton backend support

[aspctu]

Overview

This PR adds support for Triton as a backend for Truss. Specifically, this PR contains logic for

1. Building Docker images that run user code within Triton
2. Generating the appropriate configs for Triton based on a user's config.yaml and model.py
3. Handling the conversion between Triton types and Python types

Logic around testing this flow will be in a follow-up PR. It's a significant testing suite and requires running tests within the Triton docker container.

Quickstart

Quickstart repo here

1. git clone https://github.com/aspctu/bert-triton-truss
2. truss image build-context ./bert-truss-context ./bert-truss
3. cd ./bert-truss-context
4. docker build ./
5. docker run --gpus=all -p8080:8080 -p8000:8000 -it (image id)

Follow the README.md in the repo above to invoke the model.

Introduction

Triton is a high-performance model serving backend developed by NVidia. For most models (outside of LLMs), it's advantageous to use Triton as the backend server. This is due to various server features that are attractive to maximizing GPU utilization and memory.

This PR introduces a simplified developer experience to enabling users to tap into some of this functionality within Truss. It's worth nothing that there is a lot of functionality in Triton that is not supported here (such as decoupled mode or ensemble models).

To enable Triton, a user needs to do a couple things:

1. Update their config.yaml to contain the following information (automatically done if the truss is created via truss init)

build:
  model_server: TRITON

7. Define 2 Pydantic classes in their model.py that correspond to the Input and Output of their model (example below)

from pydantic import BaseModel, conlist

class Input(BaseModel):
    text: str

class Output(BaseModel):
    text: str
    embedding: conlist(float, min_length=768, max_length=768)

...

8. Update their predict function to accept a List[Input] and produce a List[Output]

def predict(inputs: List[Input]) -> List[Output]:
...

TODOs

• Parametrize GPU / CPU deployment in config.pbtxt
• Look into truss image build failing to do anything

[bolasim]

lgtm. just a few comments/questions. Lemme know when it's ready for review.

[squidarth]

Took a first pass. Overall, I like the structure that we went with here!

Nits aside, main comments here are around error-handling (and my own questions about what kind of assumptions are fair to make about the input)

[bolasim]

let's merge this and use it then iterate on smaller PRs

[amiruci]

🤩 🤩 🤩