debugging test_shelltask, reworking resolve of lazy inputs

Author: tclose

pydra/engine/helpers.py

@@ -16,11 +16,14 @@
 from filelock import SoftFileLock, Timeout
 import cloudpickle as cp
 from fileformats.core import FileSet
+from pydra.utils.typing import StateArray
+
 
 if ty.TYPE_CHECKING:
     from .specs import TaskDef, Result, WorkflowOutputs, WorkflowDef
     from .core import Task
     from pydra.design.base import Field
+    from pydra.engine.lazy import LazyField
 
 
 PYDRA_ATTR_METADATA = "__PYDRA_METADATA__"
@@ -695,3 +698,27 @@ def is_lazy(obj):
     from pydra.engine.lazy import LazyField
 
     return isinstance(obj, LazyField)
+
+
+T = ty.TypeVar("T")
+U = ty.TypeVar("U")
+
+
+def state_array_support(
+    function: ty.Callable[T, U],
+) -> ty.Callable[T | StateArray[T], U | StateArray[U]]:
+    """
+    Decorator to convert a allow a function to accept and return StateArray objects,
+    where the function is applied to each element of the StateArray.
+    """
+
+    def state_array_wrapper(
+        value: "T | StateArray[T] | LazyField[T]",
+    ) -> "U | StateArray[U] | LazyField[U]":
+        if is_lazy(value):
+            return value
+        if isinstance(value, StateArray):
+            return StateArray(function(v) for v in value)
+        return function(value)
+
+    return state_array_wrapper

pydra/engine/specs.py

@@ -16,9 +16,9 @@
 import attrs
 from attrs.converters import default_if_none
 import cloudpickle as cp
-from fileformats.generic import FileSet
+from fileformats.generic import FileSet, File
 from pydra.utils.messenger import AuditFlag, Messenger
-from pydra.utils.typing import is_optional, optional_type, MultiInputObj
+from pydra.utils.typing import is_optional, optional_type
 from .helpers import (
     attrs_fields,
     attrs_values,
@@ -28,6 +28,7 @@
     ensure_list,
     parse_format_string,
     fields_in_formatter,
+    state_array_support,
 )
 from .helpers_file import template_update, template_update_single
 from . import helpers_state as hlpst
@@ -1032,17 +1033,29 @@ def _resolve_value(
 ShellOutputsType = ty.TypeVar("OutputType", bound=ShellOutputs)
 
 
+@state_array_support
+def additional_args_converter(value: ty.Any) -> list[str]:
+    """Convert additional arguments to a list of strings."""
+    if isinstance(value, str):
+        return shlex.split(value)
+    if not isinstance(value, ty.Sequence):
+        return [value]
+    return list(value)
+
+
 @attrs.define(kw_only=True, auto_attribs=False, eq=False)
 class ShellDef(TaskDef[ShellOutputsType]):
 
     _task_type = "shell"
 
     BASE_NAMES = ["additional_args"]
 
-    additional_args: MultiInputObj[str] = shell.arg(
+    additional_args: list[str | File] = shell.arg(
         name="additional_args",
         default=attrs.Factory(list),
-        type=MultiInputObj[str],
+        converter=additional_args_converter,
+        type=list[str | File],
+        sep=" ",
         help="Additional free-form arguments to append to the end of the command.",
     )
 

pydra/engine/state.py

@@ -27,14 +27,18 @@ class StateIndex:
 
     indices: OrderedDict[str, int]
 
-    def __init__(self, indices: dict[str, int] | None = None):
+    def __init__(
+        self, indices: dict[str, int] | ty.Sequence[tuple[str, int]] | None = None
+    ):
         # We used ordered dict here to ensure the keys are always in the same order
         # while OrderedDict is not strictly necessary for CPython 3.7+, we use it to
         # signal that the order of the keys is important
         if indices is None:
             self.indices = OrderedDict()
         else:
-            self.indices = OrderedDict(sorted(indices.items()))
+            if isinstance(indices, dict):
+                indices = indices.items()
+            self.indices = OrderedDict(sorted(indices))
 
     def __len__(self) -> int:
         return len(self.indices)

pydra/engine/submitter.py

@@ -20,7 +20,7 @@
 from pydra.utils.hash import PersistentCache
 from .state import StateIndex
 from pydra.utils.typing import StateArray
-from pydra.engine.lazy import LazyField
+from pydra.engine.lazy import LazyField, LazyOutField
 from .audit import Audit
 from .core import Task
 from pydra.utils.messenger import AuditFlag, Messenger
@@ -537,6 +537,10 @@ def __init__(
         self.queued = {}
         self.running = {}  # Not used in logic, but may be useful for progress tracking
         self.unrunnable = defaultdict(list)
+        # Prepare the state to be run
+        if self.state:
+            self.state.prepare_states(self.node.state_values)
+            self.state.prepare_inputs()
         self.state_names = self.node.state.names if self.node.state else []
         self.workflow = workflow
         self.graph = None
@@ -567,6 +571,21 @@ def tasks(self) -> ty.Iterable["Task[DefType]"]:
             self._tasks = {t.state_index: t for t in self._generate_tasks()}
         return self._tasks.values()
 
+    def translate_index(self, index: StateIndex, lf: LazyOutField):
+        state_key = f"{lf._node.name}.{lf._field}"
+        try:
+            upstream_state = self.state.inner_inputs[state_key]
+        except KeyError:
+            state_index = StateIndex(index)
+        else:
+            state_index = StateIndex(
+                zip(
+                    upstream_state.keys_final,
+                    upstream_state.ind_l_final[index[state_key]],
+                )
+            )
+        return state_index
+
     def matching_jobs(self, index: StateIndex = StateIndex()) -> "StateArray[Task]":
         """Get the jobs that match a given state index.
 
@@ -702,23 +721,22 @@ def _split_definition(self) -> dict[StateIndex, "TaskDef[OutputType]"]:
         if not self.node.state:
             return {None: self.node._definition}
         split_defs = {}
-        self.state.prepare_states(self.node.state_values)
-        self.state.prepare_inputs()
         for input_ind in self.node.state.inputs_ind:
             resolved = {}
             for inpt_name in set(self.node.input_names):
                 value = getattr(self._definition, inpt_name)
+                state_key = f"{self.node.name}.{inpt_name}"
                 if isinstance(value, LazyField):
                     value = resolved[inpt_name] = value._get_value(
                         workflow=self.workflow,
                         graph=self.graph,
-                        state_index=StateIndex(input_ind),
+                        state_index=input_ind,
                     )
-                if f"{self.node.name}.{inpt_name}" in input_ind:
+                elif state_key in input_ind:
                     resolved[inpt_name] = self.node.state._get_element(
                         value=value,
                         field_name=inpt_name,
-                        ind=input_ind[f"{self.node.name}.{inpt_name}"],
+                        ind=input_ind[state_key],
                     )
             split_defs[StateIndex(input_ind)] = attrs.evolve(
                 self.node._definition, **resolved

pydra/engine/tests/test_shelltask.py

@@ -1674,9 +1674,11 @@ def test_wf_shell_cmd_3(plugin, tmp_path):
     class Shelly1(ShellDef["Shelly1.Outputs"]):
         executable = "shelly"
 
+        arg: str = shell.arg(argstr=None)
+
         class Outputs(ShellOutputs):
             file: File = shell.outarg(
-                path_template="{args}",
+                path_template="{arg}",
                 help="output file",
             )
 
@@ -1700,12 +1702,12 @@ class Outputs(ShellOutputs):
             )
 
     @workflow.define(outputs=["touch_file", "out1", "cp_file", "out2"])
-    def Workflow(cmd1, cmd2, args):
+    def Workflow(cmd1, cmd2, arg):
 
         shelly1 = workflow.add(
             Shelly1(
                 executable=cmd1,
-                additional_args=args,
+                arg=arg,
             )
         )
         shelly2 = workflow.add(
@@ -1717,9 +1719,9 @@ def Workflow(cmd1, cmd2, args):
 
         return shelly1.file, shelly1.stdout, shelly2.out_file, shelly2.stdout
 
-    wf = Workflow(cmd1="touch", cmd2="cp", args=File.mock("newfile.txt"))
+    wf = Workflow(cmd1="touch", cmd2="cp", arg="newfile.txt")
 
-    with Submitter(plugin="debug") as sub:
+    with Submitter(plugin="debug", cache_dir=tmp_path) as sub:
         res = sub(wf)
 
     assert res.outputs.out1 == ""
@@ -1738,14 +1740,19 @@ def test_wf_shell_cmd_3a(plugin, tmp_path):
 
     @shell.define
     class Shelly1(ShellDef["Shelly1.Outputs"]):
+        executable = "shelly"
+        arg: str = shell.outarg(argstr=None)
+
         class Outputs(ShellOutputs):
+
             file: File = shell.outarg(
-                path_template="{args}",
+                path_template="{arg}",
                 help="output file",
             )
 
     @shell.define
     class Shelly2(ShellDef["Shelly2.Outputs"]):
+        executable = "shelly2"
         orig_file: str = shell.arg(
             position=1,
             help="output file",
@@ -1761,12 +1768,12 @@ class Outputs(ShellOutputs):
             )
 
     @workflow.define(outputs=["touch_file", "out1", "cp_file", "out2"])
-    def Workflow(cmd1, cmd2, args):
+    def Workflow(cmd1, cmd2, arg):
 
         shelly1 = workflow.add(
             Shelly1(
                 executable=cmd1,
-                additional_args=args,
+                arg=arg,
             )
         )
         shelly2 = workflow.add(
@@ -1778,7 +1785,7 @@ def Workflow(cmd1, cmd2, args):
 
         return shelly1.file, shelly1.stdout, shelly2.out_file, shelly2.stdout
 
-    wf = Workflow(cmd1="touch", cmd2="cp", args=File.mock("newfile.txt"))
+    wf = Workflow(cmd1="touch", cmd2="cp", arg="newfile.txt")
 
     with Submitter(plugin="debug") as sub:
         res = sub(wf)
@@ -1861,14 +1868,20 @@ def test_wf_shell_cmd_ndst_1(plugin, tmp_path):
 
     @shell.define
     class Shelly1(ShellDef["Shelly1.Outputs"]):
+        executable = "shelly"
+
+        arg: str = shell.arg(argstr=None)
+
         class Outputs(ShellOutputs):
             file: File = shell.outarg(
-                path_template="{args}",
+                path_template="{arg}",
                 help="output file",
             )
 
     @shell.define
     class Shelly2(ShellDef["Shelly2.Outputs"]):
+        executable = "shelly2"
+
         orig_file: str = shell.arg(
             position=1,
             help="output file",
@@ -1889,7 +1902,7 @@ def Workflow(cmd1, cmd2, args):
         shelly1 = workflow.add(
             Shelly1(
                 executable=cmd1,
-            ).split("args", args=args)
+            ).split("arg", arg=args)
         )
         shelly2 = workflow.add(
             Shelly2(
@@ -1903,10 +1916,10 @@ def Workflow(cmd1, cmd2, args):
     wf = Workflow(
         cmd1="touch",
         cmd2="cp",
-        args=[File.mock("newfile_1.txt"), File.mock("newfile_2.txt")],
+        args=["newfile_1.txt", "newfile_2.txt"],
     )
 
-    with Submitter(plugin="debug") as sub:
+    with Submitter(plugin="debug", cache_dir=tmp_path) as sub:
         res = sub(wf)
 
     assert res.outputs.out1 == ["", ""]
@@ -3288,7 +3301,7 @@ class Outputs(ShellOutputs):
     # An exception should be raised because the second mandatory output does not exist
     with pytest.raises(
         ValueError,
-        match=r"file system path provided to mandatory field .* does not exist",
+        match=r"file system path\(s\) provided to mandatory field .* does not exist",
     ):
         shelly(cache_dir=tmp_path)
     # checking if the first output was created