Use a wrapper script to add MPI envs to singularity binary

Author: kinow

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
+from importlib.resources import files
 from subprocess import check_call, check_output, run  # nosec
+from tempfile import NamedTemporaryFile
 from typing import cast
 
 from mypy_extensions import mypyc_attr
@@ -590,23 +593,50 @@ 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",
-            "--ipc",
-            "--contain",
-        ]
         mpi_req, is_req = self.builder.get_requirement(MPIRequirementName)
-        if not mpi_req or not is_req:
-            runtime.append("--cleanenv")
-        else:
+        mpi_enabled = mpi_req and is_req
+        mpi_env_vars_reference_file_name: str | None = None
+        runtime: list[str] = []
+        if mpi_enabled:
+            # Save current environment variables. The ``cwl_singularity_wrapper.sh`` will
+            # diff it against the env vars produced by mpirun/srun/etc., and use the new
+            # env vars as Singularity ``--env`` arguments for MPI.
+            with NamedTemporaryFile(mode="wb", 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(files("cwltool") / "singularity_wrapper.sh"),
+                mpi_env_vars_reference_file_name
+            ])
+
+            # MPI implementations like OpenMPI and MPICH use shared memory.
             self.append_volume(
                 runtime,
                 runtime_context.create_tmpdir(),
                 "/dev/shm",
                 writable=True,
             )
+        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 is_apptainer_1_1_or_newer() or is_version_3_10_or_newer():
             runtime.append("--no-eval")
 

cwltool/singularity_wrapper.sh

@@ -0,0 +1,86 @@
+#!/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 command
+#   line arguments to singularity (with --env).
+#
+#   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
+#
+
+usage() {
+    cat <<'EOF' >&2
+Usage:
+  singularity_wrapper.sh <baseline-env-file> <singularity-bin> [args...]
+
+Run with --help to see full documentation.
+EOF
+    exit 1
+}
+
+if [[ "${1:-}" == "--help" ]]; then
+    grep '^#' "$0" | sed 's/^#//' >&2
+    exit 0
+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
+
+if [[ ! -x "$SINGULARITY_BIN" ]]; then
+    echo "Error: singularity binary not executable: $SINGULARITY_BIN" >&2
+    exit 3
+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.
+ENV_ARGS=()
+while IFS='=' read -r k v; do
+    [[ -n "$k" ]] || continue
+    # If the current env doesn't exist (! -v) in the given baseline env (BASE_ENV),
+    # then we want to add it as --env in singularity.
+    if [[ ! -v BASE_ENV[$k] ]]; then
+        ENV_ARGS+=(--env "$k=$v")
+    fi
+done < <(env)
+
+# Debug
+# echo "Adding env vars into Singularity command: ${#ENV_ARGS[@]}" >&2
+
+# Launch wrapped singularity command with the singularity executable, followed
+# by the new --env arguments, and finally with all the rest the CWL runner had.
+exec "$1" "${ENV_ARGS[@]}" "${@:2}"