New worker event assembly

dranspose/debug_worker.py

@@ -5,7 +5,6 @@
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator, Any
 
-import zmq
 from fastapi import FastAPI
 from starlette.responses import Response
 
@@ -29,13 +28,14 @@ async def work(self) -> None:
         while True:
             self.dequeue_task = None
             self.dequeue_task = asyncio.create_task(self.assignment_queue.get())
-            ingesterset = await self.dequeue_task
-            done = await self.poll_internals(ingesterset)
-            if set(done) != {zmq.POLLIN}:
-                self._logger.warning("not all sockets are pollIN %s", done)
-                continue
+            evn, streamset = await self.dequeue_task
 
-            event = await self.build_event(ingesterset)
+            self.new_data.clear()
+            while set(self.stream_queues.get(evn, {}).keys()) != streamset:
+                await self.new_data.wait()
+                self.new_data.clear()
+
+            event = EventData(event_number=evn, streams=self.stream_queues[evn])
             self._logger.debug("adding event %s to buffer", event)
             self.buffer.append(event)
 

dranspose/ingesters/__init__.py

@@ -10,6 +10,10 @@
     ZmqPullSingleIngester,
     ZmqPullSingleSettings,
 )
+from dranspose.ingesters.stins_parallel import (  # noqa: F401
+    StinsParallelIngester,
+    StinsParallelSettings,
+)
 from dranspose.ingesters.zmqpull_eiger_legacy import (  # noqa: F401
     ZmqPullEigerLegacyIngester,
     ZmqPullEigerLegacySettings,

dranspose/ingesters/stins_parallel.py

@@ -0,0 +1,96 @@
+import asyncio
+from typing import AsyncGenerator, Optional
+
+import zmq
+
+
+from dranspose.data.stream1 import Stream1Packet, Stream1End
+from dranspose.event import StreamData, InternalWorkerMessage
+from dranspose.ingester import Ingester, IngesterSettings
+from dranspose.protocol import (
+    StreamName,
+    ZmqUrl,
+    WorkAssignment,
+    WorkerName,
+)
+
+
+class StinsParallelSettings(IngesterSettings):
+    upstream_url: ZmqUrl
+
+
+class StinsParallelIngester(Ingester):
+    """
+    A simple ingester class to comsume a stream from the streaming-receiver repub port
+    """
+
+    def __init__(self, settings: Optional[StinsParallelSettings] = None) -> None:
+        if settings is not None:
+            self._streaming_single_settings = settings
+        else:
+            self._streaming_single_settings = StinsParallelSettings()
+
+        super().__init__(settings=self._streaming_single_settings)
+        self.in_socket: Optional[zmq._future._AsyncSocket] = None
+
+    async def work(self) -> None:
+        self._logger.info("started stins ingester work task")
+
+        if len(self.active_streams) == 0:
+            self._logger.warning("this ingester has no active streams, stopping worker")
+            return
+        sourcegen = self.run_source_part(self.active_streams[0])
+        try:
+            while True:
+                nextiwm: InternalWorkerMessage = await anext(sourcegen)
+
+                work_assignment: WorkAssignment = await self.assignment_queue.get()
+                while work_assignment.event_number < nextiwm.event_number:
+                    work_assignment = await self.assignment_queue.get()
+
+                workermessages: dict[WorkerName, InternalWorkerMessage] = {}
+                for stream, workers in work_assignment.assignments.items():
+                    for worker in workers:
+                        if worker not in workermessages:
+                            workermessages[worker] = nextiwm
+                self._logger.debug("workermessages %s", workermessages)
+                await self._send_workermessages(workermessages)
+                self.state.processed_events += 1
+        except asyncio.exceptions.CancelledError:
+            self._logger.info("stopping worker")
+            for stream in self.active_streams:
+                await self.stop_source(stream)
+
+    async def run_source_part(
+        self, stream: StreamName
+    ) -> AsyncGenerator[InternalWorkerMessage, None]:
+        self.in_socket = self.ctx.socket(zmq.PULL)
+        self.in_socket.connect(str(self._streaming_single_settings.upstream_url))
+        self._logger.info(
+            "pulling from %s", self._streaming_single_settings.upstream_url
+        )
+
+        while True:
+            parts = await self.in_socket.recv_multipart(copy=False)
+            try:
+                packet = Stream1Packet.validate_json(parts[0].bytes)
+            except Exception as e:
+                self._logger.error("packet not valid %s", e.__repr__())
+                continue
+            self._logger.debug("msg number %d", packet.msg_number)
+            yield InternalWorkerMessage(
+                event_number=packet.msg_number,
+                streams={stream: StreamData(typ="STINS", frames=parts)},
+            )
+
+            if isinstance(packet, Stream1End):
+                break
+        while True:
+            self._logger.debug("discarding messages until next run")
+            await self.in_socket.recv_multipart(copy=False)
+
+    async def stop_source(self, stream: StreamName) -> None:
+        if self.in_socket:
+            self._logger.info("closing socket without linger")
+            self.in_socket.close(linger=0)
+            self.in_socket = None

dranspose/ingesters/stream1_parallel.py

@@ -0,0 +1,55 @@
+from typing import Optional, AsyncGenerator
+import zmq
+
+from dranspose.data.eiger_legacy import (
+    EigerLegacyPacket,
+    EigerLegacyEnd,
+    EigerLegacyImage,
+    EigerLegacyHeader,
+)
+from dranspose.event import InternalWorkerMessage, StreamData
+from dranspose.ingesters.stins_parallel import StinsParallelIngester
+from dranspose.ingesters.zmqpull_eiger_legacy import ZmqPullEigerLegacySettings
+from dranspose.protocol import StreamName, EventNumber
+
+
+class Stream1ParallelIngester(StinsParallelIngester):
+    def __init__(self, settings: Optional[ZmqPullEigerLegacySettings] = None) -> None:
+        if settings is not None:
+            self._streaming_settings = settings
+        else:
+            self._streaming_settings = ZmqPullEigerLegacySettings()
+
+        super().__init__(settings=self._streaming_settings)
+        self.in_socket: Optional[zmq._future._AsyncSocket] = None
+
+    async def run_source_part(
+        self, stream: StreamName
+    ) -> AsyncGenerator[InternalWorkerMessage, None]:
+        self.in_socket = self.ctx.socket(zmq.PULL)
+        self.in_socket.connect(str(self._streaming_settings.upstream_url))
+        self._logger.info("pulling from %s", self._streaming_settings.upstream_url)
+
+        while True:
+            parts = await self.in_socket.recv_multipart(copy=False)
+            try:
+                packet = EigerLegacyPacket.validate_json(parts[0].bytes)
+            except Exception as e:
+                self._logger.error("packet not valid %s", e.__repr__())
+                continue
+            msg_number = None
+            if isinstance(packet, EigerLegacyImage):
+                msg_number = EventNumber(packet.frame + 1)
+            elif isinstance(packet, EigerLegacyHeader):
+                msg_number = EventNumber(0)
+            elif isinstance(packet, EigerLegacyEnd):
+                break
+            self._logger.debug("msg number %d", msg_number)
+            yield InternalWorkerMessage(
+                event_number=msg_number,
+                streams={stream: StreamData(typ="EIGER_LEGACY", frames=parts)},
+            )
+
+        while True:
+            self._logger.debug("discarding messages until next run")
+            await self.in_socket.recv_multipart(copy=False)