[samuelgarcia]

See #4831

Totally untested just a draft proof of concept.

@chrishalcrow
@h-mayorquin

[alejoe91]

@samuelgarcia @h-mayorquin the last commit adds some core tests and sets num_chunk_per_job as the number of chunks in a shard, so that each job is going to deal with full shards. I can benchmark memory usage with this approach with respect to the previous "full-shard" one and report back here

[alejoe91]

Hi guys, this is ready to review on my end @samuelgarcia @chrishalcrow @h-mayorquin

[alejoe91]

@samuelgarcia @h-mayorquin the last commit adds some core tests and sets num_chunk_per_job as the number of chunks in a shard, so that each job is going to deal with full shards. I can benchmark memory usage with this approach with respect to the previous "full-shard" one and report back here

Here are some numbers and RAM savings with this PR:

I think we can merge this! @h-mayorquin do you want to take a look?

[h-mayorquin]

I started reading it today but could not finish my review, I will finish it tomorrow.

[chrishalcrow]

Since this is so zarr specific, is it better to give it a more specific name for the use case, like
num_chunk_per_shard?

Also, my feeling is that this is mostly for internal use and we don't expect users to pass it? I.e. if you're using sharding, we compute num_chunk_per_batch for you. Maybe this should be private or documented as an "internal job kwargs" or something

[alejoe91]

I don't think it's zarr specific. There might be other instances where you might want jobs to loop through chunks (I still have to think of one though hehehe)

[chrishalcrow]

Ok, so a user can set num_chunk_per_batch and this does affect saving binary files etc. But for non-zarr stuff, it shouldn't change memory use?

[h-mayorquin]

I think this is zarr specific. In #4905 I argue that what we need is to be able to separate the chunk size (in time) of the different steps of a processing chain. But I do think that the way it is solved here is very zarr specific because each call still gets a single chunk. What a batch adds is that one worker owns a range of consecutive chunks and that only matters when several chunks share one storage object. Right now only zarr shards do that.

I left a general comment with a way to get that ownership inside the zarr writer without a new job kwarg.

[h-mayorquin]

I think it is important for us to decouple the chunk processing from the size of the zarr writing. Right now main forces a single worker per shard (adjust_chunks_shards_and_job_kwargs sets chunk_size to the shard length). That is fine and I think it should remain a complexity confined to zarr. But I am not so sure about the approach here, I think there are a couple of problems with it. First, as Chris mentioned, it adds a job keyword for every user of TimeSeriesChunkExecutor to solve a problem that only exists for sharded zarr. Second, there are two forces here, and this PR only serves one.

I think the core of the problem here is two things that pull in different directions. Memory wants small pieces, because computing and holding the whole shard at once is what blows up memory. In zarr writing we want large pieces, because every partial write to a shard makes zarr-python read the whole shard and write it again. With this approach and our defaults of 1 second per chunk this is too slow. In my tests with 60 s shards that made the save about 2 times slower for the phase shift chain (run together with #4904) and 13 to 15 times slower for raw data. Right now one chunk_size is shared by every step, the writer included, so we have to pick a side: main pays in memory and this PR pays in write time. If we could assign one chunk size to preprocessing and another to writing, the trade-off would dissolve, and I made the more general argument on why this is a good way to move forward here: #4905.

I don't think that we should wait for #4905 to move forward with this, but I would not like to expose a parameter which I think has a high churn risk (num_chunk_per_batch) and an approach that is slower than it could be. I think we can implement the good part of this PR internally, which is computing the chain one zarr chunk at a time inside the worker that owns the shard, but also avoid writing on every chunk by writing as many whole chunks as fit in a fixed memory budget (for example 500 MB, at least one chunk), which keeps the extra memory bounded whatever the chunk length or channel count. This will solve the problem, avoid the slowdown and will not commit us to adding one extra parameter to some very general machinery that will be hard to remove after. I tested the approach on a local branch:

	time	peak memory
main	40.8 s	22.8 GB
this PR	71.4 s	4.9 GB
inside the zarr writer, 10 chunks per write	33.4 s	4.5 GB

1 s chunks and 60 s shards, with 10 chunks per write (what a 500 MB budget gives for 384 float32 channels), so each shard is rewritten 6 times instead of 60. The data is the first 60 s of a four-shank Neuropixels 2.0 recording (384 channels, 30 kHz), so a single shard, with phase_shift, bandpass_filter to float32 and median common_reference, one worker, one run each, and peak physical footprint on my Mac.

[alejoe91]

Thanks @h-mayorquin

I agree we should not wait for #4905 since it will require a lot of discussion (my first feeling is that implementation wise the job distribution across processing steps will be very very hard!).

Can you make an alternative PR with the approach you have in your local branch? Fine for me to go with that if it does what this does, but internally in the zarr writer!