Curious Out Of Memory during looping on Cosmos Catalogue

[jecampagne]

Hello,

I was trying to extract the galaxy&psf couples from the COSMOS_23.5_training_sample catalog
to make some datasets

Please find below the code.

The problem I have is that after having processed 28_000 galaxies (and not yet 29_000) the
python code crashes due to "Out of Memory" detected by the my local SLURM system

Is there a "purge" that I have missed somewhere?

    real_galaxy_catalog = galsim.RealGalaxyCatalog(
            dir=path_catalog, sample=sample
        )
    n_total = real_galaxy_catalog.nobjects  # - 56062
    sequence = np.arange(0, n_total) 
    np.random.shuffle(sequence)

    # Loop....
    t0 = time.time()
    for i,gal_idx in enumerate(sequence):
        if i==0 or i%1000==0:
            print(f"Process {i}/{gal_idx} at t:{time.time()-t0:.1f}sec")
        # read HST galaxy and PSF
        gal_ori = galsim.RealGalaxy(real_galaxy_catalog, index=gal_idx)
        psf_hst = real_galaxy_catalog.getPSF(gal_idx)

        # get the corresponding galsim.Image
        gal_ori_image = real_galaxy_catalog.getGalImage(gal_idx)
        psf_ori_image = real_galaxy_catalog.getPSFImage(gal_idx)

        #save into fits files
        fname =  out_catalog_dir +"/"+"gal_"+str(i)+".fits"
        gal_ori_image.write(fname)
        
        fname =  out_catalog_dir +"/"+"psf_"+str(i)+".fits"
        psf_ori_image.write(fname)

Thanks for your attention.
Jean_Eric

[rmjarvis]

I finally took some time to look into this. I think the memory usage is in the real_galaxy_catalog.loaded_files. We keep the various files open, since reading them from disk is usually relatively slow, so we try to keep them in memory when possible. On my system, the above script uses about 3GB total, and indeed the memory usage goes up steadily over time at first, then plateaus.

If it's too much on your system to have all the input files open at once, there is a (not well documented) feature you can use to close all the loaded files. real_galaxy_catalog.close() will clear out these saved open files, so you can reclaim their memory. Something like:

import psutil
process = psutil.Process()

[...]
    # Somewhere in the loop where you are concerned about memory getting too high:
    if process.memory_info().rss > 2 * 1024**3:
        real_galaxy_catalog.close()

This would clear them out whenever the memory usage exceeded 2 GB. (Total, not just from these files, but maybe that would be sufficient.) Of course, a lot of the files will need to be reopened right away, especially with this random ordering you're using in this script. You're bouncing around among all the individual files, but it still manages to reclaim some of the memory that astropy.io.fits uses for its memory mapping, so it does help some.

There is also a cache for the noise file, but I think that's pretty negligible. If you want to also clear that, you can do:

        real_galaxy_catalog.saved_noise_im.clear()

[rmjarvis]

Better documentation about this added in #1340.