tool.poetry.sourceを使わずにPyTorchを高速にpoetry管理下におく

はじめに

Pythonのパッケージ管理にPoetryを使用している場合、PytorchのリポジトリURLを直接していしたり、別途pipでインストールするなど裏技的に入れる方法はありましたが、1.2のリリースによって公式にサポートされました。
PyTorchを利用するデータ分析界隈では盛り上がり、その影響でPoetryの利用者も増えたのかなと思います。

インストール自体はできるようになったのですが、この方法では速度面で問題がありました。
その回避に成功したようなので備忘録として残しておきます。

実験環境

• OS
  • Windows 10 Pro 19045.2604
  • macOS Sierra, 10.12.6
• Python 3.9.6
• poetry 1.2.1
• torch 1.12.1+cu116

(読み飛ばしOK)問題のインストール方法

1.2.0以降の書き方

1.2.0リリース直後は以下のようなpyproject.tomlを書くことで、PyTorchのインストールが簡単になりました。

1.2以降のpyproject.toml

[tool.poetry.dependencies]
python = "^3.8,<3.11"
torch = { version = "=1.12.1+cu116", source = "pytorch-cu116" }
torchvision = { version = "=0.13.1+cu116", source = "pytorch-cu116" }

[[tool.poetry.source]]
name = "pytorch-cu116"
url = "https://download.pytorch.org/whl/cu116/"
secondary = true

が、この方法には欠点がありました。
パッケージ情報(.lockファイル)の更新する際に、常にtool.poetry.sourceに書かれたパッケージが依存しているとみなしてファイルをダウンロードしてしまうというものです。
つまり、すべてのパッケージにPyTorchが依存していると理解して巨大なパッケージを常にダウンロードする為更新でも非常に時間がかかってしまいます。

URLを直接指定する

この回避のために、バージョン1.2より前で使われていたURLを直接指定する方法を使っている方もいるようです。
ですが、Whlのファイル名にはOSやPythonのバージョンなどが書かれているため、チーム内でOSが異なる場合に対応できない等開発環境の柔軟性に欠けてしまいます。

URLを直接指定するpyproject.toml

[tool.poetry.dependencies]
python = "^3.8,<3.11"
torch = { url = "https://download.pytorch.org/whl/cu116/torch-1.12.1%2Bcu116-cp310-cp310-win_amd64.whl"}
torchaudio = { url = "https://download.pytorch.org/whl/cu116/torchaudio-0.12.1%2Bcu116-cp310-cp310-win_amd64.whl"}
torchvision = { url = "https://download.pytorch.org/whl/cu116/torchvision-0.13.1%2Bcu116-cp310-cp310-win_amd64.whl"}

(本題)問題を解決する書き方

環境マーカー(Environment Marker)という、指定されたマーカーが実行環境と一致する場合(特定の環境のみ)にインストール可能とするようなフラグ機能を使用します。

• https://python-poetry.org/docs/dependency-specification#using-environment-markers
• https://peps.python.org/pep-0508/#environment-markers

すべてを解決するpyproject.toml

[tool.poetry.dependencies]
python = "^3.8,<3.11"
torch = [{ version = ">=1.12+cu116", markers = "sys_platform == 'win32'" }]

この場合、環境がWindowsの場合にtorch>=1.12+cu116をインストールするという条件分岐となります。
やっていることとしてはOSとTorchのバージョンを指定しているので、URLの直接指定とあまり変わらないのですが、非常にシンプルに書くことが出来ました。

さらに、

もっとすべてを解決するpyproject.toml

[tool.poetry.dependencies]
python = "^3.8,<3.11"
torch = [{ version = ">=1.12+cu116", markers = "sys_platform == 'win32'" },
         { version = ">=1.12+cpu", markers = "sys_platform == 'darwin'" }]

とすることで、WindowsではGPU版を、macOSではCPU版をインストールする様に制御することも可能です。
また、python_full_versionのマーカーを使うことで、3.10.0b3を使うときと3.9.6を使うときでパッケージを分けることも出来ます。
2桁目のマイナーバージョンまでで問題なければ、python_versionでもOKです。

おまけ(Poetry導入時の注意)

この記事を見て、Poetryを初めて使ってみよう！と思った方がいたら、仮想環境の作成先に注意してください。
多くの方はプロジェクトディレクトリ毎に.venv環境を作る運用をされているかと思いますが、Poetryのデフォルトはその様な挙動となっていません。

poetry config --local virtualenvs.in-project true

のコマンドを叩いておくと、poetry installした時のカレントに仮想環境が作れられる様になります。