Loading pipelines containing many nested sorting extractors are non-performant

[DradeAW]

Hi,

In my pipeline (Lussac), I have very complex sorting objects (made with like tens of UnitsSelectionSorting, tens of UnitsAggregationSorting, and multiple other curation functions.
Handling those objects is fine, fast, and memory-efficient.
Saving such objects to pickle format or JSON format is also fast and memory-efficient (although the json file is tens of thousands of lines long and can weigh more than 100 MB ^^).

The problem comes when loading.
If the file was saved with recursive=False, then the loading is fast and efficient.
But if the file was saved with recursive=True, then the loading can take minutes and more than 10 GB of RAM, whereas creating the object takes less than a second and almost no RAM.

In my case, I need to use recursive=True because I am using relative_to when saving, which can only propagate if loaded recursively.

[h-mayorquin]

Thanks for writing this @DradeAW. For anyone else reading this, @DradeAW is using PhySortingExtractor already profiled this with memray and found that the following lines are making way too many memory alocations:

https://github.com/SpikeInterface/spikeinterface/blob/e261fccfc9055d55d2cac0072dea485af6cbd54e/src/spikeinterface/extractors/phykilosortextractors.py

Something else to add to the discussion. This is in the context of serializing and unserializing recording objects with deeply nested structures (like preprocessing and curation). Me and @alejoe91 previously profiled the pickling process in this context and an important finding from this time is that when recursive=False pickling specifically is not able to identify that data in differents part of the nested structure is the same so it is loaded multiple times.

This is relevant because as @DradeAW mentioned above the problem dissappears when profiling with recursive=False.

So what to do?
It seems highly desirable that pickling and unpickling deeply nested extractor objects in a performant way is desirable. We all probably agree on this. How should we do this?
I think the easiest way to achieve this is to let pickle do its work. We keep recursive=True as an option because we wanted to support the option of writing relative paths and that is a way of doing it. We should decouple the ability of changing all the file paths to relative from the ability of recursively expanding dicts. My hunch is that this is the less complicated way of achieving the goal of having performant pipelines in highly nested pipelines. I already discussed a bit with @alejoe91 about this and he has good ideas on how to achieve this.

Am I failing to see other possibilities, needs or use cases?
This issue was meant to discuss this.

[samuelgarcia]

Salut Aurelien,
Do you save to json or pickle ?
Could you make a simple script that generate and apply a long processing chain and then load then with simple measurement with time.perf_counter() at loading time ? With json/pickle and /relative/absolut ?

[DradeAW]

I'll try to make a script this morning to try to reproduce the problem so that anyone can reproduce the problem more easily. Might take me a couple of hours but that's time I'm not spending writing so I like it 😅

[DradeAW]

Here is something I hacked really quickly but that shows that (in my opinion) there is definitely a problem:

from pathlib import Path
import time
import numpy as np
import psutil
import spikeinterface.core as si
import spikeinterface.curation as scur
import spikeinterface.postprocessing as spost


# Constants

SF = 30_000			# in Hz.
MAX_TIME = 30 * 60	# in s.
FOLDER = Path(".")
PROCESS = psutil.Process()


def generate_spike_train() -> np.ndarray:
	firing_rate = np.random.uniform(low=1.0, high=140.0)
	num_spikes = int(round(firing_rate * MAX_TIME))
	
	return np.sort(np.random.randint(low=100, high=MAX_TIME * SF - 100, size=num_spikes))


# Generating the sorting object

mem_before = PROCESS.memory_info().rss / 1024**2
spike_trains = {i: generate_spike_train() for i in range(150)}
mem_after = PROCESS.memory_info().rss / 1024**2

sorting = si.NumpySorting.from_dict(spike_trains, SF)
si.NpzSortingExtractor.write_sorting(sorting, FOLDER / "tmp" / "sorting.npz")
sorting = si.NpzSortingExtractor(FOLDER / "tmp" / "sorting.npz")


# Create complex pipeline

t0 = time.perf_counter()

firing_rates = np.array([len(sorting.get_unit_spike_train(unit_id)) / MAX_TIME for unit_id in sorting.unit_ids])
sorting.set_property("category", ["low_fr" if fr < 5.0 else "high_fr" for fr in firing_rates])

sorting = sorting.select_units([unit_id for unit_id in sorting.unit_ids if unit_id not in (3, 50, 103)])

sorting_low  = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "low_fr"])
sorting_high = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "high_fr"])
sorting_high = sorting_high.select_units([unit_id for unit_id in sorting_high.unit_ids if unit_id not in sorting_high.unit_ids[:4]])
sorting = si.UnitsAggregationSorting([sorting_low, sorting_high])

sorting_low  = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "low_fr"])
sorting_high = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "high_fr"])
sorting_low  = spost.align_sorting(sorting_low,  {unit_id: np.random.randint(-4, 5) for unit_id in sorting_low.unit_ids})
sorting_high = spost.align_sorting(sorting_high, {unit_id: np.random.randint(-2, 3) for unit_id in sorting_high.unit_ids})
sorting = si.UnitsAggregationSorting([sorting_low, sorting_high])

sorting_low  = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "low_fr"])
sorting_high = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "high_fr"])
sorting_low  = scur.remove_duplicated_spikes(sorting_low,  censored_period_ms=0.3)
sorting_high = scur.remove_duplicated_spikes(sorting_high, censored_period_ms=0.1)
sorting = si.UnitsAggregationSorting([sorting_low, sorting_high])

sorting_low  = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "low_fr"])
sorting_high = sorting.select_units([unit_id for unit_id in sorting.unit_ids if sorting.get_unit_property(unit_id, "category") == "high_fr"])
sorting_high = sorting_high.select_units([unit_id for unit_id in sorting_high.unit_ids if unit_id not in sorting_high.unit_ids[:4]])
sorting = si.UnitsAggregationSorting([sorting_low, sorting_high])


# Benchmark save / load

t1 = time.perf_counter()
sorting.dump_to_pickle(FOLDER / "tmp" / "recursiveFalse.pkl", include_properties=True, recursive=False)
t2 = time.perf_counter()
sorting.dump_to_pickle(FOLDER / "tmp" / "recursiveTrue.pkl",  include_properties=True, recursive=True, relative_to=FOLDER)
t3 = time.perf_counter()
mem1 = PROCESS.memory_info().rss / 1024**2
aze = si.load_extractor(FOLDER / "tmp" / "recursiveFalse.pkl")
mem2 = PROCESS.memory_info().rss / 1024**2
t4 = time.perf_counter()
mem3 = PROCESS.memory_info().rss / 1024**2
azer = si.load_extractor(FOLDER / "tmp" / "recursiveTrue.pkl", base_folder=FOLDER)
mem4 = PROCESS.memory_info().rss / 1024**2
t5 = time.perf_counter()

print(f"Time to generate the post-processing: {t1-t0:.1f} s")
print(f"Spike trains take up {mem_after-mem_before:.1f} MiB\n")
print("SAVING:")
print(f"\t- recursive=False: {t2-t1:.3f} s")
print(f"\t- recursive=True:  {t3-t2:.3f} s")
print("LOADING:")
print(f"\t- recursive=False: {t4-t3:.2f} s (used {mem2-mem1:.1f} MiB)")
print(f"\t- recursive=True:  {t5-t4:.2f} s (used {mem4-mem3:.1f} MiB)")

And here is the output (done several times, it is reproducible):

Time to generate the post-processing: 8.8 s
Spike trains take up 156.2 MiB

SAVING:
	- recursive=False: 0.007 s
	- recursive=True:  0.011 s
LOADING:
	- recursive=False: 4.65 s (used 293.8 MiB)
	- recursive=True:  14.81 s (used 4700.4 MiB)

As you can see, loading with recursive=False is faster than generating the post-processing, and the memory usage seems correct given the memory usage of the spikes.
However, with recursive=True, it is slower than generating it, and it takes more than 4.55 GiB of RAM!!!
I think there is clearly an underlying issue, I mean once the sorting object has been loaded the recursive stuff should be gone, and the memory impact shouldn't be this!

[DradeAW]

I wanted to add JSON to this profiling, but surprisingly loading from JSON crashes :/

In the mean time, I will document some of the discussion I had with @h-mayorquin:
To him, what I showed above is expected behaviour given how recursive=True is coded, because it loses references to other objects. I think this is clearly a problem that users are not aware of when using the save function (who would expect the save to be bigger than the actual object?).

Getting away from from recursive=True/False, I propose the following plan of moving forward:

• Having the opportunity to save a SpikeInterface object to pickle, which never calls any to_dict. You just save the object as is and let pickle to its work. This is fast and efficient, but at the cost of flexibility (you can't load half the object of modify it in place)
• Having the opportunity to save a dict object (containing nothing related to SpikeInterface objects) using the pickle format. This is a bit similar to what saving to JSON does, except that in this case we can keep the numpy arrays, complex numbers and so on ... This way you keep all the flexibility because you can load the dictionary, change it if needed, and then load the SI object from_dict (i.e. re-creating the whole pipeline)

At first I thought that this was what recursive=True/False did, but chatting with @h-mayorquin, it seems that this is not what's happening.

Tell me what you think about this @alejoe91 @samuelgarcia, but I think both options have pros and cons and that both should exist so that the user can choose what's right for him.

[h-mayorquin]

OK, so after our discussion this is how I remember things:
For pickling, we want to restrict and orient pickling to serializing quikcly for parallel processing. So unless we can come up with something clever, right now we will not be supporting the use of relative paths with pickling. This was done here: #1831. Storing with relative paths is now only relegated to serializing with json and, at the moment, is not very performant. That said, @alejoe91 mentioned some ideas on how that could be done and we will probably be working long-term for on this. This is something of a long-term goal in the project as well #1235

@DradeAW are you OK with closing this or am I missing something?

Something that would be good to sort out though is the bug that you encountered when dumping to json. I think this is better done in a dedicated issue though.

[DradeAW]

I am ok with closing this, and waiting for the long-term solution :)

I talked with Alessio about this problem: it's because the unit_id is the key in a dictionary, and in json the key can only be a string. So there is a mismatch, which can easily be fixed but not in a very elegant manner :/