Repository navigation
Description
I think there are some places where zarr would benefit immensely from some async capabilities when reading and writing data. I will try to illustrate this with the simplest example I can.
Let's consider a zarr array stored in a public S3 bucket, which we can read with fsspec's HTTPFileSystem interface (no special S3 API needed, just regular http calls).
import fsspec
url_base = 'https://mur-sst.s3.us-west-2.amazonaws.com/zarr/time'
mapper = fsspec.get_mapper(url_base)
za = zarr.open(mapper)
za.infoNote that this is a highly sub-optimal choice of chunks. The 1D array of shape (6443,) is stored in chunks of only (5,) items, resulting in over 1000 tiny chunks. Reading this data takes forever, over 5 minutes
%prun tdata = za[:] 20312192 function calls (20310903 primitive calls) in 342.624 seconds
Ordered by: internal time
ncalls tottime percall cumtime percall filename:lineno(function)
1289 139.941 0.109 140.077 0.109 {built-in method _openssl.SSL_do_handshake}
2578 99.914 0.039 99.914 0.039 {built-in method _openssl.SSL_read}
1289 68.375 0.053 68.375 0.053 {method 'connect' of '_socket.socket' objects}
1289 9.252 0.007 9.252 0.007 {built-in method _openssl.SSL_CTX_load_verify_locations}
1289 7.857 0.006 7.868 0.006 {built-in method _socket.getaddrinfo}
1289 1.619 0.001 1.828 0.001 connectionpool.py:455(close)
930658 0.980 0.000 2.103 0.000 os.py:674(__getitem__)
...
I believe fsspec is introducing some major overhead by not reusing a connectionpool. But regardless, zarr is iterating synchronously over each chunk to load the data:
Lines 1023 to 1028 in 994f244
| # iterate over chunks | |
| for chunk_coords, chunk_selection, out_selection in indexer: | |
| # load chunk selection into output array | |
| self._chunk_getitem(chunk_coords, chunk_selection, out, out_selection, | |
| drop_axes=indexer.drop_axes, fields=fields) |
As a lower bound on how fast this approach could be, we bypass zarr and fsspec and just fetch all the chunks with requests:
import requests
s = requests.Session()
def get_chunk_http(n):
r = s.get(url_base + f'/{n}')
r.raise_for_status()
return r.content
%prun all_data = [get_chunk_http(n) for n in range(za.shape[0] // za.chunks[0])] 12550435 function calls (12549147 primitive calls) in 98.508 seconds
Ordered by: internal time
ncalls tottime percall cumtime percall filename:lineno(function)
2576 87.798 0.034 87.798 0.034 {built-in method _openssl.SSL_read}
13 1.436 0.110 1.437 0.111 {built-in method _openssl.SSL_do_handshake}
929936 1.042 0.000 2.224 0.000 os.py:674(__getitem__)
As expected, reusing a connection pool sped things up, but it still takes 100 s to read the array.
Finally, we can try the same thing with asyncio
import asyncio
import aiohttp
import time
async def get_chunk_http_async(n, session):
url = url_base + f'/{n}'
async with session.get(url) as r:
r.raise_for_status()
data = await r.read()
return data
async with aiohttp.ClientSession() as session:
tic = time.time()
all_data = await asyncio.gather(*[get_chunk_http_async(n, session)
for n in range(za.shape[0] // za.chunks[0])])
print(f"{time.time() - tic} seconds")
# > 1.7969944477081299 secondsThis is a MAJOR speedup!
I am aware that using dask could possibly help me here. But I don't have big data here, and I don't want to use dask. I want zarr to support asyncio natively.
I am quite new to async programming and have no idea how hard / complicated it would be to do this. But based on this experiment, I am quite sure there are major performance benefits to be had, particularly when using zarr with remote storage protocols.
Thoughts?
cc @cgentemann

Activity
jakirkham commentedon Feb 15, 2020
Have been wondering the same thing. Thanks for this nice write-up and performing this experiment. 😄
gzuidhof commentedon Feb 17, 2020
This is supported in zarr.js, for the same reasons as you wrote here. Here's the relevant PR and issue which also explains why it was a good idea (with some timings for less pathological cases where even there the speedup can be large).
Pinging @manzt who implemented above functionality, the JS implementation follows Python pretty closely here, so a similar implementation approach may be possible
manzt commentedon Feb 17, 2020
Thanks for pinging me!
alimanfoo commentedon Feb 17, 2020
manzt commentedon Feb 17, 2020
I'm not at familiar with Python's
asynciobut I'm assuming that there is overlap. Calling to WebAPIs (likefetch) returns a Promise rather than a value so requests are non-blocking and can be made concurrently.In the example above,
futuresis a list of Python Futures (similar to promises in javascript). These are awaitable objects that represent an eventual result of an async operation. Using thesession, multiple requests are fired off and awaitingasync.gather(*futures)waits for the last future to be resolved.To handle concurrency, in zarr.js we use a dependency called
p-queueto limit the number of concurrent requests to 10 by default. This way we can dynamically add to the queue while iterating through the projections. Hopefully this looks somewhat similar to zarr.I think
asyncio.Queueexists but that's about my limit of knowledge.One concern I have here is that the API for accessing a slice in an array in zarr.js is naturally awaitable since everything in javascript is async. We need to wrap our requests in async functions and await the responses.
This in zarr-python,
is the same as this in zarr.js,
regardless of the store type. Ideally you'd hide all the async fetching from the end user and have the same API independent of whether the array is on disk or over HTTP. Someone smarter than me probably knows if this is possible.
jakirkham commentedon Feb 18, 2020
To Alistair's point, if this is just about requests, I wonder if simply a
Storeleveragingasynciointernally would be sufficient to give you the same benefits. Not to say it wouldn't be interesting to addasyncioto all of Zarr. Just curious what is needed to get the performance boost.rabernat commentedon Feb 18, 2020
This is what I had in mind. I think the main performance benefit is when reading / writing many smallish chunks concurrently. We definitely don't want to force zarr users to use async.
Maybe a store could declare itself to use async for certain operations, and then zarr would know to use asyncio when performing loops with I/O operations. Otherwise it would default to the current, synchronous behavior.
I think there are cases where asyncio would hurt us. I did a few tests with files and found that, in some cases, asyncio was slower. I assume this is because issuing serial reads is more efficient than concurrent reads for some disks.
alimanfoo commentedon Feb 18, 2020
jakirkham commentedon Feb 18, 2020
Yeah this came up in issue ( #384 ). Agree this would be a good solution to this problem (amongst others).
clbarnes commentedon Mar 3, 2020
@alimanfoo would you envision the stores' API being synchronous, and just using asyncio under the hood for multi-gets? I've found asyncio quite difficult to work with because it "poisons" the entire stack above it, forcing the entire pipeline into peppering async/await annotations everywhere. You can get the event loop and resync inside a function call but it's not pretty.
alimanfoo commentedon Mar 3, 2020
Yes that was my thought.
mzjp2 commentedon Mar 5, 2020
If we just want to parallelise network requests, we don't need to go full asyncio, right? We could do thread-based parallelism. Not fully convinced we need to buy into the whole asyncio package just to parallelise a few network requests.
jakirkham commentedon Mar 5, 2020
That's a good point. We can probably live with something as simple as
ThreadPoolExecutor'smap.TBH mapping with async, threads, or otherwise could just be specified as part of the API of this function (
multigetitem? or other appropriate name).rabernat commentedon Mar 5, 2020
Async and parallel are not the same thing. From my point of view, I use dask to handle actual parallel I/O. The example I posted here is one where I don't actually want dask involved. It's a tiny amount of data. When using xarray with zarr, the dimension coordinates are always read eagerly and without dask. That's the use case that motivated this issue.
Furthermore, I would worry about any internal multithreading within zarr conflicting with dask's own parallelism. This would have to be handled carefully.
alimanfoo commentedon Mar 5, 2020
66 remaining items
martindurant commentedon Jul 23, 2020
Update:
bilts commentedon Aug 10, 2020
What's the current status of an implementation that calls
getitemsinstead of__getitem__in the Zarr library? People seem to be mentioning making use of it, but I haven't seen an actual implementation. Is it still in the cards?martindurant commentedon Aug 10, 2020
Currently, zarr calls getitem on each thing it needs, in sequence. To make use of my getitems and any concurrency, the code would need to be changed. I'm not aware of anyone having done this.
jakirkham commentedon Aug 10, 2020
I think this is still of interest. Is this something that you would be willing to contribute, @bilts? 🙂
bilts commentedon Aug 10, 2020
@jakirkham I would be / have been willing. I have a prototype implementation in a fork mentioned above for people to try it out. It's proven way, way faster for our needs. I think making a more robust change may require someone more able to make choices about zarr-python's internal organization than I am.
martindurant commentedon Aug 10, 2020
Hurray!
jakirkham commentedon Aug 10, 2020
That's probably easier/faster to address through the normal review process in a PR 😉
rabernat commentedon Nov 3, 2020
I believe we can now close this issue, now that #630 and #606 have both been merged!
Just to close the loop, I just did a simple benchmarking exercise from Pangeo Cloud in Google Cloud.
My test environment included the following versions:
My old (pre-async) environment was
Here is a comparison of the speeds
This is a fantastic improvement, particularly for reading (over 20x)!
👏 👏 👏 for all the effort by @martindurant and the other devs who helped make this happen!
shoyer commentedon Apr 2, 2021
One thing I'll add: parallel fetching of chunks in Zarr (as is supported here by
getitems) is valuable even in the case where Dask is also used for parallelism. This is because the optimal number of threads for IO is often much higher than the optimal number of threads for computation, so it's a best practice to use a separate thread pool anyways. See pangeo-forge/pangeo-forge-recipes#89 for more details.tasansal commentedon Sep 11, 2021
@shoyer I have a question about this.
I am currently fetching 500 GB to 2 TB chunked datases (~1MB chunks). Running on a 48 physical 96 logical core machine.
Data is on Google Cloud Storage (GCS).
I am using
daskwith a larger chunksize than the underlying array. i.e. from ~1MB chunks to ~100MB chunks. So eachdasktask should pull multiple chunks?I am using a
LocalClusterclient.One thing I noticed is if I run a
dask.Clientwith ONLY threads (i.e.Client(n_workers=1, threads_per_worker=48), I get terrible read performance. It is almost equal to a single process (400MB ish).When I recreate the
Clientwith 48 workers (processes) and 1 thread per worker, I can read up to 5GB/s from GCP. However, this adds some latency because of multiple processes.Another example is 8 workers with 12 threads, and that scales up to 8x performance from 1 worker, but still much slower than 48 workers. So it only seems to scale up with processes, not threads?
I was wondering if I am doing something wrong? It sounds like the thread only client is NOT fetching multiple parts / chunks at the same time.
I am not using any
asynchronous=xxxkwargs anywhere, maybe I need to do that?gcsfs,FSSpec,zarr, anddaskis the combo I am using.jakirkham commentedon Sep 11, 2021
@tasansal can you please raise a new issue? Also please include a minimal reproducer
tasansal commentedon Sep 11, 2021
@jakirkham, of course, wasn't sure if it should be an issue of its own. Will do that.
jakirkham commentedon Sep 11, 2021
Thanks! It just makes it easier for us to keep track of these things :)