Extractor and Injector async models fixed (#881)

* Fixes to injector and extractor async models Signed-off-by: bamsumit <[email protected]> * Deletd redundant time increase Signed-off-by: bamsumit <[email protected]> --------- Signed-off-by: bamsumit <[email protected]>
lava-nc · Aug 8, 2024 · efbe872 · efbe872
1 parent d22e829
commit efbe872
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git a/src/lava/proc/io/extractor.py b/src/lava/proc/io/extractor.py
@@ -130,25 +130,22 @@ def run_spk(self) -> None:
 @requires(CPU)
 class PyLoihiExtractorModelAsync(PyAsyncProcessModel):
     in_port: PyInPort = LavaPyType(PyInPort.VEC_DENSE, float)
+    out_port: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, float)
 
     def __init__(self, proc_params: dict) -> None:
         super().__init__(proc_params=proc_params)
 
         channel_config = self.proc_params["channel_config"]
-        self._pm_to_p_src_port = self.proc_params["pm_to_p_src_port"]
-        self._pm_to_p_src_port.start()
 
         self._send = channel_config.get_send_full_function()
         self.time_step = 1
 
     def run_async(self) -> None:
         while self.time_step != self.num_steps + 1:
-            self._send(self._pm_to_p_src_port, self.in_port.recv())
+            self._send(self.out_port.csp_ports[-1],
+                       self.in_port.recv())
             self.time_step += 1
 
-    def __del__(self) -> None:
-        self._pm_to_p_src_port.join()
-
 
 class VarWire(AbstractProcess):
     """VarWire allows non-Lava code, such as a third-party Python library

diff --git a/src/lava/proc/io/injector.py b/src/lava/proc/io/injector.py
@@ -46,6 +46,7 @@ class Injector(AbstractProcess):
         buffer is full and how the dst_port behaves when the buffer is empty
         and not empty.
     """
+
     def __init__(self,
                  shape: ty.Tuple[int, ...],
                  buffer_size: ty.Optional[int] = 50,
@@ -133,16 +134,14 @@ def __del__(self) -> None:
 @requires(CPU)
 class PyLoihiInjectorModelAsync(PyAsyncProcessModel):
     """PyAsyncProcessModel for the Injector Process."""
+    in_port: PyInPort = LavaPyType(PyInPort.VEC_DENSE, float)
     out_port: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, float)
 
     def __init__(self, proc_params: dict) -> None:
         super().__init__(proc_params=proc_params)
 
         shape = self.proc_params["shape"]
         channel_config = self.proc_params["channel_config"]
-        self._p_to_pm_dst_port = self.proc_params["p_to_pm_dst_port"]
-        self._p_to_pm_dst_port.start()
-
         self._zeros = np.zeros(shape)
 
         self._receive_when_empty = channel_config.get_receive_empty_function()
@@ -153,17 +152,18 @@ def __init__(self, proc_params: dict) -> None:
     def run_async(self) -> None:
         while self.time_step != self.num_steps + 1:
             self._zeros.fill(0)
-            elements_in_buffer = self._p_to_pm_dst_port._queue.qsize()
+            elements_in_buffer = self.in_port.csp_ports[-1]._queue.qsize()
 
             if elements_in_buffer == 0:
                 data = self._receive_when_empty(
-                    self._p_to_pm_dst_port,
+                    self.in_port,
                     self._zeros)
             else:
                 data = self._receive_when_not_empty(
-                    self._p_to_pm_dst_port,
+                    self.in_port,
                     self._zeros,
                     elements_in_buffer)
+
             self.out_port.send(data)
             self.time_step += 1