Add a tool to automatically pick GPUs
Created by: PhilipVinc
Every time I have to run simulations with gpus I'm annoyed that I have to use the CUDA_VISIBLE_DEVICES
hack before launching python or importing jax.
It would be great if netket could provide a netket.tools.autogpu
incantation that could be used like
mpirun -np 4 python -m netket.tools.autogpu myscript.py
which was equivalent to
mpirun -np 4 CUDA_VISIBLE_DEVICES="$RIGHTDEVICE" python myscript.py
To do this in a portable way we probably need to do something like
import os
from mpi4py import MPI
COMM = MPI.COMM_WORLD
rank = COMM.Get_rank()
all_processor_names = COMM.allgather(MPI.Get_processor_name())
if local_processor_name not in all_processor_names:
raise RuntimeError("A bug occurred")
local_sizes = {}
local_ranks = [None for _ in COMM.Get_size()]
for (rank, address) in enumerate(all_processor_names):
local_rank = local_sizes.get(address, 0)
local_ranks[rank] = local_rank
local_sizes[address] = local_rank + 1
local_rank = local_ranks = local_rank[COMM.Get_rank()]
os.environ["CUDA_VISIBLE_DEVICES"] = f"{local_rank}"
subprocess.run("python *args....")
@inailuig could you give an eye on how to finish this? most likely calling the right subprocess incantation with cmlimd arguments and ensuring that stdout and stderr is correctly passed ?