hpcflow-new2 0.2.0a173__py3-none-any.whl → 0.2.0a175__py3-none-any.whl

hpcflow/sdk/core/workflow.py

@@ -1100,7 +1100,9 @@ class Workflow:
         # TODO: add new downstream elements?
 
     @TimeIt.decorator
-    def _add_empty_loop(self, loop: app.Loop) -> app.WorkflowLoop:
+    def _add_empty_loop(
+        self, loop: app.Loop
+    ) -> Tuple[app.WorkflowLoop, List[app.ElementIteration]]:
         """Add a new loop (zeroth iterations only) to the workflow."""
 
         new_index = self.num_loops
@@ -1111,27 +1113,32 @@ class Workflow:
         # add to the WorkflowTemplate:
         self.template._add_empty_loop(loop_c)
 
+        # all these element iterations will be initialised for the new loop:
+        iters = self.get_element_iterations_of_tasks(loop_c.task_insert_IDs)
+        iter_IDs = [i.id_ for i in iters]
+
         # create and insert a new WorkflowLoop:
-        self.loops.add_object(
-            self.app.WorkflowLoop.new_empty_loop(
-                index=new_index,
-                workflow=self,
-                template=loop_c,
-            )
+        new_loop = self.app.WorkflowLoop.new_empty_loop(
+            index=new_index,
+            workflow=self,
+            template=loop_c,
+            iterations=iters,
         )
+        self.loops.add_object(new_loop)
         wk_loop = self.loops[new_index]
 
-        loop_js, _ = loop_c.to_json_like()
+        # update any child loops of the new loop to include their new parent:
+        for chd_loop in wk_loop.get_child_loops():
+            chd_loop._update_parents(wk_loop)
 
-        # all these element iterations will be initialised for the new loop:
-        iter_IDs = [
-            i.id_ for i in self.get_element_iterations_of_tasks(loop_c.task_insert_IDs)
-        ]
+        loop_js, _ = loop_c.to_json_like()
 
         # update persistent store:
         self._store.add_loop(
             loop_template=loop_js,
             iterable_parameters=wk_loop.iterable_parameters,
+            parents=wk_loop.parents,
+            num_added_iterations=wk_loop.num_added_iterations,
             iter_IDs=iter_IDs,
         )
 
@@ -1140,18 +1147,18 @@ class Workflow:
         return wk_loop
 
     @TimeIt.decorator
-    def _add_loop(self, loop: app.Loop, parent_loop_indices: Dict = None) -> None:
+    def _add_loop(self, loop: app.Loop) -> None:
         new_wk_loop = self._add_empty_loop(loop)
         if loop.num_iterations is not None:
             # fixed number of iterations, so add remaining N > 0 iterations:
             for _ in range(loop.num_iterations - 1):
-                new_wk_loop.add_iteration(parent_loop_indices=parent_loop_indices)
+                new_wk_loop.add_iteration()
 
-    def add_loop(self, loop: app.Loop, parent_loop_indices: Dict = None) -> None:
+    def add_loop(self, loop: app.Loop) -> None:
         """Add a loop to a subset of workflow tasks."""
         with self._store.cached_load():
             with self.batch_update():
-                self._add_loop(loop, parent_loop_indices)
+                self._add_loop(loop)
 
     @property
     def creation_info(self):
@@ -1231,11 +1238,15 @@ class Workflow:
             with self._store.cached_load():
                 wk_loops = []
                 for idx, loop_dat in self._store.get_loops().items():
+                    num_add_iters = {
+                        tuple(i[0]): i[1] for i in loop_dat["num_added_iterations"]
+                    }
                     wk_loop = self.app.WorkflowLoop(
                         index=idx,
                         workflow=self,
                         template=self.template.loops[idx],
-                        num_added_iterations=loop_dat["num_added_iterations"],
+                        parents=loop_dat["parents"],
+                        num_added_iterations=num_add_iters,
                         iterable_parameters=loop_dat["iterable_parameters"],
                     )
                     wk_loops.append(wk_loop)
@@ -1472,6 +1483,7 @@ class Workflow:
 
                 for loop in self.loops:
                     loop._reset_pending_num_added_iters()
+                    loop._reset_pending_parents()
 
                 self._reject_pending()
 
@@ -1496,6 +1508,7 @@ class Workflow:
 
                     for loop in self.loops:
                         loop._accept_pending_num_added_iters()
+                        loop._accept_pending_parents()
 
                     # TODO: handle errors in commit pending?
                     self._store._pending.commit_all()
@@ -2073,27 +2086,79 @@ class Workflow:
                 yield element
 
     @TimeIt.decorator
-    def get_iteration_task_pathway(self):
+    def get_iteration_task_pathway(self, ret_iter_IDs=False, ret_data_idx=False):
         pathway = []
         for task in self.tasks:
-            loop_idx = {}
-            pathway.append((task.insert_ID, loop_idx))
-
-        for loop in self.loops:  # TODO: order by depth (inner loops first?)
-            task_subset = loop.task_insert_IDs
-            subset_idx = [idx for idx, i in enumerate(pathway) if i[0] in task_subset]
-            looped_pathway = []
-            for iter_i in range(loop.num_added_iterations):
-                for j in subset_idx:
-                    item_j = copy.deepcopy(pathway[j])
-                    item_j[1][loop.name] = iter_i
-                    looped_pathway.append(item_j)
-
-            # replaced pathway `sub_idx` items with `looped_pathway` items:
-            pathway = replace_items(
-                pathway, subset_idx[0], subset_idx[-1] + 1, looped_pathway
+            pathway.append((task.insert_ID, {}))
+
+        added_loop_names = set()
+        for _ in range(self.num_loops):
+            to_add = None
+            for loop in self.loops:
+                if loop.name in added_loop_names:
+                    continue
+                elif set(loop.parents).issubset(added_loop_names):
+                    # add a loop only once their parents have been added:
+                    to_add = loop
+                    break
+
+            if to_add is None:
+                raise RuntimeError(
+                    "Failed to find a loop whose parents have already been added to the "
+                    "iteration task pathway."
+                )
+
+            iIDs = to_add.task_insert_IDs
+            relevant_idx = [idx for idx, i in enumerate(pathway) if i[0] in iIDs]
+
+            for num_add_k, num_add in to_add.num_added_iterations.items():
+                parent_loop_idx = {
+                    to_add.parents[idx]: i for idx, i in enumerate(num_add_k)
+                }
+
+                repl = []
+                repl_idx = []
+                for i in range(num_add):
+                    for p_idx, p in enumerate(pathway):
+                        skip = False
+                        if p[0] not in iIDs:
+                            continue
+                        for k, v in parent_loop_idx.items():
+                            if p[1][k] != v:
+                                skip = True
+                                break
+                        if skip:
+                            continue
+                        p = copy.deepcopy(p)
+                        p[1].update({to_add.name: i})
+                        repl_idx.append(p_idx)
+                        repl.append(p)
+
+                if repl:
+                    repl_start, repl_stop = min(repl_idx), max(repl_idx)
+                    pathway = replace_items(pathway, repl_start, repl_stop + 1, repl)
+
+            added_loop_names.add(to_add.name)
+
+        if added_loop_names != set(i.name for i in self.loops):
+            raise RuntimeError(
+                "Not all loops have been considered in the iteration task pathway."
             )
 
+        if ret_iter_IDs or ret_data_idx:
+            all_iters = self.get_all_element_iterations()
+            for idx, i in enumerate(pathway):
+                i_iters = []
+                for iter_j in all_iters:
+                    if iter_j.task.insert_ID == i[0] and iter_j.loop_idx == i[1]:
+                        i_iters.append(iter_j)
+                new = list(i)
+                if ret_iter_IDs:
+                    new += [tuple([j.id_ for j in i_iters])]
+                if ret_data_idx:
+                    new += [tuple(j.get_data_idx() for j in i_iters)]
+                pathway[idx] = tuple(new)
+
         return pathway
 
     @TimeIt.decorator

hpcflow/sdk/persistence/base.py

@@ -934,19 +934,24 @@ class PersistentStore(ABC):
         self,
         loop_template: Dict,
         iterable_parameters,
+        parents: List[str],
+        num_added_iterations: Dict[Tuple[int], int],
         iter_IDs: List[int],
         save: bool = True,
     ):
         """Add a new loop to the workflow."""
         self.logger.debug(f"Adding store loop.")
         new_idx = self._get_num_total_loops()
+        added_iters = [[list(k), v] for k, v in num_added_iterations.items()]
         self._pending.add_loops[new_idx] = {
             "loop_template": loop_template,
             "iterable_parameters": iterable_parameters,
+            "parents": parents,
+            "num_added_iterations": added_iters,
         }
 
         for i in iter_IDs:
-            self._pending.update_loop_indices[i] = {loop_template["name"]: 0}
+            self._pending.update_loop_indices[i].update({loop_template["name"]: 0})
 
         if save:
             self.save()
@@ -1298,12 +1303,30 @@ class PersistentStore(ABC):
             self.save()
 
     def update_loop_num_iters(
-        self, index: int, num_iters: int, save: bool = True
+        self, index: int, num_added_iters: int, save: bool = True
     ) -> None:
         self.logger.debug(
-            f"Updating loop {index!r} num added iterations to {num_iters!r}."
+            f"Updating loop {index!r} num added iterations to {num_added_iters!r}."
         )
-        self._pending.update_loop_num_iters[index] = num_iters
+        num_added_iters = [[list(k), v] for k, v in num_added_iters.items()]
+        self._pending.update_loop_num_iters[index] = num_added_iters
+        if save:
+            self.save()
+
+    def update_loop_parents(
+        self,
+        index: int,
+        num_added_iters: int,
+        parents: List[str],
+        save: bool = True,
+    ) -> None:
+        self.logger.debug(
+            f"Updating loop {index!r} parents to {parents!r}, and num added iterations "
+            f"to {num_added_iters}."
+        )
+        num_added_iters = [[list(k), v] for k, v in num_added_iters.items()]
+        self._pending.update_loop_num_iters[index] = num_added_iters
+        self._pending.update_loop_parents[index] = parents
         if save:
             self.save()
 
@@ -1349,6 +1372,11 @@ class PersistentStore(ABC):
             pend_num_iters = self._pending.update_loop_num_iters.get(id_)
             if pend_num_iters:
                 loop_i["num_added_iterations"] = pend_num_iters
+            # consider pending change to parents:
+            pend_parents = self._pending.update_loop_parents.get(id_)
+            if pend_parents:
+                loop_i["parents"] = pend_parents
+
             loops_new[id_] = loop_i
         return loops_new
 

hpcflow/sdk/persistence/json.py

@@ -54,6 +54,7 @@ class JSONPersistentStore(PersistentStore):
         commit_tasks=(_meta_res,),
         commit_loops=(_meta_res,),
         commit_loop_num_iters=(_meta_res,),
+        commit_loop_parents=(_meta_res,),
         commit_submissions=(_subs_res,),
         commit_submission_parts=(_subs_res,),
         commit_js_metadata=(_subs_res,),
@@ -169,6 +170,7 @@ class JSONPersistentStore(PersistentStore):
                     {
                         "num_added_iterations": loop["num_added_iterations"],
                         "iterable_parameters": loop["iterable_parameters"],
+                        "parents": loop["parents"],
                     }
                 )
                 md["template"]["loops"].append(loop["loop_template"])
@@ -225,6 +227,10 @@ class JSONPersistentStore(PersistentStore):
         with self.using_resource("metadata", action="update") as md:
             md["loops"][index]["num_added_iterations"] = num_iters
 
+    def _update_loop_parents(self, index: int, parents: List[str]):
+        with self.using_resource("metadata", action="update") as md:
+            md["loops"][index]["parents"] = parents
+
     def _append_EARs(self, EARs: List[StoreEAR]):
         with self.using_resource("metadata", action="update") as md:
             md["runs"].extend(i.encode(self.ts_fmt) for i in EARs)

hpcflow/sdk/persistence/pending.py

@@ -69,8 +69,9 @@ class PendingChanges:
         self.set_parameters: Dict[int, AnySParameter] = None
 
         self.update_param_sources: Dict[int, Dict] = None
-        self.update_loop_indices: Dict[int, Dict] = None
+        self.update_loop_indices: Dict[int, Dict[str, int]] = None
         self.update_loop_num_iters: Dict[int, int] = None
+        self.update_loop_parents: Dict[int, List[str]] = None
 
         self.reset(is_init=True)  # set up initial data structures
 
@@ -101,6 +102,7 @@ class PendingChanges:
             or bool(self.update_param_sources)
             or bool(self.update_loop_indices)
             or bool(self.update_loop_num_iters)
+            or bool(self.update_loop_parents)
         )
 
     def where_pending(self) -> List[str]:
@@ -158,6 +160,16 @@ class PendingChanges:
             loop_ids = list(self.add_loops.keys())
             self.logger.debug(f"commit: adding pending loops with indices {loop_ids!r}")
             self.store._append_loops(loops)
+
+            # pending num_added_iters and parents that belong to pending loops are now
+            # committed:
+            self.update_loop_num_iters = {
+                k: v for k, v in self.update_loop_num_iters.items() if k not in loop_ids
+            }
+            self.update_loop_parents = {
+                k: v for k, v in self.update_loop_parents.items() if k not in loop_ids
+            }
+
         self.clear_add_loops()
 
     @TimeIt.decorator
@@ -415,6 +427,14 @@ class PendingChanges:
             self.store._update_loop_num_iters(index, num_iters)
         self.clear_update_loop_num_iters()
 
+    @TimeIt.decorator
+    def commit_loop_parents(self) -> None:
+        """Make pending update to additional loop parents."""
+        for index, parents in self.update_loop_parents.items():
+            self.logger.debug(f"commit: updating loop {index!r} parents to {parents!r}.")
+            self.store._update_loop_parents(index, parents)
+        self.clear_update_loop_parents()
+
     def clear_add_tasks(self):
         self.add_tasks = {}
 
@@ -482,11 +502,14 @@ class PendingChanges:
         self.update_param_sources = {}
 
     def clear_update_loop_indices(self):
-        self.update_loop_indices = {}
+        self.update_loop_indices = defaultdict(dict)
 
     def clear_update_loop_num_iters(self):
         self.update_loop_num_iters = {}
 
+    def clear_update_loop_parents(self):
+        self.update_loop_parents = {}
+
     def reset(self, is_init=False) -> None:
         """Clear all pending data and prepare to accept new pending data."""
 
@@ -526,6 +549,7 @@ class PendingChanges:
         self.clear_update_param_sources()
         self.clear_update_loop_indices()
         self.clear_update_loop_num_iters()
+        self.clear_update_loop_parents()
 
 
 @dataclass
@@ -561,6 +585,7 @@ class CommitResourceMap:
     commit_param_sources: Optional[Tuple[str]] = tuple()
     commit_loop_indices: Optional[Tuple[str]] = tuple()
     commit_loop_num_iters: Optional[Tuple[str]] = tuple()
+    commit_loop_parents: Optional[Tuple[str]] = tuple()
 
     def __post_init__(self):
         self.groups = self.group_by_resource()

hpcflow/sdk/persistence/zarr.py

@@ -497,6 +497,7 @@ class ZarrPersistentStore(PersistentStore):
                     {
                         "num_added_iterations": loop["num_added_iterations"],
                         "iterable_parameters": loop["iterable_parameters"],
+                        "parents": loop["parents"],
                     }
                 )
                 attrs["template"]["loops"].append(loop["loop_template"])
@@ -587,6 +588,10 @@ class ZarrPersistentStore(PersistentStore):
         with self.using_resource("attrs", action="update") as attrs:
             attrs["loops"][index]["num_added_iterations"] = num_iters
 
+    def _update_loop_parents(self, index: int, parents: List[str]):
+        with self.using_resource("attrs", action="update") as attrs:
+            attrs["loops"][index]["parents"] = parents
+
     def _append_EARs(self, EARs: List[ZarrStoreEAR]):
         arr = self._get_EARs_arr(mode="r+")
         attrs_orig = arr.attrs.asdict()

hpcflow/tests/unit/test_group.py

@@ -0,0 +1,91 @@
+import pytest
+from hpcflow.app import app as hf
+from hpcflow.sdk.core.errors import MissingElementGroup
+
+
+def test_group_simple(null_config, tmp_path):
+    s1 = hf.TaskSchema(
+        objective="t1",
+        inputs=[hf.SchemaInput("p1")],
+        outputs=[hf.SchemaOutput("p2")],
+        actions=[
+            hf.Action(
+                commands=[
+                    hf.Command(
+                        "echo $(( <<parameter:p1>> + 1 ))", stdout="<<parameter:p2>>"
+                    )
+                ]
+            )
+        ],
+    )
+    s2 = hf.TaskSchema(
+        objective="t2",
+        inputs=[hf.SchemaInput("p2", group="my_group")],
+        outputs=[hf.SchemaOutput("p3")],
+        actions=[
+            hf.Action(
+                commands=[
+                    hf.Command(
+                        "echo $(( <<parameter:p2>> + 2 ))", stdout="<<parameter:p3>>"
+                    )
+                ]
+            )
+        ],
+    )
+
+    t1 = hf.Task(
+        schema=s1,
+        sequences=[hf.ValueSequence("inputs.p1", values=[1, 2, 3])],
+        groups=[hf.ElementGroup(name="my_group")],
+    )
+    t2 = hf.Task(schema=s2)
+    wk = hf.Workflow.from_template_data(
+        template_name="test_groups",
+        path=tmp_path,
+        tasks=[t1, t2],
+    )
+    assert [task.num_elements for task in wk.tasks] == [3, 1]
+    assert len(wk.tasks.t2.elements[0].get_data_idx("inputs.p2")["inputs.p2"]) == 3
+
+
+def test_group_raise_no_elements(null_config, tmp_path):
+    s1 = hf.TaskSchema(
+        objective="t1",
+        inputs=[hf.SchemaInput("p1")],
+        outputs=[hf.SchemaOutput("p2")],
+        actions=[
+            hf.Action(
+                commands=[
+                    hf.Command(
+                        "echo $(( <<parameter:p1>> + 1 ))", stdout="<<parameter:p2>>"
+                    )
+                ]
+            )
+        ],
+    )
+    s2 = hf.TaskSchema(
+        objective="t2",
+        inputs=[hf.SchemaInput("p2", group="my_group")],
+        outputs=[hf.SchemaOutput("p3")],
+        actions=[
+            hf.Action(
+                commands=[
+                    hf.Command(
+                        "echo $(( <<parameter:p2>> + 2 ))", stdout="<<parameter:p3>>"
+                    )
+                ]
+            )
+        ],
+    )
+
+    t1 = hf.Task(
+        schema=s1,
+        sequences=[hf.ValueSequence("inputs.p1", values=[1, 2, 3])],
+    )
+    t2 = hf.Task(schema=s2)
+    with pytest.raises(MissingElementGroup):
+        hf.Workflow.from_template_data(
+            template_name="test_groups",
+            path=tmp_path,
+            tasks=[t1, t2],
+        )