From 99ea06bc91cf780740820f464f0d229fc436f089 Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Sun, 26 Jan 2025 17:26:39 +0000 Subject: [PATCH] fix: terminate on SIGINT (Ctrl+C) --- src/__main__.py | 2 +- src/{cli.py => cli/__init__.py} | 14 +++++----- src/collection/ingest.py | 46 +++++++++++++++++++++------------ 3 files changed, 38 insertions(+), 24 deletions(-) rename src/{cli.py => cli/__init__.py} (80%) diff --git a/src/__main__.py b/src/__main__.py index 71ea413..e8eab89 100644 --- a/src/__main__.py +++ b/src/__main__.py @@ -5,7 +5,7 @@ from . import cli logging.basicConfig( level=logging.INFO, - format="%(asctime)s %(name)-40s %(levelname)-8s %(message)s", + format="%(asctime)s %(name)-50s %(levelname)-8s %(message)s", ) if __name__ == "__main__": diff --git a/src/cli.py b/src/cli/__init__.py similarity index 80% rename from src/cli.py rename to src/cli/__init__.py index c0063d1..3350f18 100644 --- a/src/cli.py +++ b/src/cli/__init__.py @@ -5,10 +5,11 @@ import numpy as np import numpy.typing as npt import typer -from . import config, visualise -from .collection import ingest -from .processing.aoa import AoA -from .processing.preprocess import Preprocessor +from .. import config, visualise +from ..collection import ingest +from ..processing.aoa import AoA +from ..processing.preprocess import Preprocessor +from . import file app = typer.Typer() logger = logging.getLogger(__name__) @@ -23,7 +24,7 @@ def antennas() -> None: print the antenna identifier. When you are done, press Ctrl+C to stop the script and get the final order. """ - from .utils import antenna_order + from ..utils import antenna_order antenna_order.main() @@ -32,7 +33,8 @@ def antennas() -> None: def heatmap() -> None: preprocessor = Preprocessor() aoa = AoA() - webapp_queue: "mp.Queue[AoA]" = mp.Queue(config.SAMPLE_RATE) + manager = mp.Manager() + webapp_queue: "mp.Queue[AoA]" = manager.Queue(config.SAMPLE_RATE) # Start webapp in background process webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) webapp.start() diff --git a/src/collection/ingest.py b/src/collection/ingest.py index cad463f..588aa72 100644 --- a/src/collection/ingest.py +++ b/src/collection/ingest.py @@ -6,7 +6,7 @@ import subprocess import threading import time from datetime import datetime -from typing import Callable, NamedTuple, NoReturn +from typing import Callable, NamedTuple import numpy as np import numpy.typing as npt @@ -25,6 +25,7 @@ class FeitHost: self.logger = logging.getLogger( f"{__name__}.{self.__class__.__name__}-{self.host[0]}" ) + self.active = True self.checker = threading.Thread(target=self.check_continuous) self.checker.start() @@ -44,7 +45,7 @@ class FeitHost: self.server.send(self.command.encode()) self.logger.info(f"Connected to {self.host}") - def check_continuous(self) -> NoReturn: + def check_continuous(self) -> None: """ Repeatedly check if the FeitCSI service is running @@ -53,7 +54,7 @@ class FeitHost: and logs continuously if the service is running. """ last_status = False - while True: + while self.active: if not self.check_connection(): self.logger.error(f"FeitCSI is not running on {self.host[0]}") last_status = False @@ -62,6 +63,7 @@ class FeitHost: self.connect() last_status = True time.sleep(1) + self.logger.info("Stopping host checker") class FeitTransmitter(FeitHost): @@ -86,10 +88,10 @@ class FeitReceiver(FeitHost): ) super().__init__(host, command) - def listen(self, queue: "mp.Queue[CSI]") -> NoReturn: + def listen(self, queue: "mp.Queue[CSI]") -> None: prev_time = datetime.now() self.logger.info("Listening for CSI data") - while True: + while self.active: while not hasattr(self, "server"): time.sleep(0.1) # This is the max size of a UDP packet. The size of the actual CSI @@ -105,6 +107,7 @@ class FeitReceiver(FeitHost): queue.put(csidata) except struct.error: self.logger.error("Failed to parse CSI data") + self.logger.info("Stopping CSI receiver") class CSIProcessor: @@ -117,6 +120,7 @@ class CSIProcessor: self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") self.last_processed = datetime.now() self.connections = receiver_connections + self.active = True def add_data(self, host: Host, data: CSI) -> None: if ( @@ -164,16 +168,19 @@ class CSIProcessor: self, callback: CSICallback | None = None, pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, - ) -> NoReturn: - while True: - for ip, queue in self.connections.items(): - while not queue.empty(): - self.add_data(ip, queue.get()) - if self.is_ready(): - self.process_data(callback, pre_merge_callback=pre_merge_callback) - else: - self.logger.debug("Not all data is ready") - time.sleep(0.0005) + ) -> None: + try: + while self.active: + for ip, queue in self.connections.items(): + while not queue.empty(): + self.add_data(ip, queue.get()) + if self.is_ready(): + self.process_data(callback, pre_merge_callback=pre_merge_callback) + else: + self.logger.debug("Not all data is ready") + time.sleep(0.0005) + except KeyboardInterrupt: + self.logger.info("Exiting CSI processing") class Receiver(NamedTuple): @@ -191,7 +198,7 @@ def start_processing( for ip in config.RECEIVE_HOSTS ] - FeitTransmitter() + transmitter = FeitTransmitter() # Start injecting CSI frames processor = CSIProcessor({r.ip: r.queue for r in receivers}) @@ -208,6 +215,11 @@ def start_processing( ) processing_thread.start() try: - processing_thread.join() + while True: + time.sleep(100) except KeyboardInterrupt: + for r in receivers: + r.receiver.active = False + transmitter.active = False + processor.active = False return