Pure python Modify for PR643

Signed-off-by: Miao, Hongda <[email protected]>

Author: hongdami

src/lava/magma/runtime/message_infrastructure/py_multiprocessing.py

@@ -10,12 +10,15 @@
 
 import multiprocessing as mp
 import os
-from multiprocessing.managers import SharedMemoryManager
 import traceback
 
 from lava.magma.runtime.message_infrastructure.interfaces import ChannelType
 from lava.magma.runtime.message_infrastructure import Channel
 from lava.magma.runtime.message_infrastructure.pypychannel import PyPyChannel
+from lava.magma.runtime.message_infrastructure.shared_memory_manager import (
+    SharedMemoryManager,
+)
+
 try:
     from lava.magma.compiler.channels.cpychannel import \
         CPyChannel, PyCChannel
@@ -45,6 +48,7 @@ def __init__(self, *args, **kwargs):
         mp.Process.__init__(self, *args, **kwargs)
         self._pconn, self._cconn = mp.Pipe()
         self._exception = None
+        self._is_done = False
 
     def run(self):
         try:
@@ -54,10 +58,20 @@ def run(self):
             tb = traceback.format_exc()
             self._cconn.send((e, tb))
 
+    def join(self):
+        if not self._is_done:
+            super().join()
+            super().close()
+            if self._pconn.poll():
+                self._exception = self._pconn.recv()
+            self._cconn.close()
+            self._pconn.close()
+            self._is_done = True
+
     @property
     def exception(self):
         """Exception property."""
-        if self._pconn.poll():
+        if not self._is_done and self._pconn.poll():
             self._exception = self._pconn.recv()
         return self._exception
 

tests/lava/magma/compiler/channels/test_channel.py

@@ -6,11 +6,13 @@
 import unittest
 import time
 from multiprocessing import Process
-from multiprocessing.managers import SharedMemoryManager
 from lava.magma.runtime.message_infrastructure import (
     create_channel,
     Channel,
 )
+from lava.magma.runtime.message_infrastructure.shared_memory_manager import (
+    SharedMemoryManager
+)
 
 
 class MockInterface: