Adding barriers in all processes so the init function exits only when…

… the worker and monitor loops have started

Signed-off-by: arunjose696 <[email protected]>

unidist/core/backends/mpi/core/controller/api.py

@@ -261,6 +261,8 @@ def init():
         atexit.register(_termination_handler)
         signal.signal(signal.SIGTERM, _termination_handler)
         signal.signal(signal.SIGINT, _termination_handler)
+        # Exit the init function in root only after monitor and worker loops have started.
+        mpi_state.comm.Barrier()
         return
     elif mpi_state.is_monitor_process():
         from unidist.core.backends.mpi.core.monitor.loop import monitor_loop

unidist/core/backends/mpi/core/monitor/loop.py

@@ -170,6 +170,7 @@ def monitor_loop():
     The loop exits on special cancelation operation.
     ``unidist.core.backends.mpi.core.common.Operations`` defines a set of supported operations.
     """
+    monitor_logger.debug("Monitor loop started")
     task_counter = TaskCounter.get_instance()
     mpi_state = communication.MPIState.get_instance()
     wait_handler = WaitHandler.get_instance()
@@ -182,7 +183,8 @@ def monitor_loop():
     # Once all workers excluding ``Root`` and ``Monitor`` ranks are ready to shutdown,
     # ``Monitor` sends the shutdown signal to every worker, as well as notifies ``Root`` that
     # it can exit the program.
-
+    # Barrier in monitor process to check if monitor loop has started
+    mpi_state.comm.Barrier()
     while True:
         # Listen receive operation from any source
         operation_type, source_rank = communication.mpi_recv_operation(mpi_state.comm)

unidist/core/backends/mpi/core/worker/loop.py

@@ -84,6 +84,7 @@ async def worker_loop():
     The loop exits on special cancelation operation.
     ``unidist.core.backends.mpi.core.common.Operations`` defines a set of supported operations.
     """
+    w_logger.debug("Worker loop started")
     task_store = TaskStore.get_instance()
     local_store = LocalObjectStore.get_instance()
     request_store = RequestStore.get_instance()
@@ -94,6 +95,9 @@ async def worker_loop():
     # and the worker sends the ready to shutdown signal to ``Monitor``.
     # Once all workers excluding ``Root`` and ``Monitor`` ranks are ready to shutdown,
     # ``Monitor` sends the shutdown signal to every worker so they can exit the loop.
+
+    # Barrier in worker process to check if worker loop has started
+    mpi_state.comm.Barrier()
     while True:
         # Listen receive operation from any source
         operation_type, source_rank = await async_wrap(