Repository: napari/napari

Issue:
Async slicing: dask cache error when multiple layers share data
@andy-sweet As mentioned directly already, I've started testing this new async branch with our OME-Zarr data and ran into some issues with multi-channel data I'll report here in a bit more detail:

I'm using a 3 channel OME-Zarr dataset like this little test set here: https://zenodo.org/record/7120354
I installed this branch and the `napari-ome-zarr` plugin v0.5.2. I launch napari, then set the `_force_sync` to False from the napari console and drag & drop the Zarr plate into napari.

The first channel gets loaded correctly and can be browsed in the awesome async mode. But it then hits this issue when trying to load the additional channels in the OME-Zarr file (=> they become additional layers in napari):
```
~/mambaforge/envs/napari-dev-default-ome/lib/python3.10/site-packages/dask/cache.py:56, in Cache._posttask

     55 def _posttask(self, key, value, dsk, state, id):
---> 56     duration = default_timer() - self.starttimes[key]
        key = ('from-zarr-0424ef0e6522743e819a0000fe7dc515', 1, 0, 0, 0)
        default_timer = <built-in function perf_counter>
        self.starttimes = {}
        self = <dask.cache.Cache object at 0x1036dff70>
     57     deps = state["dependencies"][key]
     58     if deps:

KeyError: ('from-zarr-0424ef0e6522743e819a0000fe7dc515', 1, 0, 0, 0)
```
In the full stacktrace, it's downstream of self._load_slice(slice_data).

Sometimes, it loads 2 channels and fails for the third, sometimes it loads just one channel and fails for the second.

_Originally posted by @jluethi in https://github.com/napari/napari/issues/5377#issuecomment-1439869471_
            


Discussion:
From https://github.com/napari/napari/pull/5377#issuecomment-1443958307

> Just to follow up on the dask KeyError, I think this is related: [dask/distributed#2372](https://github.com/dask/distributed/issues/2372)
To clarify the above issue can be caused by the following steps when 2 layers shared the same underlying dask array:

1. Create an image from a dask array on the main thread.
2. Slice the image in (1) off the main thread, executing dask's `Cache._pretask` but not `Cache._posttask` yet (i.e. waiting for data to compute).
3. Create another image from the same dask array on the main thread, completing the dask computation and executing `Dask._finish` before (2) has reached `Cache._posttask`.
4. Now `Cache._posttask` for (2) executes, but there's a key error since (3) cleared `Cache.starttimes`.

I tried to write a test to reproduce this, but it's a little tricky in the context of napari since we need to effectively pause step (2) according to the dask callbacks.

Resolving patch:
diff --git a/napari/layers/base/base.py b/napari/layers/base/base.py
--- a/napari/layers/base/base.py
+++ b/napari/layers/base/base.py
@@ -32,6 +32,7 @@
     get_extent_world,
 )
 from napari.layers.utils.plane import ClippingPlane, ClippingPlaneList
+from napari.settings import get_settings
 from napari.utils._dask_utils import configure_dask
 from napari.utils._magicgui import (
     add_layer_to_viewer,
@@ -1125,7 +1126,7 @@ def _slice_dims(
         slice_input = self._make_slice_input(point, ndisplay, order)
         if force or (self._slice_input != slice_input):
             self._slice_input = slice_input
-            self.refresh()
+            self._refresh_sync()
 
     def _make_slice_input(
         self, point=None, ndisplay=2, order=None
@@ -1335,6 +1336,15 @@ def _set_highlight(self, force=False):
     def refresh(self, event=None):
         """Refresh all layer data based on current view slice."""
         logger.debug('Layer.refresh: %s', self)
+        # If async is enabled then emit an event that the viewer should handle.
+        if get_settings().experimental.async_:
+            self.events.reload(layer=self)
+        # Otherwise, slice immediately on the calling thread.
+        else:
+            self._refresh_sync()
+
+    def _refresh_sync(self, event=None):
+        logger.debug('Layer._refresh_sync: %s', self)
         if self.visible:
             self.set_view_slice()
             self.events.set_data()
@@ -1745,7 +1755,7 @@ def _update_draw(
             ):
                 self._data_level = level
                 self.corner_pixels = corners
-                self.events.reload(Event('reload', layer=self))
+                self.refresh()
         else:
             # The stored corner_pixels attribute must contain valid indices.
             corners = np.zeros((2, self.ndim), dtype=int)


Test patch: