From cc27972d8190ee627ef7e4e83ceb21d567bf7ebe Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Sat, 26 Apr 2025 11:28:17 +0100 Subject: [PATCH] refactor into separate CSIApplication framework --- where_fi/__main__.py | 3 +- where_fi/application.py | 187 ++++++++++++++++++++++++++ where_fi/cli/__init__.py | 60 +++++---- where_fi/collection/__init__.py | 5 + where_fi/collection/ingest.py | 84 +++++++----- where_fi/config/models.py | 3 +- where_fi/processing/preprocess.py | 65 ++++++--- where_fi/utils/antenna_order.py | 16 ++- where_fi/visualise/server/__init__.py | 69 ++++++---- 9 files changed, 377 insertions(+), 115 deletions(-) create mode 100644 where_fi/application.py diff --git a/where_fi/__main__.py b/where_fi/__main__.py index e8eab89..e2412a9 100644 --- a/where_fi/__main__.py +++ b/where_fi/__main__.py @@ -2,11 +2,10 @@ import logging from . import cli - logging.basicConfig( level=logging.INFO, format="%(asctime)s %(name)-50s %(levelname)-8s %(message)s", ) if __name__ == "__main__": - cli.app() + cli.cli() diff --git a/where_fi/application.py b/where_fi/application.py new file mode 100644 index 0000000..16d9847 --- /dev/null +++ b/where_fi/application.py @@ -0,0 +1,187 @@ +import logging +import multiprocessing as mp +import threading +import time +from typing import Any, Callable, NamedTuple + +import numpy.typing as npt + +from .collection import CSIMatrix, ingest +from .config import config +from .visualise import server as visualise + + +class Receiver(NamedTuple): + ip: ingest.Host + receiver: ingest.FeitReceiver + thread: threading.Thread + + +class Scheduler: + """ + A simple scheduler that runs a function at a given rate. + """ + + def __init__(self, rate: float, func: Callable[[], None]) -> None: + self.rate = rate + self.func = func + self.active = True + + def run(self) -> None: + """ + Run the scheduler in a loop, calling the function at the given rate. + """ + while self.active: + self.func() + time.sleep(1 / self.rate) + + def start(self) -> None: + """ + Start the scheduler in a separate thread. + """ + self.thread = threading.Thread(target=self.run) + self.thread.start() + + def stop(self) -> None: + """ + Stop the scheduler. + """ + self.active = False + + +class CSIApplication: + """ + A high-level application that manages CSI data ingestion and processing. + + This allows a simple interface to be used for the main CLI logic, without having to + deal with the threads and queues directly. + + It can be used using decortors to register callbacks for different stages of the + processing pipeline: + + - `on_pre_merge`: Called with the raw CSI data (including headers) from each + receiver, before it is merged into a single matrix. + - `on_sample`: Called once per sample with the merged CSI data. + - `on_process`: Called at the processing sample rate, without any data (this is + mostly used as a scheduler). + """ + + def create_receivers(self) -> list[Receiver]: + """ + Create a thread for each receiver, so that data can be received in parallel. + + The receivers are created with a queue that is used to store the received data, + to be processed later. The receivers are created based on the configuration + file, which contains the IP addresses of the receivers. + """ + feit_receivers = [ + ingest.FeitReceiver(ip, mp.Queue(config.collection_sample_rate)) + for ip in config.receive_hosts + ] + return [ + Receiver( + ip, + receiver, + threading.Thread(target=receiver.listen), + ) + for receiver, ip in zip(feit_receivers, config.receive_hosts, strict=True) + ] + + def __init__(self) -> None: + self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") + + self.receivers = self.create_receivers() + self.transmitter = ingest.FeitTransmitter() + + self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue( + config.processing_sample_rate + ) + self.webapp = visualise.Webapp() + self.webapp_thread = threading.Thread( + target=self.webapp.start, args=(self.webapp_queue,) + ) + + self.csi_callback = None + self.pre_merge_callback = None + self.processing_callback = None + + self.buffer = ingest.CSIAntennaJoin( + {r.ip: r.receiver.queue for r in self.receivers} + ) + + def visualise_data( + self, data: npt.NDArray[Any], dtype: visualise.figures.Figure + ) -> None: + """ + Update a visualisation with the given data. The available visualisations are as + per visualise.server.all_figures. + """ + if not self.webapp_queue.full(): + self.webapp_queue.put(visualise.VisualiserData(data, dtype)) + + def on_sample( + self, func: Callable[[CSIMatrix], None] + ) -> Callable[[CSIMatrix], None]: + """ + Decorator to register a callback for the sample preprocessing. + """ + + def decorator(data: CSIMatrix) -> None: + self.visualise_data(data, visualise.figures.Figure.RAW_CSI) + func(data) + + self.csi_callback = decorator + return decorator + + def on_pre_merge( + self, + ) -> Callable[[Callable[[dict[ingest.Host, ingest.CSI]], None]], None]: + def decorator(func: Callable[[dict[ingest.Host, ingest.CSI]], None]) -> None: + self.pre_merge_callback = func + + return decorator + + def on_process(self, func: Callable[[], None]) -> Callable[[], None]: + """ + Decorator to register a callback for the sample processing. + """ + self.processing_callback = func + return func + + def start(self) -> None: + for receiver in self.receivers: + receiver.thread.start() + self.buffer_thread = threading.Thread( + target=self.buffer.process_forever, + args=(self.csi_callback, self.pre_merge_callback), + ) + self.buffer_thread.start() + self.webapp_thread.start() + + if self.processing_callback is not None: + self.scheduler = Scheduler( + config.processing_sample_rate, self.processing_callback + ) + self.scheduler.start() + try: + self.webapp_thread.join() + except KeyboardInterrupt: + self.stop() + + def stop(self) -> None: + self.logger.info("Stopping application") + for receiver in self.receivers: + receiver.receiver.active = False + receiver.thread.join() + + self.transmitter.active = False + + if hasattr(self, "scheduler"): + self.scheduler.stop() + self.scheduler.thread.join() + self.buffer.active = False + self.buffer_thread.join() + + self.webapp.active = False + self.webapp_thread.join() + self.logger.info("Application stopped") diff --git a/where_fi/cli/__init__.py b/where_fi/cli/__init__.py index 98fe07f..efebbe6 100644 --- a/where_fi/cli/__init__.py +++ b/where_fi/cli/__init__.py @@ -1,25 +1,28 @@ import logging import multiprocessing as mp -from typing import Any, cast from queue import Queue +import threading +import time +from typing import Any, cast import numpy as np import numpy.typing as npt import torch import typer +from ..application import CSIApplication from ..config import config from ..processing.aoa import AoA from ..processing.preprocess import Preprocessor from ..visualise import server as visualise from . import file, globals -app = typer.Typer(callback=globals.main) +cli = typer.Typer(callback=globals.main) logger = logging.getLogger(__name__) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -@app.command() +@cli.command() def antennas() -> None: """Utility to help determine the order in which antennas are plugged in @@ -31,18 +34,20 @@ def antennas() -> None: from ..utils import antenna_order if not globals.is_live: - raise ValueError("This command only works with live data") + raise NotImplementedError("This command only works with live data") antenna_order.main() -@app.command() +@cli.command() def heatmap() -> None: preprocessor = Preprocessor() aoa = AoA() # Start webapp in background process - webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(config.sample_rate) + webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue( + config.processing_sample_rate + ) webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) webapp.start() @@ -64,7 +69,7 @@ def heatmap() -> None: logger.info("Finished processing CSI data") -@app.command() +@cli.command() def phase_analysis( subcarrier: int = 0, rx_antenna: int = 0, tx_antenna: int = 0 ) -> None: @@ -76,38 +81,37 @@ def phase_analysis( selected subcarrier and antenna. """ + app = CSIApplication() preprocessor = Preprocessor() - # Start webapp in background process - webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(config.sample_rate) - webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) - webapp.start() + subcarrier_phase: Queue[float] = Queue(config.collection_sample_rate) - subcarrier_phase: Queue[float] = Queue(config.sample_rate) - - def visualise_data(data: npt.NDArray[Any], dtype: visualise.figures.Figure) -> None: - if not webapp_queue.full(): - webapp_queue.put(visualise.VisualiserData(data, dtype)) - - def callback(antenna_data: npt.NDArray[np.complex64]) -> None: - logger.info(f"Got final CSI data with shape {antenna_data.shape}") - visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI) - - processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data) - visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI) + @app.on_sample + def _(antenna_data: npt.NDArray[np.complex64]) -> None: + processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data) + app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI) phase = cast(float, np.angle(processed[subcarrier, rx_antenna, tx_antenna])) - if subcarrier_phase.full(): subcarrier_phase.get() subcarrier_phase.put(phase) - visualise_data( - np.array(subcarrier_phase), visualise.figures.Figure.PHASE_ANALYSIS + @app.on_process + def _() -> None: + logger.info(f"Updating phase visualisation") + app.visualise_data( + np.array(subcarrier_phase.queue), visualise.figures.Figure.PHASE_ANALYSIS ) - globals.csi_producer(csi_callback=callback) + # globals.csi_producer(csi_callback=callback) + app.start() logger.info("Finished processing CSI data") + while threading.active_count() > 1: + names = [ + t.name for t in threading.enumerate() if t != threading.current_thread() + ] + logger.info("Waiting for threads to close: " + ",".join(names)) + time.sleep(2) -app.add_typer(file.app, name="file", help="Commands for working with CSI files") +cli.add_typer(file.app, name="file", help="Commands for working with CSI files") diff --git a/where_fi/collection/__init__.py b/where_fi/collection/__init__.py index 93a1b8c..7dc383c 100644 --- a/where_fi/collection/__init__.py +++ b/where_fi/collection/__init__.py @@ -1,3 +1,6 @@ +import numpy as np +import numpy.typing as npt + from . import file, ingest from .protocols import CSICallback, CSIProducer @@ -6,4 +9,6 @@ def noop(csi_callback: CSICallback | None = None) -> None: del csi_callback +CSIMatrix = npt.NDArray[np.complex64] + __all__ = ["file", "ingest", "noop", "CSIProducer", "CSICallback"] diff --git a/where_fi/collection/ingest.py b/where_fi/collection/ingest.py index 1d1ab0f..d6c5400 100644 --- a/where_fi/collection/ingest.py +++ b/where_fi/collection/ingest.py @@ -29,13 +29,17 @@ class FeitHost: self.checker.start() def check_connection(self) -> bool: - feitcsi_status = subprocess.run( - f"ssh root@{self.host[0]} pgrep feitcsi", - check=False, - stdout=subprocess.DEVNULL, - shell=True, - ) - return feitcsi_status.returncode == 0 + try: + feitcsi_status = subprocess.run( + f"ssh root@{self.host[0]} pgrep feitcsi", + check=False, + stdout=subprocess.DEVNULL, + shell=True, + timeout=1, + ) + return feitcsi_status.returncode == 0 + except subprocess.TimeoutExpired: + return False def connect(self) -> None: self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) @@ -54,6 +58,7 @@ class FeitHost: """ last_status = False while self.active: + self.logger.debug(f"Checking connection to {self.host[0]}") if not self.check_connection(): self.logger.error(f"FeitCSI is not running on {self.host[0]}") last_status = False @@ -72,22 +77,23 @@ class FeitTransmitter(FeitHost): f"--channel-width {config.channel_width} " f"--format {config.frame_format} " f"--mode inject -s 1 --verbose " - f"--inject-delay {1_000_000 // config.sample_rate}" + f"--inject-delay {1_000_000 // config.collection_sample_rate}" ) super().__init__(config.transmit_host, command) class FeitReceiver(FeitHost): - def __init__(self, host: Host) -> None: + def __init__(self, host: Host, queue: "mp.Queue[CSI]") -> None: command = ( f"feitcsi --frequency {config.central_freq} " f"--channel-width {config.channel_width} " f"--format {config.frame_format} " f"--mode measure" ) + self.queue = queue super().__init__(host, command) - def listen(self, queue: "mp.Queue[CSI]") -> None: + def listen(self) -> None: prev_time = datetime.now() self.logger.info("Listening for CSI data") while self.active: @@ -103,13 +109,24 @@ class FeitReceiver(FeitHost): f"Received CSI data after {datetime.now() - prev_time}" ) prev_time = datetime.now() - queue.put(csidata) + self.queue.put(csidata) except struct.error: self.logger.error("Failed to parse CSI data") self.logger.info("Stopping CSI receiver") -class CSIProcessor: +class CSIAntennaJoin: + """ + This class is used to buffer a set of CSI data from multiple receivers, and once a + set of readings is ready for each receiver, it will be merged into a single sample + using the antenna order in the configuration. + + When the data is ready, the following callbacks are called: + - `pre_merge_callback`: Called with the data before merging. This is useful for + per-antenna analysis, e.g. finding the correct antenna order + - `sample_callback`: Called with the merged data + """ + def __init__( self, receiver_connections: dict[Host, "mp.Queue[CSI]"], @@ -144,7 +161,7 @@ class CSIProcessor: def process_data( self, - callback: CSICallback | None = None, + sample_callback: CSICallback | None = None, pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, ) -> None: self.last_processed = datetime.now() @@ -160,26 +177,23 @@ class CSIProcessor: # We have data from all servers all_data = np.concat(antenna_data, axis=1) - if callback: - callback(all_data) + if sample_callback: + sample_callback(all_data) def process_forever( self, callback: CSICallback | None = None, pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, ) -> 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") + 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) class Receiver(NamedTuple): @@ -193,26 +207,26 @@ def start_processing( pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, ) -> None: receivers = [ - Receiver(ip, FeitReceiver(ip), mp.Queue(config.sample_rate)) + Receiver(ip, FeitReceiver(ip), mp.Queue(config.collection_sample_rate)) for ip in config.receive_hosts ] transmitter = FeitTransmitter() # Start injecting CSI frames - processor = CSIProcessor({r.ip: r.queue for r in receivers}) + buffer = CSIAntennaJoin({r.ip: r.queue for r in receivers}) - receiver_processes = [ + receiver_threads = [ threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers ] - for proc in receiver_processes: + for proc in receiver_threads: proc.start() - processing_thread = mp.Process( - target=processor.process_forever, args=(csi_callback, pre_merge_callback) + buffering_thread = threading.Thread( + target=buffer.process_forever, args=(csi_callback, pre_merge_callback) ) - processing_thread.start() + buffering_thread.start() try: while True: time.sleep(100) @@ -220,5 +234,5 @@ def start_processing( for r in receivers: r.receiver.active = False transmitter.active = False - processor.active = False + buffer.active = False return diff --git a/where_fi/config/models.py b/where_fi/config/models.py index f3f11d2..845b002 100644 --- a/where_fi/config/models.py +++ b/where_fi/config/models.py @@ -43,7 +43,8 @@ class Config(BaseModel): antennas: Antennas - sample_rate: int + collection_sample_rate: int + processing_sample_rate: int central_freq: int channel_width: Literal[20, 40, 80, 160] frame_format: Literal["NOHT", "HT", "VHT", "HESU"] diff --git a/where_fi/processing/preprocess.py b/where_fi/processing/preprocess.py index 2aac7d7..b7e2651 100644 --- a/where_fi/processing/preprocess.py +++ b/where_fi/processing/preprocess.py @@ -6,6 +6,8 @@ import numpy as np import numpy.typing as npt from scipy.signal import butter, correlate, sosfilt, sosfilt_zi +from where_fi.collection import CSIMatrix + from ..config import config from ..visualise import server as visualise @@ -17,33 +19,56 @@ np.seterr(invalid="ignore") class Preprocessor: def __init__(self) -> None: - self.prev_entries: Queue[npt.NDArray[np.complex64]] = Queue(maxsize=100) + self.prev_entries: Queue[CSIMatrix] = Queue(maxsize=100) self.short_term_avg = np.zeros((1,), dtype=np.complex64) self.long_term_avg = np.zeros((1,), dtype=np.complex64) - self.filter = butter( - 5, - config.preprocessing.bandpass.bounds, - fs=config.sample_rate, - btype="band", - output="sos", - ) + # self.filter = butter( + # 5, + # config.preprocessing.bandpass.bounds, + # fs=config.sample_rate, + # btype="band", + # output="sos", + # ) + + def remove_sto(self, csi: CSIMatrix) -> CSIMatrix: + """ + Remove sampling time offsets caused by: + - Sampling frequency offset + - Packet detection delay + + According to [1]: + > Conjugate multiplication and division are the only two methods to + > eliminate the SFO and PDD. + + No citation or explanation is provided, so not sure why/whether it works. + Something similar is also done in [2] without explanation. + + [1] - https://tns.thss.tsinghua.edu.cn/wst/docs/sanitization + [2] - https://doi.org/10.1109/ICC51166.2024.10623053 + """ + # Conjugate multiplication + csi_remove_sto = np.zeros_like(csi) + for antenna in range(csi.shape[1]): + antenna_nxt = (antenna + 1) % csi.shape[1] + csi_remove_sto[:, antenna, :] = np.multiply( + csi[:, antenna, :], csi[:, antenna_nxt, :].conj() + ) + return csi_remove_sto def preprocess( self, - h: npt.NDArray[np.complex64], + h: CSIMatrix, visualiser: None - | Callable[[npt.NDArray[Any], visualise.DataType], None] = None, - ) -> npt.NDArray[np.complex64]: + | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None, + ) -> CSIMatrix: # CSI data is not available for pilot subcarriers. - h_hat: npt.NDArray[np.complex64] = np.where( + h_hat: CSIMatrix = np.where( np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)), correlate(h, [[[1 / 2]], [[0]], [[1 / 2]]], mode="same"), h, ) - # Skip subcarrierss per config - h_hat = h_hat[:: config.preprocessing.subcarrier_step, :, :] - # logger.info(f"CSI shape: {h_hat.shape}") + logger.debug(f"CSI shape: {h_hat.shape}") # h_hat = np.multiply(h_hat, h_hat.conj() / abs(h_hat.conj())) h_hat = np.nan_to_num(h_hat) @@ -52,12 +77,11 @@ class Preprocessor: # h_hat = correlate(h_hat, [[[1 / 4]], [[1 / 2]], [[1 / 4]]], mode="valid") # Unwrap phase and remove linear fit - print(h_hat.shape) unwrapped = np.unwrap(np.angle(h_hat[:, :, 0]), axis=0).reshape( h_hat.shape[0], h_hat.shape[1], 1 ) if visualiser: - visualiser(unwrapped, visualise.DataType.UNWRAPPED_PHASE) + visualiser(unwrapped, visualise.figures.Figure.UNWRAPPED_PHASE) for antenna in range(h_hat.shape[1]): tau, rho = np.linalg.lstsq( np.vstack([np.arange(h_hat.shape[0]), np.ones(h_hat.shape[0])]).T, @@ -71,7 +95,11 @@ class Preprocessor: ) ) - return h_hat + # Skip subcarrierss per config + h_hat = h_hat[:: config.preprocessing.subcarrier_step, :, :] + + # h_hat = self.remove_sto(h_hat) + # Assume that all csi matrices will have the same shape if self.long_term_avg.shape != h_hat.shape: self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex64) @@ -82,6 +110,7 @@ class Preprocessor: ) # Remove long term average, to remove static paths + return h_hat h_hat -= self.long_term_avg # Apply bandpass filter to remove low and high frequency noise diff --git a/where_fi/utils/antenna_order.py b/where_fi/utils/antenna_order.py index 281b96f..254b88a 100644 --- a/where_fi/utils/antenna_order.py +++ b/where_fi/utils/antenna_order.py @@ -11,11 +11,19 @@ AntennaIdentifier = tuple[ingest.Host, int] order: list[AntennaIdentifier] = [] prev_unplugged: set[AntennaIdentifier] = set() long_antenna_average: dict[AntennaIdentifier, deque[int]] = { - (host, 0): deque(maxlen=15 * config.sample_rate) for host in config.receive_hosts -} | {(host, 1): deque(maxlen=15 * config.sample_rate) for host in config.receive_hosts} + (host, 0): deque(maxlen=15 * config.collection_sample_rate) + for host in config.receive_hosts +} | { + (host, 1): deque(maxlen=15 * config.collection_sample_rate) + for host in config.receive_hosts +} short_antenna_average: dict[AntennaIdentifier, deque[int]] = { - (host, 0): deque(maxlen=2 * config.sample_rate) for host in config.receive_hosts -} | {(host, 1): deque(maxlen=2 * config.sample_rate) for host in config.receive_hosts} + (host, 0): deque(maxlen=2 * config.collection_sample_rate) + for host in config.receive_hosts +} | { + (host, 1): deque(maxlen=2 * config.collection_sample_rate) + for host in config.receive_hosts +} RSSI_THRESHOLD = 10 diff --git a/where_fi/visualise/server/__init__.py b/where_fi/visualise/server/__init__.py index ff4adc5..ddac0e8 100644 --- a/where_fi/visualise/server/__init__.py +++ b/where_fi/visualise/server/__init__.py @@ -2,6 +2,7 @@ import logging import multiprocessing as mp import queue import threading +import time from concurrent import futures from dataclasses import dataclass from typing import Any, Generator @@ -13,11 +14,6 @@ import numpy.typing as npt from . import figures from .generated import figure_pb2, figure_pb2_grpc -logger = logging.getLogger(__name__) - -clients: dict[str, list[queue.Queue[figure_pb2.FigureData]]] = {} -clients_lock = threading.Lock() - @dataclass class VisualiserData: @@ -57,29 +53,48 @@ class FigureServer(figure_pb2_grpc.FigureServiceServicer): clients[request.uuid].remove(q) -def add_data(dtype: figures.Figure, new_data: npt.NDArray[np.complex128]) -> None: - updates = figures.all_figures[dtype].update(new_data) - for fig_id, update in updates.items(): - for client in clients.get(fig_id, []): - client.put(update) +class Webapp: + def __init__(self) -> None: + self.logger = logging.getLogger(__name__) + self.active = True + self.clients: dict[str, list[queue.Queue[figure_pb2.FigureData]]] = {} + self.clients_lock = threading.Lock() -def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None: - while True: - data = data_queue.get() - add_data(data.dtype, data.data) + def add_data( + self, dtype: figures.Figure, new_data: npt.NDArray[np.complex128] + ) -> None: + updates = figures.all_figures[dtype].update(new_data) + for fig_id, update in updates.items(): + for client in self.clients.get(fig_id, []): + client.put(update) + def listen_for_data(self, data_queue: "mp.Queue[VisualiserData]") -> None: + while self.active: + self.logger.debug("Listening for data") + try: + data = data_queue.get(timeout=0.5) + self.add_data(data.dtype, data.data) + except queue.Empty: + pass -def start(data_queue: "mp.Queue[VisualiserData]") -> None: - server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) - figure_pb2_grpc.add_FigureServiceServicer_to_server(FigureServer(), server) - server.add_insecure_port("[::]:50051") - server.add_insecure_port("0.0.0.0:50051") - logger.info("Starting server on port 50051") - server.start() - logger.info("Server started") - threading.Thread(target=listen_for_data, args=(data_queue,), daemon=True).start() - try: - server.wait_for_termination() - except KeyboardInterrupt: - logger.info("Exiting visualisation server") + while not data_queue.empty(): + data = data_queue.get() + + def start(self, data_queue: "mp.Queue[VisualiserData]") -> None: + server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + figure_pb2_grpc.add_FigureServiceServicer_to_server(FigureServer(), server) + server.add_insecure_port("[::]:50051") + server.add_insecure_port("0.0.0.0:50051") + self.logger.info("Starting server on port 50051") + server.start() + self.logger.info("Server started") + + data_thread = threading.Thread(target=self.listen_for_data, args=(data_queue,)) + data_thread.start() + + while self.active: + time.sleep(1) + self.logger.debug("Server is running") + server.stop(0.5) + data_thread.join()