Input on mpirun Launcher for JAX Support in Kubeflow Trainer V2

[mahdikhashan]

Hi everyone,

As part of a KEP to integrate JAX into Kubeflow Trainer V2, I'm currently evaluating whether using the mpirun launcher is a good fit. One of my main concerns is around its suitability for TPUs, which seem to rely on gRPC for communication rather than MPI.

More generally, I'm also interested in understanding the performance characteristics of MPI, particularly in terms of latency. Are there any known limitations or considerations when using MPI in latency-sensitive workloads?

Any guidance, experiences, or references would be greatly appreciated.

Here’s the KEP for context: kubeflow/trainer#2643

Thanks
Mahdi