Open
Description
What happened + What you expected to happen
I have a feeling that unbatched ray.wait is pretty slow -- while I expect it to be slower than batched ray.wait, the difference is pretty extreme. We can probably optimize this.
In [3]: import ray
In [4]: @ray.remote
...: def f(i):
...: return str(i)
In [8]: def g(rest):
...: while True:
...: done, rest = ray.wait(rest, num_returns=1)
...: if len(rest) == 0:
...: return
...: ray.get(done)
...:
In [9]: objs = [f.remote(i) for i in range(10000)]
In [10]: %time g(objs)
CPU times: user 7.05 s, sys: 375 ms, total: 7.42 s
Wall time: 7.52 s
In [11]: def g(rest):
...: while True:
...: done, rest = ray.wait(rest, num_returns=5)
...: if len(rest) == 0:
...: return
...: ray.get(done)
...:
In [12]: objs = [f.remote(i) for i in range(10000)]
In [13]: %time g(objs)
CPU times: user 1.48 s, sys: 99.9 ms, total: 1.58 s
Wall time: 1.58 s
In [14]: def g(rest):
...: while True:
...: done, rest = ray.wait(rest, num_returns=20)
...: if len(rest) == 0:
...: return
...: ray.get(done)
...:
In [15]: objs = [f.remote(i) for i in range(10000)]
In [16]: %time g(objs)
CPU times: user 457 ms, sys: 71 ms, total: 528 ms
Wall time: 531 msVersions / Dependencies
Ray 2.40
Reproduction script
see above