Add MPI-capable tasks

jobqueue_features/mpi_dask_worker.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+"""
+Distribution of MPI enabled tasks
+"""
+
+from jobqueue_features.mpi_wrapper import (
+    mpi_deserialize_and_execute,
+    serialize_function_and_args,
+    shutdown_mpitask_worker,
+)
+from distributed.cli.dask_worker import go
+from mpi4py import MPI
+
+
+def prepare_for_mpi_tasks():
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+
+    if rank == 0:
+        # Start dask so root reports to scheduler and accepts tasks
+        # Task distribution is part of task itself (via our wrapper)
+        go()
+
+        # As a final task, send a shutdown to the other MPI ranks
+        serialized_object = serialize_function_and_args(shutdown_mpitask_worker)
+        mpi_deserialize_and_execute(serialized_object=serialized_object)
+    else:
+        while True:
+            # Calling with no arguments means these are non-root processes
+            mpi_deserialize_and_execute()
+
+
+if __name__ == "__main__":
+    prepare_for_mpi_tasks()

jobqueue_features/mpi_wrapper.py

@@ -1,5 +1,7 @@
+from distributed.protocol import serialize, deserialize
 import shlex
 import subprocess
+import sys
 from typing import Dict  # noqa
 
 
@@ -19,6 +21,7 @@ def mpi_wrap(
     cpus_per_task=None,
     ntasks_per_node=None,
     exec_args="",
+    return_wrapped_command=False,
     **kwargs
 ):
     # type: (str, str, str, str, str, ...) -> Dict[str, str]
@@ -79,16 +82,126 @@ def get_default_mpi_params(
             exec_args=exec_args,
         )
     )
-    try:
-        proc = subprocess.Popen(
-            shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    if return_wrapped_command:
+        result = cmd
+    else:
+        try:
+            proc = subprocess.Popen(
+                shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+            out = proc.stdout.read()
+            err = proc.stderr.read()
+        except OSError as err:
+            raise OSError(
+                "OS error caused by constructed command: {cmd}\n\n{err}".format(
+                    cmd=cmd, err=err
+                )
+            )
+        result = {"cmd": cmd, "out": out, "err": err}
+
+    return result
+
+
+def shutdown_mpitask_worker():
+    from mpi4py import MPI
+
+    # Finalise MPI
+    MPI.Finalize()
+    # and then exit
+    exit()
+
+
+def deserialize_and_execute(serialized_object):
+    # Ensure the serialized object is of the expected type
+    if isinstance(serialized_object, dict):
+        # Make sure it has the expected entries
+        if not ("header" in serialized_object and "frames" in serialized_object):
+            raise RuntimeError(
+                "serialized_object dict does not have expected keys [header, frames]"
+            )
+    else:
+        raise RuntimeError("Cannot deserialize without a serialized_object")
+    func = deserialize(serialized_object["header"], serialized_object["frames"])
+    if serialized_object.get("args_header"):
+        args = deserialize(
+            serialized_object["args_header"], serialized_object["args_frames"]
         )
-        out = proc.stdout.read()
-        err = proc.stderr.read()
-    except OSError as err:
-        raise OSError(
-            "OS error caused by constructed command: {cmd}\n\n{err}".format(
-                cmd=cmd, err=err
+    else:
+        args = []
+    if serialized_object.get("kwargs_header"):
+        kwargs = deserialize(
+            serialized_object["kwargs_header"], serialized_object["kwargs_frames"]
+        )
+    else:
+        kwargs = {}
+
+    # Free memory space used by (potentially large) serialised object
+    del serialized_object
+
+    # Execute the function and return
+    return func(*args, **kwargs)
+
+
+def flush_and_abort(msg="Flushing print buffer and aborting", comm=None, error_code=1):
+    from mpi4py import MPI
+
+    if comm is None:
+        comm = MPI.COMM_WORLD
+    print(msg)
+    sys.stdout.flush()
+    if error_code == 0:
+        print("To abort correctly, we need to use a non-zero error code")
+        error_code = 1
+    comm.Abort(error_code)
+
+
+def mpi_deserialize_and_execute(serialized_object=None, root=0, comm=None):
+    from mpi4py import MPI
+
+    if comm is None:
+        comm = MPI.COMM_WORLD
+
+    # We only handle the case where root has the object and is the one who returns
+    # something
+    if serialized_object:
+        # Check we have a valid communicator
+        try:
+            rank = comm.Get_rank()
+        except AttributeError:
+            flush_and_abort(
+                msg="Looks like you did not pass a valid MPI communicator, aborting "
+                "using global communicator"
             )
+            sys.stdout.flush()
+            MPI.COMM_WORLD.Abort(1)
+        if rank != root:
+            flush_and_abort(
+                msg="Only root rank (%d) can contain a serialized object for this "
+                "call, my rank is %d...aborting!" % (root, rank),
+                comm=comm,
+            )
+        return_something = True
+    else:
+        return_something = False
+    serialized_object = comm.bcast(serialized_object, root=root)
+    result = deserialize_and_execute(serialized_object)
+
+    if return_something and result:
+        return result
+
+
+def serialize_function_and_args(func, *args, **kwargs):
+    header, frames = serialize(func)
+    serialized_object = {"header": header, "frames": frames}
+    if args:
+        args_header, args_frames = serialize(args)
+        serialized_object.update(
+            {"args_header": args_header, "args_frames": args_frames}
+        )
+    if kwargs:
+        kwargs_header, kwargs_frames = serialize(kwargs)
+        serialized_object.update(
+            {"kwargs_header": kwargs_header, "kwargs_frames": kwargs_frames}
         )
-    return {"cmd": cmd, "out": out, "err": err}
+
+    return serialized_object

mpi_pickling.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+"""
+Distribution of MPI enabled tasks
+"""
+
+from jobqueue_features.mpi_wrapper import (
+    mpi_deserialize_and_execute,
+    serialize_function_and_args,
+    shutdown_mpitask_worker,
+)
+from mpi4py import MPI
+
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+
+if rank == 0:
+    # This is the task, which is only defined on root
+    def task1(task_name):
+        from mpi4py import MPI
+
+        comm = MPI.COMM_WORLD
+        size = comm.Get_size()
+        name = MPI.Get_processor_name()
+        all_nodes = comm.gather(name, root=0)
+        if all_nodes:
+            all_nodes = set(all_nodes)
+        else:
+            all_nodes = []
+        # Since it is a return  value it will only get printed by root
+        return "Running %d tasks of type %s on nodes %s." % (size, task_name, all_nodes)
+
+    def task2(name, task_name="default"):
+        from mpi4py import MPI
+
+        comm = MPI.COMM_WORLD
+        rank = comm.Get_rank()
+        print("Hi %s, my rank is %d for task of type %s" % (name, rank, task_name))
+
+    serialized_object = serialize_function_and_args(task1, "task1")
+    result = mpi_deserialize_and_execute(serialized_object=serialized_object)
+    if result:
+        print(result)
+    serialized_object = serialize_function_and_args(task2, "alan", task_name="task2")
+    mpi_deserialize_and_execute(serialized_object=serialized_object)
+
+    # As a final task, send a shutdown to the other MPI ranks
+    serialized_object = serialize_function_and_args(shutdown_mpitask_worker)
+    mpi_deserialize_and_execute(serialized_object=serialized_object)
+
+else:
+    while True:
+        # Calling with no arguments means these are non-root processes
+        mpi_deserialize_and_execute()