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

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)]

hpcflow/sdk/core/workflow.py

@@ -25,6 +25,7 @@ from hpcflow.sdk.core import (
     ABORT_EXIT_CODE,
 )
 from hpcflow.sdk.core.actions import EARStatus
+from hpcflow.sdk.core.loop_cache import LoopCache
 from hpcflow.sdk.log import TimeIt
 from hpcflow.sdk.persistence import store_cls_from_str, DEFAULT_STORE_FORMAT
 from hpcflow.sdk.persistence.base import TEMPLATE_COMP_TYPES, AnySEAR
@@ -41,6 +42,7 @@ from hpcflow.sdk.submission.schedulers.direct import DirectScheduler
 from hpcflow.sdk.typing import PathLike
 from hpcflow.sdk.core.json_like import ChildObjectSpec, JSONLike
 from .utils import (
+    nth_key,
     read_JSON_file,
     read_JSON_string,
     read_YAML_str,
@@ -625,19 +627,28 @@ class Workflow:
             )
             with wk._store.cached_load():
                 with wk.batch_update(is_workflow_creation=True):
-                    for idx, task in enumerate(template.tasks):
+                    with wk._store.cache_ctx():
+                        for idx, task in enumerate(template.tasks):
+                            if status:
+                                status.update(
+                                    f"Adding task {idx + 1}/{len(template.tasks)} "
+                                    f"({task.name!r})..."
+                                )
+                            wk._add_task(task)
                         if status:
                             status.update(
-                                f"Adding task {idx + 1}/{len(template.tasks)} "
-                                f"({task.name!r})..."
+                                f"Preparing to add {len(template.loops)} loops..."
                             )
-                        wk._add_task(task)
-                    for idx, loop in enumerate(template.loops):
-                        if status:
-                            status.update(
-                                f"Adding loop {idx + 1}/" f"{len(template.loops)}..."
-                            )
-                        wk._add_loop(loop)
+                        if template.loops:
+                            # TODO: if loop with non-initialisable actions, will fail
+                            cache = LoopCache.build(workflow=wk, loops=template.loops)
+                            for idx, loop in enumerate(template.loops):
+                                if status:
+                                    status.update(
+                                        f"Adding loop {idx + 1}/"
+                                        f"{len(template.loops)} ({loop.name!r})"
+                                    )
+                                wk._add_loop(loop, cache=cache, status=status)
         except Exception:
             if status:
                 status.stop()
@@ -1101,7 +1112,7 @@ class Workflow:
 
     @TimeIt.decorator
     def _add_empty_loop(
-        self, loop: app.Loop
+        self, loop: app.Loop, cache: LoopCache
     ) -> Tuple[app.WorkflowLoop, List[app.ElementIteration]]:
         """Add a new loop (zeroth iterations only) to the workflow."""
 
@@ -1114,15 +1125,15 @@ class Workflow:
         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]
+        iter_IDs = cache.get_iter_IDs(loop_c)
+        iter_loop_idx = cache.get_iter_loop_indices(iter_IDs)
 
         # create and insert a new WorkflowLoop:
         new_loop = self.app.WorkflowLoop.new_empty_loop(
             index=new_index,
             workflow=self,
             template=loop_c,
-            iterations=iters,
+            iter_loop_idx=iter_loop_idx,
         )
         self.loops.add_object(new_loop)
         wk_loop = self.loops[new_index]
@@ -1144,15 +1155,28 @@ class Workflow:
 
         self._pending["loops"].append(new_index)
 
+        # update cache loop indices:
+        cache.update_loop_indices(new_loop_name=loop_c.name, iter_IDs=iter_IDs)
+
         return wk_loop
 
     @TimeIt.decorator
-    def _add_loop(self, loop: app.Loop) -> None:
-        new_wk_loop = self._add_empty_loop(loop)
+    def _add_loop(
+        self, loop: app.Loop, cache: Optional[Dict] = None, status: Optional[Any] = None
+    ) -> None:
+        if not cache:
+            cache = LoopCache.build(workflow=self, loops=[loop])
+        new_wk_loop = self._add_empty_loop(loop, cache)
         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()
+            if status:
+                status_prev = status.status
+            for iter_idx in range(loop.num_iterations - 1):
+                if status:
+                    status.update(
+                        f"{status_prev}: iteration {iter_idx + 2}/{loop.num_iterations}."
+                    )
+                new_wk_loop.add_iteration(cache=cache)
 
     def add_loop(self, loop: app.Loop) -> None:
         """Add a loop to a subset of workflow tasks."""
@@ -1326,6 +1350,7 @@ class Workflow:
                     iters.append(iter_i)
         return iters
 
+    @TimeIt.decorator
     def get_elements_from_IDs(self, id_lst: Iterable[int]) -> List[app.Element]:
         """Return element objects from a list of IDs."""
 
@@ -1334,6 +1359,7 @@ class Workflow:
         task_IDs = [i.task_ID for i in store_elems]
         store_tasks = self._store.get_tasks_by_IDs(task_IDs)
 
+        element_idx_by_task = defaultdict(set)
         index_paths = []
         for el, tk in zip(store_elems, store_tasks):
             elem_idx = tk.element_IDs.index(el.id_)
@@ -1343,15 +1369,23 @@ class Workflow:
                     "task_idx": tk.index,
                 }
             )
+            element_idx_by_task[tk.index].add(elem_idx)
+
+        elements_by_task = {}
+        for task_idx, elem_idx in element_idx_by_task.items():
+            task = self.tasks[task_idx]
+            elements_by_task[task_idx] = dict(
+                zip(elem_idx, task.elements[list(elem_idx)])
+            )
 
         objs = []
         for idx_dat in index_paths:
-            task = self.tasks[idx_dat["task_idx"]]
-            elem = task.elements[idx_dat["elem_idx"]]
+            elem = elements_by_task[idx_dat["task_idx"]][idx_dat["elem_idx"]]
             objs.append(elem)
 
         return objs
 
+    @TimeIt.decorator
     def get_element_iterations_from_IDs(
         self, id_lst: Iterable[int]
     ) -> List[app.ElementIteration]:
@@ -1365,6 +1399,8 @@ class Workflow:
         task_IDs = [i.task_ID for i in store_elems]
         store_tasks = self._store.get_tasks_by_IDs(task_IDs)
 
+        element_idx_by_task = defaultdict(set)
+
         index_paths = []
         for it, el, tk in zip(store_iters, store_elems, store_tasks):
             iter_idx = el.iteration_IDs.index(it.id_)
@@ -1376,11 +1412,18 @@ class Workflow:
                     "task_idx": tk.index,
                 }
             )
+            element_idx_by_task[tk.index].add(elem_idx)
+
+        elements_by_task = {}
+        for task_idx, elem_idx in element_idx_by_task.items():
+            task = self.tasks[task_idx]
+            elements_by_task[task_idx] = dict(
+                zip(elem_idx, task.elements[list(elem_idx)])
+            )
 
         objs = []
         for idx_dat in index_paths:
-            task = self.tasks[idx_dat["task_idx"]]
-            elem = task.elements[idx_dat["elem_idx"]]
+            elem = elements_by_task[idx_dat["task_idx"]][idx_dat["elem_idx"]]
             iter_ = elem.iterations[idx_dat["iter_idx"]]
             objs.append(iter_)
 

hpcflow/sdk/persistence/base.py

@@ -716,6 +716,11 @@ class PersistentStore(ABC):
         """Cache for number of persistent tasks."""
         return self._cache["num_tasks"]
 
+    @property
+    def num_EARs_cache(self):
+        """Cache for total number of persistent EARs."""
+        return self._cache["num_EARs"]
+
     @property
     def param_sources_cache(self):
         """Cache for persistent parameter sources."""
@@ -730,6 +735,10 @@ class PersistentStore(ABC):
     def num_tasks_cache(self, value):
         self._cache["num_tasks"] = value
 
+    @num_EARs_cache.setter
+    def num_EARs_cache(self, value):
+        self._cache["num_EARs"] = value
+
     def _reset_cache(self):
         self._cache = {
             "tasks": {},
@@ -739,6 +748,7 @@ class PersistentStore(ABC):
             "param_sources": {},
             "num_tasks": None,
             "parameters": {},
+            "num_EARs": None,
         }
 
     @contextlib.contextmanager
@@ -873,6 +883,7 @@ class PersistentStore(ABC):
         """Get the total number of persistent and pending element iterations."""
         return self._get_num_persistent_elem_iters() + len(self._pending.add_elem_iters)
 
+    @TimeIt.decorator
     def _get_num_total_EARs(self):
         """Get the total number of persistent and pending EARs."""
         return self._get_num_persistent_EARs() + len(self._pending.add_EARs)
@@ -1296,9 +1307,11 @@ class PersistentStore(ABC):
             self.save()
 
     @TimeIt.decorator
-    def update_param_source(self, param_id: int, source: Dict, save: bool = True) -> None:
-        self.logger.debug(f"Updating parameter ID {param_id!r} source to {source!r}.")
-        self._pending.update_param_sources[param_id] = source
+    def update_param_source(
+        self, param_sources: Dict[int, Dict], save: bool = True
+    ) -> None:
+        self.logger.debug(f"Updating parameter sources with {param_sources!r}.")
+        self._pending.update_param_sources.update(param_sources)
         if save:
             self.save()
 

hpcflow/sdk/persistence/json.py

@@ -303,12 +303,13 @@ class JSONPersistentStore(PersistentStore):
 
     def _get_num_persistent_tasks(self) -> int:
         """Get the number of persistent tasks."""
-        if self.num_tasks_cache is not None:
+        if self.use_cache and self.num_tasks_cache is not None:
             num = self.num_tasks_cache
         else:
             with self.using_resource("metadata", action="read") as md:
                 num = len(md["tasks"])
-                self.num_tasks_cache = num
+        if self.use_cache and self.num_tasks_cache is None:
+            self.num_tasks_cache = num
         return num
 
     def _get_num_persistent_loops(self) -> int:
@@ -333,8 +334,14 @@ class JSONPersistentStore(PersistentStore):
 
     def _get_num_persistent_EARs(self) -> int:
         """Get the number of persistent EARs."""
-        with self.using_resource("metadata", action="read") as md:
-            return len(md["runs"])
+        if self.use_cache and self.num_EARs_cache is not None:
+            num = self.num_EARs_cache
+        else:
+            with self.using_resource("metadata", action="read") as md:
+                num = len(md["runs"])
+        if self.use_cache and self.num_EARs_cache is None:
+            self.num_EARs_cache = num
+        return num
 
     def _get_num_persistent_parameters(self):
         with self.using_resource("parameters", "read") as params:

hpcflow/sdk/persistence/pending.py

@@ -275,6 +275,7 @@ class PendingChanges:
             EAR_ids = list(self.add_EARs.keys())
             self.logger.debug(f"commit: adding pending EARs with IDs: {EAR_ids!r}")
             self.store._append_EARs(EARs)
+            self.store.num_EARs_cache = None  # invalidate cache
             # pending start/end times/snapshots, submission indices, and skips that belong
             # to pending EARs are now committed (accounted for in `get_EARs` above):
             self.set_EAR_submission_indices = {
@@ -408,6 +409,7 @@ class PendingChanges:
     @TimeIt.decorator
     def commit_loop_indices(self) -> None:
         """Make pending update to element iteration loop indices persistent."""
+        # TODO: batch up
         for iter_ID, loop_idx in self.update_loop_indices.items():
             self.logger.debug(
                 f"commit: updating loop indices of iteration ID {iter_ID!r} with "

hpcflow/sdk/persistence/zarr.py

@@ -774,9 +774,16 @@ class ZarrPersistentStore(PersistentStore):
         """Get the number of persistent element iterations."""
         return len(self._get_iters_arr())
 
+    @TimeIt.decorator
     def _get_num_persistent_EARs(self) -> int:
         """Get the number of persistent EARs."""
-        return len(self._get_EARs_arr())
+        if self.use_cache and self.num_EARs_cache is not None:
+            num = self.num_EARs_cache
+        else:
+            num = len(self._get_EARs_arr())
+        if self.use_cache and self.num_EARs_cache is None:
+            self.num_EARs_cache = num
+        return num
 
     def _get_num_persistent_parameters(self):
         return len(self._get_parameter_base_array())