Sending multiple buffers together

[jakirkham]

UCX supports transmitting multiple buffers together. It would be useful to have this functionality to work with objects that Dask serializes into multiple buffers. This comes up with many cuDF objects that are composed of multiple buffers. Also comes up with sparse arrays, which contain multiple buffers.

[quasiben]

I think this happens at the UCT layer and UCX-Py is only working with the UCP parts of UCX.

[pentschev]

@quasiben is right. We do not support any of the other UCX layers rather than UCP so that is not possible today. To be more concise, I don't know if we want to support other layers in the future, but perhaps the feature of communicating multiple buffers is an interesting feature for UCX core to expose to UCP, I'm not sure if there's some technical limitation to that but would be an interesting request in the future.

[jakirkham]

@madsbk, do you have any thoughts here?

[jakirkham]

Should add am exploring this at the Dask level in PR ( dask/distributed#3732 ). However it comes with the caveat that we need ~2x as much memory for send and receive as we need to pack buffers together to send and then unpack them after receiving. If UCX has an API that would avoid this memory allocation, that would be really useful.

[jakirkham]

Apparently this is not an option for us currently ( openucx/ucx#6912 )

[jakirkham]

Even though UCX lacks this functionality, there is still value in having fewer passes through the full Python stack as we do now for multiple buffers.

Namely we could imagine a Python API that takes a sequence of objects that can be coerced into Arrays, performs that coercion on them, and then calls the appropriate UCX API to send each one. A similar API could be created for receiving.

As a result we would have one pass through the Python stack to send the buffers and one pass to receive them. If in the future UCX provides the needed APIs for IOV, we could use them to replace the individual sends/recvs we would otherwise do. However we still get to cut out a lot of overhead for ourselves in the near term with this approach.