wrap_rrdesi not running on coadds from a list?

[araichoor]

I report here an apparently expected behavior of wrap_rrdesi, which stops running after few coadds, with no error, and does not deal with the rest of the coadds from the input file.
tagging @craigwarner-ufastro (sugggestion from @sbailey)

note that I am a newbie in running redrock with gpus, so it can very be that there is no bug and that I am just mis-calling the code.

here is how I proceed:

# load the desi main environment
source /global/cfs/cdirs/desi/software/desi_environment.sh main

# grab an interactive node with gpu
salloc --nodes 1 --qos interactive --time 4:00:00 --constraint gpu --gpus-per-node=4 --account desi_g

# set environment variables (custom templates)
export RR_TEMPLATE_DIR=/global/cfs/cdirs/desi/users/raichoor/laelbg/templates/rr2023oct
export OMP_NUM_THREADS=1

# inputs/outputs
MYDIR=/global/cfs/cdirs/desi/users/raichoor/laelbg/loa/healpix/tertiary37-thru20240309-loa
MYRRDIR=$MYDIR/rr2023oct
ls $MYDIR/coadd*fits > $MYRRDIR/list_coadds.ascii

# run redrock
wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $MYRRDIR --gpu > $MYRRDIR/redrock.log 2>&1

the $MYRRDIR/list_coadds.ascii file has 18 coadds, but the code stops after dealing with the first 5 coadds, with no apparent error.
the log file is $MYRRDIR/redrock.log.

[araichoor]

for info:

• I ve moved the previous redrock run to /global/cfs/cdirs/desi/users/raichoor/laelbg/loa/healpix/tertiary37-thru20240309-loa/rr2023oct-v0
• I ve re-run with the suggested command from @sbailey nicely works, handling all; I ve run the same commands as above, except that I grab two interactive nodes, and use this calling sequence:

srun -N 2 -n 8 -c 2 --gpu-bind=map_gpu:3,2,1,0  wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $MYRRDIR --gpu > $MYRRDIR/redrock.log 2>&1

[craigwarner-ufastro]

@sbailey @araichoor I took a look last night and it looks like it was just an incorrect syntax -- using 1 node also works fine with srun -N 1 -n 4 -c 2 --gpu-bind=map_gpu:3,2,1,0 wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $MYRRDIR --gpu

[sbailey]

@craigwarner-ufastro please confirm: is wrap_rrdesi only supposed to work in MPI-mode, i.e. with an srun prefix, and it doesn't support wrapping rrdesi+multiprocessing? That could be an acceptable restriction, but if so we should make that more obvious, e.g. by renaming to wrap_rrdesi_mpi and/or auto-detecting that it isn't being run in MPI mode (I don't actually know how to do that) and exiting with an informative message. From Anand's original report, it sounds like currently without the srun prefix it runs and does some of the work and then exits without an error, i.e. silently doing the wrong thing. Let's find a way to give the user more of a hint that they aren't calling it correctly and what to do instead.

[craigwarner-ufastro]

@sbailey yes that is correct wrap_rrdesi is only supposed to work with MPI. In fact GPU mode really only works with MPI period:

            if (mpprocs > 1):
                #Force mpprocs == 1 for multiprocessing mode with GPU
                print("WARNING:  using GPU mode without MPI requires --mp 1")
                print("WARNING:  Overriding {} multiprocesses to force this.".format(mpprocs))
                print("WARNING:  Running with 1 process.")

Although wrap_rrdesi does of course support CPU only runs as well. But as far as I was aware, wrap_rrdesi was only ever supposed to deal with splitting MPI communicators to do multiple runs on multiple ranks / multiple nodes.

A quick look shows that several env vars are created when its called by srun, e.g.

SLURM_NTASKS 4
SLURM_NPROCS 4

I suggest checking one of these and if not there, we can exit with an informative error.

[sbailey]

Thanks for clarifying. It looks like $SLURM_NTASKS can appear in a job's environment if launched with sbatch -n ... or salloc -n ... even before the srun call. It looks like $SLURM_TASK_PID might only exist after the srun call though. At the same time, I don't want to break wrap_redrock on non-slurm systems, even if NERSC is the only place we currently use it. Suggestion:

• rename wrap_rrdesi to wrap_rrdesi_mpi to give a hint to the user
• if $SLURM_JOB_NAME is set and $SLURM_TASK_PID is not set, assume that we are running with a SLURM-based scheduler but srun has not been called, so print an informative error message and exit with a non-zero exit code.
• otherwise print a warning message about maybe not being in MPI mode but proceed anyway and hope for the best

[craigwarner-ufastro]

@sbailey Oh - actually if I run

salloc --nodes 1 --qos interactive --time 4:00:00 --constraint gpu --gpus-per-node=4 --account desi_g

to get an interactive node, and then run wrap_rrdesi without srun, I see those vars both set:

SLURM_JOB_NAME interactive
SLURM_TASK_PID 1149951

Here is a list of SLURM env vars when doing an interactive node running

./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite

SLURMD_DEBUG 2
SLURMD_NODENAME nid001096
SLURM_CLUSTER_NAME perlmutter
SLURM_CPUS_ON_NODE 128
SLURM_GPUS_ON_NODE 4
SLURM_GPUS_PER_NODE 4
SLURM_GTIDS 0
SLURM_JOBID 33522533
SLURM_JOB_ACCOUNT desi_g
SLURM_JOB_CPUS_PER_NODE 128
SLURM_JOB_END_TIME 1733275777
SLURM_JOB_GID 98259
SLURM_JOB_GPUS 0,1,2,3
SLURM_JOB_ID 33522533
SLURM_JOB_LICENSES cfs:1
SLURM_JOB_NAME interactive
SLURM_JOB_NODELIST nid001096
SLURM_JOB_NUM_NODES 1
SLURM_JOB_PARTITION urgent_gpu_ss11
SLURM_JOB_QOS gpu_interactive
SLURM_JOB_START_TIME 1733261377
SLURM_JOB_UID 98259
SLURM_JOB_USER cdwarner
SLURM_LAUNCH_NODE_IPADDR 128.55.64.25
SLURM_LOCALID 0
SLURM_MPI_TYPE cray_shasta
SLURM_NNODES 1
SLURM_NODEID 0
SLURM_NODELIST nid001096
SLURM_PRIO_PROCESS 0
SLURM_PROCID 0
SLURM_PTY_PORT 44263
SLURM_PTY_WIN_COL 80
SLURM_PTY_WIN_ROW 24
SLURM_SCRIPT_CONTEXT prolog_task
SLURM_SRUN_COMM_HOST 128.55.64.25
SLURM_SRUN_COMM_PORT 36851
SLURM_STEPID 4294967290
SLURM_STEP_ID 4294967290
SLURM_STEP_LAUNCHER_PORT 36851
SLURM_STEP_NODELIST nid001096
SLURM_STEP_NUM_NODES 1
SLURM_STEP_NUM_TASKS 1
SLURM_STEP_TASKS_PER_NODE 1
SLURM_SUBMIT_DIR /global/cfs/cdirs/desicollab/users/cdwarner/code/desispec/bin
SLURM_SUBMIT_HOST login16
SLURM_TASKS_PER_NODE 128
SLURM_TASK_PID 968497
SLURM_TOPOLOGY_ADDR nid001096
SLURM_TOPOLOGY_ADDR_PATTERN node

Versus running

srun -N 1 -n 4 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite

SLURMD_DEBUG 2
SLURMD_NODENAME nid001096
SLURMD_TRES_BIND gres/gpu:map_gpu:3,2,1,0
SLURMD_TRES_FREQ gpu:high,memory=high
SLURM_CLUSTER_NAME perlmutter
SLURM_CPUS_ON_NODE 8
SLURM_CPUS_PER_TASK 2
SLURM_CPU_BIND quiet,mask_cpu:0x00000000000000010000000000000001,0x00000000000000020000000000000002,0x00000000000000040000000000000004,0x00000000000000080000000000000008
SLURM_CPU_BIND_LIST 0x00000000000000010000000000000001,0x00000000000000020000000000000002,0x00000000000000040000000000000004,0x00000000000000080000000000000008
SLURM_CPU_BIND_TYPE mask_cpu:
SLURM_CPU_BIND_VERBOSE quiet
SLURM_DISTRIBUTION block
SLURM_GPUS_ON_NODE 4
SLURM_GPUS_PER_NODE 4
SLURM_GTIDS 0,1,2,3
SLURM_JOBID 33522533
SLURM_JOB_ACCOUNT desi_g
SLURM_JOB_CPUS_PER_NODE 128
SLURM_JOB_END_TIME 1733275777
SLURM_JOB_GID 98259
SLURM_JOB_GPUS 0,1,2,3
SLURM_JOB_ID 33522533
SLURM_JOB_LICENSES cfs:1
SLURM_JOB_NAME interactive
SLURM_JOB_NODELIST nid001096
SLURM_JOB_NUM_NODES 1
SLURM_JOB_PARTITION urgent_gpu_ss11
SLURM_JOB_QOS gpu_interactive
SLURM_JOB_START_TIME 1733261377
SLURM_JOB_UID 98259
SLURM_JOB_USER cdwarner
SLURM_LAUNCH_NODE_IPADDR 128.55.70.19
SLURM_LOCALID 0
SLURM_MPI_TYPE cray_shasta
SLURM_NNODES 1
SLURM_NODEID 0
SLURM_NODELIST nid001096
SLURM_NPROCS 4
SLURM_NTASKS 4
SLURM_PRIO_PROCESS 0
SLURM_PROCID 0
SLURM_PTY_PORT 44263
SLURM_PTY_WIN_COL 80
SLURM_PTY_WIN_ROW 24
SLURM_SCRIPT_CONTEXT prolog_task
SLURM_SRUN_COMM_HOST 128.55.70.19
SLURM_SRUN_COMM_PORT 45615
SLURM_STEPID 0
SLURM_STEP_GPUS 0
SLURM_STEP_ID 0
SLURM_STEP_LAUNCHER_PORT 45615
SLURM_STEP_NODELIST nid001096
SLURM_STEP_NUM_NODES 1
SLURM_STEP_NUM_TASKS 4
SLURM_STEP_RESV_PORTS 63576-63580
SLURM_STEP_TASKS_PER_NODE 4
SLURM_SUBMIT_DIR /global/cfs/cdirs/desicollab/users/cdwarner/code/desispec/bin
SLURM_SUBMIT_HOST login16
SLURM_TASKS_PER_NODE 4
SLURM_TASK_PID 968968
SLURM_TOPOLOGY_ADDR nid001096
SLURM_TOPOLOGY_ADDR_PATTERN node
SLURM_TRES_BIND gres/gpu:map_gpu:3,2,1,0
SLURM_TRES_PER_TASK cpu:2
SLURM_UMASK 0007

What about CPU_BIND? What is SLURM_DISTRIBUTION block? Let me know if any of these look good to you to differentiate based upon.

[craigwarner-ufastro]

To save you time, these are the additional env vars when running with srun from an interactive node already obtained with salloc.

SLURMD_TRES_BIND
SLURMD_TRES_FREQ
SLURM_CPUS_PER_TASK
SLURM_CPU_BIND
SLURM_CPU_BIND_LIST
SLURM_CPU_BIND_TYPE
SLURM_CPU_BIND_VERBOSE
SLURM_DISTRIBUTION
SLURM_NPROCS
SLURM_NTASKS
SLURM_STEP_GPUS
SLURM_STEP_RESV_PORTS
SLURM_TRES_BIND
SLURM_TRES_PER_TASK
SLURM_UMASK

[sbailey]

Yeah, the mess comes that someone might get a node with

salloc --nodes 1 -n 128 --qos interactive --time 4:00:00 --constraint gpu --gpus-per-node=4 --account desi_g

in which case $SLURM_NPROCS is set for the job and gets inherited by srun if it is called without -n, but then that gets in the way of telling if the actual command is prefixed with an srun or not. I suspect that is true for most of those environment variables, since SLURM has a design hierarchy of multiple ways of controlling the options with different levels overriding others (envvar -> script #SBATCH headers -> sbatch/salloc options -> srun options).

Arguably I'm way overthinking this and using a more obscure envvar like $SLURM_CPU_BIND is probably fine. I find it somewhat surprising that there isn't a more obvious way for a code to inspect whether it has been invoked with srun or not.

[dmargala]

perhaps a PMI* env variable of some sort would help? slightly complicated by various PMI libraries but may be flexible enough to work with slurm at NERSC and for non-slurm+NERSC MPI invocations.

conda-forge mpich (uses pmi)

~> conda create -p /tmp/test -c conda-forge python mpi4py
...
(test) ~> conda list | grep mpi
mpi                       1.0.1                     mpich    conda-forge
mpi4py                    4.0.1           py313h7246b6a_0    conda-forge
mpich                     4.2.3              h239ebd3_102    conda-forge
(test) ~> mpirun -np 1 env | grep '^PMI.*RANK'
PMI_RANK=0

conda-forge openmpi (uses pmix)

~> conda create -p /tmp/test2 -c conda-forge python mpi4py openmpi
...
(test2) ~> conda list | grep mpi
mpi                       1.0                     openmpi    conda-forge
mpi4py                    4.0.1           py313hffb1895_0    conda-forge
openmpi                   5.0.6              hd45feaf_100    conda-forge
(test2) ~> mpirun -np 1 env | grep '^PMI.*RANK'
PMIX_RANK=0

NERSC slurm (uses cray-pmi):

dmargala@nid001033:~> env | grep 'SLURM_NPROCS'
SLURM_NPROCS=1
dmargala@nid001033:~> env | grep '^PMI.*RANK'
dmargala@nid001033:~> srun env | grep '^PMI.*RANK'
PMI_LOCAL_RANK=0
PMI_RANK=0

there may also be a PMI2_ variant to consider (and potentially others) ?

[craigwarner-ufastro]

@sbailey I tested your

salloc --nodes 1 -n 128 --qos interactive --time 4:00:00 --constraint gpu --gpus-per-node=4 --account desi_g

And yes, you're correct that this sets SLURM_NPROCS and SLURM_NTASKS however here is a list of several candidates that are still not set with this when I call ./wrap_rrdesi but are set when I call srun ... ./wrap_rrdesi:

SLURM_CPUS_PER_TASK 2
SLURM_CPU_BIND quiet,mask_cpu:0x00000000000000010000000000000001,0x00000000000000020000000000000002,0x00000000000000040000000000000004,0x00000000000000080000000000000008
SLURM_CPU_BIND_LIST 0x00000000000000010000000000000001,0x00000000000000020000000000000002,0x00000000000000040000000000000004,0x00000000000000080000000000000008
SLURM_CPU_BIND_TYPE mask_cpu:
SLURM_CPU_BIND_VERBOSE quiet
SLURM_DISTRIBUTION block
SLURM_STEP_RESV_PORTS 63938-63942
SLURM_TRES_BIND gres/gpu:map_gpu:3,2,1,0
SLURM_TRES_PER_TASK cpu:2
SLURM_UMASK 0007

It looks to me like SLURM_CPU_BIND would be a good candidate. Maybe even better is SLURM_STEP_RESV_PORTS which I assume is a list of reserved ports for MPI inter-process communication? What is SLURM_DISTRIBUTION block and SLURM_UMASK 0007

@dmargala what do you think?

[craigwarner-ufastro]

@sbailey @dmargala I think I have everything cleared up. The root bug was that if the number of GPUs available was > the rank of the communicator, it was making a bad assumption that you wanted to use at least ngpu ranks. So when calling wrap_rrdesi directly without srun, the length of the communicator was obviously 1 but there were 4 GPUs in the node so it was splitting the input files and rank 0 was only taking 1/4 of them but there were no other ranks to run anything.

I fixed this, added informative warning messages where appropriate, and cleaned up the login node logic that had been copy/pasted from elsewhere. Here are a bunch of example test cases:

• Run on login node

cdwarner@perlmutter:login16:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
wrap_rrdesi should not be run on a login node.

The following were all run after getting an interactive node with

 salloc -N 1 -C gpu -q interactive -t 3:00:00 -A desi_g --gpus-per-node=4

• Run directly - now works with warnings:

cdwarner@nid001173:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
WARNING: Detected that wrap_rrdesi is not being run with srun command.
WARNING: Calling directly can lead to under-utilizing resources.
Recommended syntax: srun -N nodes -n tasks -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi [options]
	Ex: 8 tasks each with GPU support on 2 nodes:
		srun -N 2 -n 8 -c 2 --gpu-bind=map_gpu:3,2,1,0  wrap_rrdesi ...
	Ex: 64 tasks on 1 node and 4 GPUs - this will run on both GPU and non-GPU nodes at once:
		srun -N 1 -n 64 -c 2 --gpu-bind=map_gpu:3,2,1,0  wrap_rrdesi ...
WARNING: wrap_rrdesi was called with 4 GPUs but only 1 MPI ranks.
WARNING: Will only use 1 GPUs.
Running 18 input files on 1 GPUs and 1 total procs...

• Run with srun and n < ngpu:

cdwarner@nid001173:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 1 -n 2 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
WARNING: wrap_rrdesi was called with 4 GPUs but only 2 MPI ranks.
WARNING: Will only use 2 GPUs.
Running 18 input files on 2 GPUs and 2 total procs...

• As expected run:

cdwarner@nid001173:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 1 -n 4 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
Running 18 input files on 4 GPUs and 4 total procs...

• Run with GPU + CPU:

cdwarner@nid001173:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 1 -n 64 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
Running 18 input files on 4 GPUs and 6 total procs...

• Run with -n 64 but --gpuonly

cdwarner@nid001133:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 1 -n 64 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite --gpuonly
Running 18 input files on 4 GPUs and 4 total procs...

• Run with too many nodes requested (handled by srun):

cdwarner@nid001173:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 2 -n 8 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
srun: error: Only allocated 1 nodes asked for 2

The following were all run after getting an interactive 2 nodes with

salloc --nodes 2 --qos interactive --time 4:00:00 --constraint gpu --gpus-per-node=4 --account desi_g

• Run as expected

cdwarner@nid001048:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 2 -n 8 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
Running 18 input files on 8 GPUs and 8 total procs...

• Run with too few n

cdwarner@nid001048:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 2 -n 6 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
WARNING: wrap_rrdesi was called with 8 GPUs but only 6 MPI ranks.
WARNING: Will only use 6 GPUs.
Running 18 input files on 6 GPUs and 6 total procs...

The following were run with an interactive node obtained with -n argument:

salloc --nodes 1 -n 128 --qos interactive --time 4:00:00 --constraint gpu --gpus-per-node=4 --account desi_g

• Run directly

cdwarner@nid001133:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
WARNING: Detected that wrap_rrdesi is not being run with srun command.
WARNING: Calling directly can lead to under-utilizing resources.
Recommended syntax: srun -N nodes -n tasks -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi [options]
	Ex: 8 tasks each with GPU support on 2 nodes:
		srun -N 2 -n 8 -c 2 --gpu-bind=map_gpu:3,2,1,0  wrap_rrdesi ...
	Ex: 64 tasks on 1 node and 4 GPUs - this will run on both GPU and non-GPU nodes at once:
		srun -N 1 -n 64 -c 2 --gpu-bind=map_gpu:3,2,1,0  wrap_rrdesi ...
WARNING: wrap_rrdesi was called with 4 GPUs but only 1 MPI ranks.
WARNING: Will only use 1 GPUs.
Running 18 input files on 1 GPUs and 1 total procs...

• Run as expected

cdwarner@nid001133:/global/cfs/cdirs/desi/users/cdwarner/code/desispec/bin> srun -N 1 -n 4 -c 2 --gpu-bind=map_gpu:3,2,1,0  ./wrap_rrdesi -i $MYRRDIR/list_coadds.ascii -o $SCRATCH/wrap/ --gpu --overwrite
Running 18 input files on 4 GPUs and 4 total procs...

Finally, if MPI is not available:

    try:
        import mpi4py.MPI as MPI
    except ImportError:
        have_mpi = False
        print ("MPI not available - required to run wrap_rrdesi")
        sys.exit(0)

Look good? Still want me to rename to wrap_rrdesi_mpi?

[sbailey]

@craigwarner-ufastro those look good! Thanks for handling all the cases including the "recommended syntax" message, which is great for telling users what they should do instead of just that they are doing it wrong. Please go ahead with a pull request. OK to leave as "wrap_rrdesi" without the "_mpi" suffix.

[araichoor]

I ve not followed all the technical details, but thanks for digging into that, and glad to hear that issue led up to some improvements.