memory leak in calcPSF (possibly a POPPY issue)

[ariedel]

Some background: Pandeia uses WebbPSF to precompute a library of monochromatic PSFs which we then use to generate our 3D scenes. To do this, we create and configure an instrument object, and then iterate through a list of wavelengths to produce the suite of PSFs.

Recent experimentation (and then follow-up with memory_profiler) demonstrates that running instrument.calcPSF() generates a net gain of ~250 MB per PSF (except for rare occasions where it shrinks memory usage by 60 MB). This is something we didn't notice before because we were only generating a small number of PSFs, one at a time; now I'm trying to run large numbers of PSFs in parallel.

Questions: Is there a way to close a PSF calculation after I'm done with it? Is there a more efficient way to have WebbPSF generate a list of PSFs?

[mperrin]

I was actually thinking about this a little while ago. I suspect it's almost certainly related to the opening of the FITS files for the OPDs. The current code is a bit lax about closing those file handles and was relying on the Python garbage collector to close those after a calculation. I think for whatever reason the garbage collector is not actually closing those efficiently (but it does sometimes, hence those occasions when the memory usage shrinks?)

It ought to be straightforward to profile this and figure out what's being kept open and where, and then update the code to not do that. @robelgeda do you think you might have time to take a look at this alongside the WFI updates?

[robelgeda]

I can look into this. I have a tall task queue right now so it may take a while before I get to it.

[JarronL]

When I saw this today, I became a little concerned for my own application (although, I've never seen memory blow-ups on my machine in the past), so decided to look into it. However, I am not seeing this behavior in either Python 2.7 (webbpsf 0.5.0) or 3.5 (webbpsf 0.6.0). For instance, the following code gives me no change in memory usage after the initial run:

import webbpsf
from memory_profiler import memory_usage
from time import sleep

nc = webbpsf.NIRCam()
nc.options['output_mode'] = 'oversampled'
nc.options['parity'] = 'odd'

mem_vals = []
for w in np.arange(1,2,0.1):
    print(w)
    sleep(0.1)
    kw = {'outfile':None, 'oversample':4, 'fov_pixels':1024, 'monochromatic':w*1e-6}
    mem_max, hdul = memory_usage((nc.calc_psf,{},kw), max_usage=True, interval=0.01, retval=True)
    mem_vals.append(mem_max[0])

mem_vals = np.array(mem_vals)
diff = mem_vals[1:] - mem_vals[:-1]
print(diff)

## Output:
# [  1.28621094e+02   4.60937500e-01   2.73437500e-01   3.90625000e-01
#   4.68750000e-01   1.32812500e-01   3.63281250e-01   1.13281250e-01
#   1.95312500e-02]

Final memory usage was 283MB after 10 monochromatic runs on the same instrument configuration. I performed a similar loop with 100 wavelengths and another one without memory_profiler (just watching activity monitor) and still did not see your described behavior.

Do you have a snippet of code that would allow me to reproduce your results?

[ariedel]

I'll try creating a test case...

[ariedel]

Try the file in this gist: https://gist.github.com/ariedel/b478034bf2f96ce3bc8b2f21e6932c9d

It's already set up for use with memory_profiler

[JarronL]

Just had a look at your code. Part of the problem here is that you're creating a deepcopy of the webbpsf instrument instance for every single wavelength. There's no real need to do this as far as I can tell. A simple solution is to instead have the outer for loop within create_psf() iterate over (r,theta), while the inner loop iterates over wavelength. For each psf dictionary you create, you can use the same self.instrument reference without the need for a deepcopy. Essentially, you're using the same instrument instance to generate each successive PSF. This works for sending to multiple cores as well.

My slight rewrite of the create_psf function is below. In this case, I've opted to generate a new list of psfs for each (r,theta) and run _split_multi within the outer for loop.

    def create_psfs(self,n_cores):
        """
        Create PSFs by setting up all the names, wavelengths, and offsets, then passing them
        to the system that splits them off into parallel generation.

        Parameters
        ----------
        n_cores: int
            The number of parallel generation instances to run
        """

        wavelengths = self._get_wave_range_old()
        try:
            self.aperture_name # the only aperture that should already have this is nirspec's shutter-s200a1-s200a2-etc
        except AttributeError:
            self.aperture_name = self.aperture

        if not hasattr(self,'coronagraph_source_offsets'):
            self.coronagraph_source_offsets = [(0., 0.)]

        for r,theta in self.coronagraph_source_offsets:
            self.instrument.options[u'source_offset_r'] = r
            self.instrument.options[u'source_offset_theta'] = theta

            psfs = []
            for wave in wavelengths:
                if r > 0.001:
                    # Basename + aperture + wavelength + (if non-trivial) distance from on-axis and
                    # angle from vertical as defined in WebbPSF
                    longname = '{0:}_{1:}_{2:.4f}_{3:.3f}_{4:.0f}.fits'.format(self.insname,self.aperture_name,wave,r,theta)
                else:
                    longname = '{0:}_{1:}_{2:.4f}.fits'.format(self.insname,self.aperture_name,wave)

                psf = {'psf':self.instrument, 'wave':wave*1e-6, 'filename':longname, 'optimal_offset_r': None, 'optimal_offset_theta': None}
                psfs.append(psf)
            self._split_multi(n_cores,psfs)

With this function, each core held stable at around <1.5GB of memory throughout the entire process. Every so often, I would get a spike of 3GB or so, but it didn't persist.

[ariedel]

Thanks, that's fixed the memory leak.