moved resolved lazy inputs into NodeExecution class from TaskDef

Author: tclose

pydra/engine/lazy.py

@@ -6,10 +6,9 @@
 from . import node
 
 if ty.TYPE_CHECKING:
-    from .graph import DiGraph
     from .submitter import NodeExecution
     from .core import Task, Workflow
-    from .specs import TaskDef, WorkflowDef
+    from .specs import TaskDef
     from .state import StateIndex
 
 
@@ -46,6 +45,27 @@ def _apply_cast(self, value):
             value = self._type(value)
         return value
 
+    def _get_value(
+        self,
+        node_exec: "NodeExecution",
+        state_index: "StateIndex | None" = None,
+    ) -> ty.Any:
+        """Return the value of a lazy field.
+
+        Parameters
+        ----------
+        node_exec: NodeExecution
+            the object representing the execution state of the current node
+        state_index : StateIndex, optional
+            the state index of the field to access
+
+        Returns
+        -------
+        value : Any
+            the resolved value of the lazy-field
+        """
+        raise NotImplementedError("LazyField is an abstract class")
+
 
 @attrs.define(kw_only=True)
 class LazyInField(LazyField[T]):
@@ -70,23 +90,25 @@ def _source(self):
 
     def _get_value(
         self,
-        workflow_def: "WorkflowDef",
+        node_exec: "NodeExecution",
+        state_index: "StateIndex | None" = None,
     ) -> ty.Any:
         """Return the value of a lazy field.
 
         Parameters
         ----------
-        wf : Workflow
-            the workflow the lazy field references
-        state_index : int, optional
-            the state index of the field to access
+        node_exec: NodeExecution
+            the object representing the execution state of the current node
+        state_index : StateIndex, optional
+            the state index of the field to access (ignored, used for duck-typing with
+            LazyOutField)
 
         Returns
         -------
         value : Any
             the resolved value of the lazy-field
         """
-        value = workflow_def[self._field]
+        value = node_exec.workflow_inputs[self._field]
         value = self._apply_cast(value)
         return value
 
@@ -105,16 +127,16 @@ def __repr__(self):
 
     def _get_value(
         self,
-        graph: "DiGraph[NodeExecution]",
+        node_exec: "NodeExecution",
         state_index: "StateIndex | None" = None,
     ) -> ty.Any:
         """Return the value of a lazy field.
 
         Parameters
         ----------
-        wf : Workflow
-            the workflow the lazy field references
-        state_index : int, optional
+        node_exec: NodeExecution
+            the object representing the execution state of the current node
+        state_index : StateIndex, optional
             the state index of the field to access
 
         Returns
@@ -130,7 +152,7 @@ def _get_value(
         if state_index is None:
             state_index = StateIndex()
 
-        task = graph.node(self._node.name).task(state_index)
+        task = node_exec.graph.node(self._node.name).task(state_index)
         _, split_depth = TypeParser.strip_splits(self._type)
 
         def get_nested(task: "Task[DefType]", depth: int):

pydra/engine/node.py

@@ -121,14 +121,28 @@ def lzout(self) -> OutputType:
                 type=field.type,
             )
         outputs = self.inputs.Outputs(**lazy_fields)
-        # Flag the output lazy fields as being not typed checked (i.e. assigned to another
-        # node's inputs) yet
+
         outpt: lazy.LazyOutField
         for outpt in attrs_values(outputs).values():
-            outpt._type_checked = False
+            # Assign the current node to the lazy fields so they can access the state
             outpt._node = self
+            # If the node has a non-empty state, wrap the type of the lazy field in
+            # a combination of an optional list and a number of nested StateArrays
+            # types based on the number of states the node is split over and whether
+            # it has a combiner
+            if self._state:
+                type_, _ = TypeParser.strip_splits(outpt._type)
+                if self._state.combiner:
+                    type_ = list[type_]
+                for _ in range(self._state.depth - int(bool(self._state.combiner))):
+                    type_ = StateArray[type_]
+                outpt._type = type_
+            # Flag the output lazy fields as being not typed checked (i.e. assigned to
+            # another node's inputs) yet. This is used to prevent the user from changing
+            # the type of the output after it has been accessed by connecting it to an
+            # output of an upstream node with additional state variables.
+            outpt._type_checked = False
         self._lzout = outputs
-        self._wrap_lzout_types_in_state_arrays()
         return outputs
 
     @property
@@ -217,20 +231,6 @@ def _check_if_outputs_have_been_used(self, msg):
                 + msg
             )
 
-    def _wrap_lzout_types_in_state_arrays(self) -> None:
-        """Wraps a types of the lazy out fields in a number of nested StateArray types
-        based on the number of states the node is split over"""
-        # Unwrap StateArray types from the output types
-        if not self.state:
-            return
-        outpt_lf: lazy.LazyOutField
-        for outpt_lf in attrs_values(self.lzout).values():
-            assert not outpt_lf._type_checked
-            type_, _ = TypeParser.strip_splits(outpt_lf._type)
-            for _ in range(self._state.depth):
-                type_ = StateArray[type_]
-            outpt_lf._type = type_
-
     def _set_state(self) -> None:
         # Add node name to state's splitter, combiner and cont_dim loaded from the def
         splitter = self._definition._splitter

pydra/engine/specs.py

@@ -35,14 +35,12 @@
 from pydra.utils.typing import StateArray, MultiInputObj
 from pydra.design.base import Field, Arg, Out, RequirementSet, NO_DEFAULT
 from pydra.design import shell
-from pydra.engine.lazy import LazyInField, LazyOutField
 
 if ty.TYPE_CHECKING:
     from pydra.engine.core import Task
     from pydra.engine.graph import DiGraph
     from pydra.engine.submitter import NodeExecution
     from pydra.engine.core import Workflow
-    from pydra.engine.state import StateIndex
     from pydra.engine.environments import Environment
     from pydra.engine.workers import Worker
 
@@ -476,41 +474,6 @@ def _compute_hashes(self) -> ty.Tuple[bytes, ty.Dict[str, bytes]]:
         }
         return hash_function(sorted(field_hashes.items())), field_hashes
 
-    def _resolve_lazy_inputs(
-        self,
-        workflow_inputs: "WorkflowDef",
-        graph: "DiGraph[NodeExecution]",
-        state_index: "StateIndex | None" = None,
-    ) -> Self:
-        """Resolves lazy fields in the task definition by replacing them with their
-        actual values.
-
-        Parameters
-        ----------
-        workflow : Workflow
-            The workflow the task is part of
-        graph : DiGraph[NodeExecution]
-            The execution graph of the workflow
-        state_index : StateIndex, optional
-            The state index for the workflow, by default None
-
-        Returns
-        -------
-        Self
-            The task definition with all lazy fields resolved
-        """
-        from pydra.engine.state import StateIndex
-
-        if state_index is None:
-            state_index = StateIndex()
-        resolved = {}
-        for name, value in attrs_values(self).items():
-            if isinstance(value, LazyInField):
-                resolved[name] = value._get_value(workflow_inputs)
-            elif isinstance(value, LazyOutField):
-                resolved[name] = value._get_value(graph, state_index)
-        return attrs.evolve(self, **resolved)
-
     def _check_rules(self):
         """Check if all rules are satisfied."""
 

pydra/engine/submitter.py

@@ -9,15 +9,18 @@
 from copy import copy
 from datetime import datetime
 from collections import defaultdict
+import attrs
 from .workers import Worker, WORKERS
 from .graph import DiGraph
 from .helpers import (
     get_open_loop,
     list_fields,
+    attrs_values,
 )
 from pydra.utils.hash import PersistentCache
 from .state import StateIndex
 from pydra.utils.typing import StateArray
+from pydra.engine.lazy import LazyField
 from .audit import Audit
 from .core import Task
 from pydra.utils.messenger import AuditFlag, Messenger
@@ -607,10 +610,7 @@ def all_failed(self) -> bool:
     def _generate_tasks(self) -> ty.Iterable["Task[DefType]"]:
         if not self.node.state:
             yield Task(
-                definition=self.node._definition._resolve_lazy_inputs(
-                    workflow_inputs=self.workflow_inputs,
-                    graph=self.graph,
-                ),
+                definition=self._resolve_lazy_inputs(task_def=self.node._definition),
                 submitter=self.submitter,
                 environment=self.node._environment,
                 hooks=self.node._hooks,
@@ -619,9 +619,8 @@ def _generate_tasks(self) -> ty.Iterable["Task[DefType]"]:
         else:
             for index, split_defn in self.node._split_definition().items():
                 yield Task(
-                    definition=split_defn._resolve_lazy_inputs(
-                        workflow_inputs=self.workflow_inputs,
-                        graph=self.graph,
+                    definition=self._resolve_lazy_inputs(
+                        task_def=split_defn,
                         state_index=index,
                     ),
                     submitter=self.submitter,
@@ -631,6 +630,32 @@ def _generate_tasks(self) -> ty.Iterable["Task[DefType]"]:
                     state_index=index,
                 )
 
+    def _resolve_lazy_inputs(
+        self,
+        task_def: "TaskDef",
+        state_index: "StateIndex | None" = None,
+    ) -> "TaskDef":
+        """Resolves lazy fields in the task definition by replacing them with their
+        actual values calculated by upstream jobs.
+
+        Parameters
+        ----------
+        task_def : TaskDef
+            The definition to resolve the lazy fields of
+        state_index : StateIndex, optional
+            The state index for the workflow, by default None
+
+        Returns
+        -------
+        TaskDef
+            The task definition with all lazy fields resolved
+        """
+        resolved = {}
+        for name, value in attrs_values(self).items():
+            if isinstance(value, LazyField):
+                resolved[name] = value._get_value(self, state_index)
+        return attrs.evolve(self, **resolved)
+
     def get_runnable_tasks(self, graph: DiGraph) -> list["Task[DefType]"]:
         """For a given node, check to see which tasks have been successfully run, are ready
         to run, can't be run due to upstream errors, or are blocked on other tasks to complete.