better position.

Author: tobiasraabe

src/pytask_parallel/backends.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import enum
-import inspect
 from concurrent.futures import Future
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import ThreadPoolExecutor
@@ -27,8 +26,6 @@ def submit(  # type: ignore[override]
         self, fn: Callable[..., Any], *args: Any, **kwargs: Any  # noqa: ARG002
     ) -> Future[Any]:
         """Submit a new task."""
-        task_module = inspect.getmodule(kwargs["task"].function)
-        cloudpickle.register_pickle_by_value(task_module)
         return super().submit(
             deserialize_and_run_with_cloudpickle,
             fn=cloudpickle.dumps(fn),
@@ -68,5 +65,5 @@ class ParallelBackend(enum.Enum):  # type: ignore[no-redef]
     PARALLEL_BACKENDS = {
         ParallelBackend.PROCESSES: CloudpickleProcessPoolExecutor,
         ParallelBackend.THREADS: ThreadPoolExecutor,
-        ParallelBackend.LOKY: (get_reusable_executor),  # type: ignore[attr-defined]
+        ParallelBackend.LOKY: get_reusable_executor,  # type: ignore[attr-defined]
     }

src/pytask_parallel/execute.py

@@ -12,6 +12,7 @@
 from typing import List
 
 import attr
+import cloudpickle
 from pytask import console
 from pytask import ExecutionReport
 from pytask import get_marks
@@ -185,6 +186,14 @@ def pytask_execute_task(session: Session, task: PTask) -> Future[Any] | None:
         if session.config["n_workers"] > 1:
             kwargs = _create_kwargs_for_task(task)
 
+            # Task modules are dynamically loaded and added to `sys.modules`. Thus,
+            # cloudpickle believes the module of the task function is also importable in
+            # the child process. We have to register the module as dynamic again, so
+            # that cloudpickle will pickle it with the function. See cloudpickle#417,
+            # pytask#373 and pytask#374.
+            task_module = inspect.getmodule(task.function)
+            cloudpickle.register_pickle_by_value(task_module)
+
             return session.config["_parallel_executor"].submit(
                 _execute_task,
                 task=task,