Skip to content

ensuring one shard per worker for sharded zarr arrays. - #71

Open
yuriyzubov wants to merge 6 commits into
JaneliaSciComp:ome-devfrom
yuriyzubov:safe-write
Open

ensuring one shard per worker for sharded zarr arrays. #71
yuriyzubov wants to merge 6 commits into
JaneliaSciComp:ome-devfrom
yuriyzubov:safe-write

Conversation

@yuriyzubov

Copy link
Copy Markdown
Contributor

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_transform was being called with align_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_size exists 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_size now 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_size to avoid duplicating it between the chunk and shard checks.
  • bigstream/tools/main_local_align_pipeline.py: _align_local_data now passes a shard-sized block to distributed_apply_transform (getattr(align, 'shards', None) or align_chunk_size) instead of the raw chunk-sized align_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 for validate_processing_block_size covering 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 raises ValueError
  • tests/test_distutils.py::test_processing_block_is_shardsize -- sharded array, shard-sized block, no error

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.
@cgoina

cgoina commented Aug 8, 2026

Copy link
Copy Markdown
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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants