use get_connection_info in Deisa and Bridge

Author: benoitmartin88

src/deisa/dask/deisa.py

@@ -44,16 +44,36 @@
 from distributed import Client, Future
 
 
+def get_connection_info(dask_scheduler_address: str | Client) -> Client:
+    if isinstance(dask_scheduler_address, Client):
+        client = dask_scheduler_address
+    elif isinstance(dask_scheduler_address, str):
+        try:
+            client = Client(address=dask_scheduler_address)
+        except ValueError:
+            # try scheduler_file
+            if os.path.isfile(dask_scheduler_address):
+                client = Client(scheduler_file=dask_scheduler_address)
+            else:
+                raise ValueError(
+                    "dask_scheduler_address must be a string containing the address of the scheduler, "
+                    "or a string containing a file name to a dask scheduler file, or a Dask Client object.")
+    else:
+        raise ValueError(
+            "dask_scheduler_address must be a string containing the address of the scheduler, "
+            "or a string containing a file name to a dask scheduler file, or a Dask Client object.")
+
+    return client
+
+
 class Bridge:
-    def __init__(self, dask_scheduler_address: str | Client, mpi_comm_size: int, mpi_rank: int,
-                 arrays_metadata: dict[str, dict], **kwargs):
+    def __init__(self, mpi_comm_size: int, mpi_rank: int, arrays_metadata: dict[str, dict],
+                 get_connection_info: Callable, *args, **kwargs):
         """
         Initializes an object to manage communication between an MPI-based distributed
         system and a Dask-based framework. The class ensures proper allocation of workers
         among processes and instantiates the required communication objects like queues.
 
-        :param dask_scheduler_address: Address of the Dask Scheduler as a string or an instance of
-            Dask Client that facilitates communication with the cluster.
         :type dask_scheduler_address: str | Client
 
         :param mpi_comm_size: Total number of MPI processes involved in the computation.
@@ -75,17 +95,14 @@ def __init__(self, dask_scheduler_address: str | Client, mpi_comm_size: int, mpi
                     }
         :type arrays_metadata: dict[str, dict]
 
+        :param get_connection_info: A function that returns a connected Dask Client.
+        :type get_connection_info: Callable
+
         :param kwargs: Currently unused.
         :type kwargs: dict
         """
 
-        if isinstance(dask_scheduler_address, str):
-            self.client = Client(dask_scheduler_address)
-        elif isinstance(dask_scheduler_address, Client):
-            self.client = dask_scheduler_address
-        else:
-            raise ValueError("dask_scheduler_address must be a string or a Dask Client object.")
-
+        self.client = get_connection_info()
         self.mpi_rank = mpi_rank
         self.arrays_metadata = arrays_metadata
         self.futures = []
@@ -143,38 +160,21 @@ def publish_data(self, array_name: str, data: np.ndarray, iteration: int):
 class Deisa:
     SLIDING_WINDOW_THREAD_PREFIX = "deisa_sliding_window_callback_"
 
-    def __init__(self, dask_scheduler_address: str | Client, mpi_comm_size: int, nb_workers: int):
+    def __init__(self, mpi_comm_size, nb_workers, get_connection_info: Callable, *args, **kwargs):
         """
         Initializes the distributed processing environment and configures workers using
         a Dask scheduler. This class handles setting up a Dask client and ensures the
         specified number of workers are available for distributed computation tasks.
 
-        :param dask_scheduler_address: Instance of Dask's Client to connect to the cluster,
-            or address string of the Dask scheduler,
-            or a string containing a file name to a dask scheduler file.
         :param mpi_comm_size: Number of MPI processes for the computation.
         :param nb_workers: Expected number of workers to be synchronized with the
             Dask client.
+        :param get_connection_info: A function that returns a connected Dask Client.
+        :type get_connection_info: Callable
         """
         # dask.config.set({"distributed.deploy.lost-worker-timeout": 60, "distributed.workers.memory.spill":0.97, "distributed.workers.memory.target":0.95, "distributed.workers.memory.terminate":0.99 })
 
-        if isinstance(dask_scheduler_address, Client):
-            self.client = dask_scheduler_address
-        elif isinstance(dask_scheduler_address, str):
-            try:
-                self.client = Client(address=dask_scheduler_address)
-            except ValueError:
-                # try scheduler_file
-                if os.path.isfile(dask_scheduler_address):
-                    self.client = Client(scheduler_file=dask_scheduler_address)
-                else:
-                    raise ValueError(
-                        "dask_scheduler_address must be a string containing the address of the scheduler, "
-                        "or a string containing a file name to a dask scheduler file, or a Dask Client object.")
-        else:
-            raise ValueError(
-                "dask_scheduler_address must be a string containing the address of the scheduler, "
-                "or a string containing a file name to a dask scheduler file, or a Dask Client object.")
+        self.client = get_connection_info()
 
         # Wait for all workers to be available.
         self.workers = [w_addr for w_addr in self.client.scheduler_info()["workers"].keys()]

test/TestSimulator.py

@@ -28,10 +28,10 @@
 # =============================================================================
 
 import numpy as np
-from deisa.common import BridgeInterface
-from deisa.dask import Bridge
 from distributed import Client
 
+from deisa.dask import Bridge
+
 
 class TestSimulation:
     __test__ = False
@@ -48,8 +48,9 @@ def __init__(self, client: Client, global_grid_size: tuple, mpi_parallelism: tup
         assert global_grid_size[1] % mpi_parallelism[1] == 0, "cannot compute local grid size for y dimension"
 
         self.nb_mpi_ranks = mpi_parallelism[0] * mpi_parallelism[1]
-        self.bridges: list[BridgeInterface] = [Bridge(client, self.nb_mpi_ranks, rank, arrays_metadata) for rank in
-                                               range(self.nb_mpi_ranks)]
+        self.bridges: list[Bridge] = [
+            Bridge(self.nb_mpi_ranks, rank, arrays_metadata, get_connection_info=lambda: client)
+            for rank in range(self.nb_mpi_ranks)]
 
     def __gen_data(self, noise_level: int = 0) -> np.ndarray:
         # Create coordinate grid

test/test_deisa.py

@@ -39,7 +39,7 @@
 from distributed import Client, LocalCluster, Queue, Variable
 
 from TestSimulator import TestSimulation
-from deisa.dask import Deisa
+from deisa.dask import Deisa, get_connection_info
 
 
 @pytest.mark.parametrize('global_shape', [(32, 32), (32, 16), (16, 32)])
@@ -128,30 +128,31 @@ def env_setup_tcp_cluster(self):
     def test_deisa_ctor_client(self, env_setup_tcp_cluster):
         cluster = env_setup_tcp_cluster
         client = Client(cluster)
-        deisa = Deisa(client, mpi_comm_size=0, nb_workers=0)
+        deisa = Deisa(mpi_comm_size=0, nb_workers=0, get_connection_info=lambda: client)
         assert deisa.client is not None, "Deisa should not be None"
         assert deisa.client.scheduler.address == cluster.scheduler_address, "Client should be the same as scheduler"
         deisa.close()
 
     def test_deisa_ctor_str(self, env_setup_tcp_cluster):
         cluster = env_setup_tcp_cluster
-        deisa = Deisa('tcp://127.0.0.1:4242', mpi_comm_size=0, nb_workers=0)
+        deisa = Deisa(mpi_comm_size=0, nb_workers=0,
+                      get_connection_info=lambda: get_connection_info('tcp://127.0.0.1:4242'))
         assert deisa.client is not None, "Deisa should not be None"
         assert deisa.client.scheduler.address == cluster.scheduler_address, "Client should be the same as scheduler"
         deisa.close()
 
     def test_deisa_ctor_scheduler_file(self, env_setup_tcp_cluster):
         cluster = env_setup_tcp_cluster
         f = os.path.abspath(os.path.dirname(__file__)) + os.path.sep + 'test-scheduler.json'
-        deisa = Deisa(f, mpi_comm_size=0, nb_workers=0)
+        deisa = Deisa(mpi_comm_size=0, nb_workers=0, get_connection_info=lambda: get_connection_info(f))
         assert deisa.client is not None, "Deisa should not be None"
         assert deisa.client.scheduler.address == cluster.scheduler_address, "Client should be the same as scheduler"
         deisa.close()
 
     def test_deisa_ctor_scheduler_file_error(self):
         with pytest.raises(ValueError) as e:
             f = os.path.abspath(os.path.dirname(__file__)) + os.path.sep + 'test-scheduler-error.json'
-            deisa = Deisa(f, mpi_comm_size=0, nb_workers=0)
+            deisa = Deisa(mpi_comm_size=0, nb_workers=0, get_connection_info=lambda: get_connection_info(f))
 
 
 class TestUsingDaskCluster:
@@ -243,7 +244,7 @@ def test_get_dask_array(self, global_grid_size: tuple, mpi_parallelism: tuple, n
         nb_mpi_ranks = mpi_parallelism[0] * mpi_parallelism[1]
         nb_workers = len(cluster.workers)
 
-        deisa = Deisa(client, nb_mpi_ranks, nb_workers)
+        deisa = Deisa(nb_mpi_ranks, nb_workers, get_connection_info=lambda: client)
         sim = TestSimulation(client,
                              global_grid_size=global_grid_size,
                              mpi_parallelism=mpi_parallelism,
@@ -279,7 +280,7 @@ def test_sliding_window_callback_register(self, global_grid_size: tuple, mpi_par
         nb_mpi_ranks = mpi_parallelism[0] * mpi_parallelism[1]
         nb_workers = len(cluster.workers)
 
-        deisa = Deisa(client, nb_mpi_ranks, nb_workers)
+        deisa = Deisa(nb_mpi_ranks, nb_workers, get_connection_info=lambda: client)
         sim = TestSimulation(client,
                              global_grid_size=global_grid_size,
                              mpi_parallelism=mpi_parallelism,
@@ -332,7 +333,7 @@ def test_sliding_window_callback_unregister(self, env_setup):
         nb_mpi_ranks = mpi_parallelism[0] * mpi_parallelism[1]
         nb_workers = len(cluster.workers)
 
-        deisa = Deisa(client, nb_mpi_ranks, nb_workers)
+        deisa = Deisa(nb_mpi_ranks, nb_workers, get_connection_info=lambda: client)
         sim = TestSimulation(client,
                              global_grid_size=global_grid_size,
                              mpi_parallelism=mpi_parallelism,
@@ -379,7 +380,7 @@ def test_sliding_window_callback_throws(self, env_setup):
         nb_mpi_ranks = mpi_parallelism[0] * mpi_parallelism[1]
         nb_workers = len(cluster.workers)
 
-        deisa = Deisa(client, nb_mpi_ranks, nb_workers)
+        deisa = Deisa(nb_mpi_ranks, nb_workers, get_connection_info=lambda: client)
         sim = TestSimulation(client,
                              global_grid_size=global_grid_size,
                              mpi_parallelism=mpi_parallelism,
@@ -450,7 +451,7 @@ def test_sliding_window_map_blocks(self, env_setup):
         nb_mpi_ranks = mpi_parallelism[0] * mpi_parallelism[1]
         nb_workers = len(cluster.workers)
 
-        deisa = Deisa(client, nb_mpi_ranks, nb_workers)
+        deisa = Deisa(nb_mpi_ranks, nb_workers, get_connection_info=lambda: client)
         sim = TestSimulation(client,
                              global_grid_size=global_grid_size,
                              mpi_parallelism=mpi_parallelism,