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

hpcflow/_version.py

@@ -1 +1 @@
-__version__ = "0.2.0a175"
+__version__ = "0.2.0a177"

hpcflow/sdk/core/cache.py

@@ -0,0 +1,142 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Set, Dict
+
+from hpcflow.sdk.log import TimeIt
+
+
+@dataclass
+class DependencyCache:
+    """Class to bulk-retrieve dependencies between elements, iterations, and runs."""
+
+    run_dependencies: Dict[int, Set]
+    run_dependents: Dict[int, Set]
+    iter_run_dependencies: Dict[int, Set]
+    iter_iter_dependencies: Dict[int, Set]
+    elem_iter_dependencies: Dict[int, Set]
+    elem_elem_dependencies: Dict[int, Set]
+    elem_elem_dependents: Dict[int, Set]
+    elem_elem_dependents_rec: Dict[int, Set]
+
+    elements: Dict
+    iterations: Dict
+
+    @classmethod
+    @TimeIt.decorator
+    def build(cls, workflow):
+        num_iters = workflow.num_element_iterations
+        num_elems = workflow.num_elements
+        num_runs = workflow.num_EARs
+
+        all_store_runs = workflow._store.get_EARs(list(range(num_runs)))
+        all_store_iters = workflow._store.get_element_iterations(list(range(num_iters)))
+        all_store_elements = workflow._store.get_elements(list(range(num_elems)))
+        all_param_sources = workflow.get_all_parameter_sources()
+        all_data_idx = [
+            {
+                k: v if isinstance(v, list) else [v]
+                for k, v in i.data_idx.items()
+                if k not in ("repeats.",)
+            }
+            for i in all_store_runs
+        ]
+
+        # run dependencies and dependents
+        run_dependencies = {}
+        run_dependents = defaultdict(set)
+        for idx, i in enumerate(all_data_idx):
+            run_i_sources = set()
+            for j in i.values():
+                for k in j:
+                    run_k = all_param_sources[k].get("EAR_ID")
+                    if run_k is not None and run_k != idx:
+                        run_i_sources.add(run_k)
+            run_dependencies[idx] = run_i_sources
+            for m in run_i_sources:
+                run_dependents[m].add(idx)
+
+        # add missing:
+        for k in range(num_runs):
+            run_dependents[k]
+
+        run_dependents = dict(run_dependents)
+
+        # iteration dependencies
+        all_iter_run_IDs = {
+            i.id_: [k for j in i.EAR_IDs.values() for k in j] for i in all_store_iters
+        }
+        # for each iteration, which runs does it depend on?
+        iter_run_dependencies = {
+            k: set(j for i in v for j in run_dependencies[i])
+            for k, v in all_iter_run_IDs.items()
+        }
+
+        # for each run, which iteration does it belong to?
+        all_run_iter_IDs = {}
+        for iter_ID, run_IDs in all_iter_run_IDs.items():
+            for run_ID in run_IDs:
+                all_run_iter_IDs[run_ID] = iter_ID
+
+        # for each iteration, which iterations does it depend on?
+        iter_iter_dependencies = {
+            k: set(all_run_iter_IDs[i] for i in v)
+            for k, v in iter_run_dependencies.items()
+        }
+
+        all_elem_iter_IDs = {i.id_: i.iteration_IDs for i in all_store_elements}
+
+        elem_iter_dependencies = {
+            k: set(j for i in v for j in iter_iter_dependencies[i])
+            for k, v in all_elem_iter_IDs.items()
+        }
+
+        # for each iteration, which element does it belong to?
+        all_iter_elem_IDs = {}
+        for elem_ID, iter_IDs in all_elem_iter_IDs.items():
+            for iter_ID in iter_IDs:
+                all_iter_elem_IDs[iter_ID] = elem_ID
+
+        # element dependencies
+        elem_elem_dependencies = {
+            k: set(all_iter_elem_IDs[i] for i in v)
+            for k, v in elem_iter_dependencies.items()
+        }
+
+        # for each element, which elements depend on it (directly)?
+        elem_elem_dependents = defaultdict(set)
+        for k, v in elem_elem_dependencies.items():
+            for i in v:
+                elem_elem_dependents[i].add(k)
+
+        # for each element, which elements depend on it (recursively)?
+        elem_elem_dependents_rec = defaultdict(set)
+        for k in list(elem_elem_dependents):
+            for i in elem_elem_dependents[k]:
+                elem_elem_dependents_rec[k].add(i)
+                elem_elem_dependents_rec[k].update(
+                    {m for m in elem_elem_dependents[i] if m != k}
+                )
+
+        # add missing keys:
+        for k in range(num_elems):
+            elem_elem_dependents[k]
+            elem_elem_dependents_rec[k]
+
+        elem_elem_dependents = dict(elem_elem_dependents)
+        elem_elem_dependents_rec = dict(elem_elem_dependents_rec)
+
+        elements = workflow.get_all_elements()
+        iterations = workflow.get_all_element_iterations()
+
+        return cls(
+            run_dependencies=run_dependencies,
+            run_dependents=run_dependents,
+            iter_run_dependencies=iter_run_dependencies,
+            iter_iter_dependencies=iter_iter_dependencies,
+            elem_iter_dependencies=elem_iter_dependencies,
+            elem_elem_dependencies=elem_elem_dependencies,
+            elem_elem_dependents=elem_elem_dependents,
+            elem_elem_dependents_rec=elem_elem_dependents_rec,
+            elements=elements,
+            iterations=iterations,
+        )

hpcflow/sdk/core/element.py

@@ -675,6 +675,7 @@ class ElementIteration:
             default=default,
         )
 
+    @TimeIt.decorator
     def get_EAR_dependencies(
         self,
         as_objects: Optional[bool] = False,
@@ -708,6 +709,7 @@ class ElementIteration:
             out = self.workflow.get_EARs_from_IDs(out)
         return out
 
+    @TimeIt.decorator
     def get_element_iteration_dependencies(
         self, as_objects: bool = False
     ) -> List[Union[int, app.ElementIteration]]:
@@ -719,6 +721,7 @@ class ElementIteration:
             out = self.workflow.get_element_iterations_from_IDs(out)
         return out
 
+    @TimeIt.decorator
     def get_element_dependencies(
         self,
         as_objects: Optional[bool] = False,
@@ -769,6 +772,7 @@ class ElementIteration:
 
         return out
 
+    @TimeIt.decorator
     def get_dependent_EARs(
         self, as_objects: bool = False
     ) -> List[Union[int, app.ElementActionRun]]:
@@ -793,6 +797,7 @@ class ElementIteration:
 
         return deps
 
+    @TimeIt.decorator
     def get_dependent_element_iterations(
         self, as_objects: bool = False
     ) -> List[Union[int, app.ElementIteration]]:
@@ -816,6 +821,7 @@ class ElementIteration:
 
         return deps
 
+    @TimeIt.decorator
     def get_dependent_elements(
         self,
         as_objects: bool = False,
@@ -1246,6 +1252,7 @@ class Element:
         """Get tasks that depend on the most recent iteration of this element."""
         return self.latest_iteration.get_dependent_tasks(as_objects=as_objects)
 
+    @TimeIt.decorator
     def get_dependent_elements_recursively(self, task_insert_ID=None):
         """Get downstream elements that depend on this element, including recursive
         dependencies.

hpcflow/sdk/core/loop.py

@@ -6,9 +6,11 @@ from typing import Dict, List, Optional, Tuple, Union
 from hpcflow.sdk import app
 from hpcflow.sdk.core.errors import LoopTaskSubsetError
 from hpcflow.sdk.core.json_like import ChildObjectSpec, JSONLike
+from hpcflow.sdk.core.loop_cache import LoopCache
 from hpcflow.sdk.core.parameters import InputSourceType
 from hpcflow.sdk.core.task import WorkflowTask
-from hpcflow.sdk.core.utils import check_valid_py_identifier
+from hpcflow.sdk.core.utils import check_valid_py_identifier, nth_key, nth_value
+from hpcflow.sdk.log import TimeIt
 
 # from .parameters import Parameter
 
@@ -198,6 +200,7 @@ class WorkflowLoop:
 
         self._validate()
 
+    @TimeIt.decorator
     def _validate(self):
         # task subset must be a contiguous range of task indices:
         task_indices = self.task_indices
@@ -328,6 +331,7 @@ class WorkflowLoop:
         return self.workflow.tasks[: self.task_objects[0].index]
 
     @staticmethod
+    @TimeIt.decorator
     def _find_iterable_parameters(loop_template: app.Loop):
         all_inputs_first_idx = {}
         all_outputs_idx = {}
@@ -355,18 +359,19 @@ class WorkflowLoop:
         return iterable_params
 
     @classmethod
+    @TimeIt.decorator
     def new_empty_loop(
         cls,
         index: int,
         workflow: app.Workflow,
         template: app.Loop,
-        iterations: List[app.ElementIteration],
+        iter_loop_idx: List[Dict],
     ) -> Tuple[app.WorkflowLoop, List[Dict[str, int]]]:
         parent_loops = cls._get_parent_loops(index, workflow, template)
         parent_names = [i.name for i in parent_loops]
         num_added_iters = {}
-        for iter_i in iterations:
-            num_added_iters[tuple([iter_i.loop_idx[j] for j in parent_names])] = 1
+        for i in iter_loop_idx:
+            num_added_iters[tuple([i[j] for j in parent_names])] = 1
 
         obj = cls(
             index=index,
@@ -379,6 +384,7 @@ class WorkflowLoop:
         return obj
 
     @classmethod
+    @TimeIt.decorator
     def _get_parent_loops(
         cls,
         index: int,
@@ -399,12 +405,14 @@ class WorkflowLoop:
                 parents.append(loop_i)
         return parents
 
+    @TimeIt.decorator
     def get_parent_loops(self) -> List[app.WorkflowLoop]:
         """Get loops whose task subset is a superset of this loop's task subset. If two
         loops have identical task subsets, the first loop in the workflow loop list is
         considered the child."""
         return self._get_parent_loops(self.index, self.workflow, self.template)
 
+    @TimeIt.decorator
     def get_child_loops(self) -> List[app.WorkflowLoop]:
         """Get loops whose task subset is a subset of this loop's task subset. If two
         loops have identical task subsets, the first loop in the workflow loop list is
@@ -426,10 +434,12 @@ class WorkflowLoop:
         children = sorted(children, key=lambda x: len(next(iter(x.num_added_iterations))))
         return children
 
-    def add_iteration(self, parent_loop_indices=None):
+    @TimeIt.decorator
+    def add_iteration(self, parent_loop_indices=None, cache: Optional[LoopCache] = None):
+        if not cache:
+            cache = LoopCache.build(self.workflow)
         parent_loops = self.get_parent_loops()
         child_loops = self.get_child_loops()
-        child_loop_names = [i.name for i in child_loops]
         parent_loop_indices = parent_loop_indices or {}
         if parent_loops and not parent_loop_indices:
             parent_loop_indices = {i.name: 0 for i in parent_loops}
@@ -438,25 +448,39 @@ class WorkflowLoop:
         cur_loop_idx = self.num_added_iterations[iters_key] - 1
         all_new_data_idx = {}  # keys are (task.insert_ID and element.index)
 
+        # initialise a new `num_added_iterations` key on each child loop:
+        for child in child_loops:
+            iters_key_dct = {
+                **parent_loop_indices,
+                self.name: cur_loop_idx + 1,
+            }
+            added_iters_key_chd = tuple([iters_key_dct.get(j, 0) for j in child.parents])
+            child._initialise_pending_added_iters(added_iters_key_chd)
+
         for task in self.task_objects:
+
+            new_loop_idx = {
+                **parent_loop_indices,
+                self.name: cur_loop_idx + 1,
+                **{
+                    child.name: 0
+                    for child in child_loops
+                    if task.insert_ID in child.task_insert_IDs
+                },
+            }
+            added_iter_IDs = []
             for elem_idx in range(task.num_elements):
-                # element needs to take into account changes made in this code
-                element = task.elements[elem_idx]
-                inp_statuses = task.template.get_input_statuses(element.element_set)
+
+                elem_ID = task.element_IDs[elem_idx]
+
                 new_data_idx = {}
-                existing_inners = []
-                for iter_i in element.iterations:
-                    if iter_i.loop_idx[self.name] == cur_loop_idx:
-                        existing_inner_i = {
-                            k: v
-                            for k, v in iter_i.loop_idx.items()
-                            if k in child_loop_names
-                        }
-                        if existing_inner_i:
-                            existing_inners.append(existing_inner_i)
 
                 # copy resources from zeroth iteration:
-                for key, val in element.iterations[0].get_data_idx().items():
+                zeroth_iter_ID, zi_iter_data_idx = cache.zeroth_iters[elem_ID]
+                zi_elem_ID, zi_idx = cache.iterations[zeroth_iter_ID]
+                zi_data_idx = nth_value(cache.data_idx[zi_elem_ID], zi_idx)
+
+                for key, val in zi_data_idx.items():
                     if key.startswith("resources."):
                         new_data_idx[key] = val
 
@@ -474,41 +498,47 @@ class WorkflowLoop:
                         # identify element(s) from which this iterable input should be
                         # parametrised:
                         if task.insert_ID == iter_dat["output_tasks"][-1]:
-                            src_elem = element
+                            src_elem_ID = elem_ID
                             grouped_elems = None
                         else:
-                            src_elems = element.get_dependent_elements_recursively(
-                                task_insert_ID=iter_dat["output_tasks"][-1]
-                            )
+                            src_elem_IDs_all = cache.element_dependents[elem_ID]
+                            src_elem_IDs = {
+                                k: v
+                                for k, v in src_elem_IDs_all.items()
+                                if cache.elements[k]["task_insert_ID"]
+                                == iter_dat["output_tasks"][-1]
+                            }
                             # consider groups
                             inp_group_name = inp.single_labelled_data.get("group")
                             grouped_elems = []
-                            for i in src_elems:
+                            for src_elem_j_ID, src_elem_j_dat in src_elem_IDs.items():
                                 i_in_group = any(
-                                    j.name == inp_group_name for j in i.element_set.groups
+                                    k == inp_group_name
+                                    for k in src_elem_j_dat["group_names"]
                                 )
                                 if i_in_group:
-                                    grouped_elems.append(i)
+                                    grouped_elems.append(src_elem_j_ID)
 
-                            if not grouped_elems and len(src_elems) > 1:
+                            if not grouped_elems and len(src_elem_IDs) > 1:
                                 raise NotImplementedError(
-                                    f"Multiple elements found in the iterable parameter {inp!r}'s"
-                                    f" latest output task (insert ID: "
-                                    f"{iter_dat['output_tasks'][-1]}) that can be used to "
-                                    f"parametrise the next iteration: {src_elems!r}."
+                                    f"Multiple elements found in the iterable parameter "
+                                    f"{inp!r}'s latest output task (insert ID: "
+                                    f"{iter_dat['output_tasks'][-1]}) that can be used "
+                                    f"to parametrise the next iteration: "
+                                    f"{list(src_elem_IDs.keys())!r}."
                                 )
 
-                            elif not src_elems:
+                            elif not src_elem_IDs:
                                 # TODO: maybe OK?
                                 raise NotImplementedError(
-                                    f"No elements found in the iterable parameter {inp!r}'s"
-                                    f" latest output task (insert ID: "
-                                    f"{iter_dat['output_tasks'][-1]}) that can be used to "
-                                    f"parametrise the next iteration."
+                                    f"No elements found in the iterable parameter "
+                                    f"{inp!r}'s latest output task (insert ID: "
+                                    f"{iter_dat['output_tasks'][-1]}) that can be used "
+                                    f"to parametrise the next iteration."
                                 )
 
                             else:
-                                src_elem = src_elems[0]
+                                src_elem_ID = nth_key(src_elem_IDs, 0)
 
                         child_loop_max_iters = {}
                         parent_loop_same_iters = {
@@ -534,63 +564,69 @@ class WorkflowLoop:
 
                         # identify the ElementIteration from which this input should be
                         # parametrised:
-                        source_iter = None
+                        loop_idx_key = tuple(sorted(source_iter_loop_idx.items()))
                         if grouped_elems:
-                            source_iter = []
-                            for src_elem in grouped_elems:
-                                for iter_i in src_elem.iterations:
-                                    if iter_i.loop_idx == source_iter_loop_idx:
-                                        source_iter.append(iter_i)
-                                        break
+                            src_data_idx = []
+                            for src_elem_ID in grouped_elems:
+                                src_data_idx.append(
+                                    cache.data_idx[src_elem_ID][loop_idx_key]
+                                )
                         else:
-                            for iter_i in src_elem.iterations:
-                                if iter_i.loop_idx == source_iter_loop_idx:
-                                    source_iter = iter_i
-                                    break
+                            src_data_idx = cache.data_idx[src_elem_ID][loop_idx_key]
 
-                        if not source_iter:
+                        if not src_data_idx:
                             raise RuntimeError(
                                 f"Could not find a source iteration with loop_idx: "
                                 f"{source_iter_loop_idx!r}."
                             )
 
                         if grouped_elems:
-                            inp_dat_idx = [
-                                i.get_data_idx()[f"outputs.{inp.typ}"]
-                                for i in source_iter
-                            ]
+                            inp_dat_idx = [i[f"outputs.{inp.typ}"] for i in src_data_idx]
                         else:
-                            inp_dat_idx = source_iter.get_data_idx()[f"outputs.{inp.typ}"]
+                            inp_dat_idx = src_data_idx[f"outputs.{inp.typ}"]
                         new_data_idx[f"inputs.{inp.typ}"] = inp_dat_idx
 
                     else:
                         inp_key = f"inputs.{inp.typ}"
 
-                        orig_inp_src = element.input_sources[inp_key]
+                        orig_inp_src = cache.elements[elem_ID]["input_sources"][inp_key]
                         inp_dat_idx = None
 
                         if orig_inp_src.source_type is InputSourceType.LOCAL:
                             # keep locally defined inputs from original element
-                            inp_dat_idx = element.iterations[0].get_data_idx()[inp_key]
+                            inp_dat_idx = zi_data_idx[inp_key]
 
                         elif orig_inp_src.source_type is InputSourceType.DEFAULT:
                             # keep default value from original element
-                            inp_dat_idx_iter_0 = element.iterations[0].get_data_idx()
                             try:
-                                inp_dat_idx = inp_dat_idx_iter_0[inp_key]
+                                inp_dat_idx = zi_data_idx[inp_key]
                             except KeyError:
                                 # if this input is required by a conditional action, and
                                 # that condition is not met, then this input will not
                                 # exist in the action-run data index, so use the initial
                                 # iteration data index:
-                                inp_dat_idx = element.iterations[0].data_idx[inp_key]
+                                inp_dat_idx = zi_iter_data_idx[inp_key]
 
                         elif orig_inp_src.source_type is InputSourceType.TASK:
                             if orig_inp_src.task_ref not in self.task_insert_IDs:
-                                # source task not part of the loop; copy existing data idx:
-                                inp_dat_idx = element.iterations[0].get_data_idx()[
-                                    inp_key
-                                ]
+                                # source the data_idx from the iteration with same parent
+                                # loop indices as the new iteration to add:
+                                # src_iters = []
+                                src_data_idx = []
+                                for li_k, di_k in cache.data_idx[elem_ID].items():
+                                    skip_iter = False
+                                    li_k_dct = dict(li_k)
+                                    for p_k, p_v in parent_loop_indices.items():
+                                        if li_k_dct.get(p_k) != p_v:
+                                            skip_iter = True
+                                            break
+                                    if not skip_iter:
+                                        src_data_idx.append(di_k)
+
+                                # could be multiple, but they should all have the same
+                                # data index for this parameter:
+                                src_data_idx = src_data_idx[0]
+                                inp_dat_idx = src_data_idx[inp_key]
                             else:
                                 is_group = False
                                 if (
@@ -613,19 +649,24 @@ class WorkflowLoop:
                                         # find which element in that task `element`
                                         # depends on:
                                         task_i = self.workflow.tasks.get(insert_ID=tiID)
-                                        elem_i = task_i.elements[e_idx]
-                                        src_elems_i = (
-                                            elem_i.get_dependent_elements_recursively(
-                                                task_insert_ID=task.insert_ID
-                                            )
-                                        )
-                                        # filter src_elems_i for matching element IDs:
-                                        src_elems_i = [
-                                            i for i in src_elems_i if i.id_ == element.id_
+                                        elem_i_ID = task_i.element_IDs[e_idx]
+                                        src_elem_IDs_all = cache.element_dependents[
+                                            elem_i_ID
+                                        ]
+                                        src_elem_IDs_i = {
+                                            k: v
+                                            for k, v in src_elem_IDs_all.items()
+                                            if cache.elements[k]["task_insert_ID"]
+                                            == task.insert_ID
+                                        }
+
+                                        # filter src_elem_IDs_i for matching element IDs:
+                                        src_elem_IDs_i = [
+                                            i for i in src_elem_IDs_i if i == elem_ID
                                         ]
                                         if (
-                                            len(src_elems_i) == 1
-                                            and src_elems_i[0].id_ == element.id_
+                                            len(src_elem_IDs_i) == 1
+                                            and src_elem_IDs_i[0] == elem_ID
                                         ):
                                             new_sources.append((tiID, e_idx))
 
@@ -648,10 +689,11 @@ class WorkflowLoop:
                         new_data_idx[inp_key] = inp_dat_idx
 
                 # add any locally defined sub-parameters:
+                inp_statuses = cache.elements[elem_ID]["input_statuses"]
                 inp_status_inps = set([f"inputs.{i}" for i in inp_statuses])
                 sub_params = inp_status_inps - set(new_data_idx.keys())
                 for sub_param_i in sub_params:
-                    sub_param_data_idx_iter_0 = element.iterations[0].get_data_idx()
+                    sub_param_data_idx_iter_0 = zi_data_idx
                     try:
                         sub_param_data_idx = sub_param_data_idx_iter_0[sub_param_i]
                     except KeyError:
@@ -659,7 +701,7 @@ class WorkflowLoop:
                         # and that condition is not met, then this input will not exist in
                         # the action-run data index, so use the initial iteration data
                         # index:
-                        sub_param_data_idx = element.iterations[0].data_idx[sub_param_i]
+                        sub_param_data_idx = zi_data_idx[sub_param_i]
 
                     new_data_idx[sub_param_i] = sub_param_data_idx
 
@@ -671,30 +713,26 @@ class WorkflowLoop:
                 schema_params = set(
                     i for i in new_data_idx.keys() if len(i.split(".")) == 2
                 )
-                all_new_data_idx[(task.insert_ID, element.index)] = new_data_idx
-
-                new_loop_idx = {
-                    **parent_loop_indices,
-                    self.name: cur_loop_idx + 1,
-                    **{
-                        child.name: 0
-                        for child in child_loops
-                        if task.insert_ID in child.task_insert_IDs
-                    },
-                }
-                # increment num_added_iterations on child loop for this parent loop index:
-                for i in child_loops:
-                    added_iters_key_chd = tuple([new_loop_idx[j] for j in i.parents])
-                    i._initialise_pending_added_iters(added_iters_key_chd)
+                all_new_data_idx[(task.insert_ID, elem_idx)] = new_data_idx
 
                 iter_ID_i = self.workflow._store.add_element_iteration(
-                    element_ID=element.id_,
+                    element_ID=elem_ID,
                     data_idx=new_data_idx,
                     schema_parameters=list(schema_params),
                     loop_idx=new_loop_idx,
                 )
+                if cache:
+                    cache.add_iteration(
+                        iter_ID=iter_ID_i,
+                        task_insert_ID=task.insert_ID,
+                        element_ID=elem_ID,
+                        loop_idx=new_loop_idx,
+                        data_idx=new_data_idx,
+                    )
+
+                added_iter_IDs.append(iter_ID_i)
 
-                task.initialise_EARs()
+            task.initialise_EARs(iter_IDs=added_iter_IDs)
 
         added_iters_key = tuple(parent_loop_indices[k] for k in self.parents)
         self._increment_pending_added_iters(added_iters_key)
@@ -713,7 +751,8 @@ class WorkflowLoop:
                             **par_idx,
                             **parent_loop_indices,
                             self.name: cur_loop_idx + 1,
-                        }
+                        },
+                        cache=cache,
                     )
 
     def test_termination(self, element_iter):