We use a pretty standard tech stack of PyTorch + NCCL + MPI. We've used both OpenMPI and MPICH to varying degrees.
Kubeflow is interesting, but it solves a slightly different problem of scheduling/coordinating ML workflows on top of Kube. It doesn't get involved with how an ML job communicates within itself cross-node.
One difference is that these operators use the Kubernetes service interface for communication, generally exposing a headless service for each replica.
But more generally, MPI over ssh on a large kubernetes deployment is not a very common pattern. Any reason you chose that ?
Have you looked at Ray or Torch-Elastic (which seems to be officially supported by AWS, etc as well) https://github.com/pytorch/elastic ?
We started a discussion in https://github.com/kubeflow/mpi-operator/issues/315