hpcflow-new2 0.2.0a176__py3-none-any.whl → 0.2.0a178__py3-none-any.whl

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}
@@ -458,24 +468,19 @@ class WorkflowLoop:
                     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
 
@@ -493,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 = {
@@ -553,76 +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:
-                                # TODO: what about groups?
                                 # source the data_idx from the iteration with same parent
                                 # loop indices as the new iteration to add:
-                                src_iters = []
-                                for iter_i in element.iterations:
+                                # 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 iter_i.loop_idx.get(p_k) != p_v:
+                                        if li_k_dct.get(p_k) != p_v:
                                             skip_iter = True
                                             break
                                     if not skip_iter:
-                                        src_iters.append(iter_i)
+                                        src_data_idx.append(di_k)
 
                                 # could be multiple, but they should all have the same
                                 # data index for this parameter:
-                                src_iter = src_iters[0]
-                                inp_dat_idx = src_iter.get_data_idx()[inp_key]
+                                src_data_idx = src_data_idx[0]
+                                inp_dat_idx = src_data_idx[inp_key]
                             else:
                                 is_group = False
                                 if (
@@ -645,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))
 
@@ -680,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:
@@ -691,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
 
@@ -703,16 +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
+                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,
+                    )
 
-                task.initialise_EARs()
+                added_iter_IDs.append(iter_ID_i)
+
+            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)
@@ -731,7 +751,8 @@ class WorkflowLoop:
                             **par_idx,
                             **parent_loop_indices,
                             self.name: cur_loop_idx + 1,
-                        }
+                        },
+                        cache=cache,
                     )
 
     def test_termination(self, element_iter):

hpcflow/sdk/core/loop_cache.py

@@ -0,0 +1,140 @@
+from dataclasses import dataclass
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+from hpcflow.sdk import app
+from hpcflow.sdk.core.utils import nth_key
+from hpcflow.sdk.log import TimeIt
+from hpcflow.sdk.core.cache import DependencyCache
+
+
+@dataclass
+class LoopCache:
+    """Class to store a cache for use in `Workflow.add_empty_loop` and
+    `WorkflowLoop.add_iterations`.
+
+    Attributes
+    ----------
+    element_dependents
+        Keys are element IDs, values are dicts whose keys are element IDs that depend on
+        the key element ID (via `Element.get_dependent_elements_recursively`), and whose
+        values are dicts with keys: `group_names`, which is a tuple of the string group
+        names associated with the dependent element's element set.
+    elements
+        Keys are element IDs, values are dicts with keys: `input_statuses`,
+        `input_sources`, and `task_insert_ID`.
+    zeroth_iters
+        Keys are element IDs, values are data associated with the zeroth iteration of that
+        element, namely a tuple of iteration ID and `ElementIteration.data_idx`.
+    data_idx
+        Keys are element IDs, values are data associated with all iterations of that
+        element, namely a dict whose keys are the iteration loop index as a tuple, and
+        whose values are data indices via `ElementIteration.get_data_idx()`.
+    iterations
+        Keys are iteration IDs, values are tuples of element ID and iteration index within
+        that element.
+    task_iterations
+        Keys are task insert IDs, values are list of all iteration IDs associated with
+        that task.
+
+    """
+
+    element_dependents: Dict[int, Dict]
+    elements: Dict[int, Dict]
+    zeroth_iters: Dict[int, Tuple]
+    data_idx: Dict[int, Dict]
+    iterations: Dict[int, Tuple]
+    task_iterations: Dict[int, List[int]]
+
+    @TimeIt.decorator
+    def get_iter_IDs(self, loop: "app.Loop") -> List[int]:
+        """Retrieve a list of iteration IDs belonging to a given loop."""
+        return [j for i in loop.task_insert_IDs for j in self.task_iterations[i]]
+
+    @TimeIt.decorator
+    def get_iter_loop_indices(self, iter_IDs: List[int]) -> List[Dict[str, int]]:
+        iter_loop_idx = []
+        for i in iter_IDs:
+            elem_id, idx = self.iterations[i]
+            iter_loop_idx.append(dict(nth_key(self.data_idx[elem_id], idx)))
+        return iter_loop_idx
+
+    @TimeIt.decorator
+    def update_loop_indices(self, new_loop_name: str, iter_IDs: List[int]):
+        elem_ids = {v[0] for k, v in self.iterations.items() if k in iter_IDs}
+        for i in elem_ids:
+            new_item = {}
+            for k, v in self.data_idx[i].items():
+                new_k = dict(k)
+                new_k.update({new_loop_name: 0})
+                new_item[tuple(sorted(new_k.items()))] = v
+            self.data_idx[i] = new_item
+
+    @TimeIt.decorator
+    def add_iteration(self, iter_ID, task_insert_ID, element_ID, loop_idx, data_idx):
+        """Update the cache to include a newly added iteration."""
+        self.task_iterations[task_insert_ID].append(iter_ID)
+        new_iter_idx = len(self.data_idx[element_ID])
+        self.data_idx[element_ID][tuple(sorted(loop_idx.items()))] = data_idx
+        self.iterations[iter_ID] = (element_ID, new_iter_idx)
+
+    @classmethod
+    @TimeIt.decorator
+    def build(cls, workflow: "app.Workflow", loops: Optional[List["app.Loop"]] = None):
+        """Build a cache of data for use in adding loops and iterations."""
+
+        deps_cache = DependencyCache.build(workflow)
+
+        loops = list(workflow.template.loops) + (loops or [])
+        task_iIDs = set(j for i in loops for j in i.task_insert_IDs)
+        tasks = [workflow.tasks.get(insert_ID=i) for i in sorted(task_iIDs)]
+        elem_deps = {}
+
+        # keys: element IDs, values: dict with keys: tuple(loop_idx), values: data index
+        data_idx_cache = {}
+
+        # keys: iteration IDs, values: tuple of (element ID, integer index into values
+        # dict in `data_idx_cache` [accessed via `.keys()[index]`])
+        iters = {}
+
+        # keys: element IDs, values: dict with keys: "input_statues", "input_sources",
+        # "task_insert_ID":
+        elements = {}
+
+        zeroth_iters = {}
+        task_iterations = defaultdict(list)
+        for task in tasks:
+            for elem_idx in task.element_IDs:
+                element = deps_cache.elements[elem_idx]
+                inp_statuses = task.template.get_input_statuses(element.element_set)
+                elements[element.id_] = {
+                    "input_statuses": inp_statuses,
+                    "input_sources": element.input_sources,
+                    "task_insert_ID": task.insert_ID,
+                }
+                elem_deps[element.id_] = {
+                    i: {
+                        "group_names": tuple(
+                            j.name for j in deps_cache.elements[i].element_set.groups
+                        ),
+                    }
+                    for i in deps_cache.elem_elem_dependents_rec[element.id_]
+                }
+                elem_iters = {}
+                for idx, iter_i in enumerate(element.iterations):
+                    if idx == 0:
+                        zeroth_iters[element.id_] = (iter_i.id_, iter_i.data_idx)
+                    loop_idx_key = tuple(sorted(iter_i.loop_idx.items()))
+                    elem_iters[loop_idx_key] = iter_i.get_data_idx()
+                    task_iterations[task.insert_ID].append(iter_i.id_)
+                    iters[iter_i.id_] = (element.id_, idx)
+                data_idx_cache[element.id_] = elem_iters
+
+        return cls(
+            element_dependents=elem_deps,
+            elements=elements,
+            zeroth_iters=zeroth_iters,
+            data_idx=data_idx_cache,
+            iterations=iters,
+            task_iterations=dict(task_iterations),
+        )

hpcflow/sdk/core/task.py

@@ -2062,29 +2062,36 @@ class WorkflowTask:
         return element_dat_idx
 
     @TimeIt.decorator
-    def initialise_EARs(self) -> List[int]:
+    def initialise_EARs(self, iter_IDs: Optional[List[int]] = None) -> List[int]:
         """Try to initialise any uninitialised EARs of this task."""
+        if iter_IDs:
+            iters = self.workflow.get_element_iterations_from_IDs(iter_IDs)
+        else:
+            iters = []
+            for element in self.elements:
+                # We don't yet cache Element objects, so `element`, and also it's
+                # `ElementIterations, are transient. So there is no reason to update these
+                # objects in memory to account for the new EARs. Subsequent calls to
+                # `WorkflowTask.elements` will retrieve correct element data from the
+                # store. This might need changing once/if we start caching Element
+                # objects.
+                iters.extend(element.iterations)
+
         initialised = []
-        for element in self.elements[:]:
-            # We don't yet cache Element objects, so `element`, and also it's
-            # `ElementIterations, are transient. So there is no reason to update these
-            # objects in memory to account for the new EARs. Subsequent calls to
-            # `WorkflowTask.elements` will retrieve correct element data from the store.
-            # This might need changing once/if we start caching Element objects.
-            for iter_i in element.iterations:
-                if not iter_i.EARs_initialised:
-                    try:
-                        self._initialise_element_iter_EARs(iter_i)
-                        initialised.append(iter_i.id_)
-                    except UnsetParameterDataError:
-                        # raised by `Action.test_rules`; cannot yet initialise EARs
-                        self.app.logger.debug(
-                            f"UnsetParameterDataError raised: cannot yet initialise runs."
-                        )
-                        pass
-                    else:
-                        iter_i._EARs_initialised = True
-                        self.workflow.set_EARs_initialised(iter_i.id_)
+        for iter_i in iters:
+            if not iter_i.EARs_initialised:
+                try:
+                    self._initialise_element_iter_EARs(iter_i)
+                    initialised.append(iter_i.id_)
+                except UnsetParameterDataError:
+                    # raised by `Action.test_rules`; cannot yet initialise EARs
+                    self.app.logger.debug(
+                        f"UnsetParameterDataError raised: cannot yet initialise runs."
+                    )
+                    pass
+                else:
+                    iter_i._EARs_initialised = True
+                    self.workflow.set_EARs_initialised(iter_i.id_)
         return initialised
 
     @TimeIt.decorator
@@ -2097,7 +2104,6 @@ class WorkflowTask:
         param_src_updates = {}
 
         count = 0
-        # TODO: generator is an IO op here, can be pre-calculated/cached?
         for act_idx, action in self.template.all_schema_actions():
             log_common = (
                 f"for action {act_idx} of element iteration {element_iter.index} of "
@@ -2151,8 +2157,7 @@ class WorkflowTask:
                 metadata={},
             )
 
-        for pid, src in param_src_updates.items():
-            self.workflow._store.update_param_source(pid, src)
+        self.workflow._store.update_param_source(param_src_updates)
 
     @TimeIt.decorator
     def _add_element_set(self, element_set):

hpcflow/sdk/core/utils.py

@@ -3,7 +3,7 @@ import enum
 from functools import wraps
 import contextlib
 import hashlib
-from itertools import accumulate
+from itertools import accumulate, islice
 import json
 import keyword
 import os
@@ -871,3 +871,13 @@ def dict_values_process_flat(d, callable):
         out[k] = proc_idx_k
 
     return out
+
+
+def nth_key(dct, n):
+    it = iter(dct)
+    next(islice(it, n, n), None)
+    return next(it)
+
+
+def nth_value(dct, n):
+    return dct[nth_key(dct, n)]