WIP: Use PTDS for synchronization in UCX communication

[jakirkham]

This changes the stream we synchronize on to support the per-thread default stream.

Requires numba version 0.50.0+ ( numba/numba#5347 )

Note: This is pending other changes in RAPIDS. Just putting this up to show what changes would be needed here at this point.

cc @shwina @kkraus14 @jrhemstad (as we discussed this earlier)
cc @quasiben @pentschev @harrism (for awareness)

[harrism]

It won't wait for other threads' default per-thread default streams...

[jakirkham]

So should we leave this as-is? Or is there something else we should be doing here?

[kkraus14]

I think Mark is referring to the section of syncing the per-thread default stream will wait for other non-blocking CUDA streams as that line is incorrect. Syncing the per-thread default stream will only wait for the current stream WITHOUT blocking work on other non-blocking CUDA streams.

[jakirkham]

Right, I just wasn't sure if this implied changes beyond the text in the comment.

[harrism]

Assuming syncing only the PTDS is correct behaviour, I think only the comment is incorrect.

[kkraus14]

One concern, what if someone is doing non-synchronized work on the legacy default stream and then goes to send here? We'll synchronize only the per-thread default stream and not synchronize the legacy default stream which can leave UCX in a bad state.

[jakirkham]

So what would the counterproposal be? Leave things as-is?

[kkraus14]

So what would the counterproposal be? Leave things as-is?

I think medium term is work with the community to get the stream included in __cuda_array_interface__. Short term, I'm not sure.

[pentschev]

Probably good to close, UCX has been removed from Distributed, see #9105.