multi-GPU reconstruction not working

[apeck12]

Currently the pipeline yields the correct structure when run on a single CPU, multiple CPUs (MPI-accelerated), and a single GPU; however, the MPI mode is not working across multiple GPUs. One issue is that the gpu_device_id=0 is hardcoded in the forward_gpu and adjoint_gpu NUFFT calculations. This was addressed by retrieving the device ID and passing it to the NUFFT calculations as follows:

import pycuda.driver as drv
drv.init()
device_id = comm.rank % drv.Device.count()

With this fix, the AC solver appears to be running on different GPUS, but fails with the following error:

Traceback (most recent call last):
...  
File "/sdf/home/a/apeck/cmtip/cmtip/nufft/adjoint.py", line 70, in adjoint_gpu
    nuvect_gpu = to_gpu(nuvect.astype(np.complex64).flatten())
  File "/sdf/home/a/apeck/.local/lib/python3.7/site-packages/pycuda/gpuarray.py", line 1049, in to_gpu
    result = GPUArray(ary.shape, ary.dtype, allocator, strides=_compact_strides(ary))
  File "/sdf/home/a/apeck/.local/lib/python3.7/site-packages/pycuda/gpuarray.py", line 210, in __init__
    self.gpudata = self.allocator(self.size * self.dtype.itemsize)
pycuda._driver.LogicError: cuMemAlloc failed: invalid device context

even when requesting the same number of orientations/images that fits on a single GPU without this memory error. It's unclear if I'm correctly requesting the desired allocation (multiple nodes with one GPU per node), so this may be a SLURM issue instead. Currently I'm submitting the job with the following parameters:

#SBATCH --cpus-per-task=4
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --gres=gpu:4

On SDF, variations of

#SBATCH --cpus-per-task=1
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:4

result in job submission failures for violating accounting/QOS policy.