コンテンツにスキップ

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

PyTriton のセットアップで Python のモジュールが見つからないエラーに遭遇し、動作を確認した環境へ切り替えて回避しました。Ubuntu 22.04 と Python 3.10、Ubuntu 24.04 と Python 3.12 で動いた Dockerfile と、他の環境で出たエラーを紹介します。エラーの原因は特定できていません。

PyTriton とは

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

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 バージョンではうまくいきませんでした。deadsnakes から Python 3.8 や 3.9、3.11 をインストールして試しましたが、エラーが発生しました。同様の問題として、_ctypesモジュールが見つからないというエラーがModuleNotFoundError: No module named '_ctypes' error when run pytriton server with 0.5.0 · Issue #56で報告されています。

エラーの原因を切り分けるため、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モジュールをエラーなくインポートできることを確認しました。

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

done import

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

追試

Python のバージョンを変える以外の解決策も試したので、追試の結果を記録します。

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

Ubuntu 22.04 で apt 経由で Python をインストールすると、3.10 が入ります。Ubuntu のバージョンを変えると apt でインストールされる Python のバージョンも変わるので、Ubuntu のバージョンを変えて PyTriton のセットアップがうまくいくかを検証します。

Ubuntu 24.04 では、apt でインストールされる Python のバージョンは 3.12 になります。次の Dockerfile でビルドしたイメージを起動すると、エラーなく動くことを確認できました。

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)

まとめ

同じ PyTriton のセットアップでも、ベースイメージや Python のバージョンの組み合わせ次第で明暗が分かれました。Ubuntu 22.04 と Python 3.10、Ubuntu 24.04 と Python 3.12 の組み合わせでは動作しましたが、python:3.12-bullseyeイメージでは pyzmq のビルドに失敗しました。エラーの原因ははっきりとは特定できていません。ただ、原因を深追いするよりも、動作実績のある組み合わせへ切り替える方が早く解決できました。