debugged test_tasks

Author: tclose

pydra/engine/core.py

@@ -100,6 +100,10 @@ def __init__(
         name: str,
         environment: "Environment | None" = None,
         state_index: "state.StateIndex | None" = None,
+        pre_run: ty.Callable["Task", None] | None = None,
+        pre_run_task: ty.Callable["Task", None] | None = None,
+        post_run_task: ty.Callable["Task", None] | None = None,
+        post_run: ty.Callable["Task", None] | None = None,
     ):
         """
         Initialize a task.
@@ -142,7 +146,12 @@ def __init__(
         self.allow_cache_override = True
         self._checksum = None
         self._uid = uuid4().hex
-        self.hooks = TaskHook()
+        self.hooks = TaskHook(
+            pre_run=pre_run,
+            post_run=post_run,
+            pre_run_task=pre_run_task,
+            post_run_task=post_run_task,
+        )
         self._errored = False
         self._lzout = None
 

pydra/engine/specs.py

@@ -14,6 +14,7 @@
 from copy import deepcopy
 from typing import Self
 import attrs
+from attrs.converters import default_if_none
 import cloudpickle as cp
 from fileformats.core import FileSet
 from pydra.utils.messenger import AuditFlag, Messenger
@@ -160,6 +161,10 @@ def __call__(
         messengers: ty.Iterable[Messenger] | None = None,
         messenger_args: dict[str, ty.Any] | None = None,
         name: str | None = None,
+        pre_run: ty.Callable["Task", None] | None = None,
+        post_run: ty.Callable["Task", None] | None = None,
+        pre_run_task: ty.Callable["Task", None] | None = None,
+        post_run_task: ty.Callable["Task", None] | None = None,
         **kwargs: ty.Any,
     ) -> OutputsType:
         """Create a task from this definition and execute it to produce a result.
@@ -212,8 +217,17 @@ def __call__(
                 worker=worker,
                 **kwargs,
             ) as sub:
-                result = sub(self, name=name)
+                result = sub(
+                    self,
+                    name=name,
+                    pre_run=pre_run,
+                    post_run=post_run,
+                    pre_run_task=pre_run_task,
+                    post_run_task=post_run_task,
+                )
         except TypeError as e:
+            # Catch any inadvertent passing of task definition parameters to the
+            # execution call
             if hasattr(e, "__notes__") and WORKER_KWARG_FAIL_NOTE in e.__notes__:
                 if match := re.match(
                     r".*got an unexpected keyword argument '(\w+)'", str(e)
@@ -1248,15 +1262,18 @@ def donothing(*args: ty.Any, **kwargs: ty.Any) -> None:
 class TaskHook:
     """Callable task hooks."""
 
-    pre_run_task: ty.Callable = donothing
-    post_run_task: ty.Callable = donothing
-    pre_run: ty.Callable = donothing
-    post_run: ty.Callable = donothing
-
-    def __setattr__(self, attr, val):
-        if attr not in ["pre_run_task", "post_run_task", "pre_run", "post_run"]:
-            raise AttributeError("Cannot set unknown hook")
-        super().__setattr__(attr, val)
+    pre_run_task: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+    post_run_task: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+    pre_run: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
+    post_run: ty.Callable = attrs.field(
+        default=donothing, converter=default_if_none(donothing)
+    )
 
     def reset(self):
         for val in ["pre_run_task", "post_run_task", "pre_run", "post_run"]:

pydra/engine/submitter.py

@@ -169,6 +169,10 @@ def __call__(
         self,
         task_def: "TaskDef",
         name: str | None = "task",
+        pre_run: ty.Callable["Task", None] | None = None,
+        post_run: ty.Callable["Task", None] | None = None,
+        pre_run_task: ty.Callable["Task", None] | None = None,
+        post_run_task: ty.Callable["Task", None] | None = None,
     ):
         """Submitter run function."""
 
@@ -194,7 +198,16 @@ def Split(defn: TaskDef, output_types: dict):
                 f"Task {self} is marked for combining, but not splitting. "
                 "Use the `split` method to split the task before combining."
             )
-        task = Task(task_def, submitter=self, name=name, environment=self.environment)
+        task = Task(
+            task_def,
+            submitter=self,
+            name=name,
+            environment=self.environment,
+            pre_run=pre_run,
+            post_run=post_run,
+            pre_run_task=pre_run_task,
+            post_run_task=post_run_task,
+        )
         try:
             self.run_start_time = datetime.now()
             if self.worker.is_async:  # Only workflow tasks can be async
@@ -203,6 +216,12 @@ def Split(defn: TaskDef, output_types: dict):
                 )
             else:
                 self.worker.run(task, rerun=self.rerun)
+        except Exception as e:
+            e.add_note(
+                f"Full crash report for {type(task_def).__name__!r} task is here: "
+                + str(task.output_dir / "_error.pklz")
+            )
+            raise e
         finally:
             self.run_start_time = None
         PersistentCache().clean_up()

pydra/engine/tests/test_task.py

@@ -1198,7 +1198,7 @@ def test_taskhooks_1(tmpdir: Path, capsys):
     cache_dir.mkdir()
 
     foo = Task(
-        definition=FunAddTwo(a=1), submitter=Submitter(cache_dir=cache_dir), name="foo"
+        definition=FunAddTwo(a=1), submitter=Submitter(cache_dir=tmpdir), name="foo"
     )
     assert foo.hooks
     # ensure all hooks are defined
@@ -1209,8 +1209,7 @@ def test_taskhooks_1(tmpdir: Path, capsys):
     def myhook(task, *args):
         print("I was called")
 
-    foo.hooks.pre_run = myhook
-    foo.run()
+    FunAddTwo(a=1)(cache_dir=cache_dir, pre_run=myhook)
     captured = capsys.readouterr()
     assert "I was called\n" in captured.out
     del captured
@@ -1219,52 +1218,31 @@ def myhook(task, *args):
     with pytest.raises(AttributeError):
         foo.hooks.mid_run = myhook
 
+    # reset all hooks
+    foo.hooks.reset()
+    for attr in ("pre_run", "post_run", "pre_run_task", "post_run_task"):
+        hook = getattr(foo.hooks, attr)
+        assert hook() is None
+
     # clear cache
     shutil.rmtree(cache_dir)
     cache_dir.mkdir()
 
     # set all hooks
-    foo.hooks.post_run = myhook
-    foo.hooks.pre_run_task = myhook
-    foo.hooks.post_run_task = myhook
-    foo.run()
-    captured = capsys.readouterr()
-    assert captured.out.count("I was called\n") == 4
-    del captured
-
-    # hooks are independent across tasks by default
-    bar = Task(
-        definition=FunAddTwo(a=3), name="bar", submitter=Submitter(cache_dir=tmpdir)
+    FunAddTwo(a=1)(
+        cache_dir=cache_dir,
+        pre_run=myhook,
+        post_run=myhook,
+        pre_run_task=myhook,
+        post_run_task=myhook,
     )
-    assert bar.hooks is not foo.hooks
-    # but can be shared across tasks
-    bar.hooks = foo.hooks
-    # and workflows
-    wf_task = Task(
-        definition=BasicWorkflow(x=1),
-        submitter=Submitter(cache_dir=tmpdir, worker="cf"),
-        name="wf",
-    )
-    wf_task.hooks = bar.hooks
-    assert foo.hooks == bar.hooks == wf_task.hooks
-
-    wf_task.run()
     captured = capsys.readouterr()
     assert captured.out.count("I was called\n") == 4
     del captured
 
-    # reset all hooks
-    foo.hooks.reset()
-    for attr in ("pre_run", "post_run", "pre_run_task", "post_run_task"):
-        hook = getattr(foo.hooks, attr)
-        assert hook() is None
-
 
 def test_taskhooks_2(tmpdir, capsys):
     """checking order of the hooks; using task's attributes"""
-    foo = Task(
-        definition=FunAddTwo(a=1), name="foo", submitter=Submitter(cache_dir=tmpdir)
-    )
 
     def myhook_prerun(task, *args):
         print(f"i. prerun hook was called from {task.name}")
@@ -1278,11 +1256,13 @@ def myhook_postrun_task(task, *args):
     def myhook_postrun(task, *args):
         print(f"iv. postrun hook was called {task.name}")
 
-    foo.hooks.pre_run = myhook_prerun
-    foo.hooks.post_run = myhook_postrun
-    foo.hooks.pre_run_task = myhook_prerun_task
-    foo.hooks.post_run_task = myhook_postrun_task
-    foo.run()
+    FunAddTwo(a=1)(
+        cache_dir=tmpdir,
+        pre_run=myhook_prerun,
+        post_run=myhook_postrun,
+        pre_run_task=myhook_prerun_task,
+        post_run_task=myhook_postrun_task,
+    )
 
     captured = capsys.readouterr()
     hook_messages = captured.out.strip().split("\n")
@@ -1318,21 +1298,17 @@ def myhook_postrun(task, result, *args):
 
 def test_taskhooks_4(tmpdir, capsys):
     """task raises an error: postrun task should be called, postrun shouldn't be called"""
-    foo = Task(
-        definition=FunAddTwo(a="one"), name="foo", submitter=Submitter(cache_dir=tmpdir)
-    )
 
     def myhook_postrun_task(task, result, *args):
         print(f"postrun task hook was called, result object is {result}")
 
     def myhook_postrun(task, result, *args):
         print("postrun hook should not be called")
 
-    foo.hooks.post_run = myhook_postrun
-    foo.hooks.post_run_task = myhook_postrun_task
-
     with pytest.raises(Exception):
-        foo()
+        FunAddTwo(a="one")(
+            cache_dir=tmpdir, post_run=myhook_postrun, post_run_task=myhook_postrun_task
+        )
 
     captured = capsys.readouterr()
     hook_messages = captured.out.strip().split("\n")
@@ -1351,11 +1327,13 @@ def test_traceback(tmpdir):
     def FunError(x):
         raise Exception("Error from the function")
 
-    with pytest.raises(Exception, match="Task 'FunError' raised an error") as exinfo:
+    with pytest.raises(Exception, match="Error from the function") as exinfo:
         FunError(x=3)(worker="cf", cache_dir=tmpdir)
 
     # getting error file from the error message
-    error_file_match = str(exinfo.value).split("here: ")[-1].split("_error.pklz")[0]
+    error_file_match = (
+        str(exinfo.value.__notes__[0]).split("here: ")[-1].split("_error.pklz")[0]
+    )
     error_file = Path(error_file_match) / "_error.pklz"
     # checking if the file exists
     assert error_file.exists()
@@ -1386,7 +1364,9 @@ def Workflow(x_list):
         wf(worker="cf")
 
     # getting error file from the error message
-    error_file_match = str(exinfo.value).split("here: ")[-1].split("_error.pklz")[0]
+    error_file_match = (
+        str(exinfo.value).split("here: ")[-1].split("_error.pklz")[0].strip()
+    )
     error_file = Path(error_file_match) / "_error.pklz"
     # checking if the file exists
     assert error_file.exists()