PostgreSQLとcupyを繋ぐ～機械学習基盤としてのPG-Stromその①～
世間の機械学習屋さんは、機械学習・統計解析のライブラリにデータを食わせる時に、どうやってデータを入力しているのだろうか？
話を聞くに、データを一度CSV形式に落とし込んで、それをPythonスクリプトで読み込むというパターンが多いようではある。
ただ、ある程度大量のデータセットをCSVファイルで扱うようになると、いくつか問題点が露わになってくる。
- 解析すべきデータセットを切り替えるたびに異なるCSVファイルを用意する事になり、ファイルの取り回しが煩雑である。
- 前処理をかけた後のCSVファイルがまたできてしまい、ファイルの取り回しが更に煩雑になってくる。
- 最終的にCSVファイルの所在が誰にも分からなくなってしまい、機械学習・統計解析の元になったファイルが散逸してしまう。
- そもそも、GB単位のフラットファイルをシェル上でコピーしたり読み込ませたりするのはそれなりに時間を要する処理である。
データベース屋から言わせてもらえれば、餅は餅屋、データ管理はデータベース管理システムという事で、得意な人にやってもらった方が良いのになと思う事はある。
Gstore_fdwとInter-Process Data Collaboration
PG-StromにはGstore_fdwという機能がある。これはGPU上のデバイスメモリをPostgreSQLにとっての外部ストレージに見立てて、INSERT構文を使ってGPUデバイスメモリへデータを書き出すための機構である。
これは元々、PostgreSQLの可変長データの制限が1GBであるため、PL/CUDAユーザ定義関数が処理すべきデータが1GBを越える場合にはSQL関数の引数として与える事ができず、また、CPU側で（通常、null値を含まないベクトルor行列として表現する）メモリイメージをセットアップする時間が処理全体の支配項になってしまうために作成した機能であった。
（創薬分野における類似化合物探索ワークロードにおける化合物データベースのように）繰り返し何度も利用されるデータであればGPU上に”置きっぱなし”である方が合理的であるし、これはまた同時に、PostgreSQL可変長データの制限である1GBを回避する事にもなる。
一方、PG-Stromの実装基盤であるCUDAのAPIリファレンスを眺めてみると、興味深いAPIが提供されている事が分かる。
CUresult cuIpcGetMemHandle ( CUipcMemHandle* pHandle, CUdeviceptr dptr )- Gets an interprocess memory handle for an existing device memory allocation.
CUresult cuIpcOpenMemHandle ( CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned int Flags )- Opens an interprocess memory handle exported from another process and returns a device pointer usable in the local process.
要は、GPUデバイスメモリ上に確保した領域の一個ずつにユニークな識別子が付与されており、cuIpcGetMemHandleを利用する事でその識別子を取得する事ができる。これは64バイトの長さがあり、重複しない事が保証されている。
そして他のプロセスでこの識別子をcuIpcOpenMemHandleに渡してやると、GPUデバイスメモリ上の同じ領域を読み書きできるようになるという、GPU版のプロセス間通信の仕組みである。
そうすると、GPUへのデータロードはGstore_fdwを用いてSQLで行った上で、その後のデータ操作、データ解析はPythonスクリプトで実行するというワークフローが出来上がる。
何が良いのか？
pystrom - cupyとPG-Stromを繋ぐためのモジュール
PG-Strom側でGstore_fdwテーブルの背後に存在するGPUデバイスメモリの識別子を取得するには、以下の関数を用いる。
postgres=# select gstore_export_ipchandle('ft');
gstore_export_ipchandle
------------------------------------------------------------------------------------------------------------------------------------
\x006b730200000000601100000000000000750200000000000000200000000000000000000000000000020000000000005b000000000000002000d0c1ff00005c
(1 row)gstore_export_ipchandleは引数で指定されたGstore_fdwテーブルの識別子を bytea 型のデータとして出力する。表示はHex形式だが、64バイトある事が分かる。
次に、Python側でこの識別子を受け取り、データ操作が可能なPythonのオブジェクトとして扱えるようにするためのモジュールが必要。
そのために作成したのがpystromモジュールで、ソースコードは以下の通り*1。
pg-strom/pystrom at master · heterodb/pg-strom · GitHub
pystromクラスは唯一のメソッドipc_importを持つ。
cupy.core.ndarray pystrom.ipc_import(bytes ipc_handle [, list attrNames])
このメソッドはipc_handleで与えられた識別子でGstore_fdw外部テーブルをオープンし、その中からattrNamesで指定された列だけから成る２次元のcupy.core.ndarrayを生成する。
attrNamesで指定する列は全て同一のデータ型を持っている必要があり、現状、bool、smallint、int、bigint、realおよびfloatだけが対応している。
また、デバイスメモリ上のレイアウトは行優先の２次元配列となっている（DB的に言えば列指向データ構造）。
では、実際に使ってみる事にする。
Gstore_fdwテーブルの作成とテストデータの挿入
postgres=# CREATE FOREIGN TABLE ft (id int, x real, y real, z real)
SERVER gstore_fdw OPTIONS (pinning '0');
CREATE FOREIGN TABLE
postgres=# INSERT INTO ft (SELECT x, random(), random(), random() FROM generate_series(1,100000) x);
LOG: alloc: preserved memory 1601024 bytes
INSERT 0 100000CREATE FOREIGN TABLE構文を使ってGstore_fdwテーブルを作成する。
SERVER gstore_fdwを指定している事と、デバイスメモリを確保すべきGPUのデバイス番号をpinning '0'オプションで指定している。
次に、INSERTを使ってランダムなデータを10万件ほど挿入してみた（データ自体に意味はない）。ログに1.6MBほどGPUデバイスメモリを確保した旨が表示されている。
pystromモジュールによるPythonスクリプトへのインポート
続いて、Python側のスクリプトに移る。上記のデータをインポートするには、大きく以下の３つの手順を踏む事になる。
- PythonスクリプトからPostgreSQLに接続する（psycopg2モジュール）
- gstore_export_ipchandle関数を用いて、Gstore_fdwの識別子を取得する
- この識別子を用いてGstore_fdwをオープンし、cupy.core.ndarrayオブジェクトを作成する（pystromモジュール）
順を追って説明する。
先ずは、PostgreSQLへの接続とSQL関数の実行であるが、これはアシスト様の以下のブログを参考にさせていただいた。
www.ashisuto.co.jp
ローカル接続でtrust認証なら何も難しい事は考えず、以下の通りである。
import psycopg2
conn = psycopg2.connect("host=localhost dbname=postgres")そして、この conn オブジェクトを用いてクエリを発行し、その結果を取得する。
curr = conn.cursor()
curr.execute("select gstore_export_ipchandle('ft')::bytea")
row = curr.fetchone()
conn.close()rowは実行結果の先頭行であり、この場合、結果を一個だけ含むタプルであるため、row[0]がGstore_fdwの識別子となる。
次に、pystromを用いて上記の識別子をオープンする。
import pystrom X = pystrom.ipc_import(row[0], ['x','y','z'])
以上で完了である。
cupy.core.ndarrayクラスのオブジェクトであるXは、real型（32bit浮動小数点型）であるx列、y列、z列の内容を持つ3列x10000行の行列である。
>>> X
array([[0.05267062, 0.15842682, 0.95535886],
[0.8110889 , 0.75173104, 0.09625155],
[0.0950045 , 0.71161145, 0.6916123 ],
...,
[0.32576588, 0.8340051 , 0.82255083],
[0.12769088, 0.23999453, 0.28765103],
[0.07242639, 0.14565416, 0.7454422 ]], dtype=float32)SQLで上記のftテーブルを出力してみると以下のようになる。先頭３行分のデータを見比べていただきたい。
postgres=# SELECT * FROM ft LIMIT 5; id | x | y | z ----+-----------+----------+----------- 1 | 0.0526706 | 0.158427 | 0.955359 2 | 0.811089 | 0.751731 | 0.0962516 3 | 0.0950045 | 0.711611 | 0.691612 4 | 0.051835 | 0.405314 | 0.0207166 5 | 0.598073 | 0.4739 | 0.492226 (5 rows)
もちろん、一度Python側でcupyのデータオブジェクトとしてインポートした行列はPG-Stromの機能とは関係無しに操作する事ができるため、cupyの提供する多種多様な演算を利用する事ができる。
例えばドット積を計算してみる。
>>> cupy.dot(X[:,0],X[:,1]) array(24943.564, dtype=float32)
今回の例は、単純なランダムデータをGPUに投入してデータ交換を行っただけであるが、データ管理をデータベース管理システム（PostgreSQL）に任せてしまう事で、統計解析・機械学習の前処理にSQLというパワフルなツールを使う事ができる。また、データのエクスポートとその管理という煩雑な作業からデータサイエンティストを解放する事ができるほか、フラットファイルの場合と異なりデータファイルの散逸という問題が起こる事もない。
さらに、PostgreSQLにはリモートDBと連携して動作するための機能も備わっている。
～100GB程度のデータは個々のエンジニアが自分のワークステーションへ取り込んで作業する一方、データウェアハウス／データレイクという巨大なデータセットから適宜必要なデータを取り出し、それを他のエンジニアと共有し、容易に再現可能な形にする事こそが当面の目標である。
とりあえず、今回の記事は『とりあえず動きました』以上のものではないが、以下のような点はこれからの宿題である。
*1:将来的には pip install で簡単に入れられるようにしたいが、これはもう少し宿題事項
