Map GPU operation to GPU ID

[berceanu]

The flow project contains one GPU operation per job. The goal is to track the GPU memory consumption during a run, on a per-job, per-GPU basis.

Workaround to SLURM not yet having a --gpus-per-task option:

{% set cmd_suffix = cmd_suffix|default('') ~ (' &' if parallel else '') %}
{% for operation in operations %}
export CUDA_VISIBLE_DEVICES={{ loop.index0 }}
{{ operation.cmd }}{{ cmd_suffix }}
{% endfor %}

In this way each GPU operation executes on a distinct GPU. I would like to have a flow operation which launches simulataneously with each GPU operation and tracks its memory usage.

Current approach: periodically get the PIDs of the running GPU compute processes, and map them to signac jobs IDs by grepping ps aux output for PID | -j <job_id>.

[vyasr]

This issue is closely related to #115 regarding the best way to handle GPU assignment and management within signac-flow.

[berceanu]

Indeed, it;s the same problem, but now using SLURM.
However, I did find a solution: use srun instead of mpirun and specify --gres=gpu:x for each srun command.
Example SLURM script here. The --mpi=pmi2 flag was needed for srun to pick up MPICH which was installed by mpi4py in the conda environment.