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.10 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 バージョンだとうまくいかなかった。deadsnake から 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 のバージョンを変更して解決することにした。

追試¶

追試 1 ubuntu のバージョンを変える¶

ubuntu22.04 だと apt 経由で Python を setup すると 3.10 がインストールされる。ubuntu のバージョンを変えると apt でインストールされる Python のバージョンが変わるので、ubuntu のバージョンを変えて pytriton の setup をするとうまくいくのか検証

次のように ubnutu24 だと apt でインストールされる Python のバージョンは 3.12 になる。これを実行してみるとエラーなく動くことを確認できた。

FROM ubuntu:24.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

RUN apt install -y --fix-missing python3 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 python3-venv -y && python3 -m venv /opt/venv

# Enable venv
ENV PATH="/opt/venv/bin:$PATH"
# Install nvidia-pytriton using pip
RUN python3 -m pip install nvidia-pytriton

WORKDIR /workspace

追試 2 base image を Python:3.12-bullseye に変更¶

base image をpython:3.12-bullseyeに変更して、PyTriton のセットアップをやってみた。pytriton は pyzmq に依存しているが、pyzmq でエラーが発生し、PyTriton のセットアップが失敗した。

1
2
3

...
70.45 Failed to build pyzmq
70.56 ERROR: ERROR: Failed to build installable wheels for some pyproject.toml based projects (pyzmq)