Investigate Rust-Based Python-Accessed Stores

[ilan-gold]

Currently a store can be written in rust, wrapped in pyo3, and then sent to our library. But we don't use that store. We should investigate whether it's possible to somehow "work around" and grab the underlying rust store without compromising performance instead of re-implementing/re-instantiating it in Rust within our package.

[kylebarron]

You could access it via Python calls, but you wouldn't be able to access the underlying Rust store directly. pyo3 doesn't support that sort of dynamic linking (yet?) because it's hard to do in a stable way across pyo3, python, and Rust versions.

[ilan-gold]

You could access it via Python calls

Are you aware of the performance of this?

Re: #44 (comment) this then returns a copy? Am I understanding that right, when you go from python into the rust?

[kylebarron]

Well the proposal in #44 (comment) would have you compile your own stores on the Rust side (just reusing the Rust code from pyo3-object_store), so you wouldn't have any overhead.

But in principle, it's possible to also call the Python functions exported by obstore. And then when obstore.get or obstore.get_range returned, you could extract the returned result into a PyBuffer<u8> buffer protocol object. But you'd have to acquire the GIL whenever you call the Python function. So the pure-rust approach above is better.

[kylebarron]

This PR in async-tiff contains an example of reusing pyo3-object_store. Once object_store 0.12 is released (should be in the next few days), we can release pyo3-object_store v0.1, and then we can prototype integrating that here.

[ilan-gold]

Very cool @kylebarron !!!

[kylebarron]

I've now published pyo3-object_store 0.1, which uses pyo3 0.23, and pyo3-object_store 0.2, which uses pyo3 0.24.

It looks like you'd want to add object_store support into your StoreConfig?
https://github.com/ilan-gold/zarrs-python/blob/03af8d0fdedebf847d13f88768b0cc7171238df6/src/store.rs#L26-L30

Is there a reason why you require Hash, PartialEq, Eq, PartialOrd, Ord on StoreConfig? Those are a lot of bounds.

And it looks super messy to take a LocalStore and FsspecStore and access their underlying HTTP config settings.

Instead of storing a StoreConfig, why don't you store a Store directly?

[kylebarron]

Also, at least for object_store, it's bad to be recreating a store from a config all the time, because that prevents connection pooling from taking place.

[LDeakin]

Is there a reason why you require Hash, PartialEq, Eq, PartialOrd, Ord on StoreConfig? Those are a lot of bounds.

I think it is related to storing the configs in the StoreManager.

Also, at least for object_store, it's bad to be recreating a store from a config all the time, because that prevents connection pooling from taking place.

The StoreManager is caching stores with the same config.

I think this is all somewhat because we don't get any info about the store when an array is created, and it just starts getting thrown at us with each chunk. It would be fantastic to get an upstream change into zarr-python to get this info when the codec pipeline is created.

[kylebarron]

I think this is all somewhat because we don't get any info about the store when an array is created, and it just starts getting thrown at us with each chunk. It would be fantastic to get an upstream change into zarr-python to get this info when the codec pipeline is created.

Indeed, that does sound like a really helpful update. It would be possible to have a similar config caching approach for object_store, but it wouldn't be my first approach.

[JackKelly]

@kylebarron's awesome obstore-based Store has recently been merged into zarr-python. Does that make it any easier to directly call object_store from zarrs-python without going via Python? (I'm guessing not???)

Why do I ask? To zoom out for a moment: I'm really excited about being able to train machine learning models for forecasting renewable power generation directly from public Zarr datasets of numerical weather predictions, like those on dynamical.org. Modern GPUs are shockingly fast, so the bottleneck when training these ML models is often the IO. The dream would be for these datasets to be published as sharded Zarrs with pretty tiny chunks. And then, to train our ML models, we'd spin up a virtual machine with a fast network interface card (100 or 200 gigabits per second) near to the data.

Crucially, we'd need a Zarr library that could efficiently load (at least) thousands of chunks per second. Zarrs feels like a great candidate to be able to do this. I haven't tried yet, but my guess is that zarrs + object_store can already achieve this level of performace today, for Rust developers.

But most ML researchers use Python. So I'm looking into ways to achieve this level of performance for Python users. TBH, I'm nervous that we'll bump into Python's overhead when trying to load thousands of chunks per second.

So, I guess I have several questions:

1. Assuming that zarrs-python has to go via Python to call obstore, could zarrs-python + obstore load thousands of Zarr chunks per second from object storage?
2. If not, is there a clean way to allow zarrs-python to directly call object_store from Rust, without going via Python?
3. If not, would a useful contibution be a minimal Python wrapper around zarrs and obstore, that does not use zarr-python? (This work would probably include building a simple xarray backend).

I can dedicate some time to helping build things if needs be.

[d-v-b]

Crucially, we'd need a Zarr library that could efficiently load (at least) thousands of chunks per second.

I'm pretty sure we can do this with Zarr-python after performance optimization of our IO layer. At a minimum, it would be very valuable to know what things we would need to change in order to achieve this throughput.

[ilan-gold]

@JackKelly we always welcome contributions. The main blocker here IMO would be the fact that our code is very sync-focused and not very async-focused. As far as I can tell, we would need a fair amount of duplication of current methods to achieve an async API. Doable, but since none of us three maintainers work with cloud data as far as I can tell, not a priority for us personally. io_uring has been on our collective radar but no one has taken it up, so having the async code in place would be awesome for that as a first step!

As for what @d-v-b said, I tend to agree. I don't think io_uring in python is happening any time soon for a variety of reasons, but in Rust it seems feasible. So for that our library is a good fit and something I look forward to. Aside from that, while our performance here is good, there is definitely more juice to be squeezed out of the python implementation as benchmarks shows (which does make this library very useful for knowing that as well!), but when the python implementation is "optimized," we can be quite close. For example, our read benchmarks (both chunk-by-chunk and full read) are quite competitive outside sharding, so I don't see why sharding is a fundamental bottle neck, although if it is, I would love to understand why as @d-v-b pointed out.

Of course contributions are welcome! Happy to provide any assistance!