JavaScriptを有効にしてください

【ray】@remoteがメモリを大量に食う時はray.put()を使おう

 ·  ☕ 2 min read

はじめに


ray.remoteは怖い

  • @ray.remoteが付された関数が使用するオブジェクトは, 呼び出される都度内部でray.put()が呼ばれ, データがメモリ上に展開される
  • 若干この仕様が厄介で, ray.get()なんかを使うと, 同じオブジェクトを何度もメモリ上に展開してしまう
    • なので大量にremoteの関数を呼んだりすると, RAMを64GB積んでてもメモリ不足になったりする

solution

  • そういうわけで, サイズの大きいオブジェクトを渡したい場合は, 明示的にray.put()を使って先にメモリ上に展開しておく

    • 特に, サーバを自動で建ててソケット通信するクラスなんかを扱っていると, 大量にプロセスが乱立するので注意
  • また, class内のオブジェクトに対してもray.put()で明示確保すべき

    • というのも, どうやらrayではself.hogehoge系統も何度もメモリ展開してしまうらしい
    • 例えばこういう感じにすべき👇
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class XXX():
    def __init__(self, size=8, num_cpus=None) -> None:
        ray.init(num_cpus=num_cpus or size)
        self.yy = [YYY() for _ in range(size)]
        self.yy_id = [ray.put(self.yy[i]) for i in range(size)]
        self.size = size

    def __call__(self, batch_references, batch_candidate) -> Any:
        assert len(batch_references) == len(batch_candidate)
        assert len(batch_references) <= self.size

        @ray.remote
        def run(x, references, candidate):
            return x(references, candidate)

        process = [run.remote(self.yy_id[i], hoge, fuga) for i in range(N)]
        batch_results = ray.get(process)
        return batch_results

example

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
import time
import numpy as np
import ray

ray.init(num_cpus = 4)

@ray.remote
def no_work(a):
    return

start = time.time()
a = np.zeros((5000, 5000))
result_ids = [no_work.remote(a) for x in range(10)]
results = ray.get(result_ids)
print("duration =", time.time() - start) # duration = 1.0837509632110596

👇👇👇

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
import time
import numpy as np
import ray

ray.init(num_cpus = 4)

@ray.remote
def no_work(a):
    return

start = time.time()
a_id = ray.put(np.zeros((5000, 5000)))  # ここ
result_ids = [no_work.remote(a_id) for x in range(10)]
results = ray.get(result_ids)
print("duration =", time.time() - start) # duration = 0.132796049118042 🎉

https://docs.ray.io/en/latest/ray-core/tips-for-first-time.html#tip-3-avoid-passing-same-object-repeatedly-to-remote-tasks

共有

YuWd (Yuiga Wada)
著者
YuWd (Yuiga Wada)
機械学習・競プロ・iOS・Web