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

sierra/plugins/execenv/hpc/pbs/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2021 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Container module for the PBS execution environment.
+
+See :ref:`plugins/execenv/hpc/pbs`.
+"""
+
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+
+
+def sierra_plugin_type() -> str:
+    return "pipeline"

sierra/plugins/execenv/hpc/pbs/cmdline.py

@@ -0,0 +1,30 @@
+#
+# Copyright 2025 John Harwell, All rights reserved.
+#
+# SPDX-License Identifier: MIT
+#
+"""Command line definitions for the :ref:`plugins/execenv/hpc/PBS`."""
+
+# Core packages
+import typing as tp
+import argparse
+
+# 3rd party packages
+
+# Project packages
+from sierra.plugins.execenv import hpc
+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 ``hpc.pbs`` execution environment.
+    """
+    return hpc.cmdline.HPCCmdline(parents, stages)
+
+
+def to_cmdopts(args: argparse.Namespace) -> types.Cmdopts:
+    return hpc.cmdline.to_cmdopts(args)

sierra/plugins/execenv/hpc/pbs/plugin.py

@@ -0,0 +1,121 @@
+# Copyright 2020 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+HPC plugin for running SIERRA on HPC clusters using the TORQUE-PBS scheduler.
+"""
+
+# Core packages
+import os
+import typing as tp
+import argparse
+import shutil
+import pathlib
+
+# 3rd party packages
+import implements
+
+# Project packages
+from sierra.core import types
+from sierra.core.experiment import bindings
+
+
+def cmdline_postparse_configure(args: argparse.Namespace) -> argparse.Namespace:
+    """
+    Configure SIERRA for PBS HPC.
+
+    Uses the following environment variables (if any of them are not defined an
+    assertion will be triggered):
+
+    - :envvar:`PBS_NUM_PPN`
+
+    - :envvar:`PBS_NODEFILE`
+
+    - :envvar:`PBS_JOBID`
+    """
+
+    keys = ["PBS_NUM_PPN", "PBS_NODEFILE", "PBS_JOBID"]
+
+    for k in keys:
+        assert k in os.environ, f"Non-PBS environment detected: '{k}' not found"
+
+    assert (
+        args.exec_jobs_per_node is not None
+    ), "--exec-jobs-per-node is required (can't be computed from PBS)"
+
+    assert not args.engine_vc, "Engine visual capture not supported on PBS"
+
+    return args
+
+
+@implements.implements(bindings.IExpShellCmdsGenerator)
+class ExpShellCmdsGenerator:
+    """Generate the cmd to invoke GNU Parallel on PBS HPC."""
+
+    def __init__(self, cmdopts: types.Cmdopts, exp_num: int) -> None:
+        self.cmdopts = cmdopts
+
+    def pre_exp_cmds(self) -> tp.List[types.ShellCmdSpec]:
+        shell = shutil.which("bash")
+
+        return [
+            # Since parallel doesn't export any envvars to child processes by
+            # default, we add some common ones.
+            types.ShellCmdSpec(
+                cmd='export PARALLEL="${PARALLEL} --env LD_LIBRARY_PATH"',
+                shell=True,
+                wait=True,
+                env=True,
+            ),
+            # Make sure GNU parallel uses the right shell, because it seems to
+            # defaults to /bin/sh since all cmds are run in a python shell which
+            # does not have $SHELL set.
+            types.ShellCmdSpec(
+                cmd=f"export PARALLEL_SHELL={shell}", shell=True, wait=True, env=True
+            ),
+        ]
+
+    def post_exp_cmds(self) -> tp.List[types.ShellCmdSpec]:
+        return []
+
+    def exec_exp_cmds(self, exec_opts: types.StrDict) -> tp.List[types.ShellCmdSpec]:
+        resume = ""
+        jobid = os.environ["PBS_JOBID"]
+        nodelist = pathlib.Path(exec_opts["exp_input_root"], f"{jobid}-nodelist.txt")
+
+        resume = ""
+        # This can't be --resume, because then GNU parallel looks at the results
+        # directory, and if there is stuff in it, (apparently) assumes that the
+        # job finished...
+        if exec_opts["exec_resume"]:
+            resume = "--resume-failed"
+
+        unique_nodes = types.ShellCmdSpec(
+            cmd=f"sort -u $PBS_NODEFILE > {nodelist}", shell=True, wait=True
+        )
+
+        parallel = (
+            "parallel {2} "
+            "--jobs {1} "
+            "--results {4} "
+            "--joblog {3} "
+            "--sshloginfile {0} "
+            '--workdir {4} < "{5}"'
+        )
+
+        log = pathlib.Path(exec_opts["exp_scratch_root"], "parallel.log")
+        parallel = parallel.format(
+            nodelist,
+            exec_opts["n_jobs"],
+            resume,
+            log,
+            exec_opts["exp_scratch_root"],
+            exec_opts["cmdfile_stem_path"] + exec_opts["cmdfile_ext"],
+        )
+
+        parallel_spec = types.ShellCmdSpec(cmd=parallel, shell=True, wait=True)
+
+        return [unique_nodes, parallel_spec]
+
+
+__all__ = ["cmdline_postparse_configure"]

sierra/plugins/execenv/hpc/slurm/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2021 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Container module for the SLURM execution environment.
+
+See :ref:`plugins/execenv/hpc/slurm`.
+"""
+
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+
+
+def sierra_plugin_type() -> str:
+    return "pipeline"

sierra/plugins/execenv/hpc/slurm/cmdline.py

@@ -0,0 +1,30 @@
+#
+# Copyright 2025 John Harwell, All rights reserved.
+#
+# SPDX-License Identifier: MIT
+#
+"""Command line definitions for the :ref:`plugins/execenv/hpc/PBS`."""
+
+# Core packages
+import typing as tp
+import argparse
+
+# 3rd party packages
+
+# Project packages
+from sierra.plugins.execenv import hpc
+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 ``hpc.slurm`` execution environment.
+    """
+    return hpc.cmdline.HPCCmdline(parents, stages)
+
+
+def to_cmdopts(args: argparse.Namespace) -> types.Cmdopts:
+    return hpc.cmdline.to_cmdopts(args)

sierra/plugins/execenv/hpc/slurm/plugin.py

@@ -0,0 +1,133 @@
+# Copyright 2020 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+HPC plugin for running SIERRA on HPC clusters using the SLURM scheduler.
+"""
+
+# Core packages
+import typing as tp
+import argparse
+import shutil
+import pathlib
+import os
+
+# 3rd party packages
+import implements
+
+# Project packages
+from sierra.core import types
+from sierra.core.experiment import bindings
+
+
+def cmdline_postparse_configure(args: argparse.Namespace) -> argparse.Namespace:
+    """
+    Configure SIERRA for SLURM HPC.
+
+    Uses the following environment variables (if any of them are not defined an
+    assertion will be triggered):
+
+    - :envvar:`SLURM_CPUS_PER_TASK`
+
+    - :envvar:`SLURM_TASKS_PER_NODE`
+
+    - :envvar:`SLURM_JOB_NODELIST`
+
+    - :envvar:`SLURM_JOB_ID`
+
+    """
+
+    keys = [
+        "SLURM_CPUS_PER_TASK",
+        "SLURM_TASKS_PER_NODE",
+        "SLURM_JOB_NODELIST",
+        "SLURM_JOB_ID",
+    ]
+
+    for k in keys:
+        assert k in os.environ, f"Non-SLURM environment detected: '{k}' not found"
+
+    assert not args.engine_vc, "Engine visual capture not supported on SLURM"
+
+    # SLURM_TASKS_PER_NODE can be set to things like '1(x32),3', indicating
+    # that not all nodes will run the same # of tasks. SIERRA expects all
+    # nodes to have the same # tasks allocated to each (i.e., a homogeneous
+    # allocation), so we check for this.
+    assert (
+        "," not in os.environ["SLURM_TASKS_PER_NODE"]
+    ), "SLURM_TASKS_PER_NODE not homogeneous"
+
+    return args
+
+
+@implements.implements(bindings.IExpShellCmdsGenerator)
+class ExpShellCmdsGenerator:
+    """Generate the cmd to correctly invoke GNU Parallel on SLURM HPC."""
+
+    def __init__(self, cmdopts: types.Cmdopts, exp_num: int) -> None:
+        self.cmdopts = cmdopts
+
+    def pre_exp_cmds(self) -> tp.List[types.ShellCmdSpec]:
+        shell = shutil.which("bash")
+
+        return [
+            # Since parallel doesn't export any envvars to child processes by
+            # default, we add some common ones.
+            types.ShellCmdSpec(
+                cmd='export PARALLEL="${PARALLEL} --env LD_LIBRARY_PATH"',
+                shell=True,
+                wait=True,
+                env=True,
+            ),
+            # Make sure GNU parallel uses the right shell, because it seems to
+            # defaults to /bin/sh since all cmds are run in a python shell which
+            # does not have $SHELL set.
+            types.ShellCmdSpec(
+                cmd=f"export PARALLEL_SHELL={shell}", shell=True, wait=True, env=True
+            ),
+        ]
+
+    def post_exp_cmds(self) -> tp.List[types.ShellCmdSpec]:
+        return []
+
+    def exec_exp_cmds(self, exec_opts: types.StrDict) -> tp.List[types.ShellCmdSpec]:
+        jobid = os.environ["SLURM_JOB_ID"]
+        nodelist = pathlib.Path(exec_opts["exp_input_root"], f"{jobid}-nodelist.txt")
+
+        resume = ""
+        # This can't be --resume, because then GNU parallel looks at the results
+        # directory, and if there is stuff in it, (apparently) assumes that the
+        # job finished...
+        if exec_opts["exec_resume"]:
+            resume = "--resume-failed"
+
+        unique_nodes = types.ShellCmdSpec(
+            cmd=f"scontrol show hostnames $SLURM_JOB_NODELIST > {nodelist}",
+            shell=True,
+            wait=True,
+        )
+
+        parallel = (
+            "parallel {2} "
+            "--jobs {1} "
+            "--results {4} "
+            "--joblog {3} "
+            "--sshloginfile {0} "
+            '--workdir {4} < "{5}"'
+        )
+
+        log = pathlib.Path(exec_opts["exp_scratch_root"], "parallel.log")
+        parallel = parallel.format(
+            nodelist,
+            exec_opts["n_jobs"],
+            resume,
+            log,
+            exec_opts["exp_scratch_root"],
+            exec_opts["cmdfile_stem_path"] + exec_opts["cmdfile_ext"],
+        )
+        parallel_spec = types.ShellCmdSpec(cmd=parallel, shell=True, wait=True)
+
+        return [unique_nodes, parallel_spec]
+
+
+__all__ = ["cmdline_postparse_configure", "ExpShellCmdsGenerator"]

sierra/plugins/execenv/prefectserver/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Container module for plugins related to prefect execution environments.
+
+Driven by ``--execenv``.
+"""
+
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+from . import cmdline
+
+
+__all__ = ["cmdline"]

sierra/plugins/execenv/prefectserver/cmdline.py

@@ -0,0 +1,66 @@
+# Copyright 2020 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Common cmdline classes for the various :term:`Prefect` plugins.
+"""
+
+# Core packages
+import argparse
+import typing as tp
+
+# 3rd party packages
+
+# Project packages
+from sierra.core import types
+from sierra.plugins.execenv import hpc
+
+
+class PrefectCmdline(hpc.cmdline.HPCCmdline):
+    """
+    Define the common :term:`Prefect` cmdline.
+
+    This class decorates the HPC cmdline with some additional options; it
+    *mostly* makes sense to consider HPC as a more general form of execution
+    environment and prefect a specialization of it. Mostly.
+    """
+
+    def __init__(
+        self, parents: tp.List[argparse.ArgumentParser], stages: tp.List[int]
+    ) -> None:
+        super().__init__(parents, stages)
+
+    def init_stage2(self) -> None:
+        """Add Prefect cmdline options."""
+        super().init_stage2()
+
+        self.stage2.add_argument(
+            "--work-pool",
+            default="sierra-pool",
+            help="""
+                 Name of the prefect worker pool to use.
+                 """
+            + self.stage_usage_doc([2]),
+        )
+        self.stage2.add_argument(
+            "--work-queue",
+            default="sierra-queue",
+            help="""
+                 Name of the prefect work queue to use.
+                 """
+            + self.stage_usage_doc([2]),
+        )
+
+
+def to_cmdopts(args: argparse.Namespace) -> types.Cmdopts:
+    """Update cmdopts dictionary with the prefect-specific cmdline options."""
+    opts = hpc.cmdline.to_cmdopts(args)
+    updates = {
+        "work_pool": args.work_pool,
+        "work_queue": args.work_queue,
+    }
+    opts |= updates
+    return opts
+
+
+__all__ = ["to_cmdopts"]

sierra/plugins/execenv/prefectserver/dockerremote/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2021 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+Container module for the :term:`Prefect` remote execution environment.
+
+See :ref:`plugins/execenv/prefectserver/dockerremote`.
+"""
+
+# Core packages
+
+# 3rd party packages
+
+# Project packages
+
+
+def sierra_plugin_type() -> str:
+    return "pipeline"

sierra/plugins/execenv/prefectserver/dockerremote/cmdline.py

@@ -0,0 +1,66 @@
+#
+# 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.plugins.execenv import prefectserver
+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 for the ``prefectserver.dockerremote`` execution environment.
+    """
+    cmdline = prefectserver.cmdline.PrefectCmdline(parents, stages)
+
+    cmdline.stage2.add_argument(
+        "--docker-extra-mounts",
+        nargs="+",
+        help="""
+             Add extra mounts in the usual docker format, space separated.
+             """
+        + cmdline.stage_usage_doc([2]),
+    )
+    cmdline.stage2.add_argument(
+        "--docker-image",
+        help="""
+             Path to the docker image to use.
+             """
+        + cmdline.stage_usage_doc([2]),
+    )
+    cmdline.stage2.add_argument(
+        "--docker-is-host-user",
+        action="store_true",
+        default=False,
+        help="""
+             Direct prefect -> docker to run things as the host user within
+             docker containers.  Said user obviously has to exist in the
+             container for this to work.
+             """
+        + cmdline.stage_usage_doc([2]),
+    )
+    return cmdline
+
+
+def to_cmdopts(args: argparse.Namespace) -> types.Cmdopts:
+    opts = prefectserver.cmdline.to_cmdopts(args)
+    updates = {
+        # stage 2
+        "docker_extra_mounts": args.docker_extra_mounts,
+        "docker_image": args.docker_image,
+        "docker_is_host_user": args.docker_is_host_user,
+    }
+    opts |= updates
+
+    return opts

sierra/plugins/execenv/prefectserver/dockerremote/plugin.py

@@ -0,0 +1,132 @@
+# Copyright 2020 John Harwell, All rights reserved.
+#
+#  SPDX-License-Identifier: MIT
+"""
+HPC plugin for running SIERRA on HPC clusters using the SLURM scheduler.
+"""
+
+# Core packages
+import typing as tp
+import argparse
+import pathlib
+import os
+
+# 3rd party packages
+import implements
+import json
+
+# Project packages
+from sierra.core import types
+from sierra.core.experiment import bindings
+
+
+@implements.implements(bindings.IBatchShellCmdsGenerator)
+class BatchShellCmdsGenerator:
+    """
+    Generate commands to invoke :term:`Prefect` for remote computing using docker.
+    """
+
+    def __init__(self, cmdopts: types.Cmdopts) -> None:
+        self.cmdopts = cmdopts
+        self.api_url = os.environ["PREFECT_API_URL"]
+
+    def pre_batch_cmds(self) -> tp.List[types.ShellCmdSpec]:
+        flow_path = pathlib.Path(__file__).parent / "../flow.py"
+        no_prompt = "PREFECT_CLI_PROMPT=false"
+
+        # --sierra-root is always mounted; users can specify whatever other dirs
+        # from the host system are needed in the container to actually run
+        # things. We mount all extra dirs read/write because we can't know what
+        # the engine will be using them for. For --sierra-root, the scratch dir
+        # under there gets written, so rw is appropriate is well.
+        volumes = ["{0}:{0}:rw".format(self.cmdopts["sierra_root"])] + [
+            f"{m}:rw" for m in self.cmdopts["docker_extra_mounts"]
+        ]
+        job_vars = {
+            "image": self.cmdopts["docker_image"],
+            # Only pull the image if it doesn't exist locally
+            "image_pull_policy": "IfNotPresent",
+            "volumes": volumes,
+        }
+
+        # Run as your user, to avoid issues with files in mounted volumes
+        # being owned by root when the container exits.
+        if self.cmdopts["docker_is_host_user"]:
+            job_vars["user"] = f"{os.getuid()}:{os.getgid()}"
+
+        ret = [
+            # Make sure that the image has prefect installed and accessible.
+            types.ShellCmdSpec(
+                cmd="which prefect",
+                shell=True,
+                wait=True,
+            ),
+            types.ShellCmdSpec(
+                cmd="prefect config set PREFECT_API_URL={0}".format(self.api_url),
+                shell=True,
+                wait=True,
+            ),
+        ]
+
+        build = (
+            "{0} prefect deploy {1}:sierra "
+            "--name sierra/dockerremote "
+            "--job-variable '{2}' "
+            "--pool {3} "
+            "--work-queue {4}"
+        ).format(
+            no_prompt,
+            flow_path,
+            json.dumps(job_vars),
+            self.cmdopts["work_pool"],
+            self.cmdopts["work_queue"],
+        )
+        ret.append(
+            types.ShellCmdSpec(
+                cmd=build,
+                shell=True,
+                wait=True,
+            ),
+        )
+        return ret
+
+    def exec_batch_cmds(self, exec_opts: types.StrDict) -> tp.List[types.ShellCmdSpec]:
+        # The needed flow is already present on the server as a deployment, so
+        # we can just run it directly as many times as we want. You don't need
+        # to set the queue/pool--it is inferred from the flow you are running.
+        flow_cmd = (
+            "prefect deployment run sierra/dockerremote "
+            '--params=\'{{"input_root": "{0}","scratch_root": "{1}"}}\' --watch'
+        )
+        # Since this execenv is batch level, so are the outputs and scratch
+        # dirs.
+        flow = flow_cmd.format(exec_opts["batch_root"], exec_opts["batch_scratch_root"])
+        spec = types.ShellCmdSpec(cmd=flow, shell=True, wait=True)
+
+        return [spec]
+
+    def post_batch_cmds(self) -> tp.List[types.ShellCmdSpec]:
+        return []
+
+
+def cmdline_postparse_configure(args: argparse.Namespace) -> argparse.Namespace:
+    """
+    Configure SIERRA to use :term:`Prefect` to run remotely using docker.
+
+    Uses the following environment variables (if any of them are not defined an
+    assertion will be triggered):
+
+    - :envvar:`PREFECT_API_URL`
+    """
+
+    keys = [
+        "PREFECT_API_URL",
+    ]
+
+    for k in keys:
+        assert k in os.environ, f"Non-Prefect environment detected: '{k}' not found"
+
+    return args
+
+
+__all__ = ["cmdline_postparse_configure", "BatchShellCmdsGenerator"]