Repository: napari/napari

Issue:
Add ability to force sync slicing
# Description

The `LayerSlicer` should perform asynchronous slicing when all layers being sliced are asynchronously sliceable. In some cases, we might want to temporarily force slicing to be synchronous. For example, we might consider usage like

```python
with viewer.force_sync():
    viewer.dims.current_step = (5, 0, 0)
image = viewer.screenshot()
```

where we need to wait for slicing to finish before taking a screenshot.

This task should define public ways to allow synchronous slicing to be forced. This should include a temporary force with the context manager above, and may also include forcing sync slicing for a particular run/session of napari (so that plugins can easily force it once). We may also consider adding a setting that causes sync slicing to always be forced.

## Update

Kim: I started to add this to the initial base slicer layer PR, but ended up taking it out. Its not much code, but might be useful for reference. https://github.com/napari/napari/pull/5170/commits/56d4244112d0bc001e5f82ae31a7e5058cca1e51


Discussion:

Resolving patch:
diff --git a/napari/components/_layer_slicer.py b/napari/components/_layer_slicer.py
--- a/napari/components/_layer_slicer.py
+++ b/napari/components/_layer_slicer.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import logging
-from concurrent.futures import Executor, Future, ThreadPoolExecutor
+from concurrent.futures import Executor, Future, ThreadPoolExecutor, wait
+from contextlib import contextmanager
 from threading import RLock
 from typing import (
     Callable,
@@ -60,6 +61,8 @@ def __init__(self):
         ----------
         _executor : concurrent.futures.ThreadPoolExecutor
             manager for the slicing threading
+        _force_sync: bool
+            if true, forces slicing to execute synchronously
         _layers_to_task : dict
             task storage for cancellation logic
         _lock_layers_to_task : threading.RLock
@@ -68,9 +71,49 @@ def __init__(self):
         """
         self.events = EmitterGroup(source=self, ready=Event)
         self._executor: Executor = ThreadPoolExecutor(max_workers=1)
+        self._force_sync = False
         self._layers_to_task: Dict[Tuple[Layer], Future] = {}
         self._lock_layers_to_task = RLock()
 
+    @contextmanager
+    def force_sync(self):
+        """Context manager to temporarily force slicing to be synchronous.
+        This should only be used from the main thread.
+
+        >>> layer_slicer = _LayerSlicer()
+        >>> layer = Image(...)  # an async-ready layer
+        >>> with layer_slice.force_sync():
+        >>>     layer_slicer.slice_layers_async(layers=[layer], dims=Dims())
+        """
+        prev = self._force_sync
+        self._force_sync = True
+        try:
+            yield None
+        finally:
+            self._force_sync = prev
+
+    def wait_until_idle(self, timeout: Optional[float] = None) -> None:
+        """Wait for all slicing tasks to complete before returning.
+
+        Attributes
+        ----------
+        timeout: float or None
+            (Optional) time in seconds to wait before raising TimeoutError. If set as None,
+            there is no limit to the wait time. Defaults to None
+
+        Raises
+        ------
+        TimeoutError: when the timeout limit has been exceeded and the task is
+            not yet complete
+        """
+        futures = self._layers_to_task.values()
+        _, not_done_futures = wait(futures, timeout=timeout)
+
+        if len(not_done_futures) > 0:
+            raise TimeoutError(
+                f'Slicing {len(not_done_futures)} tasks did not complete within timeout ({timeout}s).'
+            )
+
     def slice_layers_async(
         self, layers: Iterable[Layer], dims: Dims
     ) -> Future[dict]:
@@ -101,14 +144,20 @@ def slice_layers_async(
         # term as we develop, and also in the long term if there are cases
         # when we want to perform sync slicing anyway.
         requests = {}
+        sync_layers = []
         for layer in layers:
-            if isinstance(layer, _AsyncSliceable):
+            if isinstance(layer, _AsyncSliceable) and not self._force_sync:
                 requests[layer] = layer._make_slice_request(dims)
             else:
-                layer._slice_dims(dims.point, dims.ndisplay, dims.order)
+                sync_layers.append(layer)
+
         # create task for slicing of each request/layer
         task = self._executor.submit(self._slice_layers, requests)
 
+        # submit the sync layers (purposefully placed after async submission)
+        for layer in sync_layers:
+            layer._slice_dims(dims.point, dims.ndisplay, dims.order)
+
         # store task for cancellation logic
         # this is purposefully done before adding the done callback to ensure
         # that the task is added before the done callback can be executed


Test patch:
diff --git a/napari/components/_tests/test_layer_slicer.py b/napari/components/_tests/test_layer_slicer.py
--- a/napari/components/_tests/test_layer_slicer.py
+++ b/napari/components/_tests/test_layer_slicer.py
@@ -1,3 +1,5 @@
+import time
+from concurrent.futures import Future, wait
 from dataclasses import dataclass
 from threading import RLock, current_thread, main_thread
 from typing import Tuple, Union
@@ -64,6 +66,9 @@ def _make_slice_request(self, dims) -> FakeSliceRequest:
         self.slice_count += 1
         return FakeSliceRequest(id=self.slice_count, lock=self.lock)
 
+    def _slice_dims(self, *args, **kwargs) -> None:
+        self.slice_count += 1
+
 
 class FakeSyncLayer:
     def __init__(self):
@@ -271,6 +276,94 @@ def test_slice_layers_async_task_to_layers_lock(layer_slicer):
     assert task not in layer_slicer._layers_to_task
 
 
+def test_slice_layers_exception_main_thread(layer_slicer):
+    """Exception is raised on the main thread from an error on the main
+    thread immediately when the task is created."""
+
+    class FakeAsyncLayerError(FakeAsyncLayer):
+        def _make_slice_request(self, dims) -> FakeSliceRequest:
+            raise RuntimeError('_make_slice_request')
+
+    layer = FakeAsyncLayerError()
+    with pytest.raises(RuntimeError, match='_make_slice_request'):
+        layer_slicer.slice_layers_async(layers=[layer], dims=Dims())
+
+
+def test_slice_layers_exception_subthread_on_result(layer_slicer):
+    """Exception is raised on the main thread from an error on a subthread
+    only after result is called, not upon submission of the task."""
+
+    @dataclass(frozen=True)
+    class FakeSliceRequestError(FakeSliceRequest):
+        def __call__(self) -> FakeSliceResponse:
+            assert current_thread() != main_thread()
+            raise RuntimeError('FakeSliceRequestError')
+
+    class FakeAsyncLayerError(FakeAsyncLayer):
+        def _make_slice_request(self, dims) -> FakeSliceRequestError:
+            return FakeSliceRequestError(id=self.slice_count, lock=self.lock)
+
+    layer = FakeAsyncLayerError()
+    future = layer_slicer.slice_layers_async(layers=[layer], dims=Dims())
+
+    done, _ = wait([future], timeout=5)
+    assert done, 'Test future did not complete within timeout.'
+    with pytest.raises(RuntimeError, match='FakeSliceRequestError'):
+        future.result()
+
+
+def test_wait_until_idle(layer_slicer, single_threaded_executor):
+    dims = Dims()
+    layer = FakeAsyncLayer()
+
+    with layer.lock:
+        slice_future = layer_slicer.slice_layers_async(
+            layers=[layer], dims=dims
+        )
+        _wait_until_running(slice_future)
+        # The slice task has started, but has not finished yet
+        # because we are holding the layer's slicing lock.
+        assert len(layer_slicer._layers_to_task) > 0
+        # We can't call wait_until_idle on this thread because we're
+        # holding the layer's slice lock, so submit it to be executed
+        # on another thread and also wait for it to start.
+        wait_future = single_threaded_executor.submit(
+            layer_slicer.wait_until_idle
+        )
+        _wait_until_running(wait_future)
+
+    wait_future.result()
+    assert len(layer_slicer._layers_to_task) == 0
+
+
+def _wait_until_running(future: Future):
+    while not future.running():
+        time.sleep(0.01)
+
+
+def test_layer_slicer_force_sync_on_sync_layer(layer_slicer):
+    layer = FakeSyncLayer()
+
+    with layer_slicer.force_sync():
+        assert layer_slicer._force_sync
+        future = layer_slicer.slice_layers_async(layers=[layer], dims=Dims())
+
+    assert layer.slice_count == 1
+    assert future.result() == {}
+    assert not layer_slicer._force_sync
+
+
+def test_layer_slicer_force_sync_on_async_layer(layer_slicer):
+    layer = FakeAsyncLayer()
+
+    with layer_slicer.force_sync():
+        assert layer_slicer._force_sync
+        future = layer_slicer.slice_layers_async(layers=[layer], dims=Dims())
+
+    assert layer.slice_count == 1
+    assert future.result() == {}
+
+
 def test_slice_layers_async_with_one_3d_image(layer_slicer):
     np.random.seed(0)
     data = np.random.rand(8, 7, 6)