Detect MPI with Singularity

README.rst

@@ -663,41 +663,46 @@ The host-specific parameters are configured in a simple YAML file
 (specified with the ``--mpi-config-file`` flag). The allowed keys are
 given in the following table; all are optional.
 
-+----------------+------------------+----------+------------------------------+
-| Key            | Type             | Default  | Description                  |
-+================+==================+==========+==============================+
-| runner         | str              | "mpirun" | The primary command to use.  |
-+----------------+------------------+----------+------------------------------+
-| nproc_flag     | str              | "-n"     | Flag to set number of        |
-|                |                  |          | processes to start.          |
-+----------------+------------------+----------+------------------------------+
-| default_nproc  | int              | 1        | Default number of processes. |
-+----------------+------------------+----------+------------------------------+
-| extra_flags    | List[str]        | []       | A list of any other flags to |
-|                |                  |          | be added to the runner's     |
-|                |                  |          | command line before          |
-|                |                  |          | the ``baseCommand``.         |
-+----------------+------------------+----------+------------------------------+
-| env_pass       | List[str]        | []       | A list of environment        |
-|                |                  |          | variables that should be     |
-|                |                  |          | passed from the host         |
-|                |                  |          | environment through to the   |
-|                |                  |          | tool (e.g., giving the       |
-|                |                  |          | node list as set by your     |
-|                |                  |          | scheduler).                  |
-+----------------+------------------+----------+------------------------------+
-| env_pass_regex | List[str]        | []       | A list of python regular     |
-|                |                  |          | expressions that will be     |
-|                |                  |          | matched against the host's   |
-|                |                  |          | environment. Those that match|
-|                |                  |          | will be passed through.      |
-+----------------+------------------+----------+------------------------------+
-| env_set        | Mapping[str,str] | {}       | A dictionary whose keys are  |
-|                |                  |          | the environment variables set|
-|                |                  |          | and the values being the     |
-|                |                  |          | values.                      |
-+----------------+------------------+----------+------------------------------+
-
++----------------+------------------+------------+------------------------------+
+| Key            | Type             | Default    | Description                  |
++================+==================+============+==============================+
+| runner         | str              | "mpirun"   | The primary command to use.  |
++----------------+------------------+------------+------------------------------+
+| nproc_flag     | str              | "-n"       | Flag to set number of        |
+|                |                  |            | processes to start.          |
++----------------+------------------+------------+------------------------------+
+| default_nproc  | int              | 1          | Default number of processes. |
++----------------+------------------+------------+------------------------------+
+| extra_flags    | List[str]        | []         | A list of any other flags to |
+|                |                  |            | be added to the runner's     |
+|                |                  |            | command line before          |
+|                |                  |            | the ``baseCommand``.         |
++----------------+------------------+------------+------------------------------+
+| env_pass       | List[str]        | []         | A list of environment        |
+|                |                  |            | variables that should be     |
+|                |                  |            | passed from the host         |
+|                |                  |            | environment through to the   |
+|                |                  |            | tool (e.g., giving the       |
+|                |                  |            | node list as set by your     |
+|                |                  |            | scheduler).                  |
++----------------+------------------+------------+------------------------------+
+| env_pass_regex | List[str]        | []         | A list of python regular     |
+|                |                  |            | expressions that will be     |
+|                |                  |            | matched against the host's   |
+|                |                  |            | environment. Those that match|
+|                |                  |            | will be passed through.      |
++----------------+------------------+------------+------------------------------+
+| env_set        | Mapping[str,str] | {}         | A dictionary whose keys are  |
+|                |                  |            | the environment variables set|
+|                |                  |            | and the values being the     |
+|                |                  |            | values.                      |
++----------------+------------------+------------+------------------------------+
+| shm_enabled    | bool             | True       | Flag to control whether      |
+|                |                  |            | shared memory is used or not.|
++----------------+------------------+------------+------------------------------+
+| shm_dir        | str              | "/dev/shm" | Location to use for shared   |
+|                |                  |            | memory.                      |
++----------------+------------------+------------+------------------------------+
 
 Enabling Fast Parser (experimental)
 ===================================

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,10 @@
 import sys
 import threading
 from collections.abc import Callable, MutableMapping, MutableSequence
+from contextlib import suppress
+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 +33,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
@@ -203,7 +208,7 @@ def __init__(
         hints: list[CWLObjectType],
         name: str,
     ) -> None:
-        """Builder for invoking the Singularty software container engine."""
+        """Builder for invoking the Singularity software container engine."""
         super().__init__(builder, joborder, make_path_mapper, requirements, hints, name)
 
     @staticmethod
@@ -592,14 +597,55 @@ 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."""
+                with suppress(FileNotFoundError):  # pragma: no cover
+                    os.remove(mpi_env_vars_reference_file_name)  # pragma: no cover
+
+            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")
 
@@ -665,4 +711,4 @@ def create_runtime(
         if container_HOME:
             # Restore HOME if we removed it above.
             self.environment["HOME"] = container_HOME
-        return (runtime, None)
+        return runtime, None

cwltool/singularity_wrapper.sh

@@ -0,0 +1,83 @@
+#!/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 a variable.
+BASELINE_CONTENT=$'\n'"$(cat "$BASELINE_FILE")"$'\n'
+
+# 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.
+    # Check if the key exists in the BASELINE_CONTENT string in the
+    # form \n$KEY= (that's why we start the BASELINE and end it with \n).
+    if [[ ! "$BASELINE_CONTENT" == *$'\n'"$k"=* ]]; 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" "${@}"