Open eaplatanios opened 1 month ago
PJRT is device runtime interface. IFRT is distributed runtime interface. IFRT contains a global view of arrays and computations that span devices belong to different hosts, while PJRT only has local view to single host.
Global information in PJRT such as global topology is in the process of being moved to IFRT to make the separation more clear.
TFRT usually refers to one possible implementation of PJRT.
@skye @hawkinsp Any other information you want to add? Thanks!
Thanks for clearing up the difference! In that case, I have a couple follow-up questions:
xla/pjrt/distributed/
going away in the future and I should not bother generating bindings for that part of the code?xla/python/ifrt
and xla/python/pjrt_ifrt
? Should I assume that these will move outside of the python
package at some point since it sounds like they're not specific to the Python bindings?Thanks!
xla/pjrt/distributed/
is not going away. in_memory_key_value_store.* are still being used (it is just an implementation of the key-value store interface. https://github.com/openxla/xla/blob/main/xla/tsl/distributed_runtime/coordination/coordination_service.h is used for the implementations for DistributedRuntimeService so some of the old DistributedRuntimeService are no longer used. Otherwise most of the files are interface/client which are still relevant (I think).
I am not sure why ifrt
and pjrt_ifrt
are located in xla/python
and whether they will be moved out :)
JAX mostly only interacts with IFRT (and IFRT wraps PJRT client, as in xla/python/pjrt_ifrt
). Therefore, even it is a single-node multi-GPU programs, JAX still talks to IFRT.
@eaplatanios I'd like to ask you something. I'm also working on rust binding. Do you have a repo or something available?
I'm currently building a library in Rust that leverages XLA and PJRT and while going through the APIs and the implementation of the Python bindings, I'm a little confused about the terminology and the various interfaces. Specifically, I ran into
get_tfrt_cpu_client
and its implementation in the Python bindings and I have a couple of questions related to it:PjRtClient
fromxla/pjrt
into anifrt::PjRtClient
fromxla/python/pjrt_ifrt
when a distributed client is provided. What is the difference between the two? Specifically, the CPUPjRtClient
fromxla/pjrt
seems to support a distributed setting so what do the Python bindings do differently and is it necessary or is it a by-product of an older implementation?Thank you!