Automatically using Dask for large datasets when necessary

[Zethson]

Coming out of a discussion with @ilan-gold.

Idea

We keep noticing that many (technical even) users are not aware of how to perform analyses of larger-than-memory datasets. There's multiple ways to improve this including workshops, more prominent tutorials, yadayada, but there might also be a technical solution. This is inspired a bit by @timtreis spatialdata-plot fallback where large datasets (automatically) use datashader to get the job done.

Analogously, we could calculate the dataset size & estimated memory requirements. If we suspect that memory is insufficient, we could automatically (and transparently) create a Dask array to perform OOC computation. In such cases, we'd likely need to "return" a Dask array which could confuse users. I'd suggest we enable this behavior by default with an option to opt-out but we could of course start with it being opt-in while we try it out. I argue that many users don't understand and don't want to deal with ways to make "large dataset processing possible" and if we can do it for them, this might be a cheap win.

I'll omit an API draft for now because this would just be a new setting for the settings class.

Primarily suggesting this for scanpy but if this is worthwhile considering, we should harmonize the behavior across all of scverse. Therefore, I'm kidnly asking for feedback from all of @scverse/core-devs .

caveats

1. Ideally, we should of course complete the Dask support across all of scanpy's algorithms.
2. @ilan-gold mentioned that we don't retain the file handle of anndata objects but this could be changed.
3. I don't know how easy it is to estimate available memory on Slurm jobs. There might be some hidden complexity here. I also don't know how difficult it is to estimate the required peak memory for the respective algorithms.
4. We might need to also estimate chunk size then.

Maybe, better catching OOM and pointing users to our dask tutorials is already good enough but I think this is worthwhile discussing

[ilan-gold]

Ideally, we should of course complete the Dask support across all of scanpy's algorithms.

I'd have to check, but I think this is mostly done, at least for algorithms where it makes sense. Rank genes is probably the only major algorithm missing that would benefit (leiden probably fits in memory, and for neighbors we're moving towards hannoy: scikit-learn-contrib/sklearn-ann#123)

@ilan-gold mentioned that we don't retain the file handle of anndata objects but this could be changed.

This could be changed, yeah, in theory.

I don't know how easy it is to estimate available memory on Slurm jobs. There might be some hidden complexity here. I also don't know how difficult it is to estimate the required peak memory for the respective algorithms.

Very true. It's possible some dask algorithms are just always faster in whcih case that might alleviate this burden.

We might need to also estimate chunk size then.

Probably the biggest issue maybe - "how sparse is my data" will be an approximation at best for per-chunk. Also we'd have to decide what the default cluster would be (processes? threads? not sure)

[ilia-kats]

Large datasets will usually be sparse, and sparse support in dask isn't really where it would need to be for this, I think (see here and here.

[ilan-gold]

@ilia-kats That's true probably to end users who might want to do something with the dask array outside scanpy, but we have shims in fau for simple operations and scanpy's internal functions seem fairly robust. So maybe

In such cases, we'd likely need to "return" a Dask array which could confuse users.

doesn't have to happen so as not to confuse users in this other, different sense (potentially unexpected breaking behvaior)

[timtreis]

Not properly quote-citing because I'm on mobile:

We chose to implement it the way we did in spatialdata-plot because I don't think, complexity should be passed on to the user. I think Ilia once actually said something along the tiles "easy things should be easy to do, but complex things should be possible" (back when SpatialData was still spatialmuon).

Even amateur users might get increasingly exposed to data that requires this, so, ideally, everything works automagically under the hood (imho) and experts users can make adjustments if they want/need. But I still think that a biologists doesn't need to know what Dask even is, even if we use it for everything under the hood.

Re: large data will always be sparse
Only if we stick to a purely transcriptomics world view. In other modalities you can easily imagine similarly sized but dense matrices.

Re: SLURM jobs
At least on our Helmholtz HPC, it's fairly easy to find out which ressources are actually allocated to you, even on shared notes. I think the os library has some "affine ressources" logic or sth? No idea about peak memory though

[zboldyga]

I have also had this experience! 😄 More broadly than just Dask, there are a number of small friction points in the single cell space that I feel are slowing down progress in the space. It's not the fault of code maintainers, btw! Just seems like natural refinement of tools.

For instance, a major area I feel is neglected is how data is stored and shared. I ran a benchmark that included over 30 datasets, so I had to go around and download and normalize these. There are many cases of no compression being used in h5ad files when sharing, no sorting of rows within the files (e.g. by group), chunking dissonances, and in some cases the wrong dtype is used (I think some were not integers, and some not sparse when they should have been). And also share points like figshare with quite limited download bandwidth. And then when processing files, different compressions can be applied depending on your disk speed for optimal performance (e.g. in my case, uncompressed on disk was best).

Anndata also has flexibility and some people stuff large objects into uns. I also found cases of common operations like loading anndata files which were accidentally using up a lot of extra memory loading uns objects that were unused / you have to be careful about how you load files.

And I do think there were a few basic ops that didn't seem to have Dask support (I'll go back and surface those sometime soon / raise issue/PR).

To get a pipeline working that did common operations across these files, I wound up having to write quite a lot of code: https://github.com/Virtual-Cell-Research-Community/drf/tree/main/data/preprocessing (this was initially derived from the Miller DRF paper code, but I basically rewrote the whole thing).

So while I'm drifting from your point on Dask, I say this because I think there are a few solutions that address this broader group of friction points:

1. I'm working on a comprehensive benchmark for end-to-end sc rna seq pipelines. Snippets of code for common operations under various conditions so there is a good reference for what works vs. what doesn't, or how much more performant something can be. This stems from realizing wilcoxon DE could actually be > 100x faster than it currently is in scanpy... This probably is broader scope than scanpy. But I think it would also help determine what should be in scope of scanpy / whether there's anything we could bring to this layer to keep users on a happy path with no cognitive load.

2. While I'm still unsure on the exact framing, I think there might be a place in scanpy for a higher-level abstraction, e.g. like Keras was for machine learning, that gives you a few calls to do whole single cell workflows. That's basiclly what the code I wrote does. But right now it's specific to processing lots of files simultaneously. And it's limited to perturb-seq. And needs some refinement, etc.

(I'm also chipping in on scanpy, will raise issues/PRs as I work through items)

I'll think on this more and share when I have more tangibles, but curious on initial thoughts!

[ilan-gold]

There are many cases of no compression being used in h5ad files when sharing, no sorting of rows within the files (e.g. by group), chunking dissonances, and in some cases the wrong dtype is used (I think some were not integers, and some not sparse when they should have been).

Just on this point, https://annbatch.readthedocs.io/en/stable/ became a funneling point for observing and controlling for most of these things in real time. My general advice is to use zarr as fast as humanly possible! I will respond a bit more fully but just want to point this bit out about zarr because it is a massive win.

If you have suggestions on how to get more people using zarr, I'm all ears!

[zboldyga]

ah cool I didn't realize you were on the annbatch team, I had read that paper recently -- nice work (I'm seeing you on there when I pulled it up again)!

also agree RE zarr. Tbh I thought the single cell space was kinda locked into an hdf5 mindset haha. I will think on this more.

This does make me feel even more so that proper tool usage is confusing right now. Like for the foreseeable future, you still have to wrangle with accessing the old datasets in h5ad. And there's loads of old code floating around (and AI is probably going to generate mostly h5ad code).

I will factor this into the benchmarking tooling

But perhaps also a modernized sharepoint with zarr versions of the public datasets would be something I can help with. Idk about running something like this long term, the costs could get hefty. Maybe there's also a p2p opportunity here amongst academic institutions?

[zboldyga]

I ran into like half a dozen 'directory' type tools when trying to find the perturb seq datasets, but I felt most of them lacked important things. Like provenance thru a checksum (I want to be absolutely sure the data is the true data -- some had done normalization, for instance). And they weren't maintained. And bandwidth was very low on most of them.

[Zethson]

This stems from realizing wilcoxon DE could actually be > 100x faster than it currently is in scanpy...

@ilan-gold is currently investing substantial time into making it faster - see #4038 . Whether Wilcoxon is the right tool in the first place is a different question. It's only useful for say marker gene detection but any serious DE should not be done using Wilcoxon.

While I'm still unsure on the exact framing, I think there might be a place in scanpy for a higher-level abstraction, e.g. like Keras was for machine learning, that gives you a few calls to do whole single cell workflows.

Well, some people write pipelines like https://github.com/nf-core/scdownstream/ for mass processing of datasets. However, our field is very much reliant on the right parameter and algorithm choices. It's always possible to ask for more abstraction & automation but I think scanpy hits it very well.

But perhaps also a modernized sharepoint with zarr versions of the public datasets would be something I can help with. Idk about running something like this long term, the costs could get hefty. Maybe there's also a p2p opportunity here amongst academic institutions?

Something I mentioned recently @ilan-gold . If we want more zarr usage, we must show it off a lot more in all basic tutorials and we should reach out to our friends from e.g. CxG, hubmap, ... to discuss switching to zarr for future versions. We should also understand any worries they might have.

[ilan-gold]

Something I mentioned recently @ilan-gold . If we want more zarr usage, we must show it off a lot more in all basic tutorials and we should reach out to our friends from e.g. CxG, hubmap, ... to discuss switching to zarr for future versions. We should also understand any worries they might have.

HuBMAP was not going to happen with the current setup but now the project is running out. I may have some good news on this front soon though, as control over the "backend" data may be more realistic. CxG I have tried to talk to @ivirshup about but ultimately we have no control over, which makes this a bit of a chicken and the egg. So annbatch is sort of my "marketing" tool. I would absoltuely love to see literally anyone just make a bunch of zarr data available on s3 (and foot the bill haha)