Fix pickling errors by importing task modules in workers

Author: tobiasraabe

requirements-dev.lock

@@ -12,14 +12,55 @@
 from typing import Any
 from typing import ClassVar
 
+import os
+import sys
 import cloudpickle
 from attrs import define
 from loky import get_reusable_executor
 
 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.
+
+    """
+    global _WORKER_ROOT
+    root = os.fspath(path)
+    _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,16 @@ 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,13 @@
 
 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
 
 if TYPE_CHECKING:
     from concurrent.futures import Future
@@ -57,6 +59,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"]
     )
@@ -208,7 +211,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,
@@ -230,7 +234,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,7 +2,9 @@
 
 from __future__ import annotations
 
+import importlib.util
 import inspect
+from pathlib import Path
 from functools import partial
 from typing import TYPE_CHECKING
 from typing import Any
@@ -39,6 +41,7 @@ class CoiledFunction: ...
     "create_kwargs_for_task",
     "get_module",
     "parse_future_result",
+    "should_pickle_module_by_value",
 ]
 
 
@@ -150,3 +153,30 @@ 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 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

@@ -364,3 +364,56 @@ def task_create_file(
     )
     assert result.exit_code == ExitCode.OK
     assert tmp_path.joinpath("file.txt").read_text() == "This is the text."
+
+
+@pytest.mark.parametrize(
+    "parallel_backend",
+    [
+        ParallelBackend.PROCESSES,
+        pytest.param(ParallelBackend.LOKY, marks=skip_if_deadlock),
+    ],
+)
+def test_parallel_execution_with_mark_import(runner, tmp_path, parallel_backend):
+    source = """
+    from pytask import mark, task
+
+    @task
+    def task_assert_math():
+        assert 2 + 2 == 4
+    """
+    tmp_path.joinpath("task_mark.py").write_text(textwrap.dedent(source))
+    result = runner.invoke(
+        cli, [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend]
+    )
+    assert result.exit_code == ExitCode.OK
+
+
+@pytest.mark.parametrize(
+    "parallel_backend",
+    [
+        ParallelBackend.PROCESSES,
+        pytest.param(ParallelBackend.LOKY, marks=skip_if_deadlock),
+    ],
+)
+def test_parallel_execution_with_closed_file_handle(
+    runner, tmp_path, parallel_backend
+):
+    source = """
+    from pathlib import Path
+    from pytask import task
+
+    data_path = Path(__file__).parent / "data.txt"
+    data_path.write_text("hello", encoding="utf-8")
+
+    with data_path.open(encoding="utf-8") as f:
+        content = f.read()
+
+    @task
+    def task_assert_math():
+        assert content == "hello"
+    """
+    tmp_path.joinpath("task_file.py").write_text(textwrap.dedent(source))
+    result = runner.invoke(
+        cli, [tmp_path.as_posix(), "-n", "2", "--parallel-backend", parallel_backend]
+    )
+    assert result.exit_code == ExitCode.OK