Strip local from task.

Author: tobiasraabe

src/pytask_parallel/execute.py

@@ -32,6 +32,7 @@
 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 strip_annotation_locals
 from pytask_parallel.utils import should_pickle_module_by_value
 
 if TYPE_CHECKING:
@@ -198,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,
@@ -225,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
 

src/pytask_parallel/utils.py

@@ -40,6 +40,7 @@ class CoiledFunction: ...
     "create_kwargs_for_task",
     "get_module",
     "parse_future_result",
+    "strip_annotation_locals",
     "should_pickle_module_by_value",
 ]
 
@@ -154,6 +155,19 @@ def get_module(func: Callable[..., Any], path: Path | None) -> ModuleType:
     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.