Skip to content

Warn if running under mpi but mpi4py/mpi4jax not installed

Vicentini Filipporequested to merge
mpi-warn into master

Created by: PhilipVinc

This PR attempts to detect if netket is running under mpi when no mpi dependencies are installed, and raise a warning, as some users might not realise (aka, they don't read the documentation) that even if mpi is installed in the cluster, they need to install mpi4py in the environment.

The detection of MPI is very arbitrary. It checks a few known env variables. I suspect it will work reliably for OpenMPI, probably work for MPICH and maybe but who knows work for intel MPI. Apparently there does not exist a way to detect this reliably.

Example output:

➜ mpirun -np 2 python Examples/Ising1d/ising1d_sr.py
/Users/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/utils/mpi.py:69: UserWarning:
MPI WARNING: It seems you might be running Python with MPI, but dependencies required
by NetKet to enable MPI support are missing or cannot be loaded, so MPI support is
disabled.

NetKet will not be taking advantage of MPI, and every MPI rank will execute the
same code of the others.

MPI dependencies are:
  - mpi4py>=3.0.1     ....... available
  - mpi4jax>=0.2.11   ....... missing

To enable MPI support, install the missing dependencies.
To learn more about MPI and NetKet consult the documentation at
https://www.netket.org/docs/getting_started.html

  warnings.warn(dedent(
  0%|          | 0/300 [00:00<?, ?it/s]/Users/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/utils/mpi.py:69: UserWarning:
MPI WARNING: It seems you might be running Python with MPI, but dependencies required
by NetKet to enable MPI support are missing or cannot be loaded, so MPI support is
disabled.

NetKet will not be taking advantage of MPI, and every MPI rank will execute the
same code of the others.

MPI dependencies are:
  - mpi4py>=3.0.1     ....... available
  - mpi4jax>=0.2.11   ....... missing

To enable MPI support, install the missing dependencies.
To learn more about MPI and NetKet consult the documentation at
https://www.netket.org/docs/getting_started.html

  warnings.warn(dedent(
  5%|▌         | 16/300 [00:28<09:23,  1.98s/it, Energy=-50.7913 ± 0.0087 [σ²=0.4800, R̂=1.0009]]^C[mpiexec@MacBook-Pro-di-Filippo.local] Sending Ctrl-C to processes as requested
[mpiexec@MacBook-Pro-di-Filippo.local] Press Ctrl-C again to force abort
  5%|▌         | 16/300 [00:30<09:09,  1.94s/it, Energy=-50.7913 ± 0.0087 [σ²=0.4800, R̂=1.0009]]No output specified (out=[apath|nk.logging.JsonLogger(...)]).Running the optimization but not saving the output.

Merge request reports