Preventing Deadlocks When Reading Metadata Concurrently via asyncio.gather

[dgegen]

As described in #3196, I encountered issues opening Zarr v3 arrays stored over SFTP using fsspec. Specifically, python would freeze opening zarr arrays.

Root Cause

The issue stems from the use of asyncio.gather in zarr.core.array.get_array_metadata, which attempts to read multiple metadata files (e.g., .zarray, .zattrs, zarr.json) concurrently. This works well for truly asynchronous filesystems, but breaks when using systems like SFTPFileSystem, which does not seem to be concurrency-safe in async contexts (potentially relying on blocking I/O internally or managing connection states using global locks) leading to deadlocks or indefinite hangs when asyncio.gather is used to perform multiple reads simultaneously.

Solution

To address this, I’ve implemented a fallback to sequential reads for filesystems that are not concurrency-safe. The logic is as follows: For non asynchronous file systems, the user sets store.fs.asynchronous=False. The helper function is_concurrency_safe(store_path: StorePath) -> bool, checks this getattr(fs, "asynchronous", True). If True asyncio.gather is used, else we fall back to sequential await. This Preserves the performance benefit of concurrent reads for safe filesystems (e.g., local disk, S3, GCS), while preventing deadlocks and improved robustness when using backends like SFTP.

These changes may not address all scenarios not asynchronous file systems could cause issues, as there are several other instances of asyncio.gather in zarr.core.array and zarr.core.group. However, I opted to focus on this specific problem first, as enabling the opening of arrays and groups is likely the highest priority, and I wanted to discuss this approach before making too many changes.

I look forward to hearing your thoughts and seeing this issue resolved!

TODO:

• Add unit tests and/or doctests in docstrings
• Add docstrings and API docs for any new/modified user-facing classes and functions
• New/modified features documented in docs/user-guide/*.rst
• Changes documented as a new file in changes/
• GitHub Actions have all passed
• Test coverage is 100% (Codecov passes)

[d-v-b]

Good detective work here! I think the ideal solution would keep store implementation details confined to the store classes themselves. So instead of the solution here, what if we override the get_many method on the fsspec store to include the logic you have added here, and then use that method instead of multiple gets

[dgegen]

Very good point! Perhaps something along these lines?

class StorePath:
    # ...
    async def _is_concurrency_save(self):
        fs = getattr(self.store, "fs", None)
        return getattr(fs, "asynchronous", True)

    async def get_many(
        self,
        *suffixes : str,
        prototype: BufferPrototype | None = None,
        byte_range: ByteRequest | None = None,
    ):
        tasks = [
            (self / suffix).get(prototype=prototype, byte_range=byte_range) for suffix in suffixes
        ]
        if await self._is_concurrency_save():
            return await gather(*tasks)
        else:
            results = []
            for task in tasks:
                result = await task
                results.append(result)
            return results
            
class FsspecStore:
    # ...
    async def _get_many(
        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
    ) -> AsyncGenerator[tuple[str, Buffer | None], None]:
        if getattr(self.fs, "asynchronous", True):
            async for result in super()._get_many(requests=requests):
                yield result
        else:
            for key, prototype, byte_range in requests:
                value = await self.get(key, prototype, byte_range)
                yield (key, value)
                

[codecov]

Codecov Report

❌ Patch coverage is 86.84211% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 60.76%. Comparing base (a0c56fb) to head (8046f94).
⚠️ Report is 351 commits behind head on main.

Files with missing lines	Patch %	Lines
src/zarr/storage/_common.py	66.66%	2 Missing ⚠️
src/zarr/storage/_fsspec.py	75.00%	2 Missing ⚠️
src/zarr/abc/store.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3207      +/-   ##
==========================================
+ Coverage   60.68%   60.76%   +0.07%     
==========================================
  Files          78       78              
  Lines        9356     9384      +28     
==========================================
+ Hits         5678     5702      +24     
- Misses       3678     3682       +4     

Files with missing lines	Coverage Δ
src/zarr/core/array.py	69.76% <100.00%> (+0.13%)	⬆️
src/zarr/core/group.py	70.36% <100.00%> (+0.16%)	⬆️
src/zarr/abc/store.py	37.16% <87.50%> (+2.19%)	⬆️
src/zarr/storage/_common.py	69.19% <66.66%> (-0.08%)	⬇️
src/zarr/storage/_fsspec.py	68.08% <75.00%> (+0.30%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[d-v-b]

this should be a method specific to the fsspec store. the fsspec store should call it inside _get_many in order to chose which implementation to use.

[dgegen]

Done!

[d-v-b]

lets bind _join_paths([path, X]) to a variable so that we don't call the _join_paths function so many times. for example:

zarray_path = _join_paths([path, ZARRAY_JSON])
...

[dgegen]

Thanks to the feedback I received and some additional testing, I have substantially improved my initial solution. I look forward to hearing your thoughts on it! :)

PS: I'm also not sure why two of the tests failed because they don't seem to relate to any of the changes I've made.

[d-v-b]

I'm not sure I see the use for this method. If store users want fetches to happen in a specific order, then users can call get in a loop. If users only want the results of their fetches to be ordered, they can re-order the results after receiving them.

[dgegen]

I introduced the method _get_many_ordered because it is the only used case in the current implementation. Not using it leads to repetitive boilerplate code because we would always find ourselves awaiting all fetches and then sorting the results afterwards. Using _get_many_ordered encapsulates this logic, resulting in code that is less verbose, less repetitive, and easier-to-read.

Example

Compare for example

zarray_bytes, zgroup_bytes, zattrs_bytes = await store._get_many_ordered(
    [
        (_join_paths([path, ZARRAY_JSON]), default_buffer_prototype(), None),
        (_join_paths([path, ZGROUP_JSON]), default_buffer_prototype(), None),
        (_join_paths([path, ZATTRS_JSON]), default_buffer_prototype(), None),
    ]
)

with

ordered_keys = [
    (_join_paths([path, ZARRAY_JSON]), default_buffer_prototype(), None),
    (_join_paths([path, ZGROUP_JSON]), default_buffer_prototype(), None),
    (_join_paths([path, ZATTRS_JSON]), default_buffer_prototype(), None),
]

retrieved_objects = {}
async for key, value in store._get_many(ordered_keys):
    retrieved_objects[key] = value

zarray_bytes, zgroup_bytes, zattrs_bytes = tuple(retrieved_objects.get(key[0]) for key in ordered_keys)

The first block can be read much faster.

At the same time, we cannot use the original,

zarray_bytes, zgroup_bytes, zattrs_bytes = await asyncio.gather(
    store.get(_join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()),
    store.get(_join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype()),
    store.get(_join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype()),

as this would lead to deadlocks in synchronous file systems, and we of course also don't want to call .get sequentially because it reduces performance in asynchronous systems.

[d-v-b]

.get sequentially because it reduces performance in asynchronous systems.

get is async. caling it returns a coroutine. these can be scheduled together with asyncio.gather, which will preserve order. So calling get sequentially is not a performance problem.

[dgegen]

Of course, but it is, if we were to await them sequentially, e.g. if we instead used

zarray_bytes = await store.get(_join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()) 
zgroup_bytes = await store.get(_join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype())
zattrs_bytes = await store.get(_join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype())

This would assure that synchronous file systems don't run into deadlocks. But it would not be a good alternative for asynchronous systems.

[d-v-b]

i'm confused, aren't you doing all this exactly because the fsspec sftp backend is not async? so then sequential awaiting (inside the logic of _get_many) is exactly what's we expect to happen, no?

[dgegen]

Yes, but in the general case, the implementation should be asynchronous. It must be implemented, however, in such a way that we can make it synchronous by overwriting store methods in the FSSpec store, making what is generally asynchronous synchronous if store.fs.asynchronous==False.

Note also, that originally, we were not using await.gather wrapping multiple get statements, so this was not possible. And a general sequential solution is not desirable in an I/O-limited system.

[d-v-b]

what is the advantage of this new implementation? the previous implementation was extremely simple, which I think is good for an abc.

[dgegen]

The claim in the docstring is incorrect given the previous implementation.

This loop is sequential: it awaits each self.get(*req) and yields it before moving on to the next. Each request is handled one at a time, in the exact order provided. Therefore, results are always yielded in the same order as the input requests.

It is thus not fully concurrent which would be desirable in an I/O-limited system and, at least as I understand, kind of defeats the purpose of having an asynchronous _get_many method yielding results in the first place. Because if we stick to the order, we might as well await all results and simply replace the implementation of _get_many with that of _get_many_ordered, making it faster and arguable more easy to used in the asynchronous case. If we want to give the extra flexibility of not awaiting all at once, but still requesting all at the same time, the new implementation would be the right one.

[d-v-b]

the point of the default implementation is to be the simplest possible implementation of _get_many that any child class can safely support, given an implementation of get. But child classes should also be able to override this with more efficient methods where applicable, and in these cases the order of results is not guaranteed. hence the type annotation in the original method.

[dgegen]

I find this somewhat confusing, as I would have expected the standard implementation to be fully asynchronous. However, if the goal is to maximize simplicity, then having an asynchronous implementation that runs synchronously might be the way to go.

That being said, if we revert this to the original, we would only have to also remove the FsspecStore._get_many from my current solution. Unless you think we should not have a _get_many_ordered method and use the _get_many method instead and then always sort the values locally, as they could be of a different order in other implementations.

[dgegen]

Resolving this issue would be really useful for my current research project, so I am wondering what I should change specifically to get this pull request approved :)
Should I revert the _get_many method to its synchronous version, even though this would mean that the metadata would no longer be retrieved asynchronously? And should I remove the _get_many_ordered method and sort the output of the _get_many method locally whenever it is called, or can we keep it as it does make the code less verbose, less repetitive, and more readable?

[d-v-b]

Should I revert the _get_many method to its synchronous version, even though this would mean that the metadata would no longer be retrieved asynchronously? And should I remove the _get_many_ordered method and sort the output of the _get_many method locally whenever it is called, or can we keep it as it does make the code less verbose, less repetitive, and more readable?

What if we make the default implementation of _get_many async (it could either by a gather or something fancier using as_completed), and remove the _get_many_sorted routines for simplicity. I'm not convinced that the burden of sorting results post-hoc is worth the extra method at this time.

[dgegen]

Ok. Just to clarify, as you were hinting at potentially using asyncio.gather in the _get_many method: We still want _get_many be a generator, right? In other words, we think that scenarios where an asynchronous generator is indispensable are common enough to warrant its inclusion in the base class.

Assuming this is the case, I will provide a solution with _get_many remaining an asynchronous generator.

[jhamman]

@dgegen - thanks for working on this. Curious if we've ruled out a potential fix upstream in fsspec?

cc @martindurant

[martindurant]

I thought zarr wrapped all non-asynchronous fsspec FSs in the async wrapper, so that you should not have to do anything like this. Then, zarr should only ever be calling asynchronous methods. Is that wrong?

[dgegen]

I thought zarr wrapped all non-asynchronous fsspec FSs in the async wrapper, so that you should not have to do anything like this. Then, zarr should only ever be calling asynchronous methods. Is that wrong?

Zarr does indeed wrap everything in an AsyncFileSystemWrapper(fs, asynchronous=True) if fs.async_impl is false. However, this does not solve the current problem, which occurs when the file system wrapped by AsyncFileSystemWrapper does not actually support concurrent requests. This is the case for fsspec.filesystem('sftp', **storage_options), where we encountered deadlocks when trying to read different metadata files concurrently (as described in detail at the beginning of this PR).