Lmod-environment

.github/workflows/ci-cd.yml

@@ -143,6 +143,50 @@ jobs:
           fail_ci_if_error: true
           token: ${{ secrets.CODECOV_TOKEN }}
 
+  test-lmod:
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        python-version: ['3.11', '3.12', '3.13']
+        dependencies: [latest]
+      fail-fast: False
+
+    env:
+      DEPENDS: ${{ matrix.dependencies }}
+
+    steps:
+      - name: Install Lmod
+        run: sudo apt-get install -y lmod
+      - name: Set env
+        run: |
+          echo "RELEASE_VERSION=v3.7.1" >> $GITHUB_ENV
+          echo "NO_ET=TRUE" >> $GITHUB_ENV
+      - uses: actions/checkout@v5
+        with:
+          fetch-depth: 0
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install tox
+        run: |
+          uv tool install tox --with=tox-uv --with=tox-gh-actions
+      - name: Show tox config
+        run: tox c
+      - name: Run tox
+        # Run test files with singularity tests; re-add the overridable "-n auto"
+        run: |
+          source /etc/profile.d/lmod.sh
+          export
+          tox -v --exit-and-dump-after 1200 -- -n auto \
+            pydra/environments/tests/test_lmod.py
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          fail_ci_if_error: true
+          token: ${{ secrets.CODECOV_TOKEN }}
+
+
   test-slurm:
     strategy:
       matrix:
@@ -334,7 +378,7 @@ jobs:
         path: docs/build/html
 
   deploy:
-    needs: [build, build-docs, test, test-singularity, test-slurm]
+    needs: [build, build-docs, test, test-singularity, test-slurm, test-lmod]
     runs-on: ubuntu-latest
     if: github.event_name == 'release'
     permissions:

docs/source/explanation/environments.rst

@@ -7,16 +7,18 @@ construction, and allows tasks to be run in environments that are isolated from
 host system, and that have specific software dependencies.
 
 The environment a task runs within is specified by the ``environment`` argument passed
-to the execution call (e.g. ``my_task(worker="cf", environment="docker")``) or in the
+to the execution call (e.g. ``my_task(worker="cf", environment=Docker("brainlife/fsl")``) or in the
 ``workflow.add()`` call in workflow constructors.
 
 Specifying at execution
 -----------------------
 
 The environment for a task can be specified at execution time by passing the ``environment`` argument to the task call.
 This can be an instance of `pydra.environments.native.Environment` (for the host system),
-`pydra.environments.docker.Environment` (for Docker containers), or
-`pydra.environments.singularity.Environment` (for Singularity containers), or a custom environment.
+`pydra.environments.docker.Environment` (for Docker containers),
+`pydra.environments.singularity.Environment` (for Singularity containers),
+`pydra.environments.lmod.Lmod` (for lmod environment modules, e.g. ``module load fsl/6.0.7``)
+or a your own custom environment.
 
 Example:
 

docs/source/reference/api.rst

@@ -32,3 +32,21 @@ Engine classes
    :members:
    :undoc-members:
    :show-inheritance:
+
+
+Environments
+------------
+
+.. automodule:: pydra.environments.docker.Docker
+   :members:
+   :show-inheritance:
+
+
+.. automodule:: pydra.environments.singularity.Singularity
+   :members:
+   :show-inheritance:
+
+
+.. automodule:: pydra.environments.lmod.Lmod
+   :members:
+   :show-inheritance:

docs/source/tutorial/2-advanced-execution.ipynb

@@ -312,7 +312,7 @@
     "## Environments and hooks\n",
     "\n",
     "For shell tasks, it is possible to specify that the command runs within a specific\n",
-    "software environment, such as those provided by software containers (e.g. Docker or Singularity/Apptainer).\n",
+    "software environment, such as those provided by software containers (e.g. Docker, Singularity/Apptainer or Lmod).\n",
     "This is down by providing the environment to the submitter/execution call,"
    ]
   },

pydra/engine/tests/utils.py

@@ -1,4 +1,5 @@
 # Tasks for testing
+import os
 import time
 import sys
 import shutil
@@ -25,6 +26,9 @@
 need_singularity = pytest.mark.skipif(
     shutil.which("singularity") is None, reason="no singularity available"
 )
+need_lmod = pytest.mark.skipif(
+    "MODULEPATH" not in os.environ, reason="modules not available"
+)
 no_win = pytest.mark.skipif(
     sys.platform.startswith("win"),
     reason="docker command not adjusted for windows docker",

pydra/environments/base.py

@@ -200,7 +200,11 @@ def map_path(fileset: os.PathLike | FileSet) -> Path:
         return bindings, values
 
 
-def execute(cmd, strip=False):
+def execute(
+    cmd: ty.Sequence[str],
+    strip: bool = False,
+    **kwargs: ty.Any,
+) -> tuple[int, str, str]:
     """
     Run the event loop with coroutine.
 
@@ -213,26 +217,23 @@ def execute(cmd, strip=False):
     cmd : :obj:`list` or :obj:`tuple`
         The command line to be executed.
     strip : :obj:`bool`
-        TODO
+        Whether to strip the output strings. Default is ``False``.
+    kwargs : keyword arguments
+        Additional keyword arguments passed to the subprocess call.
 
     """
-    rc, stdout, stderr = read_and_display(*cmd, strip=strip)
-    """
-    loop = get_open_loop()
-    if loop.is_running():
-        rc, stdout, stderr = read_and_display(*cmd, strip=strip)
-    else:
-        rc, stdout, stderr = loop.run_until_complete(
-            read_and_display_async(*cmd, strip=strip)
-        )
-    """
+    rc, stdout, stderr = read_and_display(*cmd, strip=strip, **kwargs)
     return rc, stdout, stderr
 
 
-def read_and_display(*cmd, strip=False, hide_display=False):
+def read_and_display(
+    *cmd: str, strip: bool = False, **kwargs: ty.Any
+) -> tuple[int, str, str]:
     """Capture a process' standard output."""
     try:
-        process = sp.run(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
+        process: sp.CompletedProcess = sp.run(
+            cmd, stdout=sp.PIPE, stderr=sp.PIPE, **kwargs
+        )
     except Exception:
         # TODO editing some tracing?
         raise

pydra/environments/lmod.py

@@ -0,0 +1,83 @@
+import os
+import typing as ty
+import logging
+from pathlib import Path
+import re
+import subprocess as sp
+
+import attrs
+from pydra.compose import shell
+from pydra.environments import base
+from pydra.utils.general import ensure_list
+
+
+logger = logging.getLogger("pydra")
+
+if ty.TYPE_CHECKING:
+    from pydra.engine.job import Job
+
+
+@attrs.define
+class Lmod(base.Environment):
+    """Lmod environment."""
+
+    modules: list[str] = attrs.field(converter=ensure_list)
+
+    @modules.validator
+    def _validate_modules(self, _, value: ty.Any) -> None:
+        if not value:
+            raise ValueError("At least one module must be specified")
+        if not all(isinstance(v, str) for v in value):
+            raise ValueError("All module names must be strings")
+
+    def execute(self, job: "Job[shell.Task]") -> dict[str, int | str]:
+        env_src = self.run_lmod_cmd("python", "load", *self.modules)
+        env = {}
+        for key, value in re.findall(
+            r"""os\.environ\[['"](.*?)['"]\]\s*=\s*['"](.*?)['"]""", env_src
+        ):
+            env[key] = value
+        cmd_args = job.task._command_args(values=job.inputs)
+        values = base.execute(cmd_args, env=env)
+        return_code, stdout, stderr = values
+        if return_code:
+            msg = f"Error running '{job.name}' job with {cmd_args}:"
+            if stderr:
+                msg += "\n\nstderr:\n" + stderr
+            if stdout:
+                msg += "\n\nstdout:\n" + stdout
+            raise RuntimeError(msg)
+        return {"return_code": return_code, "stdout": stdout, "stderr": stderr}
+
+    @classmethod
+    def modules_are_installed(cls) -> bool:
+        return "MODULESHOME" in os.environ
+
+    @classmethod
+    def run_lmod_cmd(cls, *args: str) -> str:
+        if not cls.modules_are_installed():
+            raise RuntimeError(
+                "Could not find Lmod installation, please ensure it is installed and MODULESHOME is set"
+            )
+        lmod_exec = Path(os.environ["MODULESHOME"]) / "libexec" / "lmod"
+
+        try:
+            output_bytes, error_bytes = sp.Popen(
+                [str(lmod_exec)] + list(args),
+                stdout=sp.PIPE,
+                stderr=sp.PIPE,
+            ).communicate()
+        except (sp.CalledProcessError, OSError) as e:
+            raise RuntimeError(f"Error running 'lmod': {e}")
+
+        output = output_bytes.decode("utf-8")
+        error = error_bytes.decode("utf-8")
+
+        if output == "_mlstatus = False\n":
+            raise RuntimeError(f"Error running module cmd '{' '.join(args)}':\n{error}")
+
+        return output
+
+
+# Alias so it can be referred to as lmod.Environment
+Environment = Lmod