fix: remove hard-coded scheduler from da.store - #2584
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2584 +/- ##
==========================================
- Coverage 87.66% 85.83% -1.84%
==========================================
Files 49 49
Lines 8121 8123 +2
==========================================
- Hits 7119 6972 -147
- Misses 1002 1151 +149
|
flying-sheep
left a comment
flying-sheep
left a comment
There was a problem hiding this comment.
The change you’re reverting was made intentionally here: #2183
You can see there which dask versions have the bug it circumvents. Makes sense to not hard-code it for dask versions where it’s fixed, but you should make it so that threads is still hardcoded for afflicted dask version.
|
The bug was caused by the anndata lock, not by the scheduler so it should be safe to remove the scheduler as the lock has been removed. |
|
Not only that, there’s more, see test failures. You could probably make this work by leaving the hardcoded scheduler in place only for Then when using modern dask and zarr, the user scheduler is respected. |
…arr-store-scheduler # Conflicts: # src/anndata/_io/specs/methods.py
|
Sorry, I didn't realize this code block was used for saving h5 files as well. I updated the pull request to use threaded scheduler with h5. Thanks. |
|
Added dask version gate. Thanks. |
flying-sheep
changed the title
da.store
|
OK, looks good to me now, but @ilan-gold said that the threading solution seems safer. So now that we no longer support zarr 2, maybe that’s obsolete, but let’s ask Ilan. |
| if isinstance(f, h5py.Group) or Version(version("dask")) <= Version("2025.3.0"): | ||
| da.store(elem, g, scheduler="threads") | ||
| else: | ||
| da.store(elem, g) |
There was a problem hiding this comment.
Maybe the xarray developers would have some insight? I checked their code and they do not hardcode the scheduler.
There was a problem hiding this comment.
Doesn't https://github.com/dask/dask/blob/ba5045e001d67decd6a83a0f50f666a5e3845250/dask/array/core.py#L1196 ensure safety across processes or threads?
There was a problem hiding this comment.
Also see dask/dask#12105 and zarr-developers/zarr-python#3514.
There was a problem hiding this comment.
Not certain why I only see a notification just now. To reply https://github.com/dask/dask/blob/ba5045e001d67decd6a83a0f50f666a5e3845250/dask/array/core.py#L1196, as far as I am aware this actually locks the whole array and not individual blocks. I ran into this when implementing the direct writing of sharding with da.to_zarr. Whether it is safe depends on whether shards and chunks are aligned, but I believe it should throw an error in any case when that is not the case.
There was a problem hiding this comment.
I have no experience here. @melonora could you please check? If your belief being true is all that’s needed here we’re close to wrapping this up.
There was a problem hiding this comment.
Any updated here @melonora ? Maybe @joshua-gould could just post a script a verifying one way or the other that this safety check exists i.e., the error is thrown when not aligned?
There was a problem hiding this comment.
Here's an example that shows when chunks or shards are not aligned, data is corrupted:
import dask.array as da
a = da.arange(100).reshape((10, 10)).rechunk((2, 2))
da.to_zarr(a, "test1.zarr", overwrite=True)
print((a == da.from_zarr("test1.zarr")).all().compute()) # works
da.to_zarr(a, "test2.zarr", chunks=(4, 4), overwrite=True)
print((a == da.from_zarr("test2.zarr")).all().compute()) # corrupted
da.to_zarr(a, "test3.zarr", shards=(4, 4), overwrite=True)
print((a == da.from_zarr("test3.zarr")).all().compute()) # corrupted
4 participants
test_dask_distributed_writetests writing to zarr store using dask so no new tests added