ensuring one shard per worker for sharded zarr arrays. - #71
Open
yuriyzubov wants to merge 6 commits into
Open
Conversation
Baseline case: block size matching chunk size on an unsharded array should pass validation.
Proves the gap: validate_processing_block_size only checks chunk size, so a chunk-sized block on a sharded output passes validation even though a shard, not the chunk, is the real atomic write unit for zarr v3.
Shard-sized block on a sharded array must not be rejected: this is the safe configuration each worker owning a whole shard.
Extracts the compare/log/raise logic shared by the chunk-size and shard-size checks. Not wired in yet.
A zarr v3 shard, not the chunk, is the atomic write unit: writing part of a shard reads, merges, and rewrites the whole shard, so a block smaller than the shard risks the same race a chunk-size check alone can't catch.
Was passing align_chunksize (the zarr chunk size) as the dask blocksize, so when the output was sharded, multiple chunk-sized blocks could land in the same shard with no synchronization between workers. Use the shard size instead when the output is sharded, so each worker owns a whole shard.
Collaborator
|
Similar change needs to be done for distributed_invert_displacement_vector_field - This code actually has not been used in a very long time because the pipeline only computes the deformation field |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes a data-loss race condition in the local alignment pipeline's transform-application step when writing to a sharded zarr v3 output array.
distributed_apply_transformwas being called withalign_chunksize(the zarr chunk size) as its dask blocksize. For an unsharded output that's fine -- one worker per chunk-file, no overlap. But once sharding groups several chunks into one shard, a shard (not the chunk) becomes the real atomic write unit: writing part of a shard requires zarr to read the whole shard, merge in the new chunk, and rewrite the whole shard (ShardingCodec._encode_partial_single). With multiple chunk-sized dask blocks landing in the same shard and no synchronization between workers, one worker's write can silently overwrite another's already-written chunk.validate_processing_block_sizeexists specifically to catch "block too small" configurations like this, but it only ever checked the output's chunk size, not its shard size, so this exact configuration passed validation unflagged.Changes
bigstream/distutils.py:validate_processing_block_sizenow also checks the output array's shard size (when sharded), not just its chunk size. Extracted the shared compare/log/raise logic into_check_storage_unit_sizeto avoid duplicating it between the chunk and shard checks.bigstream/tools/main_local_align_pipeline.py:_align_local_datanow passes a shard-sized block todistributed_apply_transform(getattr(align, 'shards', None) or align_chunk_size) instead of the raw chunk-sizedalign_chunksize, so each worker owns a whole shard and no two workers can ever write into the same one.tests/test_distutils.py(new): unit tests forvalidate_processing_block_sizecovering the unsharded case, the unsafe chunk-sized-block-on-sharded-array case, and the safe shard-sized-block case.Test plan
tests/test_distutils.py::test_unsharded_block-- unsharded array, block == chunk size, no error (existing behavior preserved)tests/test_distutils.py::test_processing_block_is_chunksize-- sharded array, chunk-sized block, now correctly raisesValueErrortests/test_distutils.py::test_processing_block_is_shardsize-- sharded array, shard-sized block, no error