Detect MPI with Singularity.

It uses a wrapper script to detect environment variables added
by an MPI launcher program such as mpirun or srun, and exports
them as SINGULARITYENV_$KEY=$VALUE.

Updates the MpiConfig of the MPIRequirement extension to add
the shared memory directory, and a flag to enable or disable
shared memory with Singularity (on by default). When enabled,
it maps a volume for the directory used (default /dev/shm).

Author: kinow

cwltool/mpi.py

@@ -23,6 +23,8 @@ def __init__(
         env_pass: list[str] | None = None,
         env_pass_regex: list[str] | None = None,
         env_set: Mapping[str, str] | None = None,
+        shm_enabled: bool = True,
+        shm_dir: str = "/dev/shm",  # nosec B108 - required for MPI/shared memory in containers
     ) -> None:
         """
         Initialize from the argument mapping.
@@ -35,6 +37,8 @@ def __init__(
         env_pass: []
         env_pass_regex: []
         env_set: {}
+        shm_enabled: True
+        shm_dir: "/dev/shm
 
         Any unknown keys will result in an exception.
         """
@@ -45,6 +49,11 @@ def __init__(
         self.env_pass = env_pass or []
         self.env_pass_regex = env_pass_regex or []
         self.env_set = env_set or {}
+        self.shm_enabled = shm_enabled
+        # POSIX only contains functions to handle shared memory, but it does not
+        # specify the directory to be used, nor if a directory needs to be used
+        # at all -- ref: https://pubs.opengroup.org/onlinepubs/9699919799/
+        self.shm_dir = shm_dir
 
     @classmethod
     def load(cls: type[MpiConfigT], config_file_name: str) -> MpiConfigT:

cwltool/singularity.py

@@ -1,5 +1,6 @@
 """Support for executing Docker format containers using Singularity {2,3}.x or Apptainer 1.x."""
 
+import atexit
 import copy
 import hashlib
 import json
@@ -11,7 +12,9 @@
 import sys
 import threading
 from collections.abc import Callable, MutableMapping, MutableSequence
+from importlib.resources import files as resource_files
 from subprocess import check_call, check_output, run  # nosec
+from tempfile import NamedTemporaryFile
 from typing import cast
 
 from cwl_utils.types import CWLDirectoryType, CWLFileType, CWLObjectType
@@ -29,6 +32,7 @@
 from .errors import WorkflowException
 from .job import ContainerCommandLineJob
 from .loghandler import _logger
+from .mpi import MPIRequirementName
 from .pathmapper import MapperEnt, PathMapper
 from .singularity_utils import singularity_supports_userns
 from .utils import create_tmp_dir, ensure_non_writable, ensure_writable
@@ -592,14 +596,57 @@ def create_runtime(
         """Return the Singularity runtime list of commands and options."""
         any_path_okay = self.builder.get_requirement("DockerRequirement")[1] or False
 
-        runtime = [
-            "singularity",
-            "--quiet",
-            "run" if (is_apptainer_1_1_or_newer() or is_version_3_10_or_newer()) else "exec",
-            "--contain",
-            "--ipc",
-            "--cleanenv",
-        ]
+        mpi_req, is_req = self.builder.get_requirement(MPIRequirementName)
+        mpi_enabled = mpi_req and is_req
+        mpi_config = runtime_context.mpi_config
+        mpi_env_vars_reference_file_name: str | None = None
+        runtime: list[str] = []
+        if mpi_enabled:
+            # Save current environment variables. The ``singularity_wrapper.sh`` will
+            # diff it against the env vars produced by mpirun/srun/etc., and use the new
+            # env vars as SINGULARITYENV_... for Singularity.
+            with NamedTemporaryFile(mode="w+", delete=False) as f:
+                for k, v in os.environ.items():
+                    f.write(f"{k}={v}\n")
+                mpi_env_vars_reference_file_name = f.name
+
+            def delete_mpi_baseline_env() -> None:
+                """Clean up the MPI baseline environment variables file at exit."""
+                try:
+                    os.remove(mpi_env_vars_reference_file_name)
+                except FileNotFoundError:
+                    pass
+
+            atexit.register(delete_mpi_baseline_env)
+
+            runtime.extend(
+                [
+                    str(resource_files("cwltool") / "singularity_wrapper.sh"),
+                    mpi_env_vars_reference_file_name,
+                    "singularity",
+                ]
+            )
+        else:
+            runtime.append("singularity")
+
+        runtime.extend(
+            [
+                "--quiet",
+                "run" if (is_apptainer_1_1_or_newer() or is_version_3_10_or_newer()) else "exec",
+                "--contain",
+                "--ipc",
+                "--cleanenv",
+            ]
+        )
+        if mpi_enabled and mpi_config.shm_enabled:
+            # MPI implementations like OpenMPI and MPICH use shared memory.
+            self.append_volume(
+                runtime,
+                runtime_context.create_tmpdir(),
+                mpi_config.shm_dir,
+                writable=True,
+            )
+
         if is_apptainer_1_1_or_newer() or is_version_3_10_or_newer():
             runtime.append("--no-eval")
 

cwltool/singularity_wrapper.sh

@@ -0,0 +1,85 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# singularity_wrapper.sh
+#
+# DESCRIPTION
+#   Wrapper around Singularity/Apptainer for CWL + MPI + Singularity.
+#
+#   This script identifies environment variables added by an MPI launcher
+#   (e.g. srun, mpirun) and adds these environment variables as Singularity
+#   environment variables using the format ``SINGULARITYENV_$KEY=$VALUE``.
+#
+#   This allows CWL (which uses ``--cleanenv``) to launch MPI + Singularity.
+#
+# USAGE
+#   singularity_wrapper.sh <baseline-env-file> <singularity-bin> <args>
+#
+# ARGUMENTS
+#   <baseline-env-file>
+#       Path to the file containing KEY=VALUE pairs with the baseline env.
+#
+#   <singularity-bin>
+#       Path to singularity/apptainer executable.
+#
+#   [args...]
+#       Arguments passed to the singularity binary.
+#
+# EXAMPLE
+#   singularity_wrapper.sh env.txt singularity --cleanenv exec image.sif
+#
+# DEPENDENCIES
+#   It uses the following binaries:
+#   - printenv
+
+usage() {
+    cat >&2 <<EOF
+singularity_wrapper.sh
+
+Wrapper around Singularity/Apptainer for CWL + MPI + Singularity.
+
+USAGE:
+  singularity_wrapper.sh <baseline-env-file> <singularity-bin> [args...]
+EOF
+    exit 1
+}
+
+if [[ "${1:-}" == "--help" ]]; then
+    usage
+fi
+
+[[ $# -ge 2 ]] || usage
+
+BASELINE_FILE="$1"
+SINGULARITY_BIN="$2"
+shift 2
+
+if [[ ! -f "$BASELINE_FILE" ]]; then
+    echo "Error: baseline env file not found: $BASELINE_FILE" >&2
+    exit 2
+fi
+
+# Read baseline env into an array.
+declare -A BASE_ENV
+while IFS='=' read -r k v; do
+    [[ -n "$k" && -n "$v" ]] || continue
+    BASE_ENV["$k"]="$v"
+done < "$BASELINE_FILE"
+
+# Build new environment variables for Singularity (i.e. ``SINGULARITYENV_KEY=VALUE``).
+# Excludes empty variables and variables whose name do not follow POSIX (e.g. some
+# Bash environments on HPC clusters such as BSC MareNostrum5, ``BASH_FUNC_module%%=``).
+while IFS='=' read -r k v; do
+    [[ -n "$k" ]] || continue
+    [[ "$k" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] || continue
+    # If the current env doesn't exist (``! -z``) in the given baseline env (``BASE_ENV``),
+    # then we want to add it as ``--env`` in singularity.
+    if [[ -z "${BASE_ENV[$k]+x}" ]]; then
+        # Debug
+        # echo "Adding env var for Singularity command: SINGULARITYENV_$k=$v" >&2
+        export "SINGULARITYENV_$k=$v"
+    fi
+done < <(printenv)
+
+# Launch the Singularity binary.
+exec "$SINGULARITY_BIN" "${@}"

tests/test_singularity_wrapper.py

@@ -0,0 +1,85 @@
+"""Tests for the Shell wrapper of the Singularity command.
+
+This script tests a Shell script. This script does not contribute to the
+project test coverage (although kcov, or bats+kcov could be used in the
+future).
+"""
+
+import os
+import subprocess
+from importlib.resources import files
+from textwrap import dedent
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+@pytest.mark.parametrize(
+    "args,expected_return_code",
+    [(["--help"], 1), ([""], 1), (["singularity"], 1)],
+    ids=[
+        "Print usage because user passed --help",
+        "Print usage because of missing args",
+        "Print usage because not all args provided",
+    ],
+)
+def test_wrapper_usage(args: list[str], expected_return_code: int) -> None:
+    """Test the usage of the Singularity wrapper is printed."""
+    wrapper = str(files("cwltool") / "singularity_wrapper.sh")
+    command: list[str] = [wrapper] + args
+    result = subprocess.run(command, capture_output=True, text=True)
+
+    assert result.returncode == expected_return_code
+    assert "Wrapper around Singularity/Apptainer for CWL + MPI + Singularity" in result.stderr
+
+
+def test_wrapper_invalid_baseline_env_file() -> None:
+    """Test the script fails if the given file is not valid."""
+    wrapper = str(files("cwltool") / "singularity_wrapper.sh")
+    command: list[str] = [wrapper, "parangaricutirimicuaro.dat", "foo"]
+    result = subprocess.run(command, capture_output=True, text=True)
+
+    assert result.returncode == 2
+    assert "file not found" in result.stderr
+
+
+def test_wrapper_env_vars(tmp_path: "Path") -> None:
+    """Test that the wrapper script adds the new environment variables."""
+    fake_singularity = tmp_path / "fake_singularity"
+    fake_singularity.write_text(dedent("""\
+    #!/bin/bash
+    echo "Fake Singularity script"
+    env
+    """))
+    fake_singularity.chmod(0o755)
+
+    new_env_var = "TEST_WRAPPER_ENV_VARS_INJECTED_VAR"
+
+    # Create the baseline environment variables file.
+    baseline_env = os.environ
+    assert new_env_var not in baseline_env, "The test needs a new env var!"
+    baseline = tmp_path / "baseline.env"
+    baseline.write_text("A=1\nB=2\n")
+    for k, v in baseline_env.items():
+        baseline.write_text(f"{k}={v}")
+
+    # Now pretend we are mpirun, and we are adding a new env var.
+    new_env = os.environ.copy()
+    new_env[new_env_var] = "42"
+
+    wrapper = str(files("cwltool") / "singularity_wrapper.sh")
+    command: list[str] = [wrapper, str(baseline), str(fake_singularity), "--cleanenv"]
+
+    result = subprocess.run(command, capture_output=True, text=True, env=new_env)
+
+    assert result.returncode == 0
+    # There, now the wrapper just runs `env`, and the output must
+    # contain the new environment variable. We know the wrapper
+    # must have worked because we have thew new variable in the
+    # output...
+    assert new_env_var in result.stdout
+    # And also because we have the new SINGULARITYENV_{new_env_var}!
+    assert f"SINGULARITYENV_{new_env_var}" in result.stdout