Zarr v2 -> Zarr v3 migration across computational imaging team's frameworks

[mattersoflight]

We have decided to upgrade our image processing and analysis pipelines from zarr v2 (ome-zarr-0.4) to zarr v3 (ome-zarr-0.5) format.
The key reason for the switch is that the support for sharding and small chunks in zarr v3 datasets will enable real-time inference for smart microscopy, efficient model training through parallel/distributed data loading, and efficient visualization.

This issue is focused on identifying any disruption in processing that computational imaging team and collaborators may encounter due to the change in format, across iohub, biahub, waveorder, viscy, and shrimpy. We are centralizing the discussion in biahub repo.

iohub has been updated to support reading and writing with both versions: When we write the zarrs during the processing steps we just need to set the version flag in iohub.ngff.open_ome_zarr to v0.5.

We need to make a few key decisions:

1. Should the refactored CLI across the repos "silently" migrate from zarr v2 to zarr v3, or propagate the format? For silent migration, the CLI can read zarr v2/v3 data, and write only zarr v3 data. This is currently preferred to migrate to a better format and reduce data duplication. If the conversion can break existing analysis workflows and the fix is not straightforward, we can propagate the format.

2. Which CLI need to be updated in what order so that silent migration doesn't break existing analysis workflows? We need key users's inputs to decide if migration to zarr v3 will break their data analysis pipelines.

@Soorya19Pradeep, @tayllatheodoro, @gav-sturm can you please add a sub-issue to clarify which CLI are you currently dependent on that could break with migration to zarr v3? We can use the sub-issue to resolve the problems.

[ieivanov]

I think the change to v0.5 is a no-brainer, when it comes to implementation, the devil may be in the details. We won't know until we try, I think we can do these tests in parallel and surface any issues that we can resolve as a group. My preference would be for a "silent" transition - our pipelines would read zarr v2/v3 and write zarr v3 only. We can create fallback cases (e.g. write zarr v2 if input was v2, if needed).

I checked and it looks like biahub, waveorder, and viscy already depend on iohub>0.3. shrimpy does not depend on iohub - it uses acquire-zarr to write data. Note that shrimpy writes zarr v3 by default, so new data we generate will be coming in this format (I can check if there is a way to write zarr v2, but I think that would be a step in the wrong direction). I've confirmed zarr stores written with shrimpy can be read by iohub: czbiohub-sf/shrimPy#173

iohub v0.3 is still in alpha phase: https://github.com/czbiohub-sf/iohub/milestone/3. @srivarra what is needed to release a stable v0.3? Which of the issues on this milestone are critical, and which can be pushed to a v3.1 milestone? We've been using v0.3 alpha versions since last September, and I think we've already resolved a lot of problem - I don't have any complains with the current codebase.

It's great that @talonchandler already worked on transitioning waveorder to zarr v3 and surfaced a get gotchas: mehta-lab/waveorder#519

@srivarra @tayllatheodoro can you prototype the transition of biahub to zarr v3 using the deskew function? The input should be a zarr v3 raw/skewed dataset and the output should be a zarr v3 deskewed dataset. You can get the zarr v3 input by either using biahub convert to convert an old dataset, say of the argolight target, or I can share with you a new dataset acquired with the shrimpy v2 acquisition engine. Let us know what issues come up.

@srivarra @edyoshikun can you do a similar prototype for viscy? Can you confirm if the dataloaders can work with zarr v3 dataset? Can you test if virtual staining can output zarr v3 datasets? Again, surface any issues that come up.

Previous discussion on the transition to zarr v3 have centered around picking the right chunk / shard size such that we can enable easy data loading for visualization and processing, reduce the total number of files on the hpc, and preserve the ability to do parallel processing across timepoints and channels. Currently, in zarr v2, or chunk size is (1, 1, Z, Y, X), limited to 500 MB if needed (i.e. we may use chunk_z < Z). I think our consensus for chunks and shards in zarr v3 is:

• chunk size of ~ (1, 1, 16, 256, 256) - should be on the order of 1-5 MB and maybe a bit smaller than the receptive field of our average model
• shard size of (1, 1, Z, Y, X) during processing - not sharding along T and C will allow us to process these dimensions in parallel, as before
• shard size of ~(8, 1, Z, Y, X) during assembly at the end of reconstruction; shard should be ~1 GB each (chunks within these shards can still be accessed independently, without loading the entire shard).

[srivarra]

what is needed to release a stable v0.3?

Looking through the issues for this milestone, I think everything can be pushed back to v0.3.1. The only thing I'd like to resolve is: czbiohub-sf/iohub#374 (do we want to add the shard size here for the conversion step?) the issue of chunking across a dimension where the size is a prime number.

Added a bit of iohub related discussion over here: czbiohub-sf/iohub#375

[edyoshikun]

mehta-lab/VisCy#372
@alxndrkalinin @srivarra I'm starting a PR here. The easy first case was just to parametrize our fixtures for building an HCS plate, but I dont think this covers all the utils in the repo. I am 80% confident this covers all our usecases. I did a quick search on places where we call zarr directly and it only happened in the hcsdataloader when we prepare_data() for cachine, so I am confident we are good to go. We are leveraging iohub as our reader and we are letting it handle that with Tensostore in the ImageArray.

[Soorya19Pradeep]

Thank you for the PR @mattersoflight .
I have been trying out biahub CLIs on zarr v3 inputs, including biahub register from #198.
I understand that biahub CLIs dont work with zarr v3 yet. I was under the impression that biahub concatenate works with zarr v3 from the main branch. I didn't get that working either.

[talonchandler]

I'm in favor of a "silent" transition. No concerns from the waveorder side.

[gav-sturm]

All good to move ahead with silent transition from the OPS pipeline. We currently use biahub for channel registration and this exists as a separate env disconnected from all other pipeline steps. Also, we use Iohub Version: 0.2.1 for the zarr-v2 portions of the pipeline (registration, reconstruction etc ) and Iohub Version: 0.3.0a2 for zarr-v3 portions (all downstream tasks like feature extraction, model inference)... we eventually plan for the entire pipeline to generate zarr-v3 (next couple months)