CASA v5 script to test memory leak in `ia.getregion`

test_getregion_leak.py

#!/usr/bin/env python
"""
Execute tests to reproduce memory leaks from `ia.getregion`. See the functions
under the "Tests" banner near the bottom of the module for the tests to run.
Images with diameters in pixels greater than or equal to 3500 appear to cause
memory leaks. By default 50 calls are made to `ia.getregion`. The call itself
selects a bounding box encompassing the full image and selects one channel (the
second channel).

This file is meant to be run using `execfile` under CASA v5 and Python v2.7.
"""

from __future__ import (print_function, division)

import os
import gc
import resource


# The default number of trials to run in each of the tests.
N_TRIALS = 50


class ImageWrapper(object):
    # To speed up reading/writing the cube data in the execution of the test,
    # write the temporary file to shared memory on /dev/shm.
    base_directory = '/dev/shm'

    def __init__(self, name):
        """
        Wrap a test image cube.

        Parameters
        ----------
        imagename : str
        shape : Iterable(int)
        """
        path = os.path.join(self.base_directory, name)
        ia.open(path)
        shape = ia.shape()
        ia.done()
        ia.close()
        assert len(shape) == 4
        self.name = name
        self.path = path
        self.shape = shape

    @classmethod
    def create_from_shape(cls, name='test.image', shape=(3500,3500,1,5),
            overwrite=True):
        """
        Create a zero-filled test image from the given shape.

        Parameters
        ----------
        name : str
        shape : Iterable(int)
            Image shape in [RA, Dec, Stokes, Freq].
        overwrite : bool
        """
        # Create an empty CASA image file with the given shape
        path = os.path.join(cls.base_directory, name)
        success = ia.fromshape(path, list(shape), overwrite=overwrite)
        ia.done()
        ia.close()
        return cls(name)

    def remove(self):
        rmtables(self.path)

    def get_region_for_plane(self, chan_ix=1):
        """
        Parameters
        ----------
        chan_ix : int
            Channel index number. Note that unintended behavior occurs when
            selecting channel index "0" and thus for the purposes of this
            script it is not allowed. Instead of the plane for the first
            channel, the whole image cube is selected.
        """
        assert chan_ix != 0
        region_fmt = 'box[[0pix,0pix],[{0}pix,{1}pix]], range=[{2}chan,{2}chan]'
        npix0, npix1, nstokes, nchan = self.shape
        return region_fmt.format(npix0-1, npix1-1, chan_ix)


class MemoryChecker(object):
    def __init__(self):
        """
        A simple context manager that prints the resident memory of the parent
        CASA Python process when entering and exiting scope.
        """
        pass

    def __enter__(self):
        self.print_mem()
        return self

    def __exit__(self, *exec_details):
        self.print_mem()

    @property
    def mem_in_kiB(self):
        # Trigger a call to the garbage collector to ensure that all latent
        # Python variables are evicted from memory.
        gc.collect()
        # Maximum resident set size used by the parent process in kiB. For
        # further information see the man page for the Linux command line tool
        # `getrusage`.
        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss

    @property
    def mem_in_GiB(self):
        return self.mem_in_kiB / 1024**2

    def print_mem(self):
        print('-- Process resident memory: {0:.6f} GiB'.format(self.mem_in_GiB))


def run_repeated_getregion_calls(img, n_trials=N_TRIALS):
    """
    Repeatedly read a one channel plane from the test image using
    the global Image Analysis tool instance.

    Parameters
    ----------
    img : ImageWrapper
    n_trials : int
    """
    for ii in range(n_trials):
        ia.open(img.path)
        region = img.get_region_for_plane(chan_ix=1)
        data = ia.getregion(region)
        ia.done()
        ia.close()


##############################################################################
#                                  Tests
##############################################################################

def run_test(imsize=3500, nchan=5, n_trials=N_TRIALS):
    shape = [imsize, imsize, 1, nchan]
    img = ImageWrapper.create_from_shape(
            name='test.image', shape=shape, overwrite=True,
    )
    with MemoryChecker() as _:
        run_repeated_getregion_calls(img, n_trials=n_trials)
    img.remove()


def run_test_small():
    """Okay!"""
    run_test(imsize=350)


def run_test_large():
    """Okay!"""
    run_test(imsize=3499)


def run_test_critical():
    """Not okay!"""
    run_test(imsize=3500)