desihub / redrock

Redshift fitting for spectroperfectionism
BSD 3-Clause "New" or "Revised" License
21 stars 13 forks source link

point2point gather of results #310

Closed sbailey closed 1 month ago

sbailey commented 1 month ago

This PR fixes a memory/objsize problem in Jura. Healpix with an especially large number of targets were failing on the final gather:

results = targets.comm.gather(results, root=0)
...
    results = targets.comm.gather(results, root=0)
Proc 0:   File "mpi4py/MPI/Comm.pyx", line 1578, in mpi4py.MPI.Comm.gather
Proc 0:   File "mpi4py/MPI/msgpickle.pxi", line 772, in mpi4py.MPI.PyMPI_gather
Proc 0:   File "mpi4py/MPI/msgpickle.pxi", line 777, in mpi4py.MPI.PyMPI_gather
Proc 0:   File "mpi4py/MPI/msgpickle.pxi", line 190, in mpi4py.MPI.pickle_allocv
Proc 0:   File "mpi4py/MPI/msgpickle.pxi", line 181, in mpi4py.MPI.pickle_alloc
Proc 0: SystemError: Negative size passed to PyBytes_FromStringAndSize

This would occur even if the individual ranks had fairly small inputs (in particular below 2GB).

This PR works around this by doing multiple point-to-point send/recv transfers instead of a gather. Admittedly, this might just be buying us time as healpix continue to grow, but it works on special-other healpix 27239 with 12590 targets, the largest case in Jura that has made it to the Redrock step. If needed, the next step in bookkeeping would be to split the results dict into multiple smaller dicts and transfer them, but I haven't tackled that yet.

In addition to testing several very large healpix, I did a direct comparison of main-dark healpix 0 with main and this PR branch in /pscratch/sd/s/sjbailey/redrock/redrock-0-{main,pr}.fits, which passes fitsdiff (data are same, timestamps and dependency keywords diff). They also take the same amount of time so that multi send/recv isn't hurting performance.

coveralls commented 1 month ago

Coverage Status

coverage: 39.013% (-0.009%) from 39.022% when pulling 72f96a4f4c0de70e0a4dbcb2db3872e364abbe83 on mpi_big_mem into ca7838158ab7da085ae14b6da909a32bbdcf9dba on main.