📝

Triton - Python Backend 実装の備忘録

2024/05/24に公開

Triton Inference Server の Python Backend の実装をした際の備忘録．発見があれば順次追記していく．

Triton Python Backend とは？

Triton Inference Server において，Python で書かれたモデルをサービングするための環境．PyTorch や Tensorflow のモデルはもちろん，基本的に Python スクリプトであれば何でもホスティングできる．C++のコードの実装が不要．

基本的な実装 (完成系)

実行したい Python コードをmodel.py に実装し，モデルのデプロイ方法を定義するconfig.pbtxtを用意すれば，モデルをホスティングできる．ただし，ファイルを下記のフォルダ構成で配置することが必要．

model_repository
└── my_model/ ... モデル名のフォルダ. config内のモデル名と一致する必要がある.
    ├── config.pbptxt ... 設定ファイル．モデルの入出力やサービングの方法を記述．
    └── 1 ... モデルのバージョン情報
        └── model.py ... 実行したいPythonスクリプトを実装するファイル．

`model.py`

実行したい Python スクリプトを実装する場所．
関数や Pytorch/Tensorflow のモデルなど，Python で書けるならば基本的になんでも実行できる (e.g., JAX)．
任意のクラスが実行できるわけではなく，テンプレート (TritonPythonModels)に従って，3 つのメソッドを実装．
- initialize(self, args) (Optional)
  - モデルの初期化を行う関数．
  - このメソッド (+ Warmup) が終了すると，Triton 上のモデルのステータスが "Ready" になる．
- execute(self, requests) (Required)
  - 実行したい Python モデルやスクリプトは，このメソッドに実装する．
  - メソッドの入出力は，テンソルではなく list[pb_utils.InferenceRequest] / list[pb_utils.InferenceResponse]
- finalize(self) (Optional)
  - Triton からモデルを降ろす前に実行されるメソッド．
  - 不要なプロセスやファイルの削除などを実装する．

実装サンプルは以下の通り.

class TritonPythonModel:

    def initialize(self, args):
        model_config = json.loads(args["model_config"])
        self.output_tensor_configs = model_config["output"]

        # Init model
        self.model = DummyModel().to(device=_DEVICE)
        self.model.eval()

    def execute(self, requests):
        # Prepare input tensors (InferenceRequestをテンソルに変換)
        preproc_outputs: dict[str, torch.Tensor] = convert_infer_request_to_tensors(
            requests,
            tensor_names=[DUMMY_IMAGE_INPUT_TENSOR_NAME],
            device=_DEVICE,
        )

        # Inference (モデルの実行)
        model_outputs: dict[str, torch.Tensor] = self.model(
            pixel_values_input=preproc_outputs[DUMMY_IMAGE_INPUT_TENSOR_NAME]
        )

        # Return InferenceResponse objects for each request (テンソルをInferenceResponseに変換)
        response = build_inference_response(
            model_outputs,
            output_tensor_configs=self.output_tensor_configs,
        )
        return [response]

`config.pbtxt` ^[1]

大まかに以下の項目が設定できる．

name: モデル名
Backend: "python" を指定
max_batch_size: 最大のバッチサイズ．余計なオーバーヘッドをなくすために，下記例では Dynamic Batching を使用しない例．
input/output: 入出力
instance_grou: モデルの実行方法 (CPU/GPU, いくつインスタンスを立てるか etc)

サンプル (Warmup 付き)

name: "my_model"
backend: "python"
max_batch_size: 0
input [
    {
        name: "image_input"
        data_type: TYPE_FP32
        dims: [ -1, 3, 1080, 1920 ]
    }
]
output [
    {
        name: "image_output"
        data_type: TYPE_FP32
        dims: [ -1, 3,  1080, 1920 ]
    }
]
instance_group [
    // ホストするモデルの数とパターンを指定．
    // GPU実行の場合は，各GPUで実行するモデル数を指定できる (e.g., GPU:0で1モデル，GPU:1で2モデル).
    {
      count: 1
      kind: KIND_AUTO
    }
]
model_warmup [
    {
        name : "Warmup Sample"
        batch_size: 1
        inputs {
            key: "image_input"
            value: {
                data_type: TYPE_FP32
                dims: [ 8, 3, 1080, 1920 ]
                random_data: true
            }
        }
        count: 5
    }
]
// GPU実行には必須
parameters: { key: "FORCE_CPU_ONLY_INPUT_TENSORS" value: {string_value:"no"}}

設定ファイルのスキーマは Protobuf (model_config.proto) で定義される．手書きしても良いが，モデルの Signature を使って，自動生成した方が効率的ではないか？

GPU 環境での実行

GPU 実行をする場合，入出力は dlpack を用いることで，データのコピーなしで Client-Model, Model-Model 間のテンソルのやり取りが可能になる．

config.pbtxt に，FORCE_CPU_ONLY_INPUT_TENSORS の設定が必要．この記述がないと，クライアントで GPU テンソルを指定しても，モデル実行前に CPU にコピーされ，モデル内で再度 GPU にコピーされるので，無駄が発生する．(GPU テンソルを渡したのに，余計に遅くなる．) ^[2]

Datatypes ^[3]

Triton で扱えるデータ型の一覧．フレームワークによって，サポートされている方が異なるので注意．
文字列 (TYPE_STRING) も扱える．ただし，DeepStream で Triton を使用した場合 (nvinferserve)，STRING 型を使用するとメモリリークと思われるエラーが発生した．

Tips: Model Warmup

注意点

Warmup を行うと，その際の実行結果が Triton の Statistics 情報に含まれる．
- TensorRT など 1 回目の実行に時間がかかる場合，外れ値に引っ張られて統計情報がおかしくなるので注意．
GPU 上のテンソルは，Warmup では選択できない．
- モデルを GPU 実行する場合でも，model.py において，CPU 上のテンソルを GPU に転送する処理が必要

実行例

(1) ランダムな値で初期化されたテンソルを 5 回モデルに入力する.

model_warmup [
    {
        name : "Warmup Sample"
        batch_size: 1
        inputs {
            key: "image_input"
            value: {
                data_type: TYPE_FP32
                dims: [ 8, 3, 1080, 1920 ]
                random_data: true
            }
        }
        count: 5 // リクエストの送信回数
    }
]

TODO (今後調べたいこと)

TritonPythonModels は，バックエンド上でどのように読み込まれるのか？ (Deepstream のように Pybind を使っている?)

脚注

Triton Python Backend とは？

基本的な実装 (完成系)

model.py

config.pbtxt [1]