Add MPIRequirement (#1276)

Author: rupertnash

README.rst

@@ -558,6 +558,80 @@ For this example, grab the test.json (and input file) from https://github.com/Ca
 
 .. _`GA4GH Tool Registry API`: https://github.com/ga4gh/tool-registry-schemas
 
+Running MPI-based tools that need to be launched
+------------------------------------------------
+
+Cwltool supports an extension to the CWL spec
+``http://commonwl.org/cwltool#MPIRequirement``. When the tool
+definition has this in its ``requirements``/``hints`` section, and
+cwltool has been run with ``--enable-ext``, then the tool's command
+line will be extended with the commands needed to launch it with
+``mpirun`` or similar. You can specify the number of processes to
+start as either a literal integer or an expression (that will result
+in an integer). For example::
+
+  #!/usr/bin/env cwl-runner
+  cwlVersion: v1.1
+  class: CommandLineTool
+  $namespaces:
+    cwltool: "http://commonwl.org/cwltool#"
+  requirements:
+    cwltool:MPIRequirement:
+      processes: $(inputs.nproc)
+  inputs:
+    nproc:
+      type: int
+
+Interaction with containers: the MPIRequirement currently prepends its
+commands to the front of the command line that is constructed. If you
+wish to run a containerised application in parallel, for simple use
+cases this does work with Singularity, depending upon the platform
+setup. However this combination should be considered "alpha" -- please
+do report any issues you have! This does not work with Docker at the
+moment. (More precisely, you get `n` copies of the same single process
+image run at the same time that cannot communicate with each other.)
+
+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        |
+|                |                  |          | nodelist 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.                      |
++----------------+------------------+----------+------------------------------+
+
+
+
 ===========
 Development
 ===========

cwltool/argparser.py

@@ -615,6 +615,15 @@ def arg_parser() -> argparse.ArgumentParser:
         "listed targets (can provide more than once).",
     )
 
+    parser.add_argument(
+        "--mpi-config-file",
+        type=str,
+        default=None,
+        help="Platform specific configuration for MPI (parallel "
+        "launcher, its flag etc). See README section 'Running MPI-"
+        "based tools' for details of the format.",
+    )
+
     parser.add_argument(
         "workflow",
         type=str,

cwltool/command_line_tool.py

@@ -45,6 +45,7 @@
 from .flatten import flatten
 from .job import CommandLineJob, JobBase
 from .loghandler import _logger
+from .mpi import MPIRequirementName
 from .mutation import MutationManager
 from .pathmapper import PathMapper
 from .process import (
@@ -360,16 +361,21 @@ def __init__(
         self.prov_obj = loadingContext.prov_obj
 
     def make_job_runner(self, runtimeContext: RuntimeContext) -> Type[JobBase]:
-        dockerReq, _ = self.get_requirement("DockerRequirement")
+        dockerReq, dockerRequired = self.get_requirement("DockerRequirement")
+        mpiReq, mpiRequired = self.get_requirement(MPIRequirementName)
+
         if not dockerReq and runtimeContext.use_container:
             if runtimeContext.find_default_container is not None:
                 default_container = runtimeContext.find_default_container(self)
                 if default_container is not None:
-                    self.requirements.insert(
-                        0,
-                        {"class": "DockerRequirement", "dockerPull": default_container},
-                    )
-                    dockerReq = self.requirements[0]
+                    dockerReq = {"class": "DockerRequirement", "dockerPull": default_container}
+                    if mpiRequired:
+                        self.hints.insert(0, dockerReq)
+                        dockerRequired = False
+                    else:
+                        self.requirements.insert(0, dockerReq)
+                        dockerRequired = True
+
                     if (
                         default_container == windows_default_container_id
                         and runtimeContext.use_container
@@ -382,17 +388,40 @@ def make_job_runner(self, runtimeContext: RuntimeContext) -> Type[JobBase]:
                         )
 
         if dockerReq is not None and runtimeContext.use_container:
+            if mpiReq is not None:
+                _logger.warning(
+                    "MPIRequirement with containers is a beta feature"
+                )
             if runtimeContext.singularity:
                 return SingularityCommandLineJob
             elif runtimeContext.user_space_docker_cmd:
                 return UDockerCommandLineJob
+            if mpiReq is not None:
+                if mpiRequired:
+                    if dockerRequired:
+                        raise UnsupportedRequirement("No support for Docker and MPIRequirement both being required")
+                    else:
+                        _logger.warning(
+                            "MPI has been required while Docker is hinted, discarding Docker hint(s)"
+                        )
+                        self.hints = [h for h in self.hints if h["class"] != "DockerRequirement"]
+                        return CommandLineJob
+                else:
+                    if dockerRequired:
+                        _logger.warning(
+                            "Docker has been required while MPI is hinted, discarding MPI hint(s)"
+                        )
+                        self.hints = [h for h in self.hints if h["class"] != MPIRequirementName]
+                    else:
+                        raise UnsupportedRequirement(
+                            "Both Docker and MPI have been hinted - don't know what to do"
+                        )
             return DockerCommandLineJob
-        for t in reversed(self.requirements):
-            if t["class"] == "DockerRequirement":
-                raise UnsupportedRequirement(
-                    "--no-container, but this CommandLineTool has "
-                    "DockerRequirement under 'requirements'."
-                )
+        if dockerRequired:
+            raise UnsupportedRequirement(
+                "--no-container, but this CommandLineTool has "
+                "DockerRequirement under 'requirements'."
+            )
         return CommandLineJob
 
     def make_path_mapper(
@@ -839,6 +868,19 @@ def register_reader(f: CWLObjectType) -> None:
         )
         j.output_callback = output_callbacks
 
+        mpi, _ = self.get_requirement(MPIRequirementName)
+
+        if mpi is not None:
+            np = cast(  # From the schema for MPIRequirement.processes
+                Union[int, str],
+                mpi.get('processes', runtimeContext.mpi_config.default_nproc)
+            )
+            if isinstance(np, str):
+                tmp = builder.do_eval(np)
+                if not isinstance(tmp, int):
+                    raise TypeError("{} needs 'processes' to evaluate to an int, got {}".format(MPIRequirementName, type(np)))
+                np = tmp
+            j.mpi_procs = np
         yield j
 
     def collect_output_ports(

cwltool/context.py

@@ -10,6 +10,7 @@
 from typing_extensions import TYPE_CHECKING
 
 from .builder import Builder, HasReqsHints
+from .mpi import MpiConfig
 from .mutation import MutationManager
 from .pathmapper import PathMapper
 from .secrets import SecretStore
@@ -78,7 +79,7 @@ def copy(self):
 
 class RuntimeContext(ContextBase):
     def __init__(self, kwargs: Optional[Dict[str, Any]] = None) -> None:
-        """Initializet the RuntimeContext from the kwargs."""
+        """Initialize the RuntimeContext from the kwargs."""
         select_resources_callable = Callable[  # pylint: disable=unused-variable
             [Dict[str, Union[int, float]], RuntimeContext], Dict[str, Union[int, float]]
         ]
@@ -141,6 +142,8 @@ def __init__(self, kwargs: Optional[Dict[str, Any]] = None) -> None:
         self.cwl_full_name = ""  # type: str
         self.process_run_id = None  # type: Optional[str]
         self.prov_obj = None  # type: Optional[ProvenanceProfile]
+        self.mpi_config = MpiConfig()  # type: MpiConfig
+
         super(RuntimeContext, self).__init__(kwargs)
 
     def copy(self):

cwltool/extensions-v1.1.yml

@@ -44,3 +44,23 @@ $graph:
         subscope: run
       doc: |
         Specifies the process to run.
+
+- name: MPIRequirement
+  type: record
+  inVocab: false
+  extends: cwl:ProcessRequirement
+  doc: |
+    Indicates that a process requires an MPI runtime.
+  fields:
+    - name: class
+      type: string
+      doc: "Always 'MPIRequirement'"
+      jsonldPredicate:
+        "_id": "@type"
+        "_type": "@vocab"
+    - name: processes
+      type: [int, string]
+      doc: |
+        The number of MPI processes to start. If you give a string,
+        this will be evaluated as a CWL Expression and it must
+        evaluate to an integer.

cwltool/extensions.yml

@@ -154,3 +154,23 @@ $graph:
         _type: "@id"
       doc: |
         Specifies the process to run.
+
+- name: MPIRequirement
+  type: record
+  inVocab: false
+  extends: cwl:ProcessRequirement
+  doc: |
+    Indicates that a process requires an MPI runtime.
+  fields:
+    - name: class
+      type: string
+      doc: "Always 'MPIRequirement'"
+      jsonldPredicate:
+        "_id": "@type"
+        "_type": "@vocab"
+    - name: processes
+      type: [int, string]
+      doc: |
+        The number of MPI processes to start. If you give a string,
+        this will be evaluated as a CWL Expression and it must
+        evaluate to an integer.

cwltool/job.py

@@ -225,6 +225,7 @@ def __init__(
         self.parent_wf = None  # type: Optional[ProvenanceProfile]
         self.timelimit = None  # type: Optional[int]
         self.networkaccess = False  # type: bool
+        self.mpi_procs = None  # type: Optional[int]
 
     def __repr__(self):  # type: () -> str
         """Represent this Job object."""
@@ -286,6 +287,16 @@ def _execute(
         if scr is not None:
             shouldquote = neverquote
 
+        # If mpi_procs (is not None and > 0) then prepend the
+        # appropriate MPI job launch command and flags before the
+        # execution.
+        if self.mpi_procs:
+            menv = runtimeContext.mpi_config
+            mpi_runtime = [menv.runner, menv.nproc_flag, str(self.mpi_procs)] + menv.extra_flags
+            runtime = mpi_runtime + runtime
+            menv.pass_through_env_vars(env)
+            menv.set_env_vars(env)
+
         _logger.info(
             "[job %s] %s$ %s%s%s%s",
             self.name,

cwltool/main.py

@@ -99,7 +99,7 @@
     windows_default_container_id,
 )
 from .workflow import Workflow
-
+from .mpi import MpiConfig
 
 def _terminate_processes() -> None:
     """Kill all spawned processes.
@@ -632,16 +632,15 @@ def setup_schema(
     if custom_schema_callback is not None:
         custom_schema_callback()
     elif args.enable_ext:
-        res = pkg_resources.resource_stream(__name__, "extensions.yml")
-        ext10 = res.read().decode("utf-8")
-        res = pkg_resources.resource_stream(__name__, "extensions-v1.1.yml")
-        ext11 = res.read().decode("utf-8")
+        with pkg_resources.resource_stream(__name__, "extensions.yml") as res:
+            ext10 = res.read().decode("utf-8")
+        with pkg_resources.resource_stream(__name__, "extensions-v1.1.yml") as res:
+            ext11 = res.read().decode("utf-8")
         use_custom_schema("v1.0", "http://commonwl.org/cwltool", ext10)
         use_custom_schema("v1.1", "http://commonwl.org/cwltool", ext11)
         use_custom_schema("v1.2.0-dev1", "http://commonwl.org/cwltool", ext11)
         use_custom_schema("v1.2.0-dev2", "http://commonwl.org/cwltool", ext11)
         use_custom_schema("v1.2.0-dev3", "http://commonwl.org/cwltool", ext11)
-        res.close()
     else:
         use_standard_schema("v1.0")
         use_standard_schema("v1.1")
@@ -888,6 +887,9 @@ def main(
         if not args.enable_ga4gh_tool_registry:
             del ga4gh_tool_registries[:]
 
+        if args.mpi_config_file is not None:
+            runtimeContext.mpi_config = MpiConfig.load(args.mpi_config_file)
+
         setup_schema(args, custom_schema_callback)
 
         if args.provenance: