truss

项目概述

Truss 是一个旨在帮助用户简化在生产环境中部署 AI/ML 模型的工具。它为模型的编写、测试、打包和部署提供了一体化的解决方案，适用于各种 Python 框架，并且无需用户花费大量时间进行 Docker 和 Kubernetes 的配置。

为什么选择 Truss？

一次编写，随处运行

Truss 提供了一个在开发和生产环境中表现一致的模型服务器，用户可以在该服务器上打包和测试模型代码、模型权重以及相关依赖项。这样一来，用户只需要编写一次代码，就能在不同环境中运行，非常方便实用。

快速迭代

通过 Truss，开发者能够在实施模型的过程中，借助实时重载服务器获得快速反馈。Truss 内置电池式模型服务环境，大大加速了开发迭代流程，使得跳过繁琐的 Docker 和 Kubernetes 配置成为可能。

支持所有 Python 框架

无论是流行的 transformers 和 diffusers，还是主流的 PyTorch 和 TensorFlow，甚至是高性能的 TensorRT 和 Triton，Truss 均提供了多样化的框架支持。

安装指导

Truss 的安装简单快捷，用户只需在终端中执行以下命令即可：

pip install --upgrade truss

快速入门

接下来，我们以一个文本分类的管道为例，演示使用 Truss 快速打包并部署模型的过程。

创建 Truss

要开始，就在终端中输入以下命令以创建一个 Truss：

truss init text-classification

根据提示，为你的 Truss 取一个名字，例如 Text classification。然后进入新创建的目录：

cd text-classification

实现模型

Truss 中最核心的文件之一是 model/model.py，在此文件中，用户需要编写一个 Model 类，提供模型与模型服务器之间的接口。

在 Model 类中，用户必须实现两个方法：

• load() 用于将模型加载到服务器上，仅在服务器启动或更新时运行一次。
• predict() 用于处理每次服务器调用的模型推理。

以下是文本分类模型的完整 model/model.py 示例：

from transformers import pipeline

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

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

    def predict(self, model_input):
        return self._model(model_input)

添加模型依赖

在 Truss 中，另一个重要的文件是 config.yaml，用于配置模型的服务环境。首先，需要指定模型的依赖项：

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

其他配置可以保持默认，无需复杂更改。

部署与调用

Truss 提供了由 Baseten 维护的 ML 模型生产环境。用户可以将 Baseten 用作远程主机来托管模型，未来还将支持 AWS SageMaker 等其他平台。

获取 API 密钥

为了在 Baseten 上远程托管模型，需要获取一个 Baseten API 密钥。用户可以通过 注册 Baseten 账号 获取免费试用的 API 密钥。

部署模型

获得 API 密钥后，用户可以通过以下命令部署模型：

truss push

用户可以通过 Baseten 的模型仪表板实时监控模型的部署状态。

模型调用

部署完成后，可以在终端调用模型进行推理：

调用命令

truss predict -d '"Truss is awesome!"'

返回结果

[
  {
    "label": "POSITIVE",
    "score": 0.999873161315918
  }
]

贡献者

Truss 由 Baseten 支持开发，并与全球的机器学习工程师协作完成。特别感谢来自 stability.ai 的 Stephan Auerhahn 和来自 Salad Technologies 的 Daniel Sarfati 的贡献。

Truss 欢迎各界的积极贡献，有意者可以查阅我们的[贡献指南]和[行为准则]。