PERF-#367: Use std::fill to fill service shared buffer with a given v…

…alue (#373)

Co-authored-by: Igoshev, Iaroslav <[email protected]>
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.global_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/memory/_memory.pyx

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from libc.stdint cimport uint8_t
+from libc.stdint cimport uint8_t, int64_t
 
 cimport memory
 
@@ -25,3 +25,17 @@ def parallel_memcopy(const uint8_t[:] src, uint8_t[:] dst, int memcopy_threads):
                                 len(src),
                                 64,
                                 memcopy_threads)
+
+def fill(int64_t[:] buff, int64_t value):
+    """
+    Fill a given buffer with a given value.
+
+    Parameters
+    ----------
+    buff : int64_t[:]
+        Original data.
+    value : int64_t
+        Value to fill.
+    """
+    with nogil:
+        memory.fill(&buff[0], len(buff), value)

unidist/core/backends/mpi/core/memory/memory.cpp

@@ -17,6 +17,10 @@ namespace unidist {
     return reinterpret_cast<uint8_t *>(value & bits);
   }
 
+  void fill(int64_t *buff, int64_t size, int64_t value){
+    std::fill(buff, buff+size, value);
+  }
+
   void parallel_memcopy(uint8_t *dst,
                         const uint8_t *src,
                         int64_t nbytes,

unidist/core/backends/mpi/core/memory/memory.h

@@ -17,6 +17,8 @@ namespace unidist {
                         int64_t nbytes,
                         uintptr_t block_size,
                         int num_threads);
+
+    void fill(int64_t *buff, int64_t size, int64_t value);
 }  // namespace unidist
 
 #endif  // MEMORY_H

unidist/core/backends/mpi/core/memory/memory.pxd

@@ -14,3 +14,6 @@ cdef extern from "memory.h" namespace "unidist" nogil:
                             int64_t nbytes, 
                             uintptr_t block_size, 
                             int num_threads)
+
+    void fill(int64_t *buff, int64_t size, int64_t value)
+

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

@@ -177,12 +177,14 @@ def monitor_loop():
     shared_store = SharedObjectStore.get_instance()
     shm_manager = SharedMemoryManager()
 
-    workers_ready_to_shutdown = []
-    shutdown_workers = False
+    # Barrier to check if monitor process is ready to start the communication loop
+    mpi_state.global_comm.Barrier()
+    monitor_logger.debug("Monitor loop started")
     # 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.
-
+    workers_ready_to_shutdown = []
+    shutdown_workers = False
     while True:
         # Listen receive operation from any source
         operation_type, source_rank = communication.mpi_recv_operation(

unidist/core/backends/mpi/core/shared_object_store.py

@@ -9,9 +9,8 @@
 import time
 import psutil
 import weakref
-from array import array
 
-from unidist.core.backends.mpi.core._memory import parallel_memcopy
+from unidist.core.backends.mpi.core._memory import parallel_memcopy, fill
 from unidist.core.backends.mpi.utils import ImmutableDict
 
 try:
@@ -190,9 +189,7 @@ def _allocate_shared_memory(self):
         self.service_shared_buffer = memoryview(service_buffer).cast("l")
         # Set -1 to the service buffer because 0 is a valid value and may be recognized by mistake.
         if mpi_state.is_monitor_process():
-            self.service_shared_buffer[:] = array(
-                "l", [-1] * len(self.service_shared_buffer)
-            )
+            fill(self.service_shared_buffer, -1)
 
     def _parse_data_id(self, data_id):
         """

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

@@ -88,12 +88,16 @@ async def worker_loop():
     local_store = LocalObjectStore.get_instance()
     request_store = RequestStore.get_instance()
     async_operations = AsyncOperations.get_instance()
-    ready_to_shutdown_posted = False
+
+    # Barrier to check if worker process is ready to start the communication loop
+    mpi_state.global_comm.Barrier()
+    w_logger.debug("Worker loop started")
     # Once the worker receives the cancel signal from ``Root`` rank,
     # it is getting to shutdown. All pending requests and communications are cancelled,
     # 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.
+    ready_to_shutdown_posted = False
     while True:
         # Listen receive operation from any source
         operation_type, source_rank = await async_wrap(