Triton Inference Serverを色々な環境から使いたい: TypeScript (Node.js) 編

イントロ: 色々な環境から使えるハイパフォーマンスで堅牢な ML 推論サーバー、ありますか？

機械学習モデルを運用するときの推論サーバーの実装方法は色々あります。 FastAPI 等でラップして API を生やすのは初手としては楽な一方、ある程度の規模で複数のモデルを運用するとなると、推論用途の API をうまく定義したり、データモデルをクライアント・サーバーで実装したり、推論リクエストのスケジューリングだったりと、困ることが増えてくる印象です。

また、ある程度の規模のシステムになると複数のプログラミング言語で構成されることもままあるので、各言語向けにクライアントを実装する手間や、リクエストスキーマの品質をどのように保つかという悩みも出てきます。

Triton Inference Server

最近では ML 推論サーバーがいくつも提案されていますが、この記事では Triton Inference Server という機械学習向け推論サーバーの OSS 実装に注目します。

Triton Inference Server 自体の紹介や ML モデルそのものをどうやってデプロイするかという内容は別項に書くかもしれません。

Triton Inference Server の特徴のうち、先に挙げた課題に効きそうなものは以下の点でしょうか¹。

様々な機械学習フレームワークをバックエンドとして使える
複数のモデルを同時に運用できる
Auto-batching などのスケジューリング機能を備える
KServe 互換の gRPC API を備える
簡易 DAG 機能で前処理・後処理を含めてデプロイできる²

TypeScript から使いたい！

ここでは特に Node.js (TypeScript) で実装されたサービスから利用する場合を例に挙げてみます。公式の Triton Inference Server Client が提供されていない言語であるため、例としても都合がよいと思ったからです³。

ちなみに、C++, Java, Python については公式実装が存在し、また Go, Java, JavaScript (Node) については gRPC からのコード生成のサンプルが公開されています。

gRPC 定義からの TypeScript 型情報付きクライアントコード生成

ここでは gRPC 定義から TypeScript (Node.js) 向けのクライアントを生成してみます。流れは以下のようになります:

Triton Inference Server の gRPC 定義の取得
proto-loader-gen-types を利用した TypeScript 型定義付きのリフレクションベースのコード生成
自動生成コードに対する簡易ラッパーの実装

TL;DR: こちらに例をまとめてあります。
https://github.com/tuxedocat/triton-client-polyglot-example

gRPC 定義の取得

Triton Inference Server の gRPC 定義のリポジトリから取得します。 Server 実装のバージョンと合った Tag を指定してください。 Tag でのバージョン体系は r{year}.{month} となので、Releases との対応付けには注意が必要です。

# https://github.com/tuxedocat/triton-client-polyglot-example/blob/c0cb5e2e76170d8f8c3599f9c9c0d4d2acb011ea/client/typescript/package.json#L15
git clone https://github.com/triton-inference-server/common -b r22.02 ./upstream
mkdir proto && mv ./upstream/protobuf/*.proto ./proto/

gRPC 定義からのクライアントコードの生成

1. Node.js 向けの gRPC 公式チュートリアルと 2. CADDi さんの Tech Blog の記事を主に参考にしました。

参考記事のように、TypeScript 向けの gRPC クライアントのコード生成に関しては大きく２つの方法があります。一方はコード生成に grpc_tools_node_proto　と protoc-gen-ts を利用する静的に利用できるコードの方法であり、もう一方は proto-loader を利用しつつ TypeScript の型定義を与える方法です。ここでの「静的」は、生成したクライアントの実行時に gRPC 定義が不要という意味です。

ここでは proto-loader を用いる方法を採用します。理由は、先に挙げた参考記事 2 で述べられているこれら２つの特徴のうち、

static_codegen

値の指定は、Message ごとに生成される setter を使用する

コンストラクタからも指定できるが使いにくい

dynamic_codegen

値の指定は、Message のコンストラクタから値を渡す or 直接指定する

内部的に JSON から Message を作成する protobuf.js の fromObject が使われる

— static_codegen と dynamic_codegen の比較

という箇所について、ML 推論サーバーは複数のロールの人たちが触れる・メンテナンスすることになるため、なるべく gRPC を意識しないで良い API にしたかったためです。

依存パッケージは以下のようになります。Node.js と TypeScript のバージョンは、node==14.17.5, typescript==4.3.5 で試しています。

// package.json
// https://github.com/tuxedocat/triton-client-polyglot-example/blob/c0cb5e2e76170d8f8c3599f9c9c0d4d2acb011ea/client/typescript/package.json#L29-L44
  "dependencies": {
    "@grpc/grpc-js": "^1.6.2",
    "@grpc/proto-loader": "^0.6.9"
  },
  "devDependencies": {
    "@types/node": "^14.0",
    "grpc-tools": "^1.11.0",
    "ts-node": "*",
    "typescript": "^4.3.5",
  },
  ...

型定義を生成します。先の参考記事 2 と異なる点として、最近のバージョンの proto-loader には proto-loader-gen-types という TypeScript 型定義を生成するユーティリティが追加されたことです。ここではgRPC のドキュメントを参考にしました。

./node_modules/bin/proto-loader-gen-types \
  --longs=String \
  --enums=String \
  --defaults \
  --oneofs \
  --keepCase \
  --grpcLib=@grpc/grpc-js \
  --outDir=./generated \
  proto/*.proto"

./generated 内にコードが生成されます。

generated/
├── inference
│   ├── たくさん...
├── grpc_service.ts
└── model_config.ts

パッケージ化する必要がなければ、この状態で proto-loader と組み合わせて使うことができるようになります。自動生成したクライアントコードを利用する際の例としては、Triton Inference Server のリポジトリにある JavaScript 版の例も参考になります。

import { loadPackageDefinition } from '@grpc/grpc-js'
import { loadSync } from '@grpc/proto-loader'
import { ProtoGrpcType } from './generated/grpc_service'
 
const PROTO_IMPORT_PATH = __dirname + '/proto'
const PROTO_PATH = PROTO_IMPORT_PATH + '/grpc_service.proto'
const packageDefinition = loadSync(PROTO_PATH, {
  includeDirs: [PROTO_IMPORT_PATH],
  keepCase: true,
  longs: String,
  enums: String,
  defaults: true,
  oneofs: true,
})
const loadedPackageDefinition = loadPackageDefinition(
  packageDefinition,
) as unknown as ProtoGrpcType
export const triton = loadedPackageDefinition.inference

簡易ラッパーの実装

自動生成されたコードはそのままでは詳細すぎたり基礎的なデータ構造の操作に寄っているため、以下の点で推論クライアントとしては使いにくいことが多い印象です。

（静的なコード生成の場合は）例えば画像や文字列などの推論ペイロードを Triton Inference Server の Tensor 型にする際に、getter/setter 経由でセットするのはつらい
Promise 以前の API となっているのでつらい
他の箇所で必要な型を探してくるのがつらい
ただ推論実行したいだけなのに要求されるコード量が多くてつらい

gRPC を利用した経験の少ないメンバー（私とか）もいるので、このあたりは導入にあたっての障壁となるかもしれません。ここでは、必要な API に絞ってラッパーを実装することで対処します。

Core API の提供

まず、必要な基礎的な API を抜き出し、promisify でラップしてみます⁴。これにより、簡単に async/await なコード内で利用できます。ただし unary な操作に限定されるため、双方向ストリーミングのような高度な操作が必要な場合は自力で callback ベースのラッパーを書いたほうがいいと思います。

// https://github.com/tuxedocat/triton-client-polyglot-example/blob/c0cb5e2e76170d8f8c3599f9c9c0d4d2acb011ea/client/typescript/src/index.ts#L39-L77
/**
 * Promisified API subset of triton.GRPCInferenceServiceClient
 * Notes: Only unary ops are supported e.g. no `modelStreamInfer`, due to limitation of util.promisify.
 * */
export class TritonCoreAPI {
  private client: GRPCInferenceServiceClient
  constructor(endpoint: string, creds?: ChannelCredentials) {
    this.client = new triton.GRPCInferenceService(
      endpoint,
      creds ?? credentials.createInsecure(),
    )
  }
  public async serverReady(
    req: ServerReadyRequest,
  ): Promise<ServerReadyResponse> {
    return promisify<ServerReadyRequest>(
      this.client.serverReady.bind(this.client),
    )(req) as unknown as ServerReadyResponse
  }
  public async serverLive(req: ServerLiveRequest): Promise<ServerLiveResponse> {
    return promisify<ServerLiveRequest>(
      this.client.serverLive.bind(this.client),
    )(req) as unknown as ServerLiveResponse
  }
  public async modelReady(req: ModelReadyRequest): Promise<ModelReadyResponse> {
    return promisify<ModelReadyRequest>(
      this.client.modelReady.bind(this.client),
    )(req) as unknown as ModelReadyResponse
  }
  public async modelConfig(
    req: ModelConfigRequest,
  ): Promise<ModelConfigResponse> {
    return promisify<ModelConfigRequest>(
      this.client.modelConfig.bind(this.client),
    )(req) as unknown as ModelConfigResponse
  }
  public async modelMetadata(
    req: ModelMetadataRequest,
  ): Promise<ModelMetadataResponse> {
    return promisify<ModelMetadataRequest>(
      this.client.modelMetadata.bind(this.client),
    )(req) as unknown as ModelMetadataResponse
  }
  public async modelStatistics(
    req: ModelStatisticsRequest,
  ): Promise<ModelStatisticsResponse> {
    return promisify<ModelStatisticsRequest>(
      this.client.modelStatistics.bind(this.client),
    )(req) as unknown as ModelStatisticsResponse
  }
  public async modelInfer(req: ModelInferRequest): Promise<ModelInferResponse> {
    return promisify<ModelInferRequest>(
      this.client.modelInfer.bind(this.client),
    )(req) as unknown as ModelInferResponse
  }
}

こんな感じで必要な API だけが生えた Client ができました。 proto-loader を利用しているため、コンストラクタに適当なオブジェクトを与えられて楽になります。厳密性が必要な箇所ではより工夫が必要ですが…。

test('client can call liveness endpoint', async () => {
  tritonClient = new TritonCoreAPI(ENDPOINT, credentials.createInsecure())
  const serverLiveResponse = (await tritonClient.serverLive(
    {},
  )) as ServerLiveResponse
  expect(serverLiveResponse).toBeDefined()
  expect(serverLiveResponse.live).toBe(true)
})

Tensor への変換と Request の組み立て

Triton Inference Server では、モデル側の定義ファイルに書かれた型と対応するように画像やテキストなどを Request に設定します。煩雑になる箇所なので変換用のユーティリティを実装しておくと楽になります。

たとえばテキストの場合は UTF-8 の Buffer を利用します。

// Utility
const toUTF8Buffer = (s: string): Buffer => {
  return Buffer.from(new TextEncoder().encode(s).buffer)
}
 
// Tensor
const inputText = {
  name: 'TEXT',
  datatype: 'BYTES',
  shape: [1, 1],
  contents: {
    bytes_contents: [toUTF8Buffer(someString)],
  },
} as InferTensorContents
 
// Request
const inferRequest = {
  model_name: 'model_name',
  inputs: [inputText],
  outputs: [{ name: 'OUTPUT_NAME' }],
} as ModelInferRequest

画像はもう少し手間がかかります。簡単に BASE64 文字列として送ったりファイルシステムを経由したりなど色々考えられるのですが、この例では画像のピクセルデータを 1 次元の配列として送ることにします。ここでは画像処理に Sharp を利用します。

公開当初の内容では多次元配列を与えていましたが、Triton Inference Server の InferTensorContents の定義どおり、1 次元の配列を与える必要があります。行指向 (Row-major) の配列であることが前提となっています。

import { Sharp } from 'sharp'
import {
  ModelInferRequest,
  _inference_ModelInferRequest_InferInputTensor as InferInputTensor,
} from './generated/inference/ModelInferRequest'
import { InferTensorContents } from './generated/inference/InferTensorContents'
 
/**
 * Convert image buffer to unsigned int array
 */
fromImageBuffer(imgBuf: ArrayBufferLike): number[] {
    return Array.from(new Uint8ClampedArray(imgBuf))
}
 
// Tensor
const { data, info } = await img.removeAlpha().raw().toBuffer({ resolveWithObject: true })
const imgBuf = data.buffer
const inputImage = {
    name,
    datatype: 'UINT8',
    shape: [1, info.height, info.width, info.channels],
    contents: {
      uint_contents: fromImageBuffer(imgBuf),
    } as InferTensorContents,
} as InferInputTensor
 
// Request
const inferRequest = {
    model_name: 'image_model',
    inputs: [inputImage],
    outputs: [{ name: 'OUTPUT' }],
} as ModelInferRequest

推論実行はこんな感じになります。

interface TritonInferenceOutput {
  name: string
  value: number
}
getFloats(response: ModelInferResponse): TritonInferenceOutput[] {
    const outputDefs = response.outputs as _inference_ModelInferResponse_InferOutputTensor[]
    const outputBuffer = response.raw_output_contents as Buffer[]
    const parsedFp32 = outputDefs.map((e, i) => {
        return { name: e.name ?? '', value: outputBuffer[i].readFloatLE(0) }
    })
    return parsedFp32 ?? []
}
const outputs = getFloats((await tritonClient.modelInfer(inferRequest)) as ModelInferResponse)

この例ではよくある回帰問題のような数値型で出力される場合を想定しています。注意点として API 体系の元となっている KServe と異なり、Triton Inference Server ではモデルの推論結果がバイト列として raw_output_contents に格納されているため、モデル定義にある出力の型を考慮しながら変換する必要があります。

これで自動生成したクライアントから Triton Inference Server を利用することができるようになりました。

まとめ

Triton Inference Server を公式クライアントが提供されていないプログラミング言語から利用することを試みました。特に、TypeScript の型情報を（ある程度）利用できるクライアントを例として挙げました。

MLOps はここ最近話題になっていて、特にツールや SaaS 周りでの派手な話題が多い印象なのですが、実際に機械学習を利用する際には何らかの形で推論を支えなくてはならず、技術的な課題を感じています。

Python 以外の言語からも堅牢な推論サーバーを利用できる Triton Inference Server と gRPC のエコシステムは、個人的に注目しているので他におもしろいことがあれば⁵ また書こうと思います。

複数のフレームワークで実装された ML モデルの同時運用、モデル管理、簡易的な計算グラフ機能による前処理・後処理込みのパイプライン化、KServe 互換の API、C++で実装された API サーバー、コンテナ化、デバイス指定の抽象化など、色々と面白い特徴があります。 ↩
Triton における用語としては “ensemble” となっています。もうちょっと混乱の少ない名称はなかったんだろうか。 ↩
私がたまたまそういう構成のサービスを触る機会が多いということもあります。 ↩
GRPCInferenceService 全体をうまく型情報を保ったまま promisify する方法について、Issue https://github.com/grpc/grpc-node/issues/54 やそこで挙がっていた参考実装 https://gist.github.com/smnbbrv/f147fceb4c29be5ce877b6275018e294 を見ながら試行錯誤していましたが、TSC の設定のせいなのか実装のせいなのか、うまくいきませんでした。 ↩
書きたいこと: 1. Triton Inference Server のための深層学習モデルの定義ファイルの書き方や DVC を用いた重みファイル等の管理、 2. ナイーブな自前実装との推論実行パフォーマンスの比較など。 ↩

Triton Inference Serverを​色々な​環境から​使いたい​: TypeScript (Node.js) 編

Table of Contents