[core][aDAG] Multi ref output doesn't work with asyncio.gather

[rkooo567]

What happened + What you expected to happen

When we have multi output refs, we want to batch wait the result. For the sync case, I verified ray.get works. But for async case, asyncio.gather doesn't seem to work.

Versions / Dependencies

master

Reproduction script

import asyncio
import ray
from ray.dag import InputNode, MultiOutputNode

async def main():
    @ray.remote
    class A:
        def f(self, i):
            return i

    a = A.remote()
    b = A.remote()

    with InputNode() as inp:
        x = a.f.bind(inp)
        y = b.f.bind(inp)
        dag =  MultiOutputNode([x, y])

    adag = dag.experimental_compile(enable_asyncio=True)
    refs = await adag.execute_async(1)
    outputs = []
    # works
    # for ref in refs:
    #     outputs.append(await ref)
    # doesn't work.
    outputs = await asyncio.gather(*refs)
    print(outputs)

asyncio.run(main())

Issue Severity

None

[jeffreyjeffreywang]

Need to support ray.wait on CompiledDAGFuture and a list of CompiledDAGFuture as well. @rkooo567, is this blocking you from adopting the new execute_async API in vLLM? I'm trying to assess the priority of this issue.

[rkooo567]

In vLLM we don't need it (we only await on the first ref). It is mainly for beta release (aiming the mid Oct).

Also one quick note is that we don't support ray.wait on futures in Ray (we use asyncio.gather instead).

Also other thing to mention is that ray.wait may be more useful when you have multiple dags and wait concurrently.

[jeffreyjeffreywang]

Got it, will prioritize this over the deserialization issue.