implemented non-asynchronous "debug" worker that avoids async/await for easier debugging of workflows

Author: tclose

new-docs/source/tutorial/tst.py

@@ -21,10 +21,12 @@
 
     # Instantiate the task definition, "splitting" over all NIfTI files in the test directory
     # by splitting the "input" input field over all files in the directory
-    mrgrid = MrGrid(voxel=(0.5, 0.5, 0.5)).split(in_file=nifti_dir.iterdir())
+    mrgrid = MrGrid(operation="regrid", voxel=(0.5, 0.5, 0.5)).split(
+        in_file=nifti_dir.iterdir()
+    )
 
     # Run the task to resample all NIfTI files
-    outputs = mrgrid(worker="serial")
+    outputs = mrgrid(worker="cf")
 
     # Print the locations of the output files
-    print("\n".join(str(p) for p in outputs.outputs))
+    print("\n".join(str(p) for p in outputs.out_file))

pydra/design/base.py

@@ -388,8 +388,7 @@ def make_task_def(
     klass : type
         The class created using the attrs package
     """
-    from pydra.engine.specs import TaskDef, WorkflowDef
-    from pydra.engine.core import Task, WorkflowTask
+    from pydra.engine.specs import TaskDef
 
     spec_type._check_arg_refs(inputs, outputs)
 
@@ -400,7 +399,6 @@ def make_task_def(
             f"{reserved_names} are reserved and cannot be used for {spec_type} field names"
         )
     outputs_klass = make_outputs_spec(out_type, outputs, outputs_bases, name)
-    task_type = WorkflowTask if issubclass(spec_type, WorkflowDef) else Task
     if klass is None or not issubclass(klass, spec_type):
         if name is None:
             raise ValueError("name must be provided if klass is not")
@@ -419,19 +417,13 @@ def make_task_def(
             name=name,
             bases=bases,
             kwds={},
-            exec_body=lambda ns: ns.update(
-                {
-                    "Outputs": outputs_klass,
-                    "Task": task_type,
-                }
-            ),
+            exec_body=lambda ns: ns.update({"Outputs": outputs_klass}),
         )
     else:
         # Ensure that the class has it's own annotations dict so we can modify it without
         # messing up other classes
         klass.__annotations__ = copy(klass.__annotations__)
         klass.Outputs = outputs_klass
-        klass.Task = task_type
     # Now that we have saved the attributes in lists to be
     for arg in inputs.values():
         # If an outarg input then the field type should be Path not a FileSet

pydra/engine/core.py

@@ -43,6 +43,7 @@
 )
 from .helpers_file import copy_nested_files, template_update
 from pydra.utils.messenger import AuditFlag
+from pydra.engine.environments import Environment, Native
 
 logger = logging.getLogger("pydra")
 
@@ -85,7 +86,8 @@ class Task(ty.Generic[DefType]):
 
     name: str
     definition: DefType
-    submitter: "Submitter"
+    submitter: "Submitter | None"
+    environment: "Environment | None"
     state_index: state.StateIndex
 
     _inputs: dict[str, ty.Any] | None = None
@@ -95,6 +97,7 @@ def __init__(
         definition: DefType,
         submitter: "Submitter",
         name: str,
+        environment: "Environment | None" = None,
         state_index: "state.StateIndex | None" = None,
     ):
         """
@@ -121,14 +124,16 @@ def __init__(
         if state_index is None:
             state_index = state.StateIndex()
 
-        self.definition = definition
+        # Copy the definition, so lazy fields can be resolved and replaced at runtime
+        self.definition = copy(definition)
+        # We save the submitter is the definition is a workflow otherwise we don't
+        # so the task can be pickled
+        self.submitter = submitter if is_workflow(definition) else None
+        self.environment = environment if environment is not None else Native()
         self.name = name
         self.state_index = state_index
 
-        # checking if metadata is set properly
-        self.definition._check_resolved()
-        self.definition._check_rules()
-        self._output = {}
+        self.return_values = {}
         self._result = {}
         # flag that says if node finished all jobs
         self._done = False
@@ -151,6 +156,11 @@ def __init__(
     def cache_dir(self):
         return self._cache_dir
 
+    @property
+    def is_async(self) -> bool:
+        """Check to see if the task should be run asynchronously."""
+        return self.submitter.worker.is_async and is_workflow(self.definition)
+
     @cache_dir.setter
     def cache_dir(self, path: os.PathLike):
         self._cache_dir = Path(path)
@@ -315,6 +325,21 @@ def _populate_filesystem(self):
         self.output_dir.mkdir(parents=False, exist_ok=self.can_resume)
 
     def run(self, rerun: bool = False):
+        """Prepare the task working directory, execute the task definition, and save the
+        results.
+
+        Parameters
+        ----------
+        rerun : bool
+            If True, the task will be re-run even if a result already exists. Will
+            propagated to all tasks within workflow tasks.
+        """
+        # TODO: After these changes have been merged, will refactor this function and
+        # run_async to use common helper methods for pre/post run tasks
+
+        # checking if the definition is fully resolved and ready to run
+        self.definition._check_resolved()
+        self.definition._check_rules()
         self.hooks.pre_run(self)
         logger.debug(
             "'%s' is attempting to acquire lock on %s", self.name, self.lockfile
@@ -334,8 +359,8 @@ def run(self, rerun: bool = False):
                 self.audit.audit_task(task=self)
             try:
                 self.audit.monitor()
-                run_outputs = self.definition._run(self)
-                result.outputs = self.definition.Outputs.from_task(self, run_outputs)
+                self.definition._run(self)
+                result.outputs = self.definition.Outputs._from_task(self)
             except Exception:
                 etype, eval, etr = sys.exc_info()
                 traceback = format_exception(etype, eval, etr)
@@ -355,6 +380,57 @@ def run(self, rerun: bool = False):
         self._check_for_hash_changes()
         return result
 
+    async def run_async(self, rerun: bool = False):
+        """Prepare the task working directory, execute the task definition asynchronously,
+        and save the results. NB: only workflows are run asynchronously at the moment.
+
+        Parameters
+        ----------
+        rerun : bool
+            If True, the task will be re-run even if a result already exists. Will
+            propagated to all tasks within workflow tasks.
+        """
+        # checking if the definition is fully resolved and ready to run
+        self.definition._check_resolved()
+        self.definition._check_rules()
+        self.hooks.pre_run(self)
+        logger.debug(
+            "'%s' is attempting to acquire lock on %s", self.name, self.lockfile
+        )
+        async with PydraFileLock(self.lockfile):
+            if not rerun:
+                result = self.result()
+                if result is not None and not result.errored:
+                    return result
+            cwd = os.getcwd()
+            self._populate_filesystem()
+            result = Result(outputs=None, runtime=None, errored=False, task=self)
+            self.hooks.pre_run_task(self)
+            self.audit.start_audit(odir=self.output_dir)
+            try:
+                self.audit.monitor()
+                await self.definition._run(self)
+                result.outputs = self.definition.Outputs._from_task(self)
+            except Exception:
+                etype, eval, etr = sys.exc_info()
+                traceback = format_exception(etype, eval, etr)
+                record_error(self.output_dir, error=traceback)
+                result.errored = True
+                self._errored = True
+                raise
+            finally:
+                self.hooks.post_run_task(self, result)
+                self.audit.finalize_audit(result=result)
+                save(self.output_dir, result=result, task=self)
+                # removing the additional file with the checksum
+                (self.cache_dir / f"{self.uid}_info.json").unlink()
+                os.chdir(cwd)
+        self.hooks.post_run(self, result)
+        # Check for any changes to the input hashes that have occurred during the execution
+        # of the task
+        self._check_for_hash_changes()
+        return result
+
     def pickle_task(self):
         """Pickling the tasks with full inputs"""
         pkl_files = self.cache_dir / "pkl_files"
@@ -398,7 +474,7 @@ def _combined_output(self, return_inputs=False):
         else:
             return combined_results
 
-    def result(self, state_index=None, return_inputs=False):
+    def result(self, return_inputs=False):
         """
         Retrieve the outcomes of this particular task.
 
@@ -415,13 +491,9 @@ def result(self, state_index=None, return_inputs=False):
         result : Result
             the result of the task
         """
-        # TODO: check if result is available in load_result and
-        # return a future if not
         if self.errored:
             return Result(outputs=None, runtime=None, errored=True, task=self)
 
-        if state_index is not None:
-            raise ValueError("Task does not have a state")
         checksum = self.checksum
         result = load_result(checksum, self.cache_locations)
         if result and result.errored:
@@ -483,58 +555,6 @@ def _check_for_hash_changes(self):
     DEFAULT_COPY_COLLATION = FileSet.CopyCollation.any
 
 
-class WorkflowTask(Task):
-
-    def __init__(
-        self,
-        definition: DefType,
-        submitter: "Submitter",
-        name: str,
-        state_index: "state.StateIndex | None" = None,
-    ):
-        super().__init__(definition, submitter, name, state_index)
-        self.submitter = submitter
-
-    async def run(self, rerun: bool = False):
-        self.hooks.pre_run(self)
-        logger.debug(
-            "'%s' is attempting to acquire lock on %s", self.name, self.lockfile
-        )
-        async with PydraFileLock(self.lockfile):
-            if not rerun:
-                result = self.result()
-                if result is not None and not result.errored:
-                    return result
-            cwd = os.getcwd()
-            self._populate_filesystem()
-            result = Result(outputs=None, runtime=None, errored=False, task=self)
-            self.hooks.pre_run_task(self)
-            self.audit.start_audit(odir=self.output_dir)
-            try:
-                self.audit.monitor()
-                await self.submitter.expand_workflow(self)
-                result.outputs = self.definition.Outputs.from_task(self)
-            except Exception:
-                etype, eval, etr = sys.exc_info()
-                traceback = format_exception(etype, eval, etr)
-                record_error(self.output_dir, error=traceback)
-                result.errored = True
-                self._errored = True
-                raise
-            finally:
-                self.hooks.post_run_task(self, result)
-                self.audit.finalize_audit(result=result)
-                save(self.output_dir, result=result, task=self)
-                # removing the additional file with the checksum
-                (self.cache_dir / f"{self.uid}_info.json").unlink()
-                os.chdir(cwd)
-        self.hooks.post_run(self, result)
-        # Check for any changes to the input hashes that have occurred during the execution
-        # of the task
-        self._check_for_hash_changes()
-        return result
-
-
 logger = logging.getLogger("pydra")
 
 OutputsType = ty.TypeVar("OutputType", bound=TaskOutputs)
@@ -847,18 +867,12 @@ def create_dotfile(self, type="simple", export=None, name=None, output_dir=None)
             return dotfile, formatted_dot
 
 
-def is_task(obj):
-    """Check whether an object looks like a task."""
-    return hasattr(obj, "_run_task")
-
-
 def is_workflow(obj):
     """Check whether an object is a :class:`Workflow` instance."""
     from pydra.engine.specs import WorkflowDef
     from pydra.engine.core import Workflow
-    from pydra.engine.core import WorkflowTask
 
-    return isinstance(obj, (WorkflowDef, WorkflowTask, Workflow))
+    return isinstance(obj, (WorkflowDef, Workflow))
 
 
 def has_lazy(obj):

pydra/engine/environments.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 if ty.TYPE_CHECKING:
-    from pydra.engine.task import ShellTask
+    from pydra.engine.core import Task
 
 
 class Environment:
@@ -17,7 +17,7 @@ class Environment:
     def setup(self):
         pass
 
-    def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
+    def execute(self, task: "Task") -> dict[str, ty.Any]:
         """
         Execute the task in the environment.
 
@@ -42,12 +42,12 @@ class Native(Environment):
     Native environment, i.e. the tasks are executed in the current python environment.
     """
 
-    def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
+    def execute(self, task: "Task") -> dict[str, ty.Any]:
         keys = ["return_code", "stdout", "stderr"]
-        values = execute(task.definition._command_args(), strip=task.strip)
+        values = execute(task.definition._command_args())
         output = dict(zip(keys, values))
         if output["return_code"]:
-            msg = f"Error running '{task.name}' task with {task.command_args()}:"
+            msg = f"Error running '{task.name}' task with {task.definition._command_args()}:"
             if output["stderr"]:
                 msg += "\n\nstderr:\n" + output["stderr"]
             if output["stdout"]:
@@ -90,7 +90,7 @@ def bind(self, loc, mode="ro"):
 class Docker(Container):
     """Docker environment."""
 
-    def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
+    def execute(self, task: "Task") -> dict[str, ty.Any]:
         docker_img = f"{self.image}:{self.tag}"
         # mounting all input locations
         mounts = task.definition._get_bindings(root=self.root)
@@ -112,7 +112,6 @@ def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
 
         values = execute(
             docker_args + [docker_img] + task.definition._command_args(root=self.root),
-            strip=task.strip,
         )
         output = dict(zip(keys, values))
         if output["return_code"]:
@@ -126,7 +125,7 @@ def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
 class Singularity(Container):
     """Singularity environment."""
 
-    def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
+    def execute(self, task: "Task") -> dict[str, ty.Any]:
         singularity_img = f"{self.image}:{self.tag}"
         # mounting all input locations
         mounts = task.definition._get_bindings(root=self.root)
@@ -151,7 +150,6 @@ def execute(self, task: "ShellTask") -> dict[str, ty.Any]:
             singularity_args
             + [singularity_img]
             + task.definition._command_args(root=self.root),
-            strip=task.strip,
         )
         output = dict(zip(keys, values))
         if output["return_code"]: