コンテンツにスキップ

pex のオーバーヘッドなしに Pants 管理の Python ファイルを実行する

pants run my_script.pyで Pants が管理する Python ファイルを実行すると、実行ごとに sandbox 環境を構築するため、実行までに時間がかかる。特に PyTorch のような大きいパッケージを含むモジュールでは、実行までに 10 分かかることもある。本記事では、venv環境を事前に作成しておき、作成した環境で実行して待ち時間を短縮する方法を紹介する。

venv 環境の構築

まず、venv環境を構築する。Pants にはresolveという概念があり、Python パッケージ管理ツールでいう group のようなものになる。1 つのリポジトリで複数のresolveを持つことができる。実行に使うresolveの依存パッケージをvenvへインストールしておくことで、実行時に sandbox 環境を構築する時間を削減できる。 pants exportコマンドを使えば resolve に対応したvenv環境を構築できるが、多くのパッケージが含まれる場合は完了までに時間がかかる。uvを使えば、環境構築にかかる時間を短縮できる。

1
2
3
4
5
# 複数の resolve があることを想定して出力先を指定
uv venv $resolve_name
# resolveに対応するrequirements.txtを指定。
# pexにはlockfileからrequirements.txt形式でパッケージの依存関係を出力する機能がある
uv install -r requirements.txt

import path の設定

Pants には root という概念があり、Python の import path を設定する仕組みに近い。pants runコマンドで実行できていた Python ファイルが、venv環境では module が見つからずエラーになることがある。モジュールが見つからない場合は、import path を設定して解決する。

pants rootsコマンドで root を確認し、import path に追加する。

export PYTHONPATH="/path/to/module:$PYTHONPATH"

上記のように shell で設定するか、Python ファイルの冒頭に次のように記述する。こうした記述により、Python 環境で import がうまくいくようになる。

import sys
sys.path.append("/path/to/module")

venv 環境での実行

import path を設定した状態でvenv環境を使って Python ファイルを実行すると、pex のオーバーヘッドなしに実行できる。

./venv/bin/python my_script.py

resolve ごとにvenv環境を構築しておく方法により、複数の環境でプログラム実行ができる。例えば CPU/GPU 環境の切り替えや、Python のバージョン切り替えなどが即時可能になる。