fixing up state maintenance in workflow nodes

Author: tclose

pydra/design/tests/test_workflow.py

@@ -9,18 +9,36 @@
 from pydra.engine.specs import TaskSpec
 from fileformats import video, image
 
+# NB: We use PascalCase for interfaces and workflow functions as it is translated into a class
 
-def test_workflow():
 
-    # NB: We use PascalCase (i.e. class names) as it is translated into a class
+@python.define
+def Add(a, b):
+    return a + b
 
-    @python.define
-    def Add(a, b):
-        return a + b
 
-    @python.define
-    def Mul(a, b):
-        return a * b
+@python.define
+def Mul(a, b):
+    return a * b
+
+
+@python.define(outputs=["divided"])
+def Divide(x, y):
+    return x / y
+
+
+@python.define
+def Sum(x: list[float]) -> float:
+    return sum(x)
+
+
+def a_converter(value):
+    if value is attrs.NOTHING:
+        return value
+    return float(value)
+
+
+def test_workflow():
 
     @workflow.define
     def MyTestWorkflow(a, b):
@@ -109,7 +127,7 @@ def MyTestShellWorkflow(
     assert wf.inputs.input_video == input_video
     assert wf.inputs.watermark == watermark
     assert wf.outputs.output_video == LazyOutField(
-        node=wf["resize"], field="out_video", type=video.Mp4
+        node=wf["resize"], field="out_video", type=video.Mp4, type_checked=True
     )
     assert list(wf.node_names) == ["add_watermark", "resize"]
 
@@ -119,19 +137,6 @@ def test_workflow_canonical():
 
     # NB: We use PascalCase (i.e. class names) as it is translated into a class
 
-    @python.define
-    def Add(a, b):
-        return a + b
-
-    @python.define
-    def Mul(a, b):
-        return a * b
-
-    def a_converter(value):
-        if value is attrs.NOTHING:
-            return value
-        return float(value)
-
     @workflow.define
     class MyTestWorkflow(TaskSpec["MyTestWorkflow.Outputs"]):
 
@@ -220,10 +225,10 @@ def MyTestShellWorkflow(
     )
     wf = Workflow.construct(workflow_spec)
     assert wf["add_watermark"].inputs.in_video == LazyInField(
-        node=wf, field="input_video", type=video.Mp4
+        workflow=wf, field="input_video", type=video.Mp4, type_checked=True
     )
     assert wf["add_watermark"].inputs.watermark == LazyInField(
-        node=wf, field="watermark", type=image.Png
+        workflow=wf, field="watermark", type=image.Png, type_checked=True
     )
 
 
@@ -236,10 +241,6 @@ def Add(x, y):
     def Mul(x, y):
         return x * y
 
-    @python.define(outputs=["divided"])
-    def Divide(x, y):
-        return x / y
-
     @workflow.define(outputs=["out1", "out2"])
     def MyTestWorkflow(a: int, b: float) -> tuple[float, float]:
         """A test workflow demonstration a few alternative ways to set and connect nodes
@@ -279,7 +280,9 @@ def MyTestWorkflow(a: int, b: float) -> tuple[float, float]:
     wf = Workflow.construct(workflow_spec)
     assert wf.inputs.a == 1
     assert wf.inputs.b == 2.0
-    assert wf.outputs.out1 == LazyOutField(node=wf["Mul"], field="out", type=float)
+    assert wf.outputs.out1 == LazyOutField(
+        node=wf["Mul"], field="out", type=float, type_checked=True
+    )
     assert wf.outputs.out2 == LazyOutField(
         node=wf["division"], field="divided", type=ty.Any
     )
@@ -288,14 +291,6 @@ def MyTestWorkflow(a: int, b: float) -> tuple[float, float]:
 
 def test_workflow_set_outputs_directly():
 
-    @python.define
-    def Add(a, b):
-        return a + b
-
-    @python.define
-    def Mul(a, b):
-        return a * b
-
     @workflow.define(outputs={"out1": float, "out2": float})
     def MyTestWorkflow(a: int, b: float):
 
@@ -362,10 +357,6 @@ def Mul(x: float, y: float) -> float:
     def Add(x: float, y: float) -> float:
         return x + y
 
-    @python.define
-    def Sum(x: list[float]) -> float:
-        return sum(x)
-
     @workflow.define
     def MyTestWorkflow(a: list[int], b: list[float], c: float) -> list[float]:
         mul = workflow.add(Mul()).split(x=a, y=b)
@@ -387,11 +378,11 @@ def test_workflow_split_after_access_fail():
     """
 
     @python.define
-    def Add(x, y):
+    def Add(x: float, y: float) -> float:
         return x + y
 
     @python.define
-    def Mul(x, y):
+    def Mul(x: float, y: float) -> float:
         return x * y
 
     @workflow.define

pydra/engine/workflow/base.py

@@ -77,7 +77,7 @@ def construct(
                 lazy_spec,
                 lzy_inpt.name,
                 LazyInField(
-                    node=wf,
+                    workflow=wf,
                     field=lzy_inpt.name,
                     type=lzy_inpt.type,
                 ),

pydra/engine/workflow/lazy.py

@@ -1,7 +1,9 @@
 import typing as ty
+import abc
 from typing_extensions import Self
 import attrs
 from pydra.utils.typing import StateArray
+from pydra.utils.hash import hash_single
 from . import node
 
 if ty.TYPE_CHECKING:
@@ -13,91 +15,22 @@
 TypeOrAny = ty.Union[type, ty.Any]
 
 
-@attrs.define(auto_attribs=True, kw_only=True)
-class LazyField(ty.Generic[T]):
+@attrs.define(kw_only=True)
+class LazyField(ty.Generic[T], metaclass=abc.ABCMeta):
     """Lazy fields implement promises."""
 
-    node: node.Node
     field: str
     type: TypeOrAny
-    # Set of splitters that have been applied to the lazy field. Note that the splitter
-    # specifications are transformed to a tuple[tuple[str, ...], ...] form where the
-    # outer tuple is the outer product, the inner tuple are inner products (where either
-    # product can be of length==1)
-    splits: ty.FrozenSet[ty.Tuple[ty.Tuple[str, ...], ...]] = attrs.field(
-        factory=frozenset, converter=frozenset
-    )
     cast_from: ty.Optional[ty.Type[ty.Any]] = None
-    # type_checked will be set to False after it is created but defaults to True here for
-    # ease of testing
-    type_checked: bool = True
+    type_checked: bool = False
 
     def __bytes_repr__(self, cache):
-        yield type(self).__name__.encode()
-        yield self.name.encode()
-        yield self.field.encode()
-
-    def cast(self, new_type: TypeOrAny) -> Self:
-        """ "casts" the lazy field to a new type
-
-        Parameters
-        ----------
-        new_type : type
-            the type to cast the lazy-field to
-
-        Returns
-        -------
-        cast_field : LazyField
-            a copy of the lazy field with the new type
-        """
-        return type(self)[new_type](
-            name=self.name,
-            field=self.field,
-            type=new_type,
-            splits=self.splits,
-            cast_from=self.cast_from if self.cast_from else self.type,
-        )
-
-    # def split(self, splitter: Splitter) -> Self:
-    #     """ "Splits" the lazy field over an array of nodes by replacing the sequence type
-    #     of the lazy field with StateArray to signify that it will be "split" across
-
-    #     Parameters
-    #     ----------
-    #     splitter : str or ty.Tuple[str, ...] or ty.List[str]
-    #         the splitter to append to the list of splitters
-    #     """
-    #     from pydra.utils.typing import (
-    #         TypeParser,
-    #     )  # pylint: disable=import-outside-toplevel
-
-    #     splits = self.splits | set([LazyField.normalize_splitter(splitter)])
-    #     # Check to see whether the field has already been split over the given splitter
-    #     if splits == self.splits:
-    #         return self
-
-    #     # Modify the type of the lazy field to include the split across a state-array
-    #     inner_type, prev_split_depth = TypeParser.strip_splits(self.type)
-    #     assert prev_split_depth <= 1
-    #     if inner_type is ty.Any:
-    #         type_ = StateArray[ty.Any]
-    #     elif TypeParser.matches_type(inner_type, list):
-    #         item_type = TypeParser.get_item_type(inner_type)
-    #         type_ = StateArray[item_type]
-    #     else:
-    #         raise TypeError(
-    #             f"Cannot split non-sequence field {self}  of type {inner_type}"
-    #         )
-    #     if prev_split_depth:
-    #         type_ = StateArray[type_]
-    #     return type(self)[type_](
-    #         name=self.name,
-    #         field=self.field,
-    #         type=type_,
-    #         splits=splits,
-    #     )
-
-    # # def combine(self, combiner: str | list[str]) -> Self:
+        yield type(self).__name__.encode() + b"("
+        yield from bytes(hash_single(self.source, cache))
+        yield b"field=" + self.field.encode()
+        yield b"type=" + bytes(hash_single(self.type, cache))
+        yield b"cast_from=" + bytes(hash_single(self.cast_from, cache))
+        yield b")"
 
     def _apply_cast(self, value):
         """\"Casts\" the value from the retrieved type if a cast has been applied to
@@ -110,19 +43,24 @@ def _apply_cast(self, value):
         return value
 
 
-@attrs.define(auto_attribs=True, kw_only=True)
+@attrs.define(kw_only=True)
 class LazyInField(LazyField[T]):
 
+    workflow: "Workflow" = attrs.field()
+
     attr_type = "input"
 
     def __eq__(self, other):
         return (
             isinstance(other, LazyInField)
             and self.field == other.field
             and self.type == other.type
-            and self.splits == other.splits
         )
 
+    @property
+    def source(self):
+        return self.workflow
+
     def get_value(self, wf: "Workflow", state_index: ty.Optional[int] = None) -> ty.Any:
         """Return the value of a lazy field.
 
@@ -155,8 +93,31 @@ def apply_splits(obj, depth):
         value = self._apply_cast(value)
         return value
 
+    def cast(self, new_type: TypeOrAny) -> Self:
+        """ "casts" the lazy field to a new type
 
+        Parameters
+        ----------
+        new_type : type
+            the type to cast the lazy-field to
+
+        Returns
+        -------
+        cast_field : LazyInField
+            a copy of the lazy field with the new type
+        """
+        return type(self)[new_type](
+            workflow=self.workflow,
+            field=self.field,
+            type=new_type,
+            cast_from=self.cast_from if self.cast_from else self.type,
+        )
+
+
+@attrs.define(kw_only=True)
 class LazyOutField(LazyField[T]):
+
+    node: node.Node
     attr_type = "output"
 
     def get_value(self, wf: "Workflow", state_index: ty.Optional[int] = None) -> ty.Any:
@@ -178,16 +139,15 @@ def get_value(self, wf: "Workflow", state_index: ty.Optional[int] = None) -> ty.
             TypeParser,
         )  # pylint: disable=import-outside-toplevel
 
-        node = getattr(wf, self.name)
-        result = node.result(state_index=state_index)
+        result = self.node.result(state_index=state_index)
         if result is None:
             raise RuntimeError(
-                f"Could not find results of '{node.name}' node in a sub-directory "
-                f"named '{node.checksum}' in any of the cache locations.\n"
-                + "\n".join(str(p) for p in set(node.cache_locations))
+                f"Could not find results of '{self.node.name}' node in a sub-directory "
+                f"named '{self.node.checksum}' in any of the cache locations.\n"
+                + "\n".join(str(p) for p in set(self.node.cache_locations))
                 + f"\n\nThis is likely due to hash changes in '{self.name}' node inputs. "
-                f"Current values and hashes: {node.inputs}, "
-                f"{node.inputs._hashes}\n\n"
+                f"Current values and hashes: {self.node.inputs}, "
+                f"{self.node.inputs._hashes}\n\n"
                 "Set loglevel to 'debug' in order to track hash changes "
                 "throughout the execution of the workflow.\n\n "
                 "These issues may have been caused by `bytes_repr()` methods "
@@ -224,3 +184,27 @@ def get_nested_results(res, depth: int):
         value = get_nested_results(result, depth=split_depth)
         value = self._apply_cast(value)
         return value
+
+    @property
+    def source(self):
+        return self.node
+
+    def cast(self, new_type: TypeOrAny) -> Self:
+        """ "casts" the lazy field to a new type
+
+        Parameters
+        ----------
+        new_type : type
+            the type to cast the lazy-field to
+
+        Returns
+        -------
+        cast_field : LazyOutField
+            a copy of the lazy field with the new type
+        """
+        return type(self)[new_type](
+            node=self.node,
+            field=self.field,
+            type=new_type,
+            cast_from=self.cast_from if self.cast_from else self.type,
+        )