今月も Triton Inference Server のリリース內容について、概要をお屆けします。「Triton Inference Server って何?」という方は、以下の記事などをご確認ください。

• GPU に推論を: Triton Inference Server でかんたんデプロイ
• NVIDIA Triton Inference Server を使用したエッジでの AI モデルの展開の簡素化

What’s New in 2.22.0 (NGC 22.05)

リリース ノート本體は https://github.com/triton-inference-server/server/releases/tag/v2.22.0 です。今月のリリースには以下の機能や改善などが含まれています。

• In-Process API に Java binding が追加されました
  • 従來 C / C++ 向けに libtritonserver.so として実裝されていたものと同様の機構が、Java でも利用可能になります
• Python backend で decoupled API が beta release されました
  • 1 つのリクエストに対して必ず 1 つのレスポンスを返すわけ「ではない」使い方が可能になります
• Model Control API で、load エンドポイントにリクエストする際、シリアライズしたモデル ファイルを渡せるようになります
  • API を介して、動的にモデルを差し替えられるようになります
• データ型として bfloat16 がサポートに追加されました
• PyTorch backend で、文字列サポートが改善されました
• Model Control Mode EXPLICIT で、全モデルを起動時に一括ロードする機能がサポートされました
  • --load-model オプションにワイルドカードを指定すればよいようです
• gRPC クライアント利用時、channel の動作をカスタマイズするためのパラメーターを渡せるようになります
• In-Process API に、動的にモデル リポジトリを登録解除する API が追加されました
• ビルド パイプラインが改善されました
• ONNX Runtime backend の ONNX Runtime バージョンが 1.11.1 に更新されました

今月も先月に引き続き更新が多めです。「Python backend で decoupled API が beta release」、「Model Control API で、load エンドポイントにリクエストする際、シリアライズしたモデル ファイルを渡せるように」あたりが特に注目すべき更新ではないでしょうか。

実は以前から、Triton の backend では decoupled API をサポートしていました。その一例として、repeat_backend などの參考実裝も提供されており、使おうと思えば使える狀態ではあったのですが、一方で Python backend ではサポートされておらず、利用シーンが限定されていたことも事実です。今回の対応で、decoupled API が Python backend でもサポートされることになりますので、より多くの狀況での活用が期待できます。

具體的な使い方としては、通常の Python backend の実裝に対して以下の點を変更、追加する必要があります。

• ModelTransactionPolicy を追加 (config.pbtxt)
• InferenceRequest.get_response_sender() で InferenceResponseSender のオブジェクトを取得する (model.py)
• InferenceResponseSender.send() を使ってレスポンスを返す (model.py)
• TritonPythonModel.execute() の戻り値は「必ず」 None にする (model.py)

加えてクライアント側では、decoupled API が対応するプロトコルは gRPC の (bidirectional) streaming のみに限定されるという制約もあります。これはつまり、HTTP やその他の gRPC では decoupled API は使えない、ということを意味します。

その他、詳細などは公式ドキュメントやサンプル (repeat_xxx と square_xxx の 2 種類) を參照いただければと思うのですが、ちょっと勘違いしやすい點だけ補足説明しておきます。

前記の通りモデル側 (=model.py の中) では、InferenceResponseSender のオブジェクトを利用してクライアントへのレスポンスを返します。通常は、TritonPythonModel.execute() の戻り値として InferenceResponse のリストを返すことによってレスポンスを構成していたことと比較すると、クライアントとのやり取りの部分が大きく異なります。実際には、InferenceResponseSender.send() にレスポンス相當のテンソルを渡すことで、クライアントへ値を返すのですが、この関數は 1 つのリクエストに対して任意の回數呼び出すことが可能で、最後の呼び出しの際に TRITONSERVER_RESPONSE_COMPLETE_FINAL をフラグとして渡すことによって、あるリクエストに対する一連のレスポンスが完了したことを示します。例えば以下のようなイメージです。

response_sender = request.get_response_sender()
response_sender.send(
    pb_utils.InferenceResponse(output_tensors=[pb_utils.Tensor(
        "output", np.array(...)
    )]),
)
response_sender.send(
    pb_utils.InferenceResponse(output_tensors=[pb_utils.Tensor(
        "output", np.array(...)
    )]),
    flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
)

一方、ここで勘違いしやすいのは、decoupled API のドキュメントでも言及されている通り、リクエストに対してレスポンスを返さない狀況と組み合わせる場合です。レスポンスを返さないことを想定しているため、あるリクエストに対して InferenceResponseSender.send() を一切呼び出さない、たとえば以下のような実裝が正解と思ってしまうのですが、

# 音聲認識のような、一連のリクエスト列の最後に結果を一回だけ返す、という用途を想定
# 前処理や推論のような主要な処理は、これ以前の部分で実施済み
response_sender = request.get_response_sender()
if is_last_request:
    # 系列の最後のリクエストにのみデータを返す
    response_sender.send(
        pb_utils.InferenceResponse(output_tensors=[pb_utils.Tensor(
            "output", np.array(...)
        )]),
        flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
    )
else:
    # なにもしない
    pass

これは誤りで、このように実裝するとクライアントの終了処理が正常に完了せず、ブロックしてしまいます。正しくは、サーバーが受信したすべてのリクエストに対して、必ず終了フラグを設定する必要があります。

response_sender = request.get_response_sender()
if is_last_request:
    # 系列の最後のリクエストにのみデータを返す
    response_sender.send(
        pb_utils.InferenceResponse(output_tensors=[pb_utils.Tensor(
            "output", np.array(...)
        )]),
        flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
    )
else:
    # それ以外は終了フラグだけ返す
    response_sender.send(
        flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL
    )

このとおり、若干追加の実裝が必要ではありますが、クライアント  サーバー間の通信回數を減らす意味では重要な機能ですので、うまく活用いただければと思います。

続いて Model Control API の機能追加です。従來はサーバー上などに配置済みのファイルしか、load エンドポイントにリクエストする際に読み込むことはできませんでした。今回の対応で、/v2/repository/models/${MODEL_NAME}/load をコールする際に、ロードさせたいモデルのファイルを同時に送信することで、サーバー上のファイルだけではなく、動的に任意のモデルを読み込ませることが可能になります。(注: Triton Inference Server 自體に認証等の機構はついていないため、セキュリティなどには十分ご注意ください)

Triton のエンドポイントとしては、base64 エンコードされたファイルを受け付ける仕様ですが、Triton の公式クライアント ライブラリを利用する場合の例は以下のようになります。

with client_module.InferenceServerClient(server_addr) as client:
    current_model_config = client.get_model_config(modelname)

    target_files = {}
    with open("./model.pt", "rb") as f:
        filecontent = f.read()
    target_files[f"file:1/model.pt"] = filecontent
    with open("./data.json", "rb") as f:
        filecontent = f.read()
    target_files[f"file:1/data.json"] = filecontent

    client.load_model(
        modelname,
        config=json.dumps(current_model_config),
        files=target_files)

上記の例からもわかる通り、以下のような仕様となっています。

• API 経由でモデル更新する場合、config の指定が必須
  • モデル ファイル以外変更しないのであれば、一度サーバーから取得してそれをそのまま送り返す等の対応が必要
• モデル ファイルと同時に使う必要のあるファイルが存在する場合、複數ファイルを同時送信することも可能
  • たとえば上記コード上の data.json は、言語モデルであれば vocabulary のファイルや、サーバー サイドで分類モデルのクラス名へのマッピングをする場合、クラス名一覧などのリソース ファイルといったもの
  • ただし ONNX などのモデル ファイルは常時必須で、その他のリソース ファイルだけを送りたい場合でも、モデル ファイルも同時に送る必要あり
• 対象のファイル パスは以下のフォーマットで指定
  • “file:<version>/<filepath>”
  • バージョン指定が必須となっている點に注意

その他、詳細は Model Repository Extension – Load などを參照いただければと思います。

What’s New に言及されていないアップデート

今月の更新や機能追加は What’s New でほぼ言及されていたようで、それ以外の主なものは以下の一件のみでした。

• Amazon SageMaker の Multi-Model Endpoint (MME) に対応したようです (PR#4181)

こちらの更新はまだ experimental な実裝のようですが、Amazon SageMaker で Triton を利用する際に、複數モデルを単一のエンドポイントでデプロイする機能である、Multi-Model Endpoint (MME) への対応が追加されたようです。各クラウド サービス上の Managed Service との連攜も順次進められていることを伺わせます。

まとめ

今月は、特により多くのユースケースに対応するための追加機能や更新が多かったように思われます。音聲認識のような系列データ向けの decoupled API の強化や、既存の様々なアプリケーションとの連攜を容易にするための In-Process API への Java binding 追加など、まだ beta release なものも含まれますが、早めにお試しいただくことで導入までの時間を短縮できるのではないでしょうか。