async in zarr

[rabernat]

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.info

Note 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:

zarr-python/zarr/core.py

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 seconds

This 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

[jakirkham]

Have been wondering the same thing. Thanks for this nice write-up and performing this experiment. 😄

[gzuidhof]

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]

Thanks for pinging me!

[alimanfoo]

That is a crazy speedup, amazing. I'd be very grateful if someone could give me an intuition for why it's so much faster. Is it that multiple http requests can basically be opened concurrently? If so, is there an upper bound on the number of requests that can be active concurrently, and would we need to manage that ourselves?

…

On Mon, 17 Feb 2020, 04:26 Trevor Manz, ***@***.***> wrote: Thanks for pinging me! — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#536?email_source=notifications&email_token=AAFLYQTUJNQX5RGJT4LD6XTRDIGYHA5CNFSM4KVQPIV2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEL5AUNY#issuecomment-586811959>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFLYQTNPUFTZGMXVG23XD3RDIGYHANCNFSM4KVQPIVQ> .

[manzt]

I'm not at familiar with Python's asyncio but I'm assuming that there is overlap. Calling to WebAPIs (like fetch) returns a Promise rather than a value so requests are non-blocking and can be made concurrently.

In the example above,futures is a list of Python Futures (similar to promises in javascript). These are awaitable objects that represent an eventual result of an async operation. Using the session, multiple requests are fired off and awaiting async.gather(*futures) waits for the last future to be resolved.

async with aiohttp.ClientSession() as session:
    tic = time.time()
    futures = [get_chunk_http_async(n, session) for n in range(za.shape[0] // za.chunks[0])]
    all_data = await asyncio.gather(*futures)
    print(f"{time.time() - tic} seconds")

To handle concurrency, in zarr.js we use a dependency called p-queue to 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.

const queue = new PQueue({ concurrency: concurrencyLimit });
for (const proj of indexer.iter()) {
    const chunkValue = this.getChunkValue(proj, indexer, value, selectionShape); // returns promise
    queue.add(() => this.chunkSetItem(proj.chunkCoords, proj.chunkSelection, chunkValue)); // returns promise that queue resolves
}
// guarantees that all work on queue has finished
await queue.onIdle();

I think asyncio.Queue exists 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,

arr = za[:] 

is the same as this in zarr.js,

const arr = await za.get(null);

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]

To Alistair's point, if this is just about requests, I wonder if simply a Store leveraging asyncio internally would be sufficient to give you the same benefits. Not to say it wouldn't be interesting to add asyncio to all of Zarr. Just curious what is needed to get the performance boost.

[rabernat]

I wonder if simply a Store leveraging asyncio internally would be sufficient to give you the same benefits

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]

I wonder if this would be a good opportunity to add a batch retrieval method to the store API. E.g., something like store.multi_get(keys) -> iterable of values. This would be an optional method, which if not implemented would mean zarr falls back to current behaviour of calling __getitem__ multiple times. If it was implemented, then the store internally could choose to implement via asyncio if appropriate, or not. I.e., the zarr core module would not need to be aware of async, it would be entirely internal to the store implementation. Also, stores could leverage this in other ways. E.g., if a store was wrapping a database, it could batch together all the key requests into a single transaction. Or, e.g., if the store was wrapping a cloud service that provided some API for batching multiple requests into a single HTTP call, it could use that. Just a thought.

…

On Tue, 18 Feb 2020 at 05:48, Ryan Abernathey ***@***.***> wrote: I wonder if simply a Store leveraging asyncio internally would be sufficient to give you the same benefits 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. — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#536?email_source=notifications&email_token=AAFLYQSSQWPHM5IIN3PBBWDRDNZBXA5CNFSM4KVQPIV2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEMAVJDY#issuecomment-587289743>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFLYQXW2ELB23IKXTEDGWLRDNZBXANCNFSM4KVQPIVQ> .

-- Alistair Miles Head of Epidemiological Informatics Centre for Genomics and Global Health Big Data Institute Li Ka Shing Centre for Health Information and Discovery University of Oxford Old Road Campus Headington Oxford OX3 7LF United Kingdom Phone: +44 (0)1865 743596 or +44 (0)7866 541624 Email: alimanfoo@googlemail.com Web: http://a <http://purl.org/net/aliman>limanfoo.github.io/ Twitter: @alimanfoo <https://twitter.com/alimanfoo> Please feel free to resend your email and/or contact me by other means if you need an urgent reply.

[jakirkham]

Yeah this came up in issue ( #384 ). Agree this would be a good solution to this problem (amongst others).

[clbarnes]

@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]

@alimanfoo would you envision the stores' API being synchronous, and just using asyncio under the hood for multi-gets?

Yes that was my thought.

[mzjp2]

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]

That's a good point. We can probably live with something as simple as ThreadPoolExecutor's map.

TBH mapping with async, threads, or otherwise could just be specified as part of the API of this function (multigetitem? or other appropriate name).