large `total_memory` when writing binary file leads to memory error if preprocessing step converts to `float32`

[JoeZiminski]

I am trying to choose chunk size to use a large percentage of available system memory*. The determination of chunk_size in job_tools/ensure_chunk_size is based on recording.dtype. However in some preprocessing steps (e.g. phase_shift, I am unsure of others) the data is temporarily converted to float32. This then leads to a memory error.

The only fix I can think of is to check the preprocessing steps on the recording and, if any of the steps works in float32, then that should be used to compute chunk_size. Maybe there is another workaround?

*this is because, as I understand it, larger chunk size will be better to reduce batch edge effects. Please let me know if there are other issues I have not considered.

[zm711]

I'm linking a couple semi-related issues since they also have to do with a lack of guard rails currently in spikeinterface for running certain operations.

#1063. In #1063 the issue is that without sufficient storage on disk spikeinterface tries the operation rather than warns the user.

And more recently this has also come up in #2026 about potentially needing an spikeinterface kwarg optimizer.

[h-mayorquin]

Thanks for collecting those issues @zm711.

@JoeZiminski
Thanks for brining this issue. This has made me change my mind about something in #914

I think your fix looks right to me:

The only fix I can think of is to check the preprocessing steps on the recording and, if any of the steps works in float32, then that should be used to compute chunk_size. Maybe there is another workaround?

When operating with BasePreprocessor the dtype that ensure_chunk_size uses should be the "largest" within all the parent recordings. This will give you the upper limit as that is what we care about.

On the topic of an accurate estimation of your memory footprint. Other issues that you might have not considered is that some preprocessing steps have a larger memory footprint that you might think from their dtype and their size (end_frame - start_frame) alone. When I did the GeneratorRecording (now NoiseGeneratorRecording) I put a lot of effort avoiding any extra memory allocation (calls to system malloc). It is kind of a craft to avoid them in numpy and I am not sure we have been very careful avoid that.

As a general rule assume that each get_traces will require 2-3 times the memory thaty ou might expect. The best thing for you to do would be to actually form your preocessing step, estimate how much memory it should consume (multiple size by dtype) and then use a memory profiler to see if that matches:

https://github.com/bloomberg/memray

Hopefully we start optimizing memory allocation for all the preprocessing steps. This should be a long term goal of the library and we should add it as a point in the next funding round.

[JoeZiminski]

Thanks @h-mayorquin apologies for the delay reply, that's very good to know, I'll also checkout some of the techniques used in NoiseGeneratorRecording!

In leiu of full memory optimisation for job_kwargs (which as you mention is more of a long term goal), as a hacky workaround, for now I am playing around with something like:

def get_preprocessing_steps_largest_dtype_itemsize(recording):
    all_itemsizes = []
    current_recording = recording
    
    while True:
        all_itemsizes.append(current_recording.get_dtype().itemsize)
        if not hasattr(current_recording, "_parent_recording"):
            break
        current_recording = current_recording._parent_recording

    return np.max(all_itemsizes)

Unfortunately this does not work as in phase_shift for example the dtype is stored in a temporary tmp_dtype variable. And so is not accessible. One option is to force any temp dtypes used in preprocessing to be stored as a tmp_dtype attribute on the recording. I see across preprocessing the dtype is often changed within the method, so exposing this on an attribute makes sense.

However, for this to work properly, I think it makes sense to plan it under the more general optimising kwargs umbrella #2232 and accept for now ensure_chunk_size is not working optimally. @h-mayorquin @zm711 what do you think?

[zm711]

Some of these rough edges have come up in other PRs as well #2189 (with dtype in preprocessing). It would be nice to have a completely regular set of naming conventions, but since some of the code was generated before the conventions there are some big style differences. So I agree that currently there are some functions around the kwargs in general that are sub-optimal, but hopefully something like #2232 can smooth out some of these issues.

For me I run into the memory map leak issue where at the start of the session my RAM is in good shape, but at the end some of these operations fail because I run out of memory. But these are more long-term issues. Thanks for bringing this back to my attention at least!

[JoeZiminski]

Thanks @zm711 I will close this issue and move the key points to #2232