[v3] SyncError: Calling sync() from within a running loop

[jhamman]

Zarr version

'3.0.0a2.dev5+g8446607e.d20240827'

Numcodecs version

0.13.0

Python Version

3.11

Operating System

Mac

Installation

pip

Description

Indexing a zarr array with another zarr array causes our sync method to be called inside another sync call.

Steps to reproduce

a = np.arange(10)
za = zarr.array(a, chunks=2)
ix = [False, True, False, True, False, True, False, True, False, True]

zix = zarr.array(ix, chunks=2)
za = zarr.array(a, chunks=2)
za.oindex[zix]

Produces this traceback 👇

---------------------------------------------------------------------------
SyncError Traceback (most recent call last)
Cell In[102], line 7
 5 zix = zarr.array(ix, chunks=2)
 6 za = zarr.array(a, chunks=2)
----> 7 za.oindex[zix] # will not load all zix into memory

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/indexing.py:845, in OIndex.__getitem__(self, selection)
 843 new_selection = ensure_tuple(new_selection)
 844 new_selection = replace_lists(new_selection)
--> 845 return self.array.get_orthogonal_selection(
 846 cast(OrthogonalSelection, new_selection), fields=fields
 847 )

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/array.py:1338, in Array.get_orthogonal_selection(self, selection, out, fields, prototype)
 1336 prototype = default_buffer_prototype()
 1337 indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
-> 1338 return sync(
 1339 self._async_array._get_selection(
 1340 indexer=indexer, out=out, fields=fields, prototype=prototype
 1341 )
 1342 )

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/sync.py:91, in sync(coro, loop, timeout)
 88 return_result = next(iter(finished)).result()
 90 if isinstance(return_result, BaseException):
---> 91 raise return_result
 92 else:
 93 return return_result

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/sync.py:50, in _runner(coro)
 45 """
 46 Await a coroutine and return the result of running it. If awaiting the coroutine raises an
 47 exception, the exception will be returned.
 48 """
 49 try:
---> 50 return await coro
 51 except Exception as ex:
 52 return ex

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/array.py:467, in AsyncArray._get_selection(self, indexer, prototype, out, fields)
 458 out_buffer = prototype.nd_buffer.create(
 459 shape=indexer.shape,
 460 dtype=out_dtype,
 461 order=self.order,
 462 fill_value=self.metadata.fill_value,
 463 )
 464 if product(indexer.shape) > 0:
 465 # reading chunks and decoding them
 466 await self.codec_pipeline.read(
--> 467 [
 468 (
 469 self.store_path / self.metadata.encode_chunk_key(chunk_coords),
 470 self.metadata.get_chunk_spec(chunk_coords, self.order, prototype=prototype),
 471 chunk_selection,
 472 out_selection,
 473 )
 474 for chunk_coords, chunk_selection, out_selection in indexer
 475 ],
 476 out_buffer,
 477 drop_axes=indexer.drop_axes,
 478 )
 479 return out_buffer.as_ndarray_like()

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/array.py:467, in <listcomp>(.0)
 458 out_buffer = prototype.nd_buffer.create(
 459 shape=indexer.shape,
 460 dtype=out_dtype,
 461 order=self.order,
 462 fill_value=self.metadata.fill_value,
 463 )
 464 if product(indexer.shape) > 0:
 465 # reading chunks and decoding them
 466 await self.codec_pipeline.read(
--> 467 [
 468 (
 469 self.store_path / self.metadata.encode_chunk_key(chunk_coords),
 470 self.metadata.get_chunk_spec(chunk_coords, self.order, prototype=prototype),
 471 chunk_selection,
 472 out_selection,
 473 )
 474 for chunk_coords, chunk_selection, out_selection in indexer
 475 ],
 476 out_buffer,
 477 drop_axes=indexer.drop_axes,
 478 )
 479 return out_buffer.as_ndarray_like()

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/indexing.py:813, in OrthogonalIndexer.__iter__(self)
 812 def __iter__(self) -> Iterator[ChunkProjection]:
--> 813 for dim_projections in itertools.product(*self.dim_indexers):
 814 chunk_coords = tuple(p.dim_chunk_ix for p in dim_projections)
 815 chunk_selection: tuple[Selector, ...] | npt.NDArray[Any] = tuple(
 816 p.dim_chunk_sel for p in dim_projections
 817 )

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/indexing.py:542, in BoolArrayDimIndexer.__iter__(self)
 539 for dim_chunk_ix in self.dim_chunk_ixs:
 540 # find region in chunk
 541 dim_offset = dim_chunk_ix * self.dim_chunk_len
--> 542 dim_chunk_sel = self.dim_sel[dim_offset : dim_offset + self.dim_chunk_len]
 544 # pad out if final chunk
 545 if dim_chunk_sel.shape[0] < self.dim_chunk_len:

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/array.py:899, in Array.__getitem__(self, selection)
 897 return self.vindex[cast(CoordinateSelection | MaskSelection, selection)]
 898 elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
--> 899 return self.get_orthogonal_selection(pure_selection, fields=fields)
 900 else:
 901 return self.get_basic_selection(cast(BasicSelection, pure_selection), fields=fields)

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/array.py:1338, in Array.get_orthogonal_selection(self, selection, out, fields, prototype)
 1336 prototype = default_buffer_prototype()
 1337 indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
-> 1338 return sync(
 1339 self._async_array._get_selection(
 1340 indexer=indexer, out=out, fields=fields, prototype=prototype
 1341 )
 1342 )

File ~/Library/CloudStorage/Dropbox/src/zarr-python/src/zarr/core/sync.py:78, in sync(coro, loop, timeout)
 76 loop0 = asyncio.events.get_running_loop()
 77 if loop0 is loop:
---> 78 raise SyncError("Calling sync() from within a running loop")
 79 except RuntimeError:
 80 pass

SyncError: Calling sync() from within a running loop

Additional output

No response

[jhamman]

@normanrz - curious if you have ideas for how to address this?

[normanrz]

Seems like a tricky problem. We could solve this for specific use cases. For example, in the array.oindex[sel] case, we could turn the selector into a numpy array first before handing it over to the sync-ed selection.
Not sure if there is a general solution.

cc @martindurant

[martindurant]

asyncio is not re-entrant, but nest-asyncio does exist (if "archived"). It is not a solution for us, because it would patch system libraries that users might be depending on.

Conversely, the dask way to deal with this is for a (non-async) function to detect if it is being called within an active event loop or not, so that when it calls async code, it can either pass back the coroutine to the caller, or sync() it.

The fsspec way to deal with this, is to have a separate sync API: every async method has another function which calls sync() on it, so that the calling code must decide which variant to call. This is how you might end up with a sync-minded fs instance and an async one, which are otherwise identical (they used to have different init behaviour, but not any more).

[martindurant]

A simplified version of what dask does is in streamz. I haven't actually checked that dask still works this way, now that I think about it!