こちらは NEC デジタルテクノロジー開発研究所 Advent Calendar 2023 1日目の記事です。FireDucks に搭載されているインポートフック機能を開発した際の裏事情…、というか、プレスリリースや論文にはまず書かない苦労話なんかを書いてみたいと思います。
FireDucks とは？
私が所属しているグループでは、pandas と互換な API を持つ高速なデータフレームライブラリ “FireDucks”（ファイアダックス）を開発しています。概要については以下の弊社プレスリリースや公式ウェブサイトをご覧ください。現在はベータ版を公開中で、なんと無料でお使いいただけます！
NEC 研究開発/新技術 プレスリリース（2023年10月19日）
FireDucks 公式ウェブサイト
さて、この FireDucks というライブラリは「pandas っぽい API」ではなく「pandas と同じ API」を提供するように作られています。つまりインポート文を以下のように書き換えてしまえば、すぐに pandas の代わりとして使うことができます。
# import pandas as pd
import fireducks.pandas as pd
わざわざこの記事を読んでくださっている皆さんにおかれましては「どういう仕組みで高速化しているの？」「類似のデータフレームライブラリとの比較は？」などの気になる点があるかもしれませんが、この記事ではそういった空気は読まずに、私が FireDucks の開発の中で携わった部分について書いていきます。
高速化の仕組みや他のライブラリとの比較は、FireDucks 開発チームの他のメンバーが Advent Calendar の他の日に書いてくれるはず…！
インポートフック機能とは？
で、私は何をしたかというと、FireDucks にインポートフック機能というものを実装しました。使い方は以下の FireDucks 開発者ブログで紹介しています。
この機能は簡単に言うと、Python インタープリターの起動オプションでインポートフックが有効になるように指定しておくと、プログラムを実行しながらインポート文の読み替えを勝手にやってくれるというものです。これを使えば自分でインポート文を書き換えなくても、pandas の代わりに FireDucks が自動で使われるようになります。便利ですね。
$ python3 -m fireducks.imhook your_program.py
インポートフックの実装方法
ここからが本題…、とか言いつつこの記事全体が余談みたいなものなのですが。
フックする場所を選ぶ
Python にはインポート機構をフックできる箇所がいくつか用意されています。インポート機構がどのような仕組みになっているのか知りたければ、Python ドキュメントのインポートシステムの章を頑張って読めばいいのですが、いきなりこのドキュメントを読み始めるのはかなりハードルが高いと思います。まずは以下に挙げた記事が参考になるでしょう。
これらを読むとだいたいわかってくるのは、インポート処理に割り込んで何か独自の処理を仕込めそうなのは以下のどこかだということです。
- インポート文が実行されると
__import__()という Python 組み込み関数が呼ばれる - 一度インポートされたモジュールは
sys.modulesにキャッシュされており、__import__()関数への進入よりも後、下のふたつのフックよりも前にキャッシュ検索が行われる -
sys.meta_pathを参照してメタパス・ファインダーを動かす -
sys.path_hooksを参照してパスベース・ファインダーを動かす
このうち、Python 公式ドキュメントで「インポートフック」と呼ばれているものは下のふたつなのですが…、これらよりも前にキャッシュ検索が行われるということに気をつけなければいけません。一度インポート済みとなったモジュールに対しては、sys.meta_path や sys.path_hooks を使ったフックがもう発動せず、キャッシュ検索だけでインポート処理が済まされてしまうのです。したがって選択肢は上のふたつ、__import__() 関数に手を入れるかモジュールキャッシュ sys.modules に手を入れるかの二択になります。FireDucks 用のインポートフック機能の開発にあたっては、実装が手軽そうな __import__() 関数の変更で対応することにしました。
この安易な選択が後に苦悩をもたらすことになるのですが…、そのお話はまた別の機会があれば。
__import__() 関数に手を入れるということになれば、あとはそう難しい話ではありません。まずは本物の __import__() 関数を捕まえておきます。そして偽物の __import__() 関数を作って、pandas をインポートしようとしたときは代わりに FireDucks をインポートするよう本物に命令し、特に pandas と関係ないモジュールをインポートしようとしたときはそのままそのモジュールをインポートするよう本物に命令するのです。実際の実装よりかなり簡略化していますが、擬似コードで表すと以下のようになります。
# 本物を捕まえる
genuine_import = builtins.__import__
# インポスター：偽者、詐欺師
def imposter(name, *args, **kwargs):
if name == "pandas" and should_hook(name):
return genuine_import("fireducks.pandas", *args, **kwargs)
return genuine_import(name, *args, **kwargs)
# 偽物に差し替える
builtins.__import__ = imposter
オーバーヘッドの問題にぶち当たる
先ほどの疑似コードではよく見ると ... and should_hook(name) としてフックする際に追加の条件判定を課しています。ここでは FireDucks の中で部分的に本物の pandas を使用している場合など、本当に pandas を FireDucks に置き換えてもいい状況なのか否かをチェックしているのですが、そのチェックのために Python 標準モジュールの inspect を使ってコールスタックを取得しています。Python のインポート処理ではまず __import__() 関数が実行されるため、コールスタックを遡ればインポート処理の足跡も辿れるというわけです。
ところがこのコールスタックを遡る処理に罠が潜んでいました。開発チームの先輩から「外部ライブラリをインポートするだけで10秒ぐらいかかるんだけど…」という連絡が来てびっくり。結論から言うと inspect.stack() 関数がやたらと遅いことが問題で、inspect.currentframe() 関数と f_back 属性を使って自力でスタックフレームを遡るようにすることで改善しました。
inspect.stack() は便利な関数で、コールスタックに格納されているスタックフレームをまとめて取得できるのですが…。良かれと思ってなんでしょうが、どうもこの子は “便利だけど重たい” フレームオブジェクトをたくさん作成しているようです。しかもイテレーターじゃなくてリストが返ってくるんですね。1回取得するぐらいならそう気にならないものの、大きなライブラリの中ではインポート処理が何回も何回も走るため、オーバーヘッドが積み重なることで目に見えて長い時間がかかるようになってしまったのでした。
本記事の執筆時点では、高速にコールスタックを遡ることができる便利関数は inspect には用意されていないため、自力でスタックフレームを辿っていくジェネレーター関数を書くしかなさそうです。皆さんもインポート処理に限らず、コールスタックを遡りたくなったら以下のような感じで自力で辿ってみてください。
def frame_genarator(f):
while f is not None:
yield f
f = f.f_back
def get_frames():
f = inspect.currentframe()
try:
return frame_generator(f.f_back)
finally:
del f # 循環参照防止
オーバーヘッドを計測してみる
という感じでオーバーヘッドの問題は改善したのですが、では実際インポートフック機能を使用した場合にどれだけのオーバーヘッドが上乗せされてしまうのか計測してみましょう。使用したマシンの CPU は Intel Xeon Gold 5317 (Ice Lake)、OS は Ubuntu 22.04 です。
$ cat /proc/cpuinfo | grep -m 1 '^model name'
model name : Intel(R) Xeon(R) Gold 5317 CPU @ 3.00GHz
$ cat /etc/lsb-release
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=22.04
DISTRIB_CODENAME=jammy
DISTRIB_DESCRIPTION="Ubuntu 22.04.3 LTS"
Python のバージョンは 3.10 で、FireDucks 0.8.6 をインストールした仮想環境を使います。
$ source py3.10-release/bin/activate
(py3.10-release) $ python3 -V
Python 3.10.13
(py3.10-release) $ pip freeze
fireducks==0.8.6
firefw==0.8.6
numpy==1.24.4
pandas==1.5.3
pyarrow==13.0.0
python-dateutil==2.8.2
pytz==2023.3.post1
six==1.16.0
pandas をインポートするだけのスクリプト import-pandas.py と、FireDucks をインポートするだけのスクリプト import-fireducks.py を用意します。
import pandas as pd
import fireducks.pandas as pd
以下の三つのパターンをそれぞれ100回実行してみましょう。
- (A) インポートフックなし、単に pandas をインポート
- (B) インポートフックなし、単に FireDucks をインポート
- (C) インポートフックあり、pandas を FireDucks に読み替えながらインポート
(py3.10-release) $ time for ((i = 0; i < 100; i++)); do python3 import-pandas.py; done
real 0m28.573s
user 0m24.562s
sys 0m3.392s
(py3.10-release) $ time for ((i = 0; i < 100; i++)); do python3 import-fireducks.py; done
real 0m31.418s
user 0m27.284s
sys 0m3.484s
(py3.10-release) $ time for ((i = 0; i < 100; i++)); do python3 -m fireducks.imhook import-pandas.py; done
real 0m31.694s
user 0m27.416s
sys 0m3.632s
さて、ここではパターン (B)「インポートフックなしで import-fireducks.py を実行した場合」、つまり「import fireducks.pandas as pd を手書きした場合」をベースラインと考えて比較することにします。上の計測結果はそれぞれ100回分の実行時間のため、1回あたりの平均実行時間を算出して結果をまとめると以下の表の通りとなります。
| インポートフック | スクリプト | 平均実行時間（秒） | 差分（秒） | |
|---|---|---|---|---|
| (A) | なし | import-pandas.py |
0.286 | -0.028 |
| (B) | なし | import-fireducks.py |
0.314 | ベースライン |
| (C) | あり | import-pandas.py |
0.317 | 0.003 |
まずパターン (A) と (B) の比較ですが、これらはどちらもインポートフックなしでの実行時間です。pandas をインポートするより FireDucks をインポートするほうが28ミリ秒ほど時間がかかっていますね。FireDucks の中でも pandas をインポートして使っているため、追加で時間がかかるのは当然の結果なのですが、インポートするための時間に28ミリ秒の上乗せ程度で気になることはあまりないのではないでしょうか。
そしてパターン (B) と (C) の比較は、結果としてインポートされるものはどちらも FireDucks になるのですが、それをインポートフックなしで実施したかありで実施したか、という部分が違うことになります。インポートフック機能をオンにしたことによって上乗せされている時間は3ミリ秒ほどでした。この3ミリ秒という時間が積み重なって大きなオーバーヘッドになることはないのか調査するにはより詳細な検証が必要ですが…、その点は今後の課題としたいと思います。
まとめにかえて
FireDucks 開発者ブログに書いたインポートフック機能の紹介記事の続きのような形で、実装の舞台裏を書いてみました。FireDucks 開発チームではインポートフック機能の拡張や FireDucks 自体のさらなる高速化も鋭意開発進行中です。お楽しみに！