spacetelescope / jwst

Python library for science observations from the James Webb Space Telescope
https://jwst-pipeline.readthedocs.io/en/latest/
Other
558 stars 164 forks source link

Running `resample` with `in_memory=False` reads all models into memory #8480

Closed stscijgbot-jp closed 1 week ago

stscijgbot-jp commented 4 months ago

Issue JP-3621 was created on JIRA by Brett Graham:

ResampleStep when run with in_memory=False reads all models into memory. For a 33 member association the memory usage is shown in the attached graph.

One line that reads all models is: https://github.com/spacetelescope/jwst/blob/8b254ae2b69fed42c413f4ff7163ff2f983dfa59/jwst/resample/resample_step.py#L67

stscijgbot-jp commented 1 month ago

Comment by Ned Molter on JIRA:

This should be fixed as part of implementing ModelLibrary (see JP-3690).

Here is a small summary of results from profiling memory for this step on its own in my PR branch, using as input an association containing a 46-cal-file subset of the data in 

███████████████████████████████████████████ 

The total size of the input _cal files on disk is roughly 5 GiB.

Setting in_memory=True, the peak memory usage is 15.4 GiB, and I'm attaching a graph of usage over time to the ticket (resample_in_memory.png, sorry for the similar name to what Brett attached originally).

Setting in=memory=False, the peak memory usage is 10.6 GiB, and a graph is again attached (resample_on_disk.png).

This appears to be fixed.  Analysis of the memory allocation breakdown from the flamegraph indicates that all the models are never loaded into memory.  However, the most memory-intensive part of the step is resample_variance_arrays(), which requires making several large output arrays for the different types of variance that need to be resampled.  In this example, each of those arrays has shape (9348, 13432).  It is unclear to me how to improve this, and I'd say it's beyond the scope of this ticket.