Now that this works, we could even do like Jax does, and dlopen xla_bridge_cuda only when needed (essentially when an user binds a cuda primitive).

By dlopening it, we could get around the linking problem, as we could have Nvidia wheels as a compile-time dependency for the headers, but then only load the .so when needed. Then, we could do as jax does and set the rpath to the parent folder of the package, which we know should contain Nvidia headers.

This would basically solve the CUDA linking problem and make it as simple to install mpi4jax+cuda without having to use --no-build-isolation.

To support installations with 'local' cuda, the dlopen mechanism would automatically pickup LD_LIBRARY_PATH if nothing is in the RPATH (again, this is the scheme used by jax). As the user must already set LD_LIBRARY PATH, this would work out well.

On possible way to ship this would be to mimic jax's extras, and have mpi4jax be mpi4jax[cpu], and then also have mpi4jax[cuda] and mpi4jax[rocm] and mpi4jax[sycl]...

I am unsure, but surely finding a way to simplifying the installation would be great.

That sounds doable with some work, but is this really that complicated?

pip install jax[cuda] mpi4py cython
pip install mpi4jax --no-build-isolation

Well, the problem is that you cannot ship a requirements file, instead you must share a script to setup the environment.
This breaks a bit 'simple to setup' reproducibility.

It would be better if we could just do pip install jax[cuda] mpi4jax[cuda] .

I agree that it's a lot of work for not much, though. So it's just something that will stay in the backlog until who knows when...

@PhilipVinc does that work?

Yes. Tested and it works.

Well, the problem is that you cannot ship a requirements file, instead you must share a script to setup the environment. This breaks a bit 'simple to setup' reproducibility.

It would be better if we could just do pip install jax[cuda] mpi4jax[cuda] .

I agree that it's a lot of work for not much, though. So it's just something that will stay in the backlog until who knows when...

Actually there's may be another option: depend on the CUDA runtime both at build time and at runtime. That way you can link to them during the build, set the rpath to site packages and rely on them to be present (since they're also a runtime dependency). Not the pretties solution but this should work out of the box.

Yea, but this must be something the user opts in to. ideally it should be exposed as an extra. But build time dependencies can only be declared before we process extras requires.

the only solution I see is to do like Jax, and have a second package with the correct build time dependencies, which is required only if the extra is specified.

but that’s quite a bit extra infrastructure..