hpcflow-new2 0.2.0a159__py3-none-any.whl → 0.2.0a160__py3-none-any.whl

hpcflow/sdk/core/task.py

@@ -1283,6 +1283,7 @@ class WorkflowTask:
         return self.template.num_element_sets
 
     @property
+    @TimeIt.decorator
     def elements(self):
         if self._elements is None:
             self._elements = self.app.Elements(self)
@@ -1828,6 +1829,7 @@ 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 "
@@ -2604,41 +2606,38 @@ class Elements:
     def task(self):
         return self._task
 
-    def _get_selection(self, selection):
+    @TimeIt.decorator
+    def _get_selection(self, selection: Union[int, slice, List[int]]) -> List[int]:
+        """Normalise an element selection into a list of element indices."""
         if isinstance(selection, int):
-            start, stop, step = selection, selection + 1, 1
+            lst = [selection]
 
         elif isinstance(selection, slice):
-            start, stop, step = selection.start, selection.stop, selection.step
-            stop = self.task.num_elements if stop is None else stop
-            start = start or 0
-            step = 1 if step is None else step
+            lst = list(range(*selection.indices(self.task.num_elements)))
 
+        elif isinstance(selection, list):
+            lst = selection
         else:
             raise RuntimeError(
-                f"{self.__class__.__name__} selection must be an `int` or a `slice` "
-                f"object, but received type {type(selection)}."
+                f"{self.__class__.__name__} selection must be an `int`, `slice` object, "
+                f"or list of `int`s, but received type {type(selection)}."
             )
-
-        selection = slice(start, stop, step)
-        length = len(range(*selection.indices(self.task.num_elements)))
-
-        return selection, length
+        return lst
 
     def __len__(self):
         return self.task.num_elements
 
     def __iter__(self):
-        all_elems = self.task.workflow.get_task_elements(self.task, slice(None))
-        for i in all_elems:
+        for i in self.task.workflow.get_task_elements(self.task):
             yield i
 
+    @TimeIt.decorator
     def __getitem__(
         self,
-        selection: Union[int, slice],
+        selection: Union[int, slice, List[int]],
     ) -> Union[app.Element, List[app.Element]]:
-        sel_normed, _ = self._get_selection(selection)
-        elements = self.task.workflow.get_task_elements(self.task, sel_normed)
+        idx_lst = self._get_selection(selection)
+        elements = self.task.workflow.get_task_elements(self.task, idx_lst)
 
         if isinstance(selection, int):
             return elements[0]
@@ -2657,34 +2656,34 @@ class Parameters:
     raise_on_unset: Optional[bool] = False
     default: Optional[Any] = None
 
-    def _get_selection(self, selection):
+    @TimeIt.decorator
+    def _get_selection(self, selection: Union[int, slice, List[int]]) -> List[int]:
+        """Normalise an element selection into a list of element indices."""
         if isinstance(selection, int):
-            start, stop, step = selection, selection + 1, 1
+            lst = [selection]
 
         elif isinstance(selection, slice):
-            start, stop, step = selection.start, selection.stop, selection.step
-            stop = self.task.num_elements if stop is None else stop
-            start = start or 0
-            step = 1 if step is None else step
+            lst = list(range(*selection.indices(self.task.num_elements)))
 
+        elif isinstance(selection, list):
+            lst = selection
         else:
             raise RuntimeError(
-                f"{self.__class__.__name__} selection must be an `int` or a `slice` "
-                f"object, but received type {type(selection)!r}."
+                f"{self.__class__.__name__} selection must be an `int`, `slice` object, "
+                f"or list of `int`s, but received type {type(selection)}."
             )
-
-        selection = slice(start, stop, step)
-        length = len(range(*selection.indices(self.task.num_elements)))
-
-        return selection, length
+        return lst
 
     def __iter__(self):
         for i in self.__getitem__(slice(None)):
             yield i
 
-    def __getitem__(self, selection: Union[int, slice]) -> Union[Any, List[Any]]:
-        selection, length = self._get_selection(selection)
-        elements = self.task.workflow.get_task_elements(self.task, selection)
+    def __getitem__(
+        self,
+        selection: Union[int, slice, List[int]],
+    ) -> Union[Any, List[Any]]:
+        idx_lst = self._get_selection(selection)
+        elements = self.task.workflow.get_task_elements(self.task, idx_lst)
         if self.return_element_parameters:
             params = [
                 self._app.ElementParameter(
@@ -2706,7 +2705,7 @@ class Parameters:
                 for i in elements
             ]
 
-        if length == 1:
+        if isinstance(selection, int):
             return params[0]
         else:
             return params

hpcflow/sdk/core/utils.py

@@ -834,3 +834,40 @@ def linspace_rect(
 
     rect = np.hstack(stacked)
     return rect
+
+
+def dict_values_process_flat(d, callable):
+    """
+    Return a copy of a dict, where the values are processed by a callable that is to
+    be called only once, and where the values may be single items or lists of items.
+
+    Examples
+    --------
+    d = {'a': 0, 'b': [1, 2], 'c': 5}
+    >>> dict_values_process_flat(d, callable=lambda x: [i + 1 for i in x])
+    {'a': 1, 'b': [2, 3], 'c': 6}
+
+    """
+    flat = []  # values of `d`, flattened
+    is_multi = []  # whether a list, and the number of items to process
+    for i in d.values():
+        try:
+            flat.extend(i)
+            is_multi.append((True, len(i)))
+        except TypeError:
+            flat.append(i)
+            is_multi.append((False, 1))
+
+    processed = callable(flat)
+
+    out = {}
+    for idx_i, (m, k) in enumerate(zip(is_multi, d.keys())):
+
+        start_idx = sum(i[1] for i in is_multi[:idx_i])
+        end_idx = start_idx + m[1]
+        proc_idx_k = processed[slice(start_idx, end_idx)]
+        if not m[0]:
+            proc_idx_k = proc_idx_k[0]
+        out[k] = proc_idx_k
+
+    return out

hpcflow/sdk/core/workflow.py

@@ -492,6 +492,7 @@ class Workflow:
         ts_fmt: Optional[str] = None,
         ts_name_fmt: Optional[str] = None,
         store_kwargs: Optional[Dict] = None,
+        status: Optional[Any] = None,
     ) -> app.Workflow:
         """Generate from a `WorkflowTemplate` object.
 
@@ -521,22 +522,38 @@ class Workflow:
         store_kwargs
             Keyword arguments to pass to the store's `write_empty_workflow` method.
         """
-        wk = cls._write_empty_workflow(
-            template=template,
-            path=path,
-            name=name,
-            overwrite=overwrite,
-            store=store,
-            ts_fmt=ts_fmt,
-            ts_name_fmt=ts_name_fmt,
-            store_kwargs=store_kwargs,
-        )
-        with wk._store.cached_load():
-            with wk.batch_update(is_workflow_creation=True):
-                for task in template.tasks:
-                    wk._add_task(task)
-                for loop in template.loops:
-                    wk._add_loop(loop)
+        if status:
+            status.update("Generating empty workflow...")
+        try:
+            wk = cls._write_empty_workflow(
+                template=template,
+                path=path,
+                name=name,
+                overwrite=overwrite,
+                store=store,
+                ts_fmt=ts_fmt,
+                ts_name_fmt=ts_name_fmt,
+                store_kwargs=store_kwargs,
+            )
+            with wk._store.cached_load():
+                with wk.batch_update(is_workflow_creation=True):
+                    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)
+                    for idx, loop in enumerate(template.loops):
+                        if status:
+                            status.update(
+                                f"Adding loop {idx + 1}/" f"{len(template.loops)}..."
+                            )
+                        wk._add_loop(loop)
+        except Exception:
+            if status:
+                status.stop()
+            raise
         return wk
 
     @classmethod
@@ -668,6 +685,7 @@ class Workflow:
         ts_name_fmt: Optional[str] = None,
         store_kwargs: Optional[Dict] = None,
         variables: Optional[Dict[str, str]] = None,
+        status: Optional[Any] = None,
     ) -> app.Workflow:
         """Generate from a JSON file.
 
@@ -712,6 +730,7 @@ class Workflow:
             ts_fmt,
             ts_name_fmt,
             store_kwargs,
+            status,
         )
 
     @classmethod
@@ -726,6 +745,7 @@ class Workflow:
         ts_name_fmt: Optional[str] = None,
         store_kwargs: Optional[Dict] = None,
         variables: Optional[Dict[str, str]] = None,
+        status: Optional[Any] = None,
     ) -> app.Workflow:
         """Generate from a JSON string.
 
@@ -770,6 +790,7 @@ class Workflow:
             ts_fmt,
             ts_name_fmt,
             store_kwargs,
+            status,
         )
 
     @classmethod
@@ -786,6 +807,7 @@ class Workflow:
         ts_name_fmt: Optional[str] = None,
         store_kwargs: Optional[Dict] = None,
         variables: Optional[Dict[str, str]] = None,
+        status: Optional[Any] = None,
     ) -> app.Workflow:
         """Generate from either a YAML or JSON file, depending on the file extension.
 
@@ -821,11 +843,16 @@ class Workflow:
         variables
             String variables to substitute in the file given by `template_path`.
         """
-        template = cls.app.WorkflowTemplate.from_file(
-            template_path,
-            template_format,
-            variables=variables,
-        )
+        try:
+            template = cls.app.WorkflowTemplate.from_file(
+                template_path,
+                template_format,
+                variables=variables,
+            )
+        except Exception:
+            if status:
+                status.stop()
+            raise
         return cls.from_template(
             template,
             path,
@@ -835,6 +862,7 @@ class Workflow:
             ts_fmt,
             ts_name_fmt,
             store_kwargs,
+            status,
         )
 
     @classmethod
@@ -1262,7 +1290,6 @@ class Workflow:
     @TimeIt.decorator
     def get_EARs_from_IDs(self, id_lst: Iterable[int]) -> List[app.ElementActionRun]:
         """Return element action run objects from a list of IDs."""
-
         self.app.persistence_logger.debug(f"get_EARs_from_IDs: id_lst={id_lst!r}")
 
         store_EARs = self._store.get_EARs(id_lst)
@@ -1276,6 +1303,10 @@ class Workflow:
         task_IDs = [i.task_ID for i in store_elems]
         store_tasks = self._store.get_tasks_by_IDs(task_IDs)
 
+        # to allow for bulk retrieval of elements/iterations
+        element_idx_by_task = defaultdict(set)
+        iter_idx_by_task_elem = defaultdict(lambda: defaultdict(set))
+
         index_paths = []
         for rn, it, el, tk in zip(store_EARs, store_iters, store_elems, store_tasks):
             act_idx = rn.action_idx
@@ -1291,12 +1322,25 @@ class Workflow:
                     "task_idx": tk.index,
                 }
             )
+            element_idx_by_task[tk.index].add(elem_idx)
+            iter_idx_by_task_elem[tk.index][elem_idx].add(iter_idx)
+
+        # retrieve elements/iterations:
+        iters_by_task_elem = defaultdict(lambda: defaultdict(dict))
+        for task_idx, elem_idx in element_idx_by_task.items():
+            elements = self.tasks[task_idx].elements[list(elem_idx)]
+            for elem_i in elements:
+                elem_i_iters_idx = iter_idx_by_task_elem[task_idx][elem_i.index]
+                elem_iters = [elem_i.iterations[j] for j in elem_i_iters_idx]
+                iters_by_task_elem[task_idx][elem_i.index].update(
+                    dict(zip(elem_i_iters_idx, elem_iters))
+                )
 
         objs = []
         for idx_dat in index_paths:
-            task = self.tasks[idx_dat["task_idx"]]
-            elem = task.elements[idx_dat["elem_idx"]]
-            iter_ = elem.iterations[idx_dat["iter_idx"]]
+            iter_ = iters_by_task_elem[idx_dat["task_idx"]][idx_dat["elem_idx"]][
+                idx_dat["iter_idx"]
+            ]
             run = iter_.actions[idx_dat["action_idx"]].runs[idx_dat["run_idx"]]
             objs.append(run)
 
@@ -1580,6 +1624,13 @@ class Workflow:
         id_lst = list(range(num_params))
         return self._store.get_parameters(id_lst, **kwargs)
 
+    @TimeIt.decorator
+    def get_all_parameter_sources(self, **kwargs: Dict) -> List[Dict]:
+        """Retrieve all store parameters."""
+        num_params = self._store._get_num_total_parameters()
+        id_lst = list(range(num_params))
+        return self._store.get_parameter_sources(id_lst, **kwargs)
+
     @TimeIt.decorator
     def get_all_parameter_data(self, **kwargs: Dict) -> Dict[int, Any]:
         """Retrieve all workflow parameter data."""
@@ -1759,10 +1810,14 @@ class Workflow:
         return Path(self.path) / self._exec_dir_name
 
     @TimeIt.decorator
-    def get_task_elements(self, task: app.Task, selection: slice) -> List[app.Element]:
+    def get_task_elements(
+        self,
+        task: app.Task,
+        idx_lst: Optional[List[int]] = None,
+    ) -> List[app.Element]:
         return [
             self.app.Element(task=task, **{k: v for k, v in i.items() if k != "task_ID"})
-            for i in self._store.get_task_elements(task.insert_ID, selection)
+            for i in self._store.get_task_elements(task.insert_ID, idx_lst)
         ]
 
     def set_EAR_submission_index(self, EAR_ID: int, sub_idx: int) -> None:
@@ -1929,6 +1984,7 @@ class Workflow:
             for element in task.elements[:]:
                 yield element
 
+    @TimeIt.decorator
     def get_iteration_task_pathway(self):
         pathway = []
         for task in self.tasks:
@@ -1955,7 +2011,7 @@ class Workflow:
     @TimeIt.decorator
     def _submit(
         self,
-        status,
+        status: Optional[Any] = None,
         ignore_errors: Optional[bool] = False,
         JS_parallelism: Optional[bool] = None,
         print_stdout: Optional[bool] = False,
@@ -1967,10 +2023,12 @@ class Workflow:
         # generate a new submission if there are no pending submissions:
         pending = [i for i in self.submissions if i.needs_submit]
         if not pending:
-            status.update("Adding new submission...")
+            if status:
+                status.update("Adding new submission...")
             new_sub = self._add_submission(tasks=tasks, JS_parallelism=JS_parallelism)
             if not new_sub:
-                status.stop()
+                if status:
+                    status.stop()
                 raise ValueError("No pending element action runs to submit!")
             pending = [new_sub]
 
@@ -1980,7 +2038,8 @@ class Workflow:
 
         # for direct execution the submission must be persistent at submit-time, because
         # it will be read by a new instance of the app:
-        status.update("Committing to the store...")
+        if status:
+            status.update("Committing to the store...")
         self._store._pending.commit_all()
 
         # submit all pending submissions:
@@ -1988,7 +2047,8 @@ class Workflow:
         submitted_js = {}
         for sub in pending:
             try:
-                status.update(f"Preparing submission {sub.index}...")
+                if status:
+                    status.update(f"Preparing submission {sub.index}...")
                 sub_js_idx = sub.submit(
                     status=status,
                     ignore_errors=ignore_errors,
@@ -2010,6 +2070,8 @@ class Workflow:
         add_to_known: Optional[bool] = True,
         return_idx: Optional[bool] = False,
         tasks: Optional[List[int]] = None,
+        cancel: Optional[bool] = False,
+        status: Optional[bool] = True,
     ) -> Dict[int, int]:
         """Submit the workflow for execution.
 
@@ -2036,39 +2098,52 @@ class Workflow:
             List of task indices to include in the new submission if no submissions
             already exist. By default all tasks are included if a new submission is
             created.
+        cancel
+            Immediately cancel the submission. Useful for testing and benchmarking.
+        status
+            If True, display a live status to track submission progress.
         """
 
-        console = rich.console.Console()
-        status = console.status("Submitting workflow...")
-        status.start()
+        if status:
+            console = rich.console.Console()
+            status = console.status("Submitting workflow...")
+            status.start()
 
         with self._store.cached_load():
             if not self._store.is_submittable:
-                status.stop()
+                if status:
+                    status.stop()
                 raise NotImplementedError("The workflow is not submittable.")
             with self.batch_update():
                 # commit updates before raising exception:
                 try:
-                    exceptions, submitted_js = self._submit(
-                        ignore_errors=ignore_errors,
-                        JS_parallelism=JS_parallelism,
-                        print_stdout=print_stdout,
-                        status=status,
-                        add_to_known=add_to_known,
-                        tasks=tasks,
-                    )
+                    with self._store.cache_ctx():
+                        exceptions, submitted_js = self._submit(
+                            ignore_errors=ignore_errors,
+                            JS_parallelism=JS_parallelism,
+                            print_stdout=print_stdout,
+                            status=status,
+                            add_to_known=add_to_known,
+                            tasks=tasks,
+                        )
                 except Exception:
-                    status.stop()
+                    if status:
+                        status.stop()
                     raise
 
         if exceptions:
             msg = "\n" + "\n\n".join([i.message for i in exceptions])
-            status.stop()
+            if status:
+                status.stop()
             raise WorkflowSubmissionFailure(msg)
 
-        status.stop()
+        if status:
+            status.stop()
+
+        if cancel:
+            self.cancel()
 
-        if wait:
+        elif wait:
             self.wait(submitted_js)
 
         if return_idx:
@@ -2281,6 +2356,7 @@ class Workflow:
             )
         self._abort_run_ID(submission_idx, run.id_)
 
+    @TimeIt.decorator
     def cancel(self, hard=False):
         """Cancel any running jobscripts."""
         for sub in self.submissions:
@@ -2328,6 +2404,7 @@ class Workflow:
 
         return self.submissions[new_idx]
 
+    @TimeIt.decorator
     def resolve_jobscripts(
         self, tasks: Optional[List[int]] = None
     ) -> List[app.Jobscript]:
@@ -2344,6 +2421,7 @@ class Workflow:
 
         return js_objs
 
+    @TimeIt.decorator
     def _resolve_singular_jobscripts(
         self, tasks: Optional[List[int]] = None
     ) -> Tuple[Dict[int, Dict], Dict]:
@@ -2362,6 +2440,10 @@ class Workflow:
         if not tasks:
             tasks = list(range(self.num_tasks))
 
+        if self._store.use_cache:
+            # pre-cache parameter sources (used in `EAR.get_EAR_dependencies`):
+            self.get_all_parameter_sources()
+
         submission_jobscripts = {}
         all_element_deps = {}
 
@@ -2408,6 +2490,19 @@ class Workflow:
                     "resource_hash": res_hash[js_dat["resources"]],
                     "dependencies": {},
                 }
+
+                all_EAR_IDs = []
+                for js_elem_idx, (elem_idx, act_indices) in enumerate(
+                    js_dat["elements"].items()
+                ):
+                    for act_idx in act_indices:
+                        EAR_ID_i = EAR_map[act_idx, elem_idx].item()
+                        all_EAR_IDs.append(EAR_ID_i)
+                        js_act_idx = task_actions.index([task.insert_ID, act_idx, 0])
+                        js_i["EAR_ID"][js_act_idx][js_elem_idx] = EAR_ID_i
+
+                all_EAR_objs = dict(zip(all_EAR_IDs, self.get_EARs_from_IDs(all_EAR_IDs)))
+
                 for js_elem_idx, (elem_idx, act_indices) in enumerate(
                     js_dat["elements"].items()
                 ):
@@ -2419,7 +2514,7 @@ class Workflow:
                         js_i["EAR_ID"][js_act_idx][js_elem_idx] = EAR_ID_i
 
                     # get indices of EARs that this element depends on:
-                    EAR_objs = self.get_EARs_from_IDs(all_EAR_IDs)
+                    EAR_objs = [all_EAR_objs[k] for k in all_EAR_IDs]
                     EAR_deps = [i.get_EAR_dependencies() for i in EAR_objs]
                     EAR_deps_flat = [j for i in EAR_deps for j in i]
                     EAR_deps_EAR_idx = [

hpcflow/sdk/demo/cli.py

@@ -17,6 +17,9 @@ from hpcflow.sdk.cli_common import (
     add_to_known_opt,
     print_idx_opt,
     tasks_opt,
+    cancel_opt,
+    submit_status_opt,
+    make_status_opt,
 )
 
 
@@ -100,6 +103,7 @@ def get_demo_workflow_CLI(app):
     @ts_fmt_option
     @ts_name_fmt_option
     @variables_option
+    @make_status_opt
     def make_demo_workflow(
         workflow_name,
         format,
@@ -110,6 +114,7 @@ def get_demo_workflow_CLI(app):
         ts_fmt=None,
         ts_name_fmt=None,
         variables=None,
+        status=True,
     ):
         wk = app.make_demo_workflow(
             workflow_name=workflow_name,
@@ -121,6 +126,7 @@ def get_demo_workflow_CLI(app):
             ts_fmt=ts_fmt,
             ts_name_fmt=ts_name_fmt,
             variables=dict(variables),
+            status=status,
         )
         click.echo(wk.path)
 
@@ -139,6 +145,8 @@ def get_demo_workflow_CLI(app):
     @add_to_known_opt
     @print_idx_opt
     @tasks_opt
+    @cancel_opt
+    @submit_status_opt
     def make_and_submit_demo_workflow(
         workflow_name,
         format,
@@ -154,6 +162,8 @@ def get_demo_workflow_CLI(app):
         add_to_known=True,
         print_idx=False,
         tasks=None,
+        cancel=False,
+        status=True,
     ):
         out = app.make_and_submit_demo_workflow(
             workflow_name=workflow_name,
@@ -170,6 +180,8 @@ def get_demo_workflow_CLI(app):
             add_to_known=add_to_known,
             return_idx=print_idx,
             tasks=tasks,
+            cancel=cancel,
+            status=status,
         )
         if print_idx:
             click.echo(out)

hpcflow/sdk/log.py

@@ -97,7 +97,7 @@ class TimeIt:
                 max_str = f"{v['max']:10.6f}" if number > 1 else f"{f'-':^12s}"
                 stddev_str = f"({v['stddev']:8.6f})" if number > 1 else f"{f' ':^10s}"
                 out.append(
-                    f"{k_str:.<60s} {v['sum']:12.6f} "
+                    f"{k_str:.<80s} {v['sum']:12.6f} "
                     f"{v['mean']:10.6f} {stddev_str} {number:8d} "
                     f"{min_str} {max_str} "
                 )
@@ -114,7 +114,7 @@ class TimeIt:
         summary = cls.summarise()
 
         out = [
-            f"{'function':^60s} {'sum /s':^12s} {'mean (stddev) /s':^20s} {'N':^8s} "
+            f"{'function':^80s} {'sum /s':^12s} {'mean (stddev) /s':^20s} {'N':^8s} "
             f"{'min /s':^12s} {'max /s':^12s}"
         ]
         out += _format_nodes(summary)