From 6ebb2eba72bc85784479525981d7c521f3b92e57 Mon Sep 17 00:00:00 2001
From: Jinghan Yao <yao.877@buckeyemail.osu.edu>
Date: Wed, 11 Jan 2023 20:44:25 -0500
Subject: [PATCH 1/2] large scale stable MPI

---
 dask_mpi/core.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/dask_mpi/core.py b/dask_mpi/core.py
index 8ff6ccb..caacead 100644
--- a/dask_mpi/core.py
+++ b/dask_mpi/core.py
@@ -8,6 +8,7 @@
 
 
 def initialize(
+    comm,
     interface=None,
     nthreads=1,
     local_directory="",
@@ -18,7 +19,6 @@ def initialize(
     protocol=None,
     worker_class="distributed.Worker",
     worker_options=None,
-    comm=None,
     exit=True,
 ):
     """
@@ -36,6 +36,8 @@ def initialize(
 
     Parameters
     ----------
+    comm: mpi4py.MPI.Intracomm
+        Optional MPI communicator to use instead of COMM_WORLD
     interface : str
         Network interface like 'eth0' or 'ib0'
     nthreads : int
@@ -57,8 +59,6 @@ def initialize(
         Class to use when creating workers
     worker_options : dict
         Options to pass to workers
-    comm: mpi4py.MPI.Intracomm
-        Optional MPI communicator to use instead of COMM_WORLD
     exit: bool
         Whether to call sys.exit on the workers and schedulers when the event
         loop completes.
@@ -69,11 +69,8 @@ def initialize(
         Only returned if exit=False. Inidcates whether this rank should continue
         to run client code (True), or if it acts as a scheduler or worker (False).
     """
-    if comm is None:
-        from mpi4py import MPI
-
-        comm = MPI.COMM_WORLD
-
+    assert comm is not None, "MPI Comm World needs to be created before import distributed."
+    
     rank = comm.Get_rank()
 
     if not worker_options:

From 64c95c39da0be00b03a84f38daab465a347bdb60 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 12 Jan 2023 02:03:57 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dask_mpi/core.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/dask_mpi/core.py b/dask_mpi/core.py
index caacead..007268a 100644
--- a/dask_mpi/core.py
+++ b/dask_mpi/core.py
@@ -69,8 +69,10 @@ def initialize(
         Only returned if exit=False. Inidcates whether this rank should continue
         to run client code (True), or if it acts as a scheduler or worker (False).
     """
-    assert comm is not None, "MPI Comm World needs to be created before import distributed."
-    
+    assert (
+        comm is not None
+    ), "MPI Comm World needs to be created before import distributed."
+
     rank = comm.Get_rank()
 
     if not worker_options: