PyTriton のセットアップに苦労した話¶

PyTriton とは¶

PyTriton は、NVIDIA が提供する Triton Inference Server の Python ラッパー。Triton Inference Server は、機械学習モデルを Web API として提供するためのフレームワークで TensorRT や TensorFlow、ONNX Runtime などのバックエンドをサポートしている。また、Dynamic Batcher や Model Repository、複数モデル対応、Metrics などの機能を提供しており、多機能でハイパフォーマンスなのが特徴。設定ファイルは Protocol Buffers で記述するのが基本の形になるが、この設定を Python で書くことができるのが PyTriton。

PyTriton のセットアップ¶

PyTriton のセットアップでうまくいかないことがあった。最終的には解決したが、その過程を記録しておく。

うまくできた方法¶

まず、環境構築が成功した例を示す。次の環境で PyTriton をセットアップできた。

Ubuntu 22.04
Python 3.10
PyTriton 0.5.13

# Use Ubuntu 22.04 as base image
FROM ubuntu:22.04

# Install necessary packages
RUN apt update -y && apt install -y --fix-missing software-properties-common

# Add repository with various Python versions
RUN add-apt-repository ppa:deadsnakes/ppa -y

#Set timezone
RUN apt-get update && \
    apt-get install -yq tzdata && \
    ln -fs /usr/share/zoneinfo/Asia/Seoul /etc/localtime && \
    dpkg-reconfigure -f noninteractive tzdata

# Install Python 3.8 and required libraries
RUN apt install -y --fix-missing python3.10 libpython3.10 python3.10-distutils python3-pip \
     build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev libsqlite3-dev libreadline-dev \
     libffi-dev curl libbz2-dev pkg-config make

# Install nvidia-pytriton using pip
RUN python3.10 -m pip install nvidia-pytriton
WORKDIR /workspace

上記の Dockerfile を使って Docker イメージをビルドし、次のコードを実行すると問題なく PyTriton を動かすことができた。

example.py
from pytriton.decorators import batch
from pytriton.model_config import ModelConfig, Tensor
from pytriton.triton import Triton, TritonConfig
import logging
import numpy as np

logging.basicConfig(level=logging.DEBUG)

@batch
def infer_fn(**inputs: np.ndarray):
    return {"token": np.zeros((1,1), dtype=np.int64) }

MAX_BATCH_SIZE = 10

with Triton(
        config=TritonConfig(http_port=1111, grpc_port=2222, metrics_port=3333)
    ) as triton:
        triton.bind(
            model_name="Whisper",
            infer_func=infer_fn,
            inputs=[
                Tensor(name="audio_feature", dtype=np.float32, shape=(80, 3000)),
            ],
            outputs=[
                Tensor(name="token", dtype=np.int64, shape=(-1,)),
            ],
            config=ModelConfig(max_batch_size=MAX_BATCH_SIZE),
            strict=True,
        )
        triton.serve()

うまくいかなかった方法¶

上記のように Python 3.10 で PyTriton を使うことができたが、それ以外の Python バージョンだとうまくいかなかった。Python 3.8 や 3.9, 3.11 で試したがエラーが発生した。同様の問題がissueとして報告されている。

この問題の原因の切り分けとして、apt を使わずに uv を使って Python をインストールする方法を試してみたが、エラーが発生した。エラーの内容は次の通り。

1
2
3

...省略
I1220 13:00:54.033669 864 pb_stub.cc:298] Failed to initialize Python stub for auto-complete: ModuleNotFoundError: No module named '_contextvars
...

エラー文を読むと、_contextvars モジュールが見つからないというエラーが出ている。しかし、次のコマンドで _contextvars モジュールをエラーなくインストールできることを確認した。

$ uv run python -c "import _contextvars; print('done import')"

done import

uv run python example.py のように実行するとエラーが発生するが、調査に時間がかかりそうだったので、Python のバージョンを変更して解決することにした。