Node state passing

juturna/cli/commands/_juturna_remote_service.py

@@ -9,7 +9,7 @@
 
 import grpc
 
-from juturna.components import Message, Node
+from juturna.components import Message, Node, State
 from juturna.remotizer._remote_context import RequestContext
 from juturna.remotizer._remote_builder import _standalone_builder
 
@@ -73,6 +73,8 @@ def __init__(self, node: Node, remote_name: str):
         )
         self._cleanup_thread.start()
 
+        self._pipe_state_store = dict()
+
         logger.info(f'Service initialized for node {node.name}')
 
     def _increment_stat(self, stat_name: str):
@@ -145,6 +147,19 @@ def SendAndReceive(self, request: ProtoEnvelope, context):
             tracking_id = next(self._tracking_id_counter)
             sender = envelope_dict.get('sender')
             envelope_id = envelope_dict.get('id')
+            pipe_id = envelope_dict['pipe_id']
+            state = envelope_dict['state']
+
+            if self._pipe_state_store.get(pipe_id, State()) != state:
+                logger.warning('stored state mismatch with received state')
+
+                self._pipe_state_store[pipe_id] = state
+
+            if self._pipe_state_store.get(pipe_id, None) is None:
+                self._pipe_state_store[pipe_id] = State()
+
+            # pass the state reference to the node
+            self.node.link_state(self._pipe_state_store[pipe_id])
 
             if not sender:
                 raise ValueError('Missing sender in request envelope')
@@ -192,6 +207,8 @@ def SendAndReceive(self, request: ProtoEnvelope, context):
             response_envelope = create_envelope(
                 message=proto_response,
                 creator=self.remote_name,
+                pipe_id=pipe_id,
+                state=self._pipe_state_store[pipe_id],
                 configuration={},
                 metadata={
                     'processing_time': time.time() - request_context.created_at

juturna/cli/commands/basic_node.template

@@ -56,8 +56,9 @@ class $_node_class_name(Node[BasePayload, BasePayload]):
         """Destroy the node"""
         ...
 
-    def update(self, message: Message[BasePayload]):
+    def update(self, message: Message[BasePayload], **kwargs):
         """Receive data from upstream, transmit data downstream"""
+        # state = kwargs['state']
         ...
 
     # uncomment next_batch to design custom synchronisation policy

juturna/components/__init__.py

@@ -3,11 +3,13 @@
 from juturna.components._node import Node
 from juturna.components._pipeline import Pipeline
 from juturna.components._buffer import Buffer
+from juturna.components._state import State
 
 
 __all__ = [
     'Message',
     'Node',
     'Pipeline',
     'Buffer',
+    'State',
 ]

juturna/components/_node.py

@@ -21,6 +21,7 @@
 from juturna.meta import JUTURNA_TELEMETRY_BATCH_SIZE
 
 from juturna.components._buffer import Buffer
+from juturna.components._state import State
 from juturna.components._telemetry_manager import TelemetryManager
 from juturna.components._synchronisers import _SYNCHRONISERS
 
@@ -65,6 +66,7 @@ def __init__(
         )
 
         self._status: ComponentStatus | None = None
+        self._state: State | None = None
 
         self._queue = queue.Queue(maxsize=JUTURNA_MAX_QUEUE_SIZE)
         self._worker_thread: threading.Thread | None = None
@@ -130,6 +132,9 @@ def destinations(self) -> list:
     def link_telemetry(self, manager: TelemetryManager):
         self._telemetry_manager = manager
 
+    def link_state(self, state: State):
+        self._state = state
+
     def put(self, message: Message | ControlSignal):
         if self._draining.is_set():
             self.logger.debug('message received while draining, discarding...')
@@ -372,7 +377,7 @@ def join(self):
 
     def configure(self): ...
 
-    def update(self, message: Message[T_Input]): ...
+    def update(self, message: Message[T_Input], **kwargs): ...
 
     def set_on_config(self, prop: str, value: Any): ...
 
@@ -415,18 +420,22 @@ def _update(self):
                 batch.payload, ControlPayload
             ):
                 self._handle_control(batch)
+
                 if batch.payload.signal < 0:
                     break
+
                 continue
 
             self._last_data_source_evt_id = batch.id
+
             with self._pending_condition:
                 self._pending_updates += 1
             try:
-                self.update(batch)
+                self.update(batch, state=self._state)
             finally:
                 with self._pending_condition:
                     self._pending_updates -= 1
+
                     if self._pending_updates == 0:
                         self._pending_condition.notify_all()
 

juturna/components/_pipeline.py

@@ -16,6 +16,7 @@
 from juturna.payloads import ControlSignal, ControlPayload
 
 from juturna.components._dag import DAG
+from juturna.components._state import State
 from juturna.components._node_builder import _builder
 from juturna.components._telemetry_manager import TelemetryManager
 
@@ -46,6 +47,7 @@ def __init__(self, config: dict):
         self._nodes: dict[str, Node] = dict()
         self._links: list = list()
         self._dag: DAG = DAG()
+        self._node_state_store = dict()
 
         self._telemetry_manager: TelemetryManager | None = None
         self._telemetry = False
@@ -168,6 +170,9 @@ def warmup(self):
 
             self._nodes[node_name] = _node
             self._dag.add_node(node_name)
+            self._node_state_store[node_name] = State()
+
+            _node.link_state(self._node_state_store[node_name])
 
         for link in links:
             from_node = link['from']

juturna/components/_state.py

@@ -0,0 +1,6 @@
+from collections import UserDict
+
+
+class State(UserDict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

juturna/nodes/proc/_warp/warp.py

@@ -14,11 +14,12 @@
 from juturna.remotizer.utils import (
     message_to_proto,
     create_envelope,
-    deserialize_message,
+    deserialize_envelope,
 )
 
 from juturna.components import Message
 from juturna.components import Node
+from juturna.components import State
 
 
 class Warp[T_Input, T_Output](Node[T_Input, T_Output]):
@@ -82,7 +83,7 @@ def warmup(self):
 
         self.logger.info(f'warmup node: {self.name}')
 
-    def update(self, message: Message[T_Input]):
+    def update(self, message: Message[T_Input], **kwargs):
         """
         Send message via gRPC and wait for response
 
@@ -92,16 +93,22 @@ def update(self, message: Message[T_Input]):
         Parameters
         ----------
         message : Message[T_Input]
-            The message to send with input payload type
+            The message to send with input payload type.
+        kwargs : dict
+            State and more.
 
         """
+        state: State = kwargs['state']
+
         try:
             self.logger.info('converting message to protobuf...')
             message_proto = message_to_proto(message)
             self.logger.debug('creating envelope...')
             envelope = create_envelope(
                 message=message_proto,
                 creator=self.name,
+                pipe_id=self.pipe_id,
+                state=dict(state),
                 request_type=type(T_Input).__name__,
                 response_type=type(T_Output).__name__,
                 priority=0,
@@ -111,9 +118,9 @@ def update(self, message: Message[T_Input]):
             )
 
             envelope.configuration.update(self._remote_config)
-            self.logger.info(f'sending message (envelope_id={envelope.id})...')
 
-            self.logger.info(f'sending message id {message.id}...')
+            self.logger.info(f'sending message (envelope_id={envelope.id})...')
+            self.logger.info(f'sending message id {message.id}')
 
             response_envelope = self.stub.SendAndReceive(
                 envelope, timeout=self._timeout
@@ -124,13 +131,16 @@ def update(self, message: Message[T_Input]):
             )
 
             self.logger.info('converting response to Message...')
-            to_send: Message[T_Output] = deserialize_message(
-                response_envelope.message
-            )
+            envelope_dict = deserialize_envelope(response_envelope)
+
+            to_send: Message[T_Output] = envelope_dict['message']
 
             self.transmit(to_send)
             self.logger.info(f'transmit: {to_send.version}')
 
+            for k, v in envelope_dict['state'].items():
+                state[k] = v
+
         except grpc.RpcError as e:
             self.logger.error(f'gRPC error: {e.code()} - {e.details()}')
             raise

juturna/nodes/sink/_notifier_http/notifier_http.py

@@ -67,7 +67,7 @@ def set_on_config(self, prop: str, value: str):
 
             self._endpoint = value
 
-    def update(self, message: Message[ObjectPayload]):
+    def update(self, message: Message[ObjectPayload], **kwargs):
         """Receive a message, transmit a message"""
         to_send = Message[ObjectPayload](
             creator=message.creator,

juturna/nodes/sink/_notifier_udp/notifier_udp.py

@@ -83,7 +83,7 @@ def set_on_config(self, prop: str, value: typing.Any):
         elif prop == 'port':
             self._address[1] = value
 
-    def update(self, message: Message[ObjectPayload]):
+    def update(self, message: Message[ObjectPayload], **kwargs):
         """Receive a message, transmit a message"""
         chunks = self._prepare_chunks(message, message.version)
 

juturna/nodes/sink/_notifier_websocket/notifier_websocket.py

@@ -34,14 +34,13 @@ def __init__(self, endpoint: str, **kwargs):
 
         self._endpoint = endpoint
 
-        self._sent = 0
         self._t = None
 
     def warmup(self):
         """Warmup the node"""
         self.logger.info(f'[{self.name}] set to endpoint {self._endpoint}')
 
-    def update(self, message: Message[BasePayload]):
+    def update(self, message: Message[BasePayload], **kwargs):
         """Receive a message, transmit a message"""
         meta = dict(message.meta)
         to_send = Message[BasePayload](
@@ -62,8 +61,6 @@ def update(self, message: Message[BasePayload]):
 
         self._t.start()
 
-        self._sent += 1
-
     def _send_message(self, message: Message[BasePayload]):
         with connect(self._endpoint) as ws:
             try:

juturna/nodes/sink/_videostream_ffmpeg/videostream_ffmpeg.py

@@ -108,7 +108,7 @@ def stop(self):
         except Exception:
             ...
 
-    def update(self, message: Message[ImagePayload]):
+    def update(self, message: Message[ImagePayload], **kwargs):
         """Receive a message, transmit a message"""
         frame = message.payload.image
         frame_bytes = frame.tobytes()

juturna/nodes/source/_audio_file/audio_file.py

@@ -18,6 +18,7 @@
 
 from juturna.components import Node
 from juturna.components import Message
+
 from juturna.payloads import AudioPayload
 from juturna.payloads import ControlPayload
 from juturna.payloads import ControlSignal
@@ -53,7 +54,6 @@ def __init__(
 
         self._audio = None
         self._audio_chunks = None
-        self._transmitted = 0
 
     def warmup(self):  # noqa: D102
         resampler = av.audio.resampler.AudioResampler(
@@ -92,6 +92,7 @@ def _generate_chunks(self) -> Message[AudioPayload | ControlPayload]:
 
         if audio_chunk is None:
             self.logger.info('last chunk processed, stopping')
+
             return Message[ControlPayload](
                 creator=self.name,
                 payload=ControlPayload(signal=ControlSignal.STOP),
@@ -126,7 +127,7 @@ def _iter_audio_chunks(self):
             yield chunk, sample_offset
             sample_offset += wave_len
 
-    def update(self, message: Message[AudioPayload | ControlPayload]):  # noqa: D102
+    def update(self, message: Message[AudioPayload | ControlPayload], **kwargs):  # noqa: D102
         message.meta['session_id'] = self.pipe_id
 
         self.transmit(message)

juturna/nodes/source/_audio_rtp/audio_rtp.py

@@ -15,9 +15,10 @@
 
 import numpy as np
 
-from juturna.components import _resource_broker as rb
 from juturna.components import Message
 from juturna.components import Node
+from juturna.components import _resource_broker as rb
+
 from juturna.payloads import BytesPayload, AudioPayload
 from juturna.names import ComponentStatus
 
@@ -79,7 +80,6 @@ def __init__(
             * self._block_size
             * self._audio_rate
         )
-        self._abs_recv = 0
 
         # infer incoming channel by encoding_clock_chan
         self._in_channels = AudioRTP._parse_audio_channels(encoding_clock_chan)
@@ -212,7 +212,7 @@ def configuration(self) -> dict:
 
         return base_config
 
-    def update(self, message: Message[BytesPayload]):
+    def update(self, message: Message[BytesPayload], **kwargs):
         """Read a message, return a message"""
         if not self._subprocess_running:
             return
@@ -221,24 +221,27 @@ def update(self, message: Message[BytesPayload]):
             message.payload.cnt, self._in_channels
         )
 
+        state: dict = kwargs.get('state')
+        _abs_recv = state.get('abs_recv', 0)
+
         to_send = Message[AudioPayload](
             creator=self.name,
-            version=self._abs_recv,
+            version=_abs_recv,
             payload=AudioPayload(
                 audio=waveform,
                 sampling_rate=self._audio_rate,
                 channels=self._channels,
-                start=self._block_size * self._abs_recv,
-                end=self._block_size * self._abs_recv + self._block_size,
+                start=self._block_size * _abs_recv,
+                end=self._block_size * _abs_recv + self._block_size,
             ),
         )
 
-        to_send.meta['source_recv'] = self._abs_recv
+        to_send.meta['source_recv'] = _abs_recv
 
         self.transmit(to_send)
         self.logger.info(f'transmitting message {to_send.version}')
 
-        self._abs_recv += 1
+        state['abs_recv'] = _abs_recv + 1
 
     def clear_source(self):
         """Clear any source functions defined on the node"""

juturna/nodes/source/_audio_rtp_av/audio_rtp_av.py

@@ -140,7 +140,7 @@ def stop(self):
         self._t.join()
         super().stop()
 
-    def update(self, message: Message[AudioPayload]):
+    def update(self, message: Message[AudioPayload], **kwargs):
         """Receive data from upstream, transmit data downstream"""
         self.logger.debug('update method not implemented for source node')