Raise an exception if the comm world is too small (#107)

* Raise an exeption if the comm world is too small

* Expose exception at module level

* Tweak exceptions to be more readable

* Tweak language a little more

dask_mpi/__init__.py

@@ -1,5 +1,6 @@
 from ._version import get_versions
 from .core import initialize, send_close_signal
+from .exceptions import WorldTooSmallException
 
 __version__ = get_versions()["version"]
 del get_versions

dask_mpi/cli.py

@@ -6,6 +6,8 @@
 from distributed.utils import import_term
 from mpi4py import MPI
 
+from .exceptions import WorldTooSmallException
+
 
 @click.command()
 @click.argument("scheduler_address", type=str, required=False)
@@ -95,6 +97,14 @@ def main(
     name,
 ):
     comm = MPI.COMM_WORLD
+
+    world_size = comm.Get_size()
+    if scheduler and world_size < 2:
+        raise WorldTooSmallException(
+            f"Not enough MPI ranks to start cluster, found {world_size}, "
+            "needs at least 2, one each for the scheduler and a worker."
+        )
+
     rank = comm.Get_rank()
 
     try:

dask_mpi/core.py

@@ -6,6 +6,8 @@
 from distributed import Client, Nanny, Scheduler
 from distributed.utils import import_term
 
+from .exceptions import WorldTooSmallException
+
 
 def initialize(
     interface=None,
@@ -74,6 +76,13 @@ def initialize(
 
         comm = MPI.COMM_WORLD
 
+    world_size = comm.Get_size()
+    if world_size < 3:
+        raise WorldTooSmallException(
+            f"Not enough MPI ranks to start cluster, found {world_size}, "
+            "needs at least 3, one each for the scheduler, client and a worker."
+        )
+
     rank = comm.Get_rank()
 
     if not worker_options:

dask_mpi/exceptions.py

@@ -0,0 +1,2 @@
+class WorldTooSmallException(RuntimeError):
+    """Not enough MPI ranks to start all required processes."""

dask_mpi/tests/test_cli.py

@@ -51,6 +51,24 @@ def test_basic(loop, worker_class, mpirun):
                 assert c.submit(lambda x: x + 1, 10).result() == 11
 
 
+def test_small_world(mpirun):
+    with tmpfile(extension="json") as fn:
+        # Set too few processes to start cluster
+        p = subprocess.Popen(
+            mpirun
+            + [
+                "-np",
+                "1",
+                "dask-mpi",
+                "--scheduler-file",
+                fn,
+            ]
+        )
+
+        p.communicate()
+        assert p.returncode != 0
+
+
 def test_no_scheduler(loop, mpirun):
     with tmpfile(extension="json") as fn:
         cmd = mpirun + ["-np", "2", "dask-mpi", "--scheduler-file", fn]

dask_mpi/tests/test_core.py

@@ -18,3 +18,15 @@ def test_basic(mpirun):
 
     p.communicate()
     assert p.returncode == 0
+
+
+def test_small_world(mpirun):
+    script_file = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "core_basic.py"
+    )
+
+    # Set too few processes to start cluster
+    p = subprocess.Popen(mpirun + ["-np", "1", sys.executable, script_file])
+
+    p.communicate()
+    assert p.returncode != 0