sierra-research 1.3.11__py3-none-any.whl → 1.5.0__py3-none-any.whl

sierra/plugins/proc/collate/plugin.py

@@ -0,0 +1,271 @@
+# Copyright 2019 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+
+"""
+Classes for collating data within a :term:`Batch Experiment`.
+
+Collation is the process of "lifting" data from :term:`Experimental Runs
+<Experimental Run>` across all :term:`Experiment` for all experiments in a
+:term:`Batch Experiment` into a single file (a reduce operation).  This is
+needed to correctly calculate summary statistics for performance measures in
+stage 3: you can't just run the calculated stddev through the calculations
+because comparing curves of stddev is not meaningful.
+"""
+
+# Core packages
+import multiprocessing as mp
+import typing as tp
+import queue
+import logging
+import pathlib
+
+# 3rd party packages
+import pandas as pd
+import yaml
+
+# Project packages
+import sierra.core.variables.batch_criteria as bc
+import sierra.core.plugin as pm
+from sierra.core import types, storage, utils, config, batchroot
+from sierra.core.pipeline.stage3 import gather
+
+_logger = logging.getLogger(__name__)
+
+
+def proc_batch_exp(
+    main_config: dict,
+    cmdopts: types.Cmdopts,
+    pathset: batchroot.PathSet,
+    criteria: bc.XVarBatchCriteria,
+) -> None:
+    """Generate :term:`Collated Output Data` files for each experiment.
+
+    :term:`Collated Output Data` files generated from :term:`Raw Output Data`
+    files across :term:`Experimental Runs <Experimental Run>`.  Gathered in
+    parallel for each experiment for speed, unless disabled with
+    ``--processing-parallelism``.
+    """
+    pool_opts = {}
+
+    pool_opts["parallelism"] = cmdopts["processing_parallelism"]
+
+    worker_opts = {
+        "project": cmdopts["project"],
+        "template_input_leaf": pathlib.Path(cmdopts["expdef_template"]).stem,
+        "df_verify": cmdopts["df_verify"],
+        "processing_mem_limit": cmdopts["processing_mem_limit"],
+        "storage": cmdopts["storage"],
+        "df_homogenize": cmdopts["df_homogenize"],
+        "project_config_root": cmdopts["project_config_root"],
+    }
+
+    exp_to_proc = utils.exp_range_calc(
+        cmdopts["exp_range"], pathset.output_root, criteria.gen_exp_names()
+    )
+
+    with mp.Pool(processes=pool_opts["parallelism"]) as pool:
+        _execute_for_batch(
+            main_config, pathset, exp_to_proc, worker_opts, pool_opts, pool
+        )
+
+
+def _execute_for_batch(
+    main_config: types.YAMLDict,
+    pathset: batchroot.PathSet,
+    exp_to_proc: tp.List[pathlib.Path],
+    worker_opts: types.SimpleDict,
+    pool_opts: types.SimpleDict,
+    pool,
+) -> None:
+    m = mp.Manager()
+    gatherq = m.Queue()
+    processq = m.Queue()
+
+    for exp in exp_to_proc:
+        gatherq.put(exp)
+
+    _logger.debug(
+        "Starting %d gatherers, method=%s",
+        pool_opts["parallelism"],
+        mp.get_start_method(),
+    )
+
+    gathered = [
+        pool.apply_async(_gather_worker, (gatherq, processq, main_config, worker_opts))
+        for _ in range(0, pool_opts["parallelism"])
+    ]
+    _logger.debug("Waiting for gathering to finish")
+    for g in gathered:
+        g.get()
+
+    _logger.debug(
+        "Starting %d processors, method=%s",
+        pool_opts["parallelism"],
+        mp.get_start_method(),
+    )
+    processed = [
+        pool.apply_async(
+            _process_worker,
+            (processq, main_config, pathset.stat_interexp_root, worker_opts),
+        )
+        for _ in range(0, pool_opts["parallelism"])
+    ]
+
+    # To capture the otherwise silent crashes when something goes wrong in
+    # worker threads. Any assertions will show and any exceptions will be
+    # re-raised.
+    for p in processed:
+        p.get()
+
+    pool.close()
+    pool.join()
+    _logger.debug("Processing finished")
+
+
+def _gather_worker(
+    gatherq: mp.Queue,
+    processq: mp.Queue,
+    main_config: types.YAMLDict,
+    gather_opts: types.SimpleDict,
+) -> None:
+    gatherer = ExpDataGatherer(main_config, gather_opts, processq)
+    while True:
+        # Wait for 3 seconds after the queue is empty before bailing
+        try:
+            exp_output_root = gatherq.get(True, 3)
+            gatherer(exp_output_root)
+            gatherq.task_done()
+
+        except queue.Empty:
+            break
+
+
+def _process_worker(
+    processq: mp.Queue,
+    main_config: types.YAMLDict,
+    batch_stat_interexp_root: pathlib.Path,
+    process_opts: types.SimpleDict,
+) -> None:
+    while True:
+        # Wait for 3 seconds after the queue is empty before bailing
+        try:
+            spec = processq.get(True, 3)
+            _proc_single_exp(main_config, batch_stat_interexp_root, process_opts, spec)
+            processq.task_done()
+        except queue.Empty:
+            break
+
+
+class ExpDataGatherer(gather.BaseGatherer):
+    """Gather :term:`Raw Output Data` files across all runs for :term:`Data Collation`.
+
+    The configured output directory for each run is searched recursively for
+    files to gather.  To be eligible for gathering and later processing, files
+    must:
+
+        - Be non-empty
+
+        - Have a suffix which supported by the selected ``--storage`` plugin.
+
+        - Have a name (last part of absolute path, including extension) which
+          matches a configured :term:`Product` in a YAML file. E.g., a graph
+          from the :ref:`plugins/prod/graphs` plugin
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.logger = logging.getLogger(__name__)
+
+    def calc_gather_items(
+        self, run_output_root: pathlib.Path, exp_name: str
+    ) -> tp.List[gather.GatherSpec]:
+        to_gather = []
+        proj_output_root = run_output_root / str(self.run_metrics_leaf)
+        plugin = pm.pipeline.get_plugin_module(self.gather_opts["storage"])
+
+        config_path = pathlib.Path(
+            self.gather_opts["project_config_root"], config.kYAML.collate
+        )
+
+        try:
+            collate_config = yaml.load(utils.utf8open(config_path), yaml.FullLoader)
+
+        except FileNotFoundError:
+            self.logger.warning("%s does not exist!", config_path)
+            collate_config = {}
+
+        for item in proj_output_root.rglob("*"):
+            # Must be a file (duh)
+            if not item.is_file():
+                continue
+
+            # Has to be a supported suffix for storage plugin
+            if (
+                not any(s in plugin.suffixes() for s in item.suffixes)
+                or item.stat().st_size == 0
+            ):
+                continue
+
+            # Any number of perf metrics can be configured, so look for a match.
+            files = collate_config["intra-exp"]
+            perf_confs = [f for f in files if f["file"] in item.name]
+            if not perf_confs:
+                continue
+
+            # If we get a file match, then all the columns from that file should
+            # be added to the set of things to collate.
+            for conf in perf_confs:
+                for col in conf["cols"]:
+                    to_gather.append(
+                        gather.GatherSpec(
+                            exp_name=exp_name,
+                            item_stem_path=item.relative_to(proj_output_root),
+                            collate_col=col,
+                        )
+                    )
+        return to_gather
+
+
+def _proc_single_exp(
+    main_config: types.YAMLDict,
+    batch_stat_collate_root: pathlib.Path,
+    process_opts: types.SimpleDict,
+    spec: gather.ProcessSpec,
+) -> None:
+    """Collate :term:`Raw Output Data` files together (reduce operation).
+
+    :term:`Raw Output Data` files gathered from N :term:`Experimental Runs
+    <Experimental Run>` are combined together into a single :term:`Batch Summary
+    Data` file per :term:`Experiment` with 1 column per run.
+    """
+    utils.dir_create_checked(batch_stat_collate_root, exist_ok=True)
+
+    collated = {}
+
+    key = (spec.gather.item_stem_path, spec.gather.collate_col)
+    collated[key] = pd.DataFrame(index=spec.dfs[0].index, columns=spec.exp_run_names)
+    for i, df in enumerate(spec.dfs):
+        assert (
+            spec.gather.collate_col in df.columns
+        ), f"{spec.gather.collate_col} not in {df.columns}"
+
+        collate_df = df[spec.gather.collate_col]
+        collated[key][spec.exp_run_names[i]] = collate_df
+
+    for file_path, col in collated:
+        df = utils.df_fill(collated[(file_path, col)], process_opts["df_homogenize"])
+        parent = batch_stat_collate_root / spec.gather.exp_name / file_path.parent
+        utils.dir_create_checked(parent, exist_ok=True)
+
+        # This preserves the directory structure of stuff in the per-run output
+        # run; if something is in a subdir there, it will show up in a subdir in
+        # the collated outputs too.
+        fname = f"{file_path.stem}-{col}" + config.kStorageExt["csv"]
+        storage.df_write(df, parent / fname, "storage.csv", index=False)
+
+
+__all__ = [
+    "proc_batch_exp",
+    "ExpDataGatherer",
+]

sierra/plugins/proc/compress/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2021 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Container module for the processing plugin to compress data.
+
+See :ref:`plugins/proc/compress`.
+"""
+
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+
+
+def sierra_plugin_type() -> str:
+    return "pipeline"

sierra/plugins/proc/compress/cmdline.py

@@ -0,0 +1,47 @@
+#
+# Copyright 2025 John Harwell, All rights reserved.
+#
+# SPDX-License Identifier: MIT
+#
+
+# Core packages
+import typing as tp
+import argparse
+
+# 3rd party packages
+
+# Project packages
+from sierra.core import types
+from sierra.plugins import PluginCmdline
+
+
+def build(
+    parents: tp.List[argparse.ArgumentParser], stages: tp.List[int]
+) -> PluginCmdline:
+    """
+    Get a cmdline parser supporting the ``proc.compress`` processing plugin.
+    """
+    cmdline = PluginCmdline(parents, stages)
+    cmdline.stage3.add_argument(
+        "--compress-remove-after",
+        action="store_true",
+        help="""
+                 If the ``proc.compress`` plugin is run, remove the uncompressed
+                 :term:`Raw Output Data` files after compression.  This can save
+                 TONS of disk space.  No data is lost because everything output
+                 by each :term:`Experimental Run` is in the compressed archive.
+                 """
+        + cmdline.stage_usage_doc([3]),
+        default=False,
+    )
+    return cmdline
+
+
+def to_cmdopts(args: argparse.Namespace) -> types.Cmdopts:
+    return {
+        "compress_remove_after": args.compress_remove_after,
+    }
+
+
+def sphinx_cmdline_stage3():
+    return build([], [3]).parser

sierra/plugins/proc/compress/plugin.py

@@ -0,0 +1,123 @@
+# Copyright 2025 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Plugin for compressing experiment data. Currently only works with .tar.gz files.
+"""
+
+# Core packages
+import multiprocessing as mp
+import typing as tp
+import logging
+import pathlib
+import shutil
+
+# 3rd party packages
+import tarfile
+
+# Project packages
+import sierra.core.variables.batch_criteria as bc
+from sierra.core import types, utils, batchroot
+
+_logger = logging.getLogger(__name__)
+
+
+def proc_batch_exp(
+    main_config: types.YAMLDict,
+    cmdopts: types.Cmdopts,
+    pathset: batchroot.PathSet,
+    criteria: bc.XVarBatchCriteria,
+) -> None:
+    """
+    Comcompress data for each :term:`Experiment` in the :term:`Batch Experiment`.
+
+    Ideally this is done in parallel across experiments, but this can be changed
+    to serial if memory on the SIERRA host machine is limited via
+    ``--processing-parallelism``.
+    """
+
+    exp_to_proc = utils.exp_range_calc(
+        cmdopts["exp_range"], pathset.output_root, criteria.gen_exp_names()
+    )
+
+    parallelism = cmdopts["processing_parallelism"]
+
+    tasks = []
+    run_metrics_leaf = main_config["sierra"]["run"]["run_metrics_leaf"]
+
+    for exp in exp_to_proc:
+        tasks.extend(
+            _build_tasklist_for_exp(
+                pathset.output_root / exp.name,
+                run_metrics_leaf,
+                cmdopts["compress_remove_after"],
+            )
+        )
+
+    _logger.debug("Starting %s workers, method=%s", parallelism, mp.get_start_method())
+    with mp.Pool(processes=parallelism, maxtasksperchild=1) as pool:
+        processed = [pool.starmap_async(_worker, tasks)]
+        _logger.debug("Waiting for workers to finish")
+
+        for p in processed:
+            p.get()
+
+        pool.close()
+        pool.join()
+
+    _logger.debug("All workers finished")
+
+
+def _build_tasklist_for_exp(
+    exp_output_root: pathlib.Path,
+    run_metrics_leaf: pathlib.Path,
+    remove_after: bool,
+) -> tp.List[tp.Tuple[pathlib.Path, pathlib.Path, bool]]:
+    """Add root dir each experimental run to queue for processing.
+
+    Enqueueing for processing is done at the file-level rather than
+    per-experiment, so that for systems with more CPUs than experiments you
+    still get maximum throughput.
+    """
+    res = []
+    for exp in exp_output_root.iterdir():
+        res.append(
+            (
+                exp_output_root,
+                exp.relative_to(exp_output_root) / run_metrics_leaf,
+                remove_after,
+            )
+        )
+
+    return res
+
+
+def _worker(
+    exp_output_root: pathlib.Path, relpath: pathlib.Path, remove_after: bool
+) -> None:
+    """Compress the output root for a single experiment into a tarball.
+
+    Arguments:
+        exp_output_root: Output root for the :term:`Experiment`.
+
+        relpath: Path to the actual tarball relative to the experiment root.
+    """
+
+    if not (exp_output_root / relpath).exists():
+        _logger.warning(
+            "Cannot compress: %s does not exist", (exp_output_root / relpath)
+        )
+        return
+
+    with tarfile.open(
+        (exp_output_root / relpath).with_suffix(".tar.gz"), "w:gz"
+    ) as tar:
+        tar.add(
+            str(exp_output_root / relpath), arcname=relpath.relative_to(relpath.parent)
+        )
+
+    if remove_after:
+        shutil.rmtree(exp_output_root / relpath)
+
+
+__all__ = ["proc_batch_exp"]

sierra/plugins/proc/decompress/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2021 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Container module for the processing plugin to decompress data.
+
+See :ref:`plugins/proc/decompress`.
+"""
+
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+
+
+def sierra_plugin_type() -> str:
+    return "pipeline"

sierra/plugins/proc/decompress/plugin.py

@@ -0,0 +1,96 @@
+# Copyright 2025 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Plugin for decompressing experiment data. Currently only works with .tar.gz files.
+"""
+
+# Core packages
+import multiprocessing as mp
+import typing as tp
+import logging
+import pathlib
+
+# 3rd party packages
+import tarfile
+
+# Project packages
+import sierra.core.variables.batch_criteria as bc
+from sierra.core import types, utils, batchroot
+
+_logger = logging.getLogger(__name__)
+
+
+def proc_batch_exp(
+    main_config: types.YAMLDict,
+    cmdopts: types.Cmdopts,
+    pathset: batchroot.PathSet,
+    criteria: bc.XVarBatchCriteria,
+) -> None:
+    """
+    Uncomcompress data for each :term:`Experiment` in the :term:`Batch Experiment`.
+
+    Ideally this is done in parallel across experiments, but this can be changed
+    to serial if memory on the SIERRA host machine is limited via
+    ``--processing-parallelism``.
+    """
+
+    exp_to_proc = utils.exp_range_calc(
+        cmdopts["exp_range"], pathset.output_root, criteria.gen_exp_names()
+    )
+
+    parallelism = cmdopts["processing_parallelism"]
+
+    tasks = []
+    for exp in exp_to_proc:
+        tasks.extend(_build_tasklist_for_exp(pathset.output_root / exp.name))
+
+    _logger.debug("Starting %s workers, method=%s", parallelism, mp.get_start_method())
+    with mp.Pool(processes=parallelism, maxtasksperchild=1) as pool:
+        processed = [pool.starmap_async(_worker, tasks)]
+        _logger.debug("Waiting for workers to finish")
+
+        for p in processed:
+            p.get()
+
+        pool.close()
+        pool.join()
+
+    _logger.debug("All workers finished")
+
+
+def _build_tasklist_for_exp(
+    exp_output_root: pathlib.Path,
+) -> tp.List[tp.Tuple[pathlib.Path, pathlib.Path]]:
+    """Add all compressed files from experiment to queue for processing.
+
+    Enqueueing for processing is done at the file-level rather than
+    per-experiment, so that for systems with more CPUs than experiments you
+    still get maximum pthroughput.
+    """
+    res = []
+    for f in exp_output_root.rglob("*.tar.gz"):
+        res.append((exp_output_root, f.relative_to(exp_output_root)))
+
+    return res
+
+
+def _worker(exp_output_root: pathlib.Path, relpath: pathlib.Path) -> None:
+    """Decompress a single tarball from a single experiment.
+
+    Arguments:
+        exp_output_root: Output root for the :term:`Experiment`.
+
+        relpath: Path to the actual tarball relative to the experiment root.
+    """
+    if not (exp_output_root / relpath).exists():
+        _logger.warning(
+            "Cannot decompress: %s does not exist", (exp_output_root / relpath)
+        )
+        return
+
+    with tarfile.open(exp_output_root / relpath, "r:gz") as tar:
+        tar.extractall(filter="data", path=str((exp_output_root / relpath).parent))
+
+
+__all__ = ["proc_batch_exp"]

sierra/plugins/proc/imagize/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright 2025 John Harwell, All rights reserved.
+#
+# SPDX-License Identifier: MIT
+#
+"""Container module for the :term:`imagizing` data processing plugin."""
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+
+
+def sierra_plugin_type() -> str:
+    return "pipeline"

sierra/plugins/proc/imagize/cmdline.py

@@ -0,0 +1,49 @@
+#
+# Copyright 2025 John Harwell, All rights reserved.
+#
+# SPDX-License Identifier: MIT
+#
+
+# Core packages
+import typing as tp
+import argparse
+
+# 3rd party packages
+
+# Project packages
+from sierra.core import types
+from sierra.plugins import PluginCmdline
+
+
+def build(
+    parents: tp.List[argparse.ArgumentParser], stages: tp.List[int]
+) -> PluginCmdline:
+    """
+    Get a cmdline parser supporting the ``proc.imagize`` processing plugin.
+    """
+    cmdline = PluginCmdline(parents, stages)
+    cmdline.stage3.add_argument(
+        "--imagize-no-stats",
+        action="store_true",
+        help="""
+             If the ``proc.imagize`` plugin is run, don't run statistics
+             generation/assume it has already been run.  This can save TONS of
+             time for large imagizing workloads/workloads where the memory
+             limitations of the SIERRA host machine are such that you need to
+             specify different levels of ``--processing-parallelism`` for
+             statistics calculations/imagizing to avoid filling up memory.
+             """
+        + cmdline.stage_usage_doc([3]),
+        default=False,
+    )
+    return cmdline
+
+
+def to_cmdopts(args: argparse.Namespace) -> types.Cmdopts:
+    return {
+        "imagize_no_stats": args.imagize_no_stats,
+    }
+
+
+def sphinx_cmdline_stage3():
+    return build([], [3]).parser