[core][adag] Support asyncio.gather on multiple CompiledDAGFutures (#47684)

[jeffreyjeffreywang]

Why are these changes needed?

Currently, awaiting individual CompiledDAGFuture works but performing asyncio.gather on a list of CompiledDAGFuture won't because there is a race condition between CompiledDAGFutures that originate from the same execution in a aDAG. One solution is to relax the assertion in _cache_execution_results, allowing multiple CompiledDAGFuture originating from the same execution to attempt to cache the result for that execution. The other solution is to protect the critical section with asyncio.Lock. The first solution could result in awaiting on the same future multiple times, but since this overhead is lower than that of the second solution (i.e. locks introduce more overhead), we go for the first solution.

Related issue number

Resolves #47684.

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[jeffreyjeffreywang]

Hey @rkooo567 @stephanie-wang, could you please take a look once you have time? This addresses the issue of not able to invoke asyncio.gather on a list of CompiledDAGFuture. Thank you!

[stephanie-wang]

Thanks for the fix, @jeffreyjeffreywang!

Can you measure throughput before/after? A little worried about the lock overhead and dynamic lock allocation.

Also, can you explain what data/code the lock protects? Is there a lock-free design that can work?

[jeffreyjeffreywang]

Thanks for the review, @stephanie-wang! The locks protects _result_buffer from being accessed and modified by multiple CompiledDAGFutures originating from the same execution index and avoids the same underlying asyncio.Future from being awaited multiple times. This will only be a problem when we have multiple output channels for the DAG.

Imagine that there are two output channels, and execute_async returns fut1 and fut2. While awaiting on fut1 and fut2, there is a race condition when they both attempt to store execution result in _result_buffer. Since the DAG has not stored the execution result, both will await on the same underlying future and write to _result_buffer. After either fut1 or fut2 writes to _result_buffer, the latter write triggers an assertion error. We could remove the assertion, but this will not solve the issue of awaiting on the same underlying future more than once.

I can add some logic to avoid using lock when there is a single output channel so that we don't influence the performance when MultiOutputNode isn't used. Besides, is there any tool I can use to measure the throughput?

[stephanie-wang]

Thanks for the review, @stephanie-wang! The locks protects _result_buffer from being accessed and modified by multiple CompiledDAGFutures originating from the same execution index and avoids the same underlying asyncio.Future from being awaited multiple times. This will only be a problem when we have multiple output channels for the DAG.

Imagine that there are two output channels, and execute_async returns fut1 and fut2. While awaiting on fut1 and fut2, there is a race condition when they both attempt to store execution result in _result_buffer. Since the DAG has not stored the execution result, both will await on the same underlying future and write to _result_buffer. After either fut1 or fut2 writes to _result_buffer, the latter write triggers an assertion error. We could remove the assertion, but this will not solve the issue of awaiting on the same underlying future more than once.

Ah I see, thanks for the explanation! Seems like it could also be fixed by checking again if the result is still not cached after await fut?

I can add some logic to avoid using lock when there is a single output channel so that we don't influence the performance when MultiOutputNode isn't used. Besides, is there any tool I can use to measure the throughput?

Hmm unfortunately I don't think we have an asyncio microbenchmark. But you can take a look at the tests that are run on nightly for compiled graph. Actually adding an asyncio version of this microbenchmark would be useful (I'll file an issue for this if you're interested).

[jeffreyjeffreywang]

Ah I see, thanks for the explanation! Seems like it could also be fixed by checking again if the result is still not cached after await fut?

Yup, we could simply check whether the result has been cached in _cache_execution_results and don't enforce this assertion anymore. This, however, can lead to awaiting the same future multiple times, but since creating and acquiring locks is more expensive, I think we can afford this overhead.

[jeffreyjeffreywang]

Hey @stephanie-wang, this is now ready for another review.

[stephanie-wang]

Oops sorry for the delay in review. This looks great, thanks!

[rkooo567]

there are test failures!

[rkooo567]

btw is the description still relevant? I am not seeing the lock from code

[jeffreyjeffreywang]

@stephanie-wang no worries, I was busy dealing with work as well. Thanks for the review!
@rkooo567 Looks like the test failures are not related to my changes. Will rebase my branch to the latest master. I updated the description as well.

[stephanie-wang]

Hey @jeffreyjeffreywang FYI I re-merged master, not sure why there was a core test failure in your earlier merge.