Writing multiscale element to disk fails when configuring Dask to uses "processes"

[ArneDefauw]

When configuring dask to use "processes" instead of "threads", writing a multiscale element to disk fails. I get the error message:

File ~/VIB/harpy/.venv_harpy/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1177, in SpatialData.write(self, file_path, overwrite, consolidate_metadata, update_sdata_path, sdata_formats)
 1174 store.close()
 1176 for element_type, element_name, element in self.gen_elements():
-> 1177 self._write_element(
 1178 element=element,
 1179 zarr_container_path=file_path,
...
--> 120 if group.metadata.zarr_format == 3 and len(multiscales := group.metadata.attributes["ome"]["multiscales"]) != 1:
 121 len_scales = len(multiscales)
 122 raise ValueError(f"The length of multiscales metadata should be 1, found the length to be {len_scales}")

KeyError: 'ome'

Minimal example to reproduce:

import os
from spatialdata.models import Image2DModel

from spatialdata.datasets import blobs

import dask

with dask.config.set(scheduler='processes'):

 sdata = blobs()
 sdata=sdata.subset( element_names=[ "blobs_image" ] )
 sdata.write(os.path.join( os.environ.get("TMPDIR"), "sdata.zarr" ), overwrite=True ) # this works
 print( "done" )

 sdata = blobs()
 sdata=sdata.subset( element_names=[ "blobs_multiscale_image" ] )
 sdata.write(os.path.join( os.environ.get("TMPDIR"), "sdata.zarr" ), overwrite=True )
# this fails with
#--> 120 if group.metadata.zarr_format == 3 and len(multiscales := group.metadata.attributes["ome"]["multiscales"]) != 1:
# 121 len_scales = len(multiscales)
# 122 raise ValueError(f"The length of multiscales metadata should be 1, found the length to be {len_scales}")

#KeyError: 'ome'

with dask.config.set(scheduler='threads'):

 sdata = blobs()
 sdata=sdata.subset( element_names=[ "blobs_image" ] )
 sdata.write(os.path.join( os.environ.get("TMPDIR"), "sdata.zarr" ), overwrite=True )
 print( "done" )

 sdata = blobs()
 sdata=sdata.subset( element_names=[ "blobs_multiscale_image" ] )
 sdata.write(os.path.join( os.environ.get("TMPDIR"), "sdata.zarr" ), overwrite=True )
 print("done")

I tested this on MacOS and CentOS, and got the same error message.

Working with spatialdata==0.6.0, spatial_image==1.2.3 and multiscale_spatial_image==2.0.3.

I also noticed, for images and labes, that the layer in the dask graph that starts with "from-zarr-" is no longer materialized when the spatialdata object is backed by a zarr store (zarr>=3). Probably this is unrelated to the current issue, but I found this a bit strange, since this was not the case in earlier versions of spatialdata(<0.5.0), see https://github.com/saeyslab/harpy/blob/609d639c7578a4c64c3eede4c974d4e90f982910/src/harpy/_tests/test_image/test_manager.py#L47

[ArneDefauw]

Edit, working with a dask client, and using 1 worker and 1 thread also results in the same error

cluster = LocalCluster(
 n_workers=1,
 threads_per_worker=1,
 memory_limit="32GB",
 local_directory=os.environ.get( "TMPDIR" ), # folder where dask will spill to disk
)
client = Client(cluster)

sdata = blobs()
sdata=sdata.subset( element_names=[ "blobs_multiscale_image" ] )
sdata.write(os.path.join( os.environ.get("TMPDIR"), "sdata.zarr" ), overwrite=True ) # KeyError: 'ome'

[LucaMarconato]

Thanks for reporting. I loop in @melonora who has been working on Dask recently wrt read/write APIs.

Internally we currently use ome-zarr-py APIs, so I think that this bugs originates upstream of spatialdata and could be reproduced without using spatialdata code. In such case it would be better to report it in ome-zarr-py.

Wouter WDYT?

[LucaMarconato]

I also noticed, for images and labes, that the layer in the dask graph that starts with "from-zarr-" is no longer materialized when the spatialdata object is backed by a zarr store (zarr>=3).

Yes, this is a change in Dask that has already been covered. Did you find any problem related to that?

[nbonine]

I am experiencing the same issue.. Any workarounds?

[LucaMarconato]

I isolated the bug in ome-zarr-py, I opened a bug report here: ome/ome-zarr-py#580. The clean solution would be to address the bug upstream.

Anyway to help find a workaround in spatialdata I have extended the MRE from @ArneDefauw in the script below

# /// script
# requires-python = ">=3.11"
# dependencies = [
# "spatialdata @ git+https://github.com/scverse/spatialdata.git@b63a1f27410879c2c252e1f6a54801aed7d9d1ce",
# "ome-zarr==0.14.0",
# "zarr==3.1.5",
# "dask[array]==2026.1.1",
# "numpy>=1.24",
# ]
# ///
#
# Reproduces https://github.com/scverse/spatialdata/issues/1024
# (KeyError: 'ome' when writing a multiscale element with the 'processes' Dask scheduler)
#
# Run with:
# uv run reproduce_1024_spatialdata.py

import tempfile
import os

import numpy as np
import dask
import dask.array as da

from spatialdata import SpatialData
from spatialdata.models import Image2DModel


def main() -> None:
 rng = np.random.default_rng(0)
 arr = da.from_array(rng.random((3, 64, 64)).astype("float32"), chunks=(3, 32, 32))

 # Single-scale image (no scale_factors) – works under both schedulers.
 single = Image2DModel.parse(arr, dims=["c", "y", "x"])

 # Multiscale image – fails under 'processes'.
 multi = Image2DModel.parse(arr, dims=["c", "y", "x"], scale_factors=[2])

 with tempfile.TemporaryDirectory() as tmp:

 # ── threads scheduler (default) ───────────────────────────────────────
 with dask.config.set(scheduler="threads"):
 sdata = SpatialData(images={"img": single})
 sdata.write(os.path.join(tmp, "single_threads.zarr"))
 print("single-scale / threads : OK")

 sdata = SpatialData(images={"img": multi})
 sdata.write(os.path.join(tmp, "multi_threads.zarr"))
 print("multiscale / threads : OK")

 # ── processes scheduler ───────────────────────────────────────────────
 with dask.config.set(scheduler="processes"):
 sdata = SpatialData(images={"img": single})
 sdata.write(os.path.join(tmp, "single_processes.zarr"), overwrite=True)
 print("single-scale / processes : OK")

 # This fails with KeyError: 'ome'
 #
 # Call chain:
 # SpatialData.write()
 # -> _write_multiscale_raster()
 # calls write_multiscale_ngff(..., compute=False) # ome-zarr-py
 # then da.compute(*dask_delayed, optimize_graph=False)
 # then overwrite_coordinate_transformations_raster()
 #
 # ome-zarr-py's write_multiscale(..., compute=False) wraps
 # write_multiscales_metadata() in a dask.delayed task. Under the
 # 'processes' scheduler that task executes in a subprocess, writes
 # group.attrs["ome"] = {...} to disk there, but the zarr.Group
 # held in the main process never reloads its in-memory GroupMetadata
 # from disk. When overwrite_coordinate_transformations_raster()
 # subsequently accesses group.metadata.attributes["ome"] it finds
 # an empty dict and raises KeyError: 'ome'.
 # See https://github.com/ome/ome-zarr-py/issues/580
 sdata = SpatialData(images={"img": multi})
 sdata.write(os.path.join(tmp, "multi_processes.zarr"), overwrite=True)
 print("multiscale / processes : OK") # never reached


if __name__ == "__main__":
 main()

Additional notes: The reason why we use compute=False when calling the APIs from ome-zarr-py can be traced back to this PR #726.

[LucaMarconato]

The problem is fixed with ome-zarr==0.16.0. After the work on sharding from Wouter we will be able to use the most recent dask versions and therefore also 0.16.0. I'll reopen the issue so we remember to remove the workaround that is added for 0.14.0.