a two-stage concurrent fetch/decode pipeline to maximize throughput from stores (tested on lustre)

Hi @ilan-gold and @LDeakin

I have been recently profiling and benchmarking then coming up with solutions to improve my random access speed on huge files on lustre. I now have a nice codepath I'd like to suggest and split into PR's but I want to first write here before opening any PRs. I also noticed some of my work was on stale since zarrs main has changed a lot :D, so the second part of my suggestion might fit better or worse with your plans but I think the first one isn't stale.

These are the topics:

• I was thinking of proposing a two phase concurrency path.
• and a run based access path for sparse datasets for example

two phase pipeline

So my plan was, we have a large pool of threads only for being blocked on the lustre preads (not on rayon). Then after each read is done, this would trigger a job on the rayon pool which would do the actual cpu bound job on the piece it receives.

I was reading related issues and PR's and saw #177. My draft had a thread pool using crossbeam_channel but I think we can even use tokio for this. Because if we only use tokio to fetch the uncompressed bytes we don't need to borrow python numpy output. Then the lifetime issue in #44 wouldn't be a problem no? Then this uncompressed bytes being received would trigger job from a scoped rayon pool and it can decode owned bytes directly into disjoint output views as I said.

This way we can limit rayon threads to number of cores bc they'd be doing all the cpu work. While tokio pool is waiting.

concurrency based on runs not shards

Then after two stage pipeline, I had:

• shard index caching
• file handler caching
• parallelizing on runs without caring for shard grouping

Because the shard index and file handling caching overhead is so little for even the huge files I am working on, there is no reason to have per-shard concurrency depth.

As I said I had very nice results for even chunksize=1 random access in annbatch (2500 samples/sec>). But I noticed the zarrs main path has diverged a lot :). So I'd like to first ask and talk to you about how I should split up my work in PRs. I can make a PR based on tokio and this idea, which would be stacked on top of this per run fetch wiring. I haven't read the latest main much yet but it looks like per run parallelization might be easier this way.

Here is how the draft of my first PR was going to look like roughly:

# fetct_pool.rs

pub struct FetchPool {
    tx: crossbeam_channel::Sender<Box<dyn FnOnce() + Send + 'static>>,
}

impl FetchPool {
    fn new(threads: usize) -> Self {
        let (tx, rx) = crossbeam_channel::unbounded::<Box<dyn FnOnce() + Send + 'static>>();
        for index in 0..threads {
            let rx = rx.clone();
            std::thread::Builder::new()
                .name(format!("zarrs-fetch-{index}"))
                .stack_size(FETCH_THREAD_STACK_BYTES)
                .spawn(move || {
                    while let Ok(job) = rx.recv() {
                        job();
                    }
                })
                .expect("failed to spawn fetch thread");
        }
        Self { tx }
    }

    /// Queue a read. Never blocks: the queue is unbounded, so submitting all
    /// of a batch's reads up front is what puts them in flight together.
    pub fn submit(&self, job: impl FnOnce() + Send + 'static) {
        // Workers hold `rx` for as long as the pool is alive, so this can only
        // fail after the last `Arc` has dropped, which cannot happen here.
        let _ = self.tx.send(Box::new(job));
    }
}

impl CodecPipelineImpl {
    /// Issue every whole-chunk read before decoding any of them, then decode
    /// each as it lands.
    ///
    /// One read per chunk, and each is independent, so submitting the batch up
    /// front is what puts the reads in flight together. Depth is then
    /// `fetch_threads` rather than the rayon width -- the distinction that
    /// matters on a store where a read is milliseconds of waiting.
    fn retrieve_whole_chunks_into(
        &self,
        items: Vec<ChunkItem>,
        output: UnsafeCellSlice<u8>,
        codec_options: &CodecOptions,
    ) -> PyResult<()> {
        let (tx, rx) = crossbeam_channel::unbounded();
        for (index, item) in items.iter().enumerate() {
            let tx = tx.clone();
            let store = self.store.clone();
            let key = item.key.clone();
            self.fetch_pool.submit(move || {
                let _ = tx.send((index, store.get(&key)));
            });
        }
        drop(tx);

        let failure: Mutex<Option<PyErr>> = Mutex::new(None);
        let mut pending = items.into_iter().map(Some).collect::<Vec<_>>();
        rayon::scope(|scope| {
            for (index, fetched) in rx {
                let Some(item) = pending[index].take() else {
                    continue;
                };
                let failure = &failure;
                scope.spawn(move |_| {
                    let outcome = (|| -> PyResult<()> {
                        let mut output_view = unsafe {
                            // SAFETY: chunks represent disjoint array subsets
                            ArrayBytesFixedDisjointView::new(
                                output,
                                self.data_type
                                    .fixed_size()
                                    .ok_or("variable length data type not supported")
                                    .map_py_err::<PyTypeError>()?,
                                bytemuck::must_cast_slice(&item.array_shape),
                                item.subset.clone(),
                            )
                            .map_py_err::<PyRuntimeError>()?
                        };
                        let target = ArrayBytesDecodeIntoTarget::Fixed(&mut output_view);
                        match fetched.map_py_err::<PyRuntimeError>()? {
                            Some(chunk_encoded) => {
                                let chunk_encoded: Vec<u8> = chunk_encoded.into();
                                self.codec_chain.decode_into(
                                    Cow::Owned(chunk_encoded),
                                    &item.shape,
                                    &self.data_type,
                                    &self.fill_value,
                                    target,
                                    codec_options,
                                )
                            }
                            // Missing chunk: fill value.
                            None => copy_fill_value_into(
                                &self.data_type,
                                &self.fill_value,
                                target,
                            ),
                        }
                        .map_codec_err()
                    })();
                    if let Err(error) = outcome {
                        let mut failure = failure.lock().unwrap();
                        if failure.is_none() {
                            *failure = Some(error);
                        }
                    }
                });
            }
        });
        if let Some(error) = failure.into_inner().unwrap() {
            return Err(error);
        }

        Ok(())
    }

Start by reading the current main branch alongside the proposed fetct_pool.rs and CodecPipelineImpl::retrieve_whole_chunks_into entry point, then review related issues #177 and #44. The scope is not settled: done first requires maintainer agreement on the PR split and the two-stage fetch/decode design, followed by benchmarking the resulting path on Lustre.