📚

Pythonの並列処理ライブラリRayでインスタンスメソッドを並列化する

2023/02/03に公開

Rayはメソッドをデコレータで囲んであげるだけで並列処理を簡単にかける便利なライブラリである。

以下はサンプルコード。

import ray
import time

ray.init(num_cpus=4)
# wait initailization ray
time.sleep(1)

def func1(self, x):
    time.sleep(1)
    return x * 2

@ray.remote
def func2(self, x):
    x = func1(x)
    time.sleep(1)
    return x + 1

vec = [1, 2, 3, 4]
results = []

start = time.time()

actor = Actor()
for x in vec:
    res = actor.func2.remote(actor, x)
    results.append(res)

# resultsはObjectIDのリスト
print(ray.get(results))  # 出力: [3, 5, 7, 9]
end = time.time()
print("elapsed time: ", end - start)

上記のスニペットは4並列でfunc2を実行するので約2秒で終わる。

大規模なコードになってくるとクラスを使用したくなってくる。
ここでActorというクラスを用意し、func1とfunc2をActorのインスタンスメソッドとして定義してみる。

class Actor:
    def func1(self, x):
        time.sleep(1)
        return x * 2

    @ray.remote
    def func2(self, x):
        x = self.func1(x)
        time.sleep(1)
        return x + 1
	
vec = [1, 2, 3, 4]
results = []

start = time.time()

actor = Actor()
for x in vec:
    res = actor.func2.remote(x)
    results.append(res)

上のコードを実行するとfunc2を呼び出す段階で TypeError: missing a required argument: 'x' が発生してしまう。
これはfunc2がActorのインスタンスメソッドであるためである。
よって、少し気持ち悪いが以下ように引数に自身のインスタンスを追加して呼び出すとうまく動くようになる。


actor = Actor()
for x in vec:
    res = actor.func2.remote(actor, x)
    results.append(res)

サンプルコードにクラスを並列化するためにクラス全体をデコレータで包む書き方があるが、あれはインスタンスを複数作成して並列で実行するためのものなのでインスタンスメソッドを並列化したい場合は上記のように書く必要がある。

参考：

Discussion