using partitions

unidist/core/backends/mpi/core/communication.py

@@ -31,6 +31,7 @@
 
 # Logger configuration
 logger = common.get_logger("communication", "communication.log")
+logger1 = common.get_logger("communication1", "communication1.log")
 is_logger_header_printed = False
 
 
@@ -463,9 +464,19 @@ def mpi_send_buffer(comm, buffer, dest_rank, data_type=MPI.CHAR, buffer_size=Non
     """
     if buffer_size:
         comm.send(buffer_size, dest=dest_rank, tag=common.MPITag.OBJECT)
-    comm.Send([buffer, data_type], dest=dest_rank, tag=common.MPITag.BUFFER)
-
-
+    else:
+        buffer_size = len(buffer)        
+    PARTITION_SIZE = 100000
+    partitions=list(range(0, buffer_size,PARTITION_SIZE))
+    partitions.append(buffer_size)
+    #with pkl5._bigmpi() as bigmpi:
+    if True:
+        for i,_ in enumerate(partitions):
+            if i+1<len(partitions):
+                logger1.debug(f" in send i={i} current={comm.rank}, destination={dest_rank}")
+                temp=buffer[partitions[i]:partitions[i+1]]
+                comm.Send([buffer[partitions[i]:partitions[i+1]],data_type], dest=dest_rank, tag=common.MPITag.BUFFER )
+
 def mpi_isend_buffer(comm, buffer_size, buffer, dest_rank):
     """
     Send buffer object to another MPI rank in a non-blocking way.
@@ -494,8 +505,18 @@ def mpi_isend_buffer(comm, buffer_size, buffer, dest_rank):
     requests = []
     h1 = comm.isend(buffer_size, dest=dest_rank, tag=common.MPITag.OBJECT)
     requests.append((h1, None))
-    h2 = comm.Isend([buffer, MPI.CHAR], dest=dest_rank, tag=common.MPITag.BUFFER)
-    requests.append((h2, buffer))
+
+    PARTITION_SIZE = 100000
+    partitions=list(range(0, buffer_size,PARTITION_SIZE))
+    partitions.append(buffer_size)
+    #with pkl5._bigmpi() as bigmpi:
+    if True:
+        for i,_ in enumerate(partitions):
+            if i+1<len(partitions):
+                temp=buffer[partitions[i]:partitions[i+1]]
+                h2=comm.Isend([buffer[partitions[i]:partitions[i+1]],MPI.CHAR], dest=dest_rank, tag=common.MPITag.BUFFER )
+                requests.append((h2, buffer))
+
     return requests
 
 
@@ -525,7 +546,21 @@ def mpi_recv_buffer(comm, source_rank, result_buffer=None):
     if result_buffer is None:
         buf_size = comm.recv(source=source_rank, tag=common.MPITag.OBJECT)
         result_buffer = bytearray(buf_size)
-    comm.Recv(result_buffer, source=source_rank, tag=common.MPITag.BUFFER)
+    else:
+        buf_size = len(result_buffer)
+    PARTITION_SIZE = 100000
+    partitions=list(range(0, buf_size,PARTITION_SIZE))
+    partitions.append(buf_size)
+    #with pkl5._bigmpi as bigmpi:
+    if True:
+        logger1.debug(f" in recv {list(partitions)}")
+        for i,_ in enumerate(partitions):
+            if i+1<len(partitions):
+                logger1.debug(f" in recv {i}")
+                temp=bytearray(partitions[i+1]-partitions[i])
+                comm.Recv(temp,  source=source_rank, tag=common.MPITag.BUFFER )
+                result_buffer[partitions[i]:partitions[i+1]]=temp
+
     return result_buffer