Support Trio

[mahmoudhas]

Is it possible to add support for trio as opposed to just asyncio?

[pentschev]

Hi @mahmoudhas , this is certainly possible. After checking briefly, it seems the most viable paths would be:

1. Build a trio-based library on top of the UCXX "core" Python library that is completely async-agnostic; or
2. Adapt the existing UCXX async library to use AnyIO, which would provide support for both async and trio.

Personally, I think option 1 would be simpler based on my brief analysis. However, this work is not something we would have time to prioritize any time soon and thus would probably require a community contribution, is this something you would be interested in working on and contributing? We may be able to provide some guidance if this is something you or some other community member is interested in working on.

Would you also be able to share more about your use case?

[mahmoudhas]

Hi @pentschev , I'm using ucxx to send inference data (hidden states, kv cache) in the Baseten Inference Stack. We prefer trio over asyncio for a couple of reasons:

• structured concurrency makes it easier to write fault tolerant code: e.g. handling inference server disconnects, reconnects
• we use trio's guest loop mode to progress ucxx while waiting for torch.cuda.synchronize() and other cuda events without introducing Python threads.

For now, we have a workaround using trio_asyncio but we might contribute trio support later.

[pentschev]

That sounds like an interesting use case. For the short term it seems trio_asyncio is a decent option, I have never tested it myself but it seems like you're already using it, so it may be at least functional enough. If you end up feeling like contributing a proper implementation please let us know here and we can at least provide some directional help with implementation and whatnot, contributions are definitely welcome. Excited to hear more about your use case if/when you have more to share!