Support functools.partial with lithops (zarr-developers#899)

* Run doctests via pytest
* Support use of functools.partial with Lithops.
* Use mode "forkserver" with ProcessPoolExecutor if the
  default mode is "fork." This can be removed once our minimum python
  version is 3.14, as "fork" will no longer be a default mode.

Author: chuckwondo

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,9 +1,14 @@
-<!-- Feel free to remove check-list items aren't relevant to your change -->
+# What I did
+<!-- Please describe what you did and why -->
+
+Acceptance criteria:
+<!-- Feel free to remove check-list items that aren't relevant to your change -->
 
 - [ ] Closes #xxxx
 - [ ] Tests added
 - [ ] Tests passing
+- [ ] No test coverage regression
 - [ ] Full type hint coverage
-- [ ] Changes are documented in `docs/releases.rst`
-- [ ] New functions/methods are listed in `api.rst`
+- [ ] Changes are documented in `docs/releases.md`
+- [ ] New functions/methods are listed in an appropriate `*.md` file under `docs/api`
 - [ ] New functionality has documentation

.github/workflows/minimum-versions.yml

@@ -18,7 +18,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  test-upstream:
+  test-minimum-versions:
     name: minimum-versions-build
     if: |
       always()

docs/releases.md

@@ -18,10 +18,15 @@
   By [Tom Nicholas](https://github.com/TomNicholas).
 - Raise clearer error when kerchunk references have malformed codec specifications.
   ([#864](https://github.com/zarr-developers/VirtualiZarr/pull/864)).
-- Fixed warnings caused by outdated imports from `obspec_utils` ([#863](https://github.com/zarr-developers/VirtualiZarr/pull/863)).
+- Fix warnings caused by outdated imports from `obspec_utils` ([#863](https://github.com/zarr-developers/VirtualiZarr/pull/863)).
   By [Tom Nicholas](https://github.com/TomNicholas).
 - Allow `ZarrParser` to work from inside a running event loop (e.g. inside a Jupyter Notebook) ([#900](https://github.com/zarr-developers/VirtualiZarr/pull/900))
   By [Julius Busecke](https://github.com/jbusecke).
+- Fix Lithops executor to allow use of `functools.partial`, and update
+  `get_executor` function to ensure `ProcessPoolExecutor` uses `"forkserver"`
+  mode on platforms that default to `"fork"`
+  ([#899](https://github.com/zarr-developers/VirtualiZarr/pull/899)). By [Chuck
+  Daniels](https://github.com/chuckwondo).
 
 ### Documentation
 
@@ -43,10 +48,10 @@ This release moves the `ObjectStoreRegistry` to a separate package `obspec_utils
   ([#844](https://github.com/zarr-developers/VirtualiZarr/pull/844)).
   By [Max Jones](https://github.com/maxrjones).
 
-    - `ObjectStoreRegistry` has moved from `virtualizarr.registry` to `obspec_utils.registry`. The old import path still works but emits a `DeprecationWarning` and will be removed in a future release.
-    - `ObstoreReader` has been removed from `virtualizarr.utils`. This should not break user's code, as it was not part of the public/documented API. See [obspec_utils](https://obspec-utils.readthedocs.io/en/latest/api/obspec/) for public file handlers.
-    - Added `obspec_utils>=0.7.0` as a required dependency. This package provides the `ObjectStoreRegistry` that was previously part of VirtualiZarr.
-    - Minimum required version of `obstore` is now `0.7.0` (previously `0.5.1`). This was the first release to implement obspec protocols.
+  - `ObjectStoreRegistry` has moved from `virtualizarr.registry` to `obspec_utils.registry`. The old import path still works but emits a `DeprecationWarning` and will be removed in a future release.
+  - `ObstoreReader` has been removed from `virtualizarr.utils`. This should not break user's code, as it was not part of the public/documented API. See [obspec_utils](https://obspec-utils.readthedocs.io/en/latest/api/obspec/) for public file handlers.
+  - Added `obspec_utils>=0.7.0` as a required dependency. This package provides the `ObjectStoreRegistry` that was previously part of VirtualiZarr.
+  - Minimum required version of `obstore` is now `0.7.0` (previously `0.5.1`). This was the first release to implement obspec protocols.
 
 ## v2.3.0 (20th January 2026)
 

pyproject.toml

@@ -301,6 +301,7 @@ exclude_lines = [
 ]
 
 [tool.pytest.ini_options]
+addopts = ["--doctest-modules", "--ignore=examples"]
 # See https://pytest-asyncio.readthedocs.io/en/latest/concepts.html#asyncio-event-loops
 # Explicitly set asyncio_default_fixture_loop_scope to eliminate the following warning:
 #

virtualizarr/accessor.py

@@ -249,7 +249,7 @@ def rename_paths(
         ...     filename = Path(old_local_path).name
         ...     return str(new_s3_bucket_url / filename)
         >>>
-        >>> ds.vz.rename_paths(local_to_s3_url)
+        >>> ds.vz.rename_paths(local_to_s3_url)  # doctest: +SKIP
         """
 
         new_ds = self.ds.copy()

virtualizarr/manifests/array.py

@@ -266,7 +266,7 @@ def rename_paths(
         ...     filename = Path(old_local_path).name
         ...     return str(new_s3_bucket_url / filename)
         >>>
-        >>> marr.rename_paths(local_to_s3_url)
+        >>> marr.rename_paths(local_to_s3_url)  # doctest: +SKIP
         """
         renamed_manifest = self.manifest.rename_paths(new)
         return ManifestArray(metadata=self.metadata, chunkmanifest=renamed_manifest)

virtualizarr/manifests/manifest.py

@@ -459,7 +459,7 @@ def rename_paths(
         ...     filename = Path(old_local_path).name
         ...     return str(new_s3_bucket_url / filename)
         >>>
-        >>> manifest.rename_paths(local_to_s3_url)
+        >>> manifest.rename_paths(local_to_s3_url)  # doctest: +SKIP
         """
         if isinstance(new, str):
             renamed_paths = np.full_like(self._paths, fill_value=new)

virtualizarr/parallel.py

@@ -1,7 +1,18 @@
 import inspect
+import multiprocessing as mp
 import warnings
-from concurrent.futures import Executor, Future
-from typing import Any, Callable, Iterable, Iterator, Literal, TypeVar
+from concurrent.futures import Executor, Future, ProcessPoolExecutor
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    Iterator,
+    Literal,
+    ParamSpec,
+    TypeVar,
+)
 
 __all__ = [
     "SerialExecutor",
@@ -15,28 +26,38 @@
 # TODO lithops should just not require a special wrapper class, see https://github.com/lithops-cloud/lithops/issues/1427
 
 
+P = ParamSpec("P")
 # Type variable for return type
 T = TypeVar("T")
 
 
 def get_executor(
-    parallel: Literal["dask", "lithops"] | type[Executor] | Literal[False],
-) -> type[Executor]:
-    """Get an executor that follows the concurrent.futures.Executor ABC API."""
+    parallel: Literal["dask", "lithops", False] | type[Executor],
+) -> Callable[..., Executor]:
+    """Get a callable with a return type that follows the concurrent.futures.Executor ABC API."""
 
     if parallel == "dask":
         return DaskDelayedExecutor
-    elif parallel == "lithops":
+    if parallel == "lithops":
         return LithopsEagerFunctionExecutor
-    elif parallel is False:
+    if parallel is False:
         return SerialExecutor
-    elif inspect.isclass(parallel) and issubclass(parallel, Executor):
+    if parallel is ProcessPoolExecutor:
+        # TODO Once we drop support for python <3.14, we can remove this context
+        # dance because from 3.14 onward, POSIX defaults to "forkserver" rather
+        # than "fork".
+        method = mp.get_context().get_start_method()
+        context = mp.get_context("forkserver" if method == "fork" else method)
+        return partial(ProcessPoolExecutor, mp_context=context)
+    if inspect.isclass(parallel) and issubclass(parallel, Executor):
         return parallel
-    else:
-        raise ValueError(
-            f"Unrecognized argument to ``parallel``: {parallel}"
-            "Please supply either ``'dask'``, ``'lithops'``, ``False``, or a concrete subclass of ``concurrent.futures.Executor``."
-        )
+
+    raise ValueError(
+        f"Invalid value for `parallel`: {parallel}.  Please supply "
+        "either the string 'dask' or 'lithops', or a concrete subclass of "
+        "concurrent.futures.Executor.  To obtain a serial executor, specify "
+        "the boolean value `False`."
+    )
 
 
 class SerialExecutor(Executor):
@@ -111,18 +132,6 @@ def map(
         """
         return map(fn, *iterables)
 
-    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
-        """
-        Shutdown the executor.
-
-        Parameters
-        ----------
-        wait
-            Whether to wait for pending futures (always True for serial executor)
-        """
-        # In a serial executor, shutdown is a no-op
-        pass
-
 
 class DaskDelayedExecutor(Executor):
     """
@@ -211,31 +220,48 @@ def map(
         # Compute all tasks
         return iter(dask.compute(*delayed_tasks))
 
-    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
-        """
-        Shutdown the executor
-
-        Parameters
-        ----------
-        wait
-            Whether to wait for pending futures (always True for serial executor))
-        """
-        # For Dask.delayed, shutdown is essentially a no-op
-        pass
-
 
 class LithopsEagerFunctionExecutor(Executor):
     """
-    Lithops-based function executor which follows the [concurrent.futures.Executor][] API.
+    Lithops-based function executor that follows the [concurrent.futures.Executor][] API.
 
-    Only required because lithops doesn't follow the [concurrent.futures.Executor][] API, see https://github.com/lithops-cloud/lithops/issues/1427.
+    Only required because lithops doesn't follow the [concurrent.futures.Executor][] API.
+    See https://github.com/lithops-cloud/lithops/issues/1427.
     """
 
+    class compatible_callable(Generic[P, T]):
+        """Wraps a callable to make it fully compatible with Lithops.
+
+        This wrapper deals with 2 oddities in Lithops:
+
+        1. Use of `functools.partial`, which Lithops fails to recognize as being
+           callable.  This is likely due to the builtin `partial` class using
+           slots, which causes Lithops to not recognize the `__call__` method as
+           a method.  See https://github.com/lithops-cloud/lithops/issues/1428.
+        2. Use of generic function wrappers that define generic `args` and
+           `kwargs` parameters.  In this case, because of the way Lithops
+           inspects function signatures to determine how to pass arguments, it
+           does not properly "spread" arguments as normally expected.  Instead,
+           it collects all positional arguments and associates them with the
+           keyword argument `"args"`, which is utterly unhelpful.
+        """
+
+        def __init__(self, f: Callable[P, T]):
+            self.f = f
+
+        def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+            if not args and "args" in kwargs:
+                # Lithops collected all positional args into an "args" kwarg,
+                # so we're undoing that nonsense here.
+                args = kwargs.pop("args", ())  # type: ignore
+
+            return self.f(*args, **kwargs)
+
     def __init__(self, **kwargs) -> None:
         import lithops  # type: ignore[import-untyped]
 
         # Create Lithops client with optional configuration
-        self.lithops_client = lithops.FunctionExecutor(**kwargs)
+        self.lithops_client = lithops.FunctionExecutor(**kwargs).__enter__()
 
         # Track submitted futures
         self._futures: list[Future] = []
@@ -263,7 +289,11 @@ def submit(self, fn: Callable[..., T], /, *args: Any, **kwargs: Any) -> Future[T
 
         try:
             # Submit to Lithops
-            lithops_future = self.lithops_client.call_async(fn, *args, **kwargs)
+            lithops_future = self.lithops_client.call_async(
+                LithopsEagerFunctionExecutor.compatible_callable(fn),
+                *args,
+                **kwargs,
+            )
 
             # Add a callback to set the result or exception
             def _on_done(lithops_result):
@@ -294,7 +324,8 @@ def map(
         """
         Apply a function to an iterable using lithops.
 
-        Only needed because [lithops.executors.FunctionExecutor.map][lithops.executors.FunctionExecutor.map] returns futures, unlike [concurrent.futures.Executor.map][].
+        Only needed because [lithops.executors.FunctionExecutor.map][lithops.executors.FunctionExecutor.map]
+        returns futures, unlike [concurrent.futures.Executor.map][].
 
         Parameters
         ----------
@@ -309,14 +340,13 @@ def map(
         -------
         Generator of results
         """
-        import lithops  # type: ignore[import-untyped]
-
-        fexec = lithops.FunctionExecutor()
-
-        futures = fexec.map(fn, *iterables)
-        results = fexec.get_result(futures)
+        fexec = self.lithops_client
+        futures = fexec.map(
+            LithopsEagerFunctionExecutor.compatible_callable(fn),
+            list(zip(*iterables)),
+        )
 
-        return results
+        return fexec.get_result(futures)  # type: ignore
 
     def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
         """
@@ -327,5 +357,4 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
         wait
             Whether to wait for pending futures.
         """
-        # Should this call lithops .clean() method?
-        pass
+        self.lithops_client.__exit__(None, None, None)

virtualizarr/tests/test_parallel.py

@@ -0,0 +1,43 @@
+import multiprocessing as mp
+
+import pytest
+
+from virtualizarr.parallel import LithopsEagerFunctionExecutor, get_executor
+from virtualizarr.tests import requires_lithops
+
+
+@requires_lithops
+def test_lithops_executor_with_multiple_args():
+    with LithopsEagerFunctionExecutor() as exec:
+        results = exec.map(lambda x, y: x + y, (1, 2, 3), (4, 5, 6))
+
+    assert tuple(results) == (5, 7, 9)
+
+
+@requires_lithops
+def test_lithops_executor_with_partial():
+    from functools import partial
+
+    inc = partial(lambda x, y: x + y, 1)
+
+    with LithopsEagerFunctionExecutor() as exec:
+        results = exec.map(inc, (1, 2, 3))
+
+    assert tuple(results) == (2, 3, 4)
+
+
+@pytest.mark.skipif(
+    mp.get_start_method() != "fork",
+    reason="Default multiprocessing start method is not 'fork'",
+)
+def test_get_executor_process_pool_mode():
+    from concurrent.futures import ProcessPoolExecutor
+
+    executor = get_executor(ProcessPoolExecutor)()
+
+    assert isinstance(executor, ProcessPoolExecutor), "Expected a ProcessPoolExecutor"
+
+    ctx = executor._mp_context
+
+    assert ctx is not None, "Expected executor to have a multiprocessing context"
+    assert ctx.get_start_method() == "forkserver"

virtualizarr/tests/test_parsers/test_dmrpp.py

@@ -1,3 +1,4 @@
+import platform
 import textwrap
 from contextlib import nullcontext
 from pathlib import Path
@@ -486,6 +487,10 @@ def test_split_groups(hdf5_groups_file, group_path):
     assert result_tags == expected_tags
 
 
+@pytest.mark.xfail(
+    platform.system() == "Linux",
+    reason="See https://github.com/zarr-developers/VirtualiZarr/issues/904.",
+)
 @pytest.mark.parametrize(
     "group,warns",
     [