モノレポの Python バージョンを 3.9 から 3.11 に上げる

はじめに

AI Team MLOps Engineer の西原です。前回は kubeflow pipeline(kfp)のローカル環境での実行について Tech Blog を書きました。kfp は 2024 年に入ってからローカル環境の実行以外にも嬉しいアップデートがあったのでそれに少し絡めて今回の取り組みを紹介しようと思います。

今回の取り組みは、モノレポで使っている Python の最低バージョンを 3.9 から 3.11 に上げるというものです。なぜ、バージョンを上げたのか、上げる際の障壁とその対応を紹介しようと思います。

なぜ Python バージョンを上げたのか

パッケージ更新を頻繁にする理由

Python のバージョンを上げるにあたって、依存するパッケージの更新の話が関係してきます。そのため、パッケージの更新について先に触れます。

以前、モノレポに置いて依存するパッケージは高頻度でアップデートしているということを Tech Blog で紹介しました。その Tech Blog のおさらいになりますが、新しいバージョンはバグ修正や新規機能追加が含まれていることが多いので、バージョン更新することでバグを回避したり、新しい機能が使えるようになります。また、頻繁にアップデートすることで一回あたりの変更量が少なくなるため、何か問題があった時も少ない変更量から原因を特定することになり、デバッグも簡単になります。こういった背景から、私たちのモノレポでは依存するパッケージのバージョンを頻繁に上げています。

パッケージの更新ができなくなった

上記の通り、モノレポでは依存するパッケージのバージョンを頻繁に上げています。しかし、パッケージの更新ができないものが出てきました。更新できないパッケージというのは torchserve です。我々 AI Team では ML WebAPI のフレームワークとして torchserve を利用しています。torchserve 公式が提供しているコンテナイメージをベースイメージとして WebAPI を構築しています。この公式が提供しているコンテナイメージの環境が、私達に合わなくなりました。具体的には、torchserve v0.9.0 から CUDA 12 系のコンテナイメージになり、これを使って Vertex AI Endpoints にデプロイすると、CUDA の認識ができずエラーになります。そのため、Vertex AI Endpoints を使うには CUDA 11 系のv0.8のコンテナイメージに依存することになりパッケージの更新ができなくなりました。とはいえ torchserve には信頼性に関する issue がずっと立っていたりと、これからのアップデートに期待してる部分があるため torchserve 自体のアップデートを諦めるわけにはいきません。そのため、公式が提供するコンテナイメージを使うのを辞め、自分たちに適したコンテナイメージを構築することにしました。

torchserve と各ソフトウェアのバージョン

自分達に適した torchserve のコンテナイメージを構築することでコントールできるようになったものが次になります。

• CUDA のバージョン
• torch と torchserve のバージョン
• Python のバージョン

前置きが長くなりましたが、ここで Python バージョンの話になります。torchserve 公式が提供するコンテナイメージの Python バージョンは 3.9 でした。この公式のコンテナイメージへの依存を止めることで、Python のバージョンをコントロールできます。v3.10 以降の Python では型表現の強化や match 構文の追加、パフォーマンス改善などが含まれており、よりよい Python コードを書くことができるようになります。加えて Python には EOL があり、古いバージョンのサポートが終わります。3.9 の EOL はまだ先ですが、直前になって慌ててバージョンアップするよりも、早めにバージョンアップしておいた方が良いと考え、Python のバージョンを上げることにしました。

余談：CUDA のバージョンをコントールするのは、torchserve と Vertex AI Endpoints の相性が悪いというのも理由の一つですが、他にも理由があります。それは、GitHub Actions の GPU runner に向けた対応です。GitHub Actions で GPU を使えるようになるという話があり、現在はクローズドベータとして提供されているようです。このサービスが一般提供された時の利用に備えて CUDA のバージョンをコントロールすることで、導入負荷を下げることも目的の一つです。

Python のバージョンをどこまで上げるか

torchserve のコンテナイメージを自分たちで構築することで Python のバージョンを自由に選べるようになります。タイトルにもあるように Python 3.9 から 3.11 に上げたのですが、なぜ 3.11 なのかの意思決定の過程を紹介します。このタスクは 2024 年の新年最初のタスクとして取り組みました。私たちは Kubeflow Pipelines(kfp)を使って ML パイプライン開発をしていますが、2024 年に入った時点で kfp は protobuf 3 系に依存していました。モノレポで使っている 3rd party package の中には Python 3.11 以上から protobuf 4 系を求めるものがあるため、protobuf の依存の関係で 1 月はじめの時点では Python 3.10 までしか上げることができませんでした。そこで一度 Python 3.10 に上げたのですが、その数日後に kfp が protobuf 4 系に対応したため Python を 3.11 まで上げることができました。stable release として Python 3.12 が出ていますが、開発で使っているパッケージが 3.11 までしか対応していないものがあるため 3.11 までの更新にとどめました。

torchserve のコンテナイメージを自分たちで構築する

先述の通り、torchserve のコンテナイメージを自分たちで管理します。この章ではコンテナイメージの管理について紹介します。まず、torchserve 公式の Dockerfile を見ると build 時の引数として Python、 CUDA のバージョンを指定できることがわかります。公式の Dockerfile を参考に、build 時に各ソフトウェアのバージョンを指定することで自分たちに適したコンテナイメージを構築することができます。

Pants を使ってコンテナビルドの設定を記述すると次のようになります。

docker_image(
    name="caddi-torchserve",
    source="Dockerfile",
    extra_build_args=[
        "PYTHON_VERSION=3.11",
        "CUDA_VERSION=cu118",
    ],
    ...
)

Pants を使わずともコンテナイメージをビルドして registry に push できますが、Pants で管理することでベースイメージを使った build/push が効率的に行えます。Pants には差分実行機能があり、あるコンテナイメージに変更があった場合に、それに依存するコンテナイメージを検知して一括で build/push できます。今回構築した torchserve のベースイメージに依存しているコンテナイメージは二桁に及び、今後も増えると予想しています。今後、開発規模が大きくなったとしても Pants を使ってコンテナイメージを管理することで一括で build/push できるため、運用負荷を下げることができそうです。

torchserve メトリクスのエラー対応

Vertex AI Endpoints では torchserve 公式の CUDA 12 系ではうまく動きませんでしたが、今回構築した CUDA 11 系のコンテナイメージで動作することが確認できました。しかし、動作確認をする中で GPU のメトリクスを計測するプログラムでエラーが発生するようになりました。推論処理自体は問題なく動作しているため、このエラーを無視して利用することもできます。しかし、エラーによって狼少年アラートが発砲すると良くないため、エラーを解消することにしました。

原因を調査するとpynvmlというパッケージでエラーが発生していました。pynvml の v11.5 では CUDA 11 との相性が悪いため、v11.4 にダウングレードすることでエラーを解消することができました。

pyupgrade で Python コードをアップデートする

Python のバージョンを上げるにあたり、コードの構文も新しいものにアップデートしました。これまで使っていた Python 3.9 から 3.11 までの間で新しい構文が追加されています。例えば、match構文や|を使った型表現です。これらの新しい構文によって古い構文を置き換えることができます。機械的に置き換えできるものもあり、pyupgradeというツールを使って Python コードをアップデートしました。pyupgrade は Pants 公式がサポートしているため、特別な設定をすることなく使うことができます。

Ruff でも pyupgrade のように Python の古い構文を新しい構文に置き換える機能があります。Ruff も Pants 公式がサポートしているため、簡単に利用できます。

まとめ

モノレポで使っている Python のバージョンを 3.9 から 3.11 に更新する取り組みについて紹介しました。AI Team では開発で使っているソフトウェアのバージョンを更新できる仕組み作りをしています。Python のパッケージにおいても kfp や torch、pydantic などがここ一年でメジャーアップデートをしています。チームでこれらのアップデートに追従し、更新可能な古い機能で新規開発をしなくていいように努めています。また、Python に限らず Terraform や GitHub Actions などのツールについても同様の取り組みをしています。 古いバージョンのソフトウェアに依存していて生産性を低下させているが、更新が難しいというのを技術負債あるあるとしてよく耳にします。このような状況を回避するために、短いサイクルでソフトウェアを更新できる仕組みと体制が大事だと考え、チームで取り組んでいます。

参考

• pants docker