How should Xarray control asynchronous calls?

[TomNicholas]

Need for async

There are several places in Xarray where utilizing asynchronous calls would hugely boost performance. This comes up particularly when interacting with high-latency storage backends, such as Zarr on remote object storage:

• When opening data, for loading metadata for many arrays/groups Reducing IO in ZarrStore #9853
• When loading multiple variables / groups
  • For creating indexes up-front open_dataset creates default indexes sequentially, causing significant latency in cloud high-latency stores #10579
  • For on-demand (lazy) loading Add an asynchronous load method? #10326 / Support concurrent loading of variables #8965
• For awaiting lazily-loading variables Dataset.to_zarr compute=False should allow access to awaitable #6383
• When loading multiple datasets Add an asynchronous load method? #10326
• When saving multiple variables / groups DataTree.to_zarr() is very slow writing to high latency store #9455

Each of these performance issues can be solved by leveraging async code somewhere in the stack, requiring explicit control of concurrency somewhere.

Where to async?

But where in the stack should that control be implemented?

As @rabernat noted in #8965 (comment), historically we have avoided this question originally by outsourcing nearly all parallelism to dask. Now that zarr-python v3 learned how to async, some of the above issues were addressed by controlling concurrency in that layer.

But pushing everything down to zarr only gets us so far. It requires zarr-python to implement API that would be more natural to live in xarray's zarr backend, prevents other backends benefiting from async (such as PyDAP), and sometimes we do want to expose the async functions publicly to the user.

Inter-library coordination

In general there are perhaps 5 layers at which concurrency might need be controlled:

1. User code
2. Xarray
3. Dask / Cubed
4. Zarr-python
5. Icechunk / fsspec

Xarray is clearly one amongst many so we need to think about how to coordinate across layers to avoid opaque and unpredictable interactions between libraries.

Choices in Xarray

If xarray does not explicitly control concurrency but does use asyncio.gather calls to issue requests across multiple variables then it could lead to oversubscribing threads, and some of the above issues are awkward to solve with this limitation.

If xarray does explicitly control concurrency then it requires either a separate event loop, or a threadpool with some way to configure the number of threads, plus some code to get this to run inside an already-running event loop (e.g. inside ipython/jupyter). It also has other risks.

#10327 adds a new primitive (the async method xr.Variable.load_async()) that will be useful for solving many of the above issues, but faces this question of whether or not to limit concurrency (see thread). For now I'm leaning towards punting on the question by not limiting it, simply exposing the awaitables all the way to the top, but with the understanding that once we have agreement here on a general approach we should go back and add in some concurrency-limiting code.

cc @d-v-b, @TomAugspurger, @jhamman, @dcherian, @shoyer

[TomAugspurger]

Thanks Tom. One small clarification:

If xarray does not explicitly control concurrency but does use asyncio.gather calls to issue requests across multiple variables then it could lead to oversubscribing threads, and some of the above issues are awkward to solve with this limitation.

I think that, at least as far as xarray can control it, using async specifically shouldn't lead to thread oversubscription. Each task xarray creates (and runs in that asyncio.gather) will run on the same thread. Each task will cooperate with other tasks on that thread, giving up time when they hit a point where they can't progress (e.g. a Zarr chunk read waiting for data on the socket).

asyncio tasks can spawn threads and run code inside them (and zarr-python does this for some Stores and essentially all Codecs). But I think that ideally xarray itself wouldn't be spawning threads here. And thead oversubscription still shouldn't be a problem, since these calls are run in a ThreadPoolExecutor attached to the asyncio event loop (these calls just need to avoid spawning their own threads, but that shouldn't be xarray's problem).

then it requires either a separate event loop

It'd be great if xarray could avoid starting its own event loop.

I'll take a look at #10327 when I get a chance.

[rabernat]

Thanks for writing this up Tom! As [other] Tom mentioned, it's important to distinguish between two types of concurrency:

1. I/O concurrency, which is best outsourced to asyncio tasks. The concurrency "limit" here is ultimately set by the remote store. S3 can handle hundreds of concurrent request from a single node.
2. CPU concurrency for compute-bound tasks. The limit here is determined by the number of CPU cores available.

In real world workflows, these are coupled together. Managing http connections takes CPU time. And decompressing chunks requires compute.

For 1, my view is that Xarray is basically just pushing down asyncio calls to the underlying store, and it's that store's responsibility to enforce any concurrency limits.

2 is a lot more complicated because it requires several layers of software (Zarr, Xarray, maybe Dask) all figuring out how to share a finite amount of CPU.

[TomNicholas]

that store's responsibility to enforce any concurrency limits.

Of course we can only make it entirely that store's responsibility if that store has a synchronous API to do what we need (e.g. load an arbitrary set of arrays concurrently, or set metadata for an arbitrary heirarchy of arrays).