Compatibility with pip-installed openmpi

[pcuenca]

Proposed changes

When installing openmpi using pip in a virtual environment, the installation looks like this (incomplete):

.venv
├── bin
│   ├── mpirun
│   ├── ompi_info
├── lib
│   ├── libmpi.40.dylib  # <- not libmpi.dylib

However, mpi.cpp uses a hardcoded name, libmpi.dylib, to dlopen(). The installation does not provide a symlink.

On the Python side, distributed_run was already special-casing homebrew installations by explicitly passing DYLD_LIBRARY_PATH to mpirun. This PR:

• Extends the existing mechanism so DYLD_LIBRARY_PATH is sent in all cases, as long as an mpi library is found.
• Derives the library name by examining ompi_info with otool. Other common tools, such as mpirun, are not linked against the shared library.
• Sends the library name as an envvar for mpi.cpp to use.

This should result in a more robust environment detection and increased compatibility.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

Manual Tests

• macOS with pip
• macOS with homebrew
• macOS with conda
• Linux (CPU) with pip
• Linux (CPU) with conda (conda install openmpi; conda install conda-forge::openmpi did not put an mpirun executable anywhere in my environment)

[pcuenca]

If this solution is acceptable, I can tweak the docs a bit.

[madrob]

Does this suggest that instead of brew install openmpi in the CI setup we should python install it?

[pcuenca]

This works with brew, pip, and conda. Needs some adjustments for Linux (cpu backend).

[pcuenca]

As updated in the initial description, I tested the following combinations:

• macOS with pip
• macOS with homebrew
• macOS with conda
• Linux (CPU) with pip
• Linux (CPU) with conda (conda install openmpi; conda install conda-forge::openmpi did not put an mpirun executable anywhere in my environment)

[angeloskath]

Thanks! This looks great!

I 'll merge after the tests clear.

[angeloskath]

Some random thought. Perhaps we should change the environment variable to MLX_MPI_LIBNAME since I don't think it is standard and we could avoid collisions. I am fine with MPI_LIBNAME too though.

[pcuenca]

change the environment variable to MLX_MPI_LIBNAME

Good call, thank you!