Refactor FITS reader

[Kraysent]

It seems to me that there is memory leak, since graph of the memory of the big analysis looks strange.

For now, code looks like this:

def from_fits(filename: str) -> Iterator[Tuple[Particles, ScalarQuantity]]:
    hdul = fits.open(filename, memmap=True)

    for frame in range(len(hdul) - 1):
        table: BinTableHDU = hdul[frame + 1]
        number_of_particles = len(table.data[list(fields.keys())[0]])

        timestamp = table.header["TIME"] | units.Myr
        particles = Particles(number_of_particles)

        for (key, val) in fields.items():
            if val is not None:
                setattr(particles, key, table.data[key] | val)
            else:
                try:
                    data = np.array(table.data[key], dtype=np.float64)
                except KeyError:
                    continue
                setattr(particles, key, data)

        yield particles, timestamp

Want to try several changes and look how it affects memory consumption of the process. Results would be in this ticket.

[Kraysent]

Using psrecord library to look at the memory graph.

Testing things on the file of 30 snapshots with 1.5 million particles in each of them. Each snapshot itself takes about 42 MB of memory. Measurements are taken once a second. Original graph looks like this (blue line):

The odd thing about this graph is that there is no corellation with number of snapshots. There are 30 snapshots but there is no pattern with this number of iterations (for example, number of finest spikes is not 30).

[Kraysent]

Intend to elaborate these thoughts:

First of all, the hdul is never closed. I do not think that this affects memory consumption in any significant way but it is good to eliminate this first.

Second, memmap parameter seems to be True by default. Should try removing all of the parameters and use simple fits.open(filename) method. Also, there is lazy_load_hdus parameter. It is also True by default but it is worth trying it and see if anything changes. It would be more rational to set it to False and see if consumption of the memory increases.

Another point is that in the loop I count len(hdul) which is incredibly slow since it goes over the whole file at once. Not sure if there is any caching but I probably should iterate over the elements without counting of the length (since it is useless anyway).

In the latter case maybe should try using del after each iteration.

And as a last resort might try manually invoking garbage collector after each iteration. Maybe it is it that creates these large-scale spikes.

Need to keep in mind that I have no idea if this memory oddity is connected with the FITS reading. It might as well be connected with anything else, reading large files just seems to be the most likely candidate. At least I would clean up this code.

[Kraysent]

Added

    hdul.close()

to the end of the function. Seems to me that nothing changed. Leaving it that way, closing file is never redundant.

By the way, CPU graphs seems to be more consistent. It could be nothing though.

By the way x2, number of CPU graph spikes roughly equals number of snapshots. At least processor feels it.

[Kraysent]

Stopped using len() of the HDUTable. Same, nothing changed.

Code

def from_fits(filename: str) -> Iterator[Tuple[Particles, ScalarQuantity]]:
    hdul = fits.open(filename, memmap=True)

    first = True

    for table in hdul:
        if first:
            # skipping first HDU; it is required by the FITS specification.
            first = False
            continue

        number_of_particles = len(table.data[list(fields.keys())[0]])

        timestamp = table.header["TIME"] | units.Myr
        particles = Particles(number_of_particles)

        for (key, val) in fields.items():
            if val is not None:
                setattr(particles, key, table.data[key] | val)
            else:
                try:
                    data = np.array(table.data[key], dtype=np.float64)
                except KeyError:
                    continue
                setattr(particles, key, data)

        yield particles, timestamp

    hdul.close()

Memory graph:

Well, time of the modelling certainly decreases.

[Kraysent]

Changed

hdul = fits.open(filename, memmap=True)

to

hdul = fits.open(filename)

Still no effect:

[Kraysent]

hdul = fits.open(filename, lazy_load_hdus=False)

Still no effect. Will try to do same thing with True.

Time is still going down though.

[Kraysent]

hdul = fits.open(filename, lazy_load_hdus=True)

Maybe it is not because of FITS reading after all.

[Kraysent]

hdul = fits.open(filename, memmap=False)

Nope

Though overall memory consuption peak is lower than it was. Not sufficient though.

[Kraysent]

Added

        del particles
        del table

after the yield statement. Still nothin'.

Garbage collector time I guess.

[Kraysent]

Finally. Added this after each iteration:

gc.collect()

This seems to make the whole thing much better (without even compromising the time):

It is strange, though, that memory gap is about 1.5 GB while snapshot takes about 42 MB. Something is wrong.

This still does not mean that the problem was in the FITS reader, maybe this garbage collector just collects a lot of other things apart from FITS remnants.

Memory still goes upwards though.

[Kraysent]

Next step would be to inspect which task (if it is task) creates these memory leaks. Need to gradually decreas number of components in the simulation during tests.

But that is the story for another time (#148).

[Kraysent]

analysis1.txt analysis2.txt analysis3.txt analysis4.txt analysis5.txt analysis6.txt analysis7.txt analysis8.txt analysis9.txt