Fix pickling errors in parallel workers (#137)

* Fix pickling errors by importing task modules in workers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Document fix for pickling errors

* Fix typing, lint, and docs checks

* Allow global worker root with ruff ignore

* Add regression test for mark import loop

* Strip local from task.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Release 0.5.2

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: tobiasraabe

docs/source/changes.md

@@ -5,12 +5,14 @@ chronological order. Releases follow [semantic versioning](https://semver.org/)
 releases are available on [PyPI](https://pypi.org/project/pytask-parallel) and
 [Anaconda.org](https://anaconda.org/conda-forge/pytask-parallel).
 
-## Unreleased
+## 0.5.2 - 2026-02-05
 
 - {pull}`129` drops support for Python 3.8 and 3.9 and adds support for Python 3.14.
 - {pull}`130` switches type checking to ty.
 - {pull}`131` updates pre-commit hooks.
 - {pull}`132` removes the tox configuration in favor of uv and just.
+- {pull}`137` fixes pickling errors in parallel workers when task modules contain
+  non-picklable globals. Fixes {issue}`136`.
 
 ## 0.5.1 - 2025-03-09
 

pyproject.toml

@@ -33,7 +33,7 @@ docs = [
     "matplotlib",
     "myst-parser",
     "nbsphinx",
-    "sphinx",
+    "sphinx<9",
     "sphinx-autobuild",
     "sphinx-click",
     "sphinx-copybutton",

requirements-dev.lock

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import os
+import sys
 import warnings
 from concurrent.futures import Executor
 from concurrent.futures import Future
@@ -19,7 +21,46 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-__all__ = ["ParallelBackend", "ParallelBackendRegistry", "WorkerType", "registry"]
+__all__ = [
+    "ParallelBackend",
+    "ParallelBackendRegistry",
+    "WorkerType",
+    "registry",
+    "set_worker_root",
+]
+
+_WORKER_ROOT: str | None = None
+
+
+def set_worker_root(path: os.PathLike[str] | str) -> None:
+    """Configure the root path for worker processes.
+
+    Spawned workers (notably on Windows) start with a clean interpreter and may not
+    inherit the parent's import path. We set both ``sys.path`` and ``PYTHONPATH`` so
+    task modules are importable by reference, which avoids pickling module globals.
+
+    """
+    root = os.fspath(path)
+    global _WORKER_ROOT  # noqa: PLW0603
+    _WORKER_ROOT = root
+    if root not in sys.path:
+        sys.path.insert(0, root)
+    # Ensure custom process backends can import task modules by reference.
+    separator = os.pathsep
+    current = os.environ.get("PYTHONPATH", "")
+    parts = [p for p in current.split(separator) if p] if current else []
+    if root not in parts:
+        parts.insert(0, root)
+        os.environ["PYTHONPATH"] = separator.join(parts)
+
+
+def _configure_worker(root: str | None) -> None:
+    """Set cwd and sys.path for worker processes."""
+    if not root:
+        return
+    os.chdir(root)
+    if root not in sys.path:
+        sys.path.insert(0, root)
 
 
 def _deserialize_and_run_with_cloudpickle(fn: bytes, kwargs: bytes) -> Any:
@@ -75,12 +116,20 @@ def _get_dask_executor(n_workers: int) -> Executor:
 
 def _get_loky_executor(n_workers: int) -> Executor:
     """Get a loky executor."""
-    return get_reusable_executor(max_workers=n_workers)
+    return get_reusable_executor(
+        max_workers=n_workers,
+        initializer=_configure_worker,
+        initargs=(_WORKER_ROOT,),
+    )
 
 
 def _get_process_pool_executor(n_workers: int) -> Executor:
     """Get a process pool executor."""
-    return _CloudpickleProcessPoolExecutor(max_workers=n_workers)
+    return _CloudpickleProcessPoolExecutor(
+        max_workers=n_workers,
+        initializer=_configure_worker,
+        initargs=(_WORKER_ROOT,),
+    )
 
 
 def _get_thread_pool_executor(n_workers: int) -> Executor:

requirements.lock

@@ -26,11 +26,14 @@
 
 from pytask_parallel.backends import WorkerType
 from pytask_parallel.backends import registry
+from pytask_parallel.backends import set_worker_root
 from pytask_parallel.typing import CarryOverPath
 from pytask_parallel.typing import is_coiled_function
 from pytask_parallel.utils import create_kwargs_for_task
 from pytask_parallel.utils import get_module
 from pytask_parallel.utils import parse_future_result
+from pytask_parallel.utils import should_pickle_module_by_value
+from pytask_parallel.utils import strip_annotation_locals
 
 if TYPE_CHECKING:
     from concurrent.futures import Future
@@ -57,6 +60,7 @@ def pytask_execute_build(session: Session) -> bool | None:  # noqa: C901, PLR091
 
     # The executor can only be created after the collection to give users the
     # possibility to inject their own executors.
+    set_worker_root(session.config["root"])
     session.config["_parallel_executor"] = registry.get_parallel_backend(
         session.config["parallel_backend"], n_workers=session.config["n_workers"]
     )
@@ -195,6 +199,7 @@ def pytask_execute_task(session: Session, task: PTask) -> Future[WrapperResult]:
     kwargs = create_kwargs_for_task(task, remote=remote)
 
     if is_coiled_function(task):
+        strip_annotation_locals(task)
         # Prevent circular import for coiled backend.
         from pytask_parallel.wrappers import (  # noqa: PLC0415
             rewrap_task_with_coiled_function,
@@ -208,7 +213,8 @@ def pytask_execute_task(session: Session, task: PTask) -> Future[WrapperResult]:
         # cloudpickle will pickle it with the function. See cloudpickle#417, pytask#373
         # and pytask#374.
         task_module = get_module(task.function, getattr(task, "path", None))
-        cloudpickle.register_pickle_by_value(task_module)
+        if should_pickle_module_by_value(task_module):
+            cloudpickle.register_pickle_by_value(task_module)
 
         return cast("Any", wrapper_func).submit(
             task=task,
@@ -221,6 +227,7 @@ def pytask_execute_task(session: Session, task: PTask) -> Future[WrapperResult]:
         )
 
     if worker_type == WorkerType.PROCESSES:
+        strip_annotation_locals(task)
         # Prevent circular import for loky backend.
         from pytask_parallel.wrappers import wrap_task_in_process  # noqa: PLC0415
 
@@ -230,7 +237,8 @@ def pytask_execute_task(session: Session, task: PTask) -> Future[WrapperResult]:
         # cloudpickle will pickle it with the function. See cloudpickle#417, pytask#373
         # and pytask#374.
         task_module = get_module(task.function, getattr(task, "path", None))
-        cloudpickle.register_pickle_by_value(task_module)
+        if should_pickle_module_by_value(task_module):
+            cloudpickle.register_pickle_by_value(task_module)
 
         return session.config["_parallel_executor"].submit(
             wrap_task_in_process,

src/pytask_parallel/backends.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import importlib.util
 import inspect
 from functools import partial
+from pathlib import Path
 from typing import TYPE_CHECKING
 from typing import Any
 
@@ -20,7 +22,6 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
     from concurrent.futures import Future
-    from pathlib import Path
     from types import ModuleType
     from types import TracebackType
 
@@ -39,6 +40,8 @@ class CoiledFunction: ...
     "create_kwargs_for_task",
     "get_module",
     "parse_future_result",
+    "should_pickle_module_by_value",
+    "strip_annotation_locals",
 ]
 
 
@@ -150,3 +153,43 @@ def get_module(func: Callable[..., Any], path: Path | None) -> ModuleType:
     if path:
         return inspect.getmodule(func, path.as_posix())  # type: ignore[return-value]
     return inspect.getmodule(func)  # type: ignore[return-value]
+
+
+def strip_annotation_locals(task: PTask) -> None:
+    """Remove annotation locals from task functions before pickling.
+
+    The locals snapshot is only needed during collection to evaluate annotations.
+    Keeping it around for execution can break pickling when it contains non-serializable
+    objects (for example, when importing ``pytask.mark`` in loop-generated tasks).
+
+    """
+    meta = getattr(task.function, "pytask_meta", None)
+    if meta is not None and getattr(meta, "annotation_locals", None) is not None:
+        meta.annotation_locals = None
+
+
+def should_pickle_module_by_value(module: ModuleType) -> bool:
+    """Return whether a module should be pickled by value.
+
+    We only pickle by value when the module is not importable by name in the worker.
+    This avoids serializing all module globals, which can fail for non-picklable
+    objects (e.g., closed file handles or locks stored at module scope).
+
+    """
+    module_name = getattr(module, "__name__", None)
+    module_file = getattr(module, "__file__", None)
+    if not module_name or module_name == "__main__" or module_file is None:
+        return True
+
+    try:
+        spec = importlib.util.find_spec(module_name)
+    except (ImportError, ValueError, AttributeError):
+        return True
+
+    if spec is None or spec.origin is None:
+        return True
+
+    try:
+        return Path(spec.origin).resolve() != Path(module_file).resolve()
+    except OSError:
+        return True

src/pytask_parallel/execute.py

@@ -217,7 +217,7 @@ def _render_traceback_to_string(
     traceback = Traceback(exc_info, show_locals=show_locals)
     segments = console.render(cast("Any", traceback), options=console_options)
     text = "".join(segment.text for segment in segments)
-    return (*exc_info[:2], text)  # ty: ignore[invalid-return-type]
+    return (*exc_info[:2], text)
 
 
 def _handle_function_products(