機械学習エンジニアの吉田です。
夏ですね。7月はLayerXエンジニアブログを活発にしよう月間です。

昨年バクラクOCRの機械学習モデルの検証から本番投入までの取り組みについて記事を書きました。 tech.layerx.co.jp

その後、運用する中で新たな課題が生まれたり、負債を解消するために当初の開発環境を見直しアップデートしてきました。今回は機械学習周辺の技術スタックに焦点を当ててその変遷について紹介したいと思います。
MLチームでは各サービスからのリクエストを処理するAPIやデータ基盤、社内のアノテーションツールなどの開発も行っており、これらは主にGo, TypeScriptで開発されていますが今回は対象外としています。

技術スタックの変遷

本番リリース時と現在の主な技術スタックの比較です。

	リリース時	現在
言語	Python	Python
パッケージ管理	pip	Poetry
API	Flask	FastAPI
静的解析	mypy, flake8	mypy, Ruff
Formatter	black, isort	black, Ruff
テスト	unittest	pytest
仮想環境	Docker	Docker, Poetry
ML	Lightning, transformers	Lightning, transformers

モデルの検証時点では本番に投入できるかはわからない状態だったので、あまり開発環境にはこだわらず最低限以下の項目を重視していました。

前処理や後処理などの決定的な処理はユニットテストを書く
型アノテーションは必須
アドホックな分析以外ではJupyterは使わない
Linter, FormatterやフレームワークとしてPyTorch Lightningを採用することでコードの自由度を制限し、個人の好みによるコードの差異を極力減らす
開発環境はDockerで再現可能な状態を維持する

以下アップデートされているものを中心に良かった点とイマイチな点を紹介していきます。

Poetry

github.com

Good

pyproject.tomlで一元管理できる

最初はデータセット作成やモデル開発といった用途毎にDockerfileを用意してrequirements.txtで管理していました。検証時はあまり依存関係に悩まずに並行してそれぞれの開発を進められるという点で良かったのですが、次第にそれぞれの環境でパッケージの重複やパージョンを揃えるのが大変になってきました。

Poetryでは用途毎にパッケージを管理することができる Dependency groups という機能があるので、pyproject.tomlですべてのパッケージの依存関係を解決しつつ環境によってインストールするパッケージを取捨選択することができます。

以下のように本番の推論APIで必要なパッケージはprd, 開発時だけ必要なパッケージはdev, すべての環境で必要なパッケージはmainという感じで管理することができます。

[tool.poetry.dependencies]
python = "3.11.3"
numpy = "^1.24.2"

[tool.poetry.group.dev.dependencies]
black = "^23.1.0"
mypy = "^1.0.1"

[tool.poetry.group.prd.dependencies]
gunicorn = "^20.1.0"
fastapi = "^0.100.0"

マルチプラットフォームに対応できる

MLチームではMacBook Pro (M1) を使って開発をしています。モデルの学習や推論を動かすようなGPUが必要なときはクラウドを使いますが、コードを書くときはMacBookを使うことが多いです。さらに、MacBookで環境構築するときにもDockerを使いたい人もいれば使いたくない人もいます。また、CIではGitHub Actionsを使っていますが、機械学習で使うDockerイメージはサイズが大きくなりがちでビルドに時間がかかることと、場合によってはディスクが足りずにエラーとなってしまうこともあるため極力Dockerを使いたくありません。

このようにOSやCPUアーキテクチャが異なる場合、環境によってインストールするパッケージが異なる場合があります。 PyTorchを例にするとpyproject.tomlに以下のように記述することで環境に対応したパッケージをインストールすることができます。

torch = [
    # M1 Mac
    { url = "https://download.pytorch.org/whl/cpu/torch-2.0.1-cp311-none-macosx_11_0_arm64.whl", markers = "sys_platform == 'darwin'" },
    # Docker on M1 Mac
    { url = "https://download.pytorch.org/whl/torch-2.0.1-cp311-cp311-manylinux2014_aarch64.whl", markers = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
    # Linux
    { url = "https://download.pytorch.org/whl/cu118/torch-2.0.1%2Bcu118-cp311-cp311-linux_x86_64.whl", markers = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
]

ちなみにwheelのURLを直接指定せずにsource, versionを指定することも可能ですが、不要な環境のwheelもすべて取得するためダウンロードに時間がかかり、また依存関係の解決がいつまでも終わらない状態になってしまったのでwheelのURLを明示的に指定するようにしています。

少し単純化していますが、ディレクトリ構成は以下のようになっています。DockerfileをCPU/GPUで分けており、環境に応じて使い分けています。 Dockerfileは基本的にベースイメージの違いだけで、CPUは ubuntu を、GPUは nvidia/cuda を使っています。上述したようにプラットフォームの違いはpyproject.tomlで吸収できるので、pyproject.tomlとpoetry.lockはどの環境でも共通のものになります。

├── docker-compose.yml
├── dockerfile
│   ├── Dockerfile.cpu
│   └── Dockerfile.gpu
├── poetry.lock
├── pyproject.toml
└── src
    ├── common
    ├── dataset
    ├── api
    └── model

Bad

依存関係の解決が遅い

Poetryの難点は依存関係の解決が遅いことです。特にPyTorchが入ってくると極端に遅くなります。私の環境では50個程度のパッケージが含まれるpyproject.tomlでpoetry lockした場合、PyTorch無しだと5分以内に完了するものが、PyTorchを入れるだけで約16分かかるようになりました。また、上にも記載しましたがPyTorchのwheelを明示的に指定しなかった場合、何時間経っても完了しないという事象も発生しました。

補足

最近RyeというRust製のパッケージ管理ツールも出てきましたが、残念ながらPoetryのようにマルチプラットフォームに対応することができませんでした。今後に期待です。 github.com

Ruff

github.com

Good

flake8, isortをRuffひとつで代替

現在Ruffは500以上のlintルールに対応しているようです。
最初すべてのルールを適用してみたところ大量のエラーが出てしまったので、以下のようにルールを絞って導入しました。徐々に適用するルールを増やしていけるのもよいですね。

[tool.ruff]
select = [
    "F",  # Pyflakes
    "E",  # pycodestyle error
    "W",  # pycodestyle warning
    "I",  # isort
    "B",  # flake8-bugbear
]

Ruffひとつでflake8, isortを削除することができるのですっきりしてよいです。

速い

flake8, isort, Ruffの実行時間を計測してみました。(flake8の設定をpyproject.tomlで管理しているために実際にはpflake8を使っています)

結果は以下の通りとなり、flake8 & isortに比べて約3倍速い結果となりました。1万行程度のコードなので元々の実行時間もそこまで遅くはありませんが、手元で実行すると速さを実感できます。コードの行数が増えるとさらに差は顕著になるようです。

time	pflake8 . && isort -c .	ruff check .
real	1.39s	0.42s
user	2.6s	0.18s
sys	0.82s	0.09s

Bad

いまのところ困ったことはないです。速さは正義。

FastAPI

github.com

Good

Pydantic, OpenAPIとのインテグレーション、コードファースト

リリース時点では推論APIはFlaskで実装されていましたが、APIのドキュメントが無く明確なスキーマ定義やバリデーションも無い状態でした。推論APIを使うのは社内からのみであることやエンドポイントもひとつだけで頻繁に改修が入る箇所ではなかったのでリリースからしばらくはそれでも良かったのですがAPIのインターフェースを変えようと思ったときにやはりドキュメントやバリデーションが欲しくなります。

最初はバリデーションやスキーマの型定義にPydanticを、ドキュメントにOpenAPIを導入し、OpenAPIのスキーマ定義からコードを自動生成することを考えていました。コードからの自動生成に関しては datamodel-code-generator を使うことでOpenAPIからPydanticのBaseModelを継承したclassを生成してくれるのでよさそうでした。

一方で特にFlaskに依存したコードも無いので、FastAPIへの切り替えも同時に検討していました。 FastAPIは標準でPydanticやOpenAPIとインテグレーションされており、コードからOpenAPIのドキュメントを自動生成してくれます。また、ちょうど検討していたタイミングでFastAPIの 0.100.0 がリリースされPydantic v2がサポートされたのも大きかったです。

当初はスキーマファーストでコードを自動生成することを考えていました。スキーマファーストの嬉しさとしては、サーバ / クライアントから切り離して管理できること、最初に明確にスキーマを定義することでサーバとクライアントが同時に実装に移れることなどがあると思います。

しかし、機械学習の推論APIは通常のWebアプリケーションほど複雑なスキーマや多くのエンドポイントを持つことがないのでそれらの恩恵はあまり受けられず逆に実装の手間が増えてしまうと考えました。また、機械学習エンジニアでOpenAPIの経験が無かったとしても学習コスト無くコードを書くだけでドキュメントを更新できるのも良いと思いました。

Bad

コードとドキュメントが密結合している

Goodでコードからドキュメントの自動生成をあげましたが、一方でFastAPIのレールから外れたときに少しつらいな、と思うこともあります。一例をあげると同じエンドポイントでAPIバージョンによってリクエストのスキーマを切り替えたいようなケースです。

通常FastAPIでは以下のようにPydanticのBaseModelを継承したclassをinvocationsの引数の型に指定することで、バリデーションを行いドキュメントにもスキーマを自動生成してくれます。

from fastapi import FastAPI, Response
from pydantic import BaseModel


class InvocationsRequestV1(BaseModel):
    first_name: str

app = FastAPI()


@app.post("/invocations")
def invocations(request: InvocationsRequestV1, response: Response) -> InvocationsResponse:

ここでリクエストヘッダからversionを取得し、versionによってスキーマを切り替えたい場合は以下のようになります。この場合、引数の型からはスキーマが分からないのでデコレータ内でスキーマを定義する必要があります。

from fastapi import FastAPI, Request, Response
from pydantic import BaseModel

class InvocationsRequestV1(BaseModel):
    first_name: str


class InvocationsRequestV2(BaseModel):
    first_name: str
    last_name: str


app = FastAPI()


@app.post(
    "/invocations",
    openapi_extra={
        "parameters": [
            {
                "in": "header",
                "name": "version",
                "schema": {
                    "type": "string",
                },
                "required": True,
            }
        ],
        "requestBody": {"content": {"application/json": {"schema": InvocationsRequestV2.model_json_schema()}}, "required": True},
    },
)
async def invocations(request: Request, response: Response) -> InvocationsResponse:
    api_version = int(request.headers["version"])
    params = await request.json()
    match api_version:
        case 1:
            InvocationsRequestV1.model_validate(params)
        case 2:
            InvocationsRequestV2.model_validate(params)

他にも不要なステータスコードを勝手に自動生成したりするので削除したい場合は以下のようなコードを実装する必要があります。

# ステータスコード200をドキュメントから削除
def custom_openapi() -> Any:
    if app.openapi_schema:
        return app.openapi_schema
    openapi_schema = get_openapi(title=app.title, version=app.version, routes=app.routes)
    for method_item in openapi_schema.get("paths").values():
        for param in method_item.values():
            responses = param.get("responses")
            if "200" in responses:
                del responses["200"]
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app = FastAPI()
app.openapi = custom_openapi