Fix wrapper (passing on MN5) add tests use MpiConfig

Author: kinow

cwltool/mpi.py

@@ -23,6 +23,7 @@ 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
     ) -> None:
         """
         Initialize from the argument mapping.
@@ -35,6 +36,7 @@ def __init__(
         env_pass: []
         env_pass_regex: []
         env_set: {}
+        shm_enabled: True
 
         Any unknown keys will result in an exception.
         """
@@ -45,6 +47,7 @@ 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
 
     @classmethod
     def load(cls: type[MpiConfigT], config_file_name: str) -> MpiConfigT:

cwltool/singularity.py

@@ -598,13 +598,14 @@ def create_runtime(
 
         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 ``cwl_singularity_wrapper.sh`` will
+            # 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 Singularity ``--env`` arguments for MPI.
-            with NamedTemporaryFile(mode="wb", delete=False) as f:
+            # 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
@@ -619,16 +620,9 @@ def delete_mpi_baseline_env() -> None:
 
             runtime.extend([
                 str(files("cwltool") / "singularity_wrapper.sh"),
-                mpi_env_vars_reference_file_name
+                mpi_env_vars_reference_file_name,
+                "singularity"
             ])
-
-            # 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")
 
@@ -639,6 +633,15 @@ def delete_mpi_baseline_env() -> None:
             "--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(),
+                "/dev/shm",
+                writable=True,
+            )
+
         if is_apptainer_1_1_or_newer() or is_version_3_10_or_newer():
             runtime.append("--no-eval")
 

cwltool/singularity_wrapper.sh

@@ -7,10 +7,10 @@ set -euo pipefail
 #   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).
+#   (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.
+#   This allows CWL (which uses ``--cleanenv``) to launch MPI + Singularity.
 #
 # USAGE
 #   singularity_wrapper.sh <baseline-env-file> <singularity-bin> <args>
@@ -28,20 +28,23 @@ set -euo pipefail
 # EXAMPLE
 #   singularity_wrapper.sh env.txt singularity --cleanenv exec image.sif
 #
+# DEPENDENCIES
+#   It uses the following binaries:
+#   - printenv
 
 usage() {
-    cat <<'EOF' >&2
-Usage:
-  singularity_wrapper.sh <baseline-env-file> <singularity-bin> [args...]
-
-Run with --help to see full documentation.
-EOF
+    # Print usage using the top comment of this file, below the shebang.
+    # Regexes:
+    # - /^#!/d    -> Delete (/d) a comment starting with #! (the shebang, above);
+    # - /^#/,$ {} -> Capture the line that starts with # until the end of the file and run {...};
+    #   - /^#/p;    -> Print (/p) lines starting with # (any comment after the shebang);
+    #   - /^[^#]/q  -> And quit (/q) as soon as the lines doesn't start with # (^[^#]).
+    sed -n '/^#!/d; /^#/,$ { /^#/p; /^[^#]/q }' "$0" | sed 's/^# \{0,1\}//' >&2
     exit 1
 }
 
 if [[ "${1:-}" == "--help" ]]; then
-    grep '^#' "$0" | sed 's/^#//' >&2
-    exit 0
+    usage
 fi
 
 [[ $# -ge 2 ]] || usage
@@ -55,32 +58,27 @@ if [[ ! -f "$BASELINE_FILE" ]]; then
     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=()
+# 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
-    # 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")
+    [[ "$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 < <(env)
-
-# Debug
-# echo "Adding env vars into Singularity command: ${#ENV_ARGS[@]}" >&2
+done < <(printenv)
 
-# 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}"
+# Launch the Singularity binary.
+exec "$SINGULARITY_BIN" "${@}"

tests/test_singularity_wrapper.py

@@ -0,0 +1,108 @@
+"""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(f"""\
+    #!/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