どこから見てもメンダコ

軟体動物門頭足綱八腕類メンダコ科

Pythonの分散並列処理ライブラリRayの使い方

いままでありがとうmultiprocessing。そしてこんにちはRay

関連:

horomary.hatenablog.com

horomary.hatenablog.com


Rayとは

Rayは分散アプリケーションを構築および実行するための高速でシンプルなフレームワークです。
What is Ray? — Ray v1.6.0

github.com

Rayライブラリを使用することで、既存のライブラリよりもはるかにシンプルでPythonicなコードで分散並列処理を実装できます。

分散並列処理は計算科学に欠かせない要素にもかかわらず、これまでのPythonでの分散並列処理の実装は直感的とは言い難いものでした。しかし、Rayライブラリの登場によってPython言語は真に"科学者のための言語"となったのではないでしょうか。


基本的な使い方

Rayによる並列化処理は非常にシンプルに記述できるためにまずはコードを見た方が早いでしょう。

並列化していないコード

5秒間 sleepする関数を3回実行するので実行時間は当然5×3=15秒程度です


Rayによる並列化コード

ほとんどコードを変更せずに並列化できることがわかります。
また、並列化によって実行時間が15秒から7秒程度まで短縮されていることを確認できます。

実行時間が5秒ではなく7秒であることからわかるように、rayはmultiprocessingなどと比較してプロセス分岐時のオーバーヘッドがやや大きいので、重いタスクの並列処理には向いていますが軽量なタスクの実行のためにサブプロセスが頻繁に生成されるような用途には向いていません。


待ち合わせ処理: ray.get

分岐したプロセスの戻り値はray.getによって取得できます。

各プロセスを格納したリストである変数result に対して ray.get することで、すべてのプロセスが終了して戻り値を返すのを待つ、待ち合わせ処理となります。

このコードからは ray.get する前のリスト result には 、ObjectRef というオブジェクトが格納されていることがわかります。

ObjectRefはプロセスの戻り値に対するプレースホルダと考えるとよいでしょう。ray.get(ObjectRef)することでプロセスの戻り値を取得できます。また、ray.get(ObjectRef)時点でプロセスが終了していない場合は対象のプロセスが終了するまで待機します。

そしてこのサンプルコードのように、ObjectRefが格納されたリストに対してray.getした場合はリスト内のすべてのサブプロセスが戻り値を返すまで待つので待ち合わせ処理となります。


逐次処理: ray.wait

同期処理はray.getで簡単に実装することができることがわかりました。では非同期処理はどうでしょう?

具体的には同時に多数の並列タスクを開始するが、各プロセスそれぞれの終了時間がまちまちであるために終了したタスクから順に次の処理を行いたい、というケースを考えます。これはray.waitを使うことで実装できます。

デフォルトでは、ray.wait はその時点ですでに終了しているタスクのリストと終了していないタスクのリストを返します。

ただし、引数 num_returns が指定されている場合は、すでに終了しているタスクがnum_returns個 格納されたリストと、それ以外のタスクが格納されたリストを返します。 ray.waitが実行された時点で終了しているタスクがnum_returns個に満たなかった場合はnum_returns個のタスクが終了するまで待機します。

このサンプルコードではray.waitを使用して、

  1. 終了済みタスクのObjectRefを1つ取得する
  2. ray.get(ObjectRef)を出力する

という流れで、終了したタスクから逐次戻り値を出力しています。


クラス単位でのサブプロセス化

Rayではクラス単位での並列化も簡単にできます。

これはmultiprocessingモジュールで実装するのはなかなか面倒だった、状態を保持するサブプロセスの実装が簡単にできるということです。


リソースの設定: ray.init

※ドキュメントに記載されている情報であり動作未検証です:The Ray API — Ray 0.3.0 documentation

ここまでのサンプルコードすべてで最初に記述されている ray.init() ではリソース配分を設定できます。

ray.init()のように引数を指定しないとローカルマシン内の全CPUが利用可能になるのに対して、ray.init(num_cpus=20, num_gpus=2) というように使用するリソースを明示することもできます。とくに子プロセスでGPUを使用したい場合はnum_gpusを明示する必要があるようです。

また、ray.init(local_mode=True) とするとコードそのままでもサブプロセス化されなくなり、printが効くようになるので開発時に便利です。


クラスターでの分散並列処理

rayを使えばマルチノードでの分散並列化も楽にできます

horomary.hatenablog.com


分散強化学習での利用

rayのユースケースのひとつは分散強化学習でありrayのサブプロジェクトとして強化学習フレームワークRLLIBが開発されています。が、もしスクラッチ実装に興味がある場合は場合はぜひ下記の記事も参照ください。

horomary.hatenablog.com

horomary.hatenablog.com