🤖

Rayシリーズ：リモート関数戻り値の返し方のアンチパターンについて

2025/09/09に公開

今回はRayでリモート関数を利用するにあたり、その戻り値としてアンチパターンとされているものを紹介しようと思います。情報源は以下になります。

単一の値を返す場合

リモート関数から値が返される場合、その値が大きいのか小さいかに関わらず、その値を直接returnするのがいいとのことです。アンチパターンとしては、ray.put()を利用して参照を作成して返すことみたいです。これがなぜかというと、リモート関数の戻り値は自動的にオブジェクトストアに登録されて参照として呼び出し元に返されますが、参照を戻り値にしてしまうと、受け取り側は参照の参照を受け取ることになります。結果として値を取得するにはray.getを二回呼び出す必要があり、このような理由からアンチパターンとなります。

import ray


@ray.remote
def task_with_single_small_return_value_bad():
    small_return_value = 1
    small_return_value_ref = ray.put(small_return_value)
    return small_return_value_ref


@ray.remote
def task_with_single_small_return_value_good():
    small_return_value = 1
    return small_return_value


bad_example = ray.get(
  ray.get(task_with_single_small_return_value_bad.remote())
)
good_example = ray.get(task_with_single_small_return_value_good.remote())
print(f"{bad_example=}, {bgood_example=}")
print(f"{bad_example == good_example}")

これを実行すると以下のようになります。

(raylet) warning: `VIRTUAL_ENV=/Users/user/Documents/Blog/blog_materials/ray/antipattern/remote_return_object_ref/.venv` does not match the project environment path `.venv` and will be ignored; use `--active` to target the active environment instead
(raylet) Using CPython 3.12.9
(raylet) Creating virtual environment at: .venv
(raylet) Installed 27 packages in 292ms
True
bad_example=1, good_example=1

結果は同じですが、参照を返している方は二重にray.getを実行しているためかなり遠回しになっていますね。

上記の例では整数値を取り扱っていましたが、取り扱うデータが大きなものになっても基本的に利用方法は変わりません。例えばNumPy配列を戻り値としたい場合も、データが大きいからといって特別別の方法を取る必要はありません。

import ray
import numpy as np


@ray.remote
def task_with_single_large_return_value_bad():
    large_return_value = np.zeros(1024 ** 2)
    large_return_value_ref = ray.put(large_return_value)
    return large_return_value_ref


@ray.remote
def task_with_single_large_return_value_good():
    large_return_value = np.zeros(1024 ** 2)
    return large_return_value


bad_example = ray.get(
  ray.get(task_with_single_large_return_value_bad.remote())
)
good_example = ray.get(task_with_single_large_return_value_good.remote())
print(f"{np.all(bad_example == good_example)=}")

こちらを実行すると、以下のように結果は同じになりました。

(raylet) warning: `VIRTUAL_ENV=/Users/user/Documents/Blog/blog_materials/ray/antipattern/remote_return_object_ref/.venv` does not match the project environment path `.venv` and will be ignored; use `--active` to target the active environment instead
(raylet) Using CPython 3.12.9
(raylet) Creating virtual environment at: .venv
(raylet) Installed 27 packages in 247ms
np.all(bad_example == good_example)=np.True_

戻り値の個数がわかっている場合

実装するリモート関数の戻り値が事前にわかっている場合は、@ray.remoteでnum_returnsオプションを指定します。例えば二つの値を返す例をみてみましょう。

import ray


@ray.remote(num_returns=1)
def task_with_two_return_values_as_bad1():
    value1, value2 = 1, 2
    value1_ref, value2_ref = ray.put(value1), ray.put(value2)
    return value1_ref, value2_ref


@ray.remote(num_returns=2)
def task_with_two_return_values_as_bad2():
    value1, value2 = 1, 2
    value1_ref, value2_ref = ray.put(value1), ray.put(value2)
    return value1_ref, value2_ref


@ray.remote(num_returns=2)
def task_with_two_return_values_as_good():
    value1, value2 = 1, 2
    return value1, value2


bad1_example_ref = ray.get(
    task_with_two_return_values_as_bad1.remote()
)
bad1_example_value1 = ray.get(bad1_example_ref[0])
bad1_example_value2 = ray.get(bad1_example_ref[1])
print(bad1_example_value1, bad1_example_value2)

bad2_example_nested_refs = task_with_two_return_values_as_bad2.remote()
bad2_example_refs = ray.get(bad2_example_nested_refs)
bad2_example_values = ray.get(bad2_example_refs)
print(bad2_example_values)

good_example_ref = task_with_two_return_values_as_good.remote()
good_example_values = ray.get(good_example_ref)
print(good_example_values)

まず最初の例ですが、num_returns=1としているため、戻り値を二つにしていますが、戻り値は二つを一つにまとめた単一の参照になります。そのため、それぞれの値への参照を取得するためにまずray.getを呼び出し、次にその結果を分割してそれぞれをray.getに与える必要があります。

二つ目の例ですが、num_returns=2としているので、戻り値はそれぞれのリファレンスとして分割することができます。しかし、最初に提示したように戻り値自体を参照として返すことはアンチパターンですので、値を参照するにはray.getを二重で呼び出す必要があります。

最後の例は、最初のアンチパターンに当てはまらないベストプラクティスにのっとり、戻り値は値そのものとしており、戻り値の個数もnum_returns=2に設定しており正しい内容となっています。

それではこれを実行してみましょう。

(raylet) warning: `VIRTUAL_ENV=/Users/user/Documents/Blog/blog_materials/ray/antipattern/remote_return_object_ref/.venv` does not match the project environment path `.venv` and will be ignored; use `--active` to target the active environment instead
(raylet) Using CPython 3.12.9
(raylet) Creating virtual environment at: .venv
(raylet) Installed 27 packages in 223ms
1 2
[1, 2]
[1, 2]

こちらの結果のように、取得できている値は同じですが、手間のかかり具合が全然違うことが確認できます。

ちなみに、戻り値の件数が事前にわからない場合（引数によって結果が変わる場合）はnum_returns='dynamic'と設定することで明示できます。例えば、サンプルで提示されているコードを利用してみましょう。

import ray
import numpy as np

@ray.remote(num_returns=1)
def task_with_dynamic_returns_bad(n):
    return_value_refs = []
    for i in range(n):
        return_value_refs.append(ray.put(np.zeros(i * 1024 * 1024)))
    return return_value_refs


@ray.remote(num_returns="dynamic")
def task_with_dynamic_returns_good(n):
    for i in range(n):
        yield np.zeros(i * 1024 * 1024)



print(np.array_equal(
    ray.get(ray.get(task_with_dynamic_returns_bad.remote(2))[0]),
    ray.get(next(iter(ray.get(task_with_dynamic_returns_good.remote(2))))),
))

こちらはtask_with_dynamic_returns_goodはジェネレータを戻り値としており、その結果を利用して結果を取得するようになっています。実行すると、取得される値は同じものであることが確認できます。

(raylet) warning: `VIRTUAL_ENV=/Users/user/Documents/Blog/blog_materials/ray/antipattern/remote_return_object_ref/.venv` does not match the project environment path `.venv` and will be ignored; use `--active` to target the active environment instead
(raylet) Using CPython 3.12.9
(raylet) Creating virtual environment at: .venv
(raylet) Installed 27 packages in 251ms
True

まとめ

今回はリモート関数の戻り値のアンチパターンについてみてみました。Rayではリモートオブジェクトを多用した実装にするため、効率よくデータのやり取りをするためにも、こちらのアンチパターンにならないような実装を心がけましょう。

単一の値を返す場合

戻り値の個数がわかっている場合

まとめ

Discussion