[sdk/python] - Await all async tasks

[komalali]

Currently, only RPC futures are awaited in the python SDK. When RPCs are complete, we collect all remaining tasks and cancel them. This means that code within .apply() and/or un-exported outputs are not always awaited and may be canceled in-flight.

This PR awaits the remaining outstanding tasks until either an exception is raised or all futures are resolved. Once asyncio.wait returns we cancel any pending tasks if they exist and check the done pile for any futures that raised an exception, re-raising if so.

We then repeat the cycle if the "extra async work" has itself scheduled any RPCs, only breaking when there are no remaining RPCs at the end of the cycle.

Python part of #3991

[komalali]

Is this actually true? I haven't found any instances where we start tasks "forever" in practice. And I don't really know why we would have a run-forever task within the pulumi_func especially.

The solution here depends on this statement being false and an assumption that all futures will resolve unless an exception is thrown.

[lblackstone]

LGTM with the caveat that I'm not very familiar with this part of the codebase.

[komalali]

Had a chat with @pgavlin about this PR and he brought up the fact that the "other async tasks" could themselves schedule new RPCs so I need to do a little more work here, at the very least to add tests for that particular scenario.

[pgavlin]

"other async tasks" could themselves schedule new RPCs so I need to do a little more work here, at the very least to add tests for that particular scenario.

One concrete example is the creation of a resource inside of an Apply callback: the callback will not run until the corresponding RPC has finished, and running the callback will add work to the RPC queue.

[komalali]

@pgavlin I've added the "create a resource inside apply" test and fixed up the code accordingly. Should be ready for your review 🙏🏽