From 581aa64ee358cdfe9dd0d59ca73ba3d56f80b46e Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Thu, 1 May 2025 14:22:00 +0100 Subject: [PATCH] refactor csi producers --- where_fi/application.py | 64 +++++------- where_fi/cli/__init__.py | 55 +++++----- where_fi/cli/globals.py | 9 +- where_fi/collection/__init__.py | 18 +++- where_fi/collection/file.py | 71 ++++++++----- where_fi/collection/ingest.py | 174 +++++++++++++++---------------- where_fi/collection/protocols.py | 18 +++- 7 files changed, 214 insertions(+), 195 deletions(-) diff --git a/where_fi/application.py b/where_fi/application.py index 16d9847..7940003 100644 --- a/where_fi/application.py +++ b/where_fi/application.py @@ -6,6 +6,8 @@ from typing import Any, Callable, NamedTuple import numpy.typing as npt +from where_fi.collection.protocols import CSIProducer + from .collection import CSIMatrix, ingest from .config import config from .visualise import server as visualise @@ -66,32 +68,10 @@ class CSIApplication: 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: + def __init__(self, producer: CSIProducer) -> None: self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") - self.receivers = self.create_receivers() - self.transmitter = ingest.FeitTransmitter() + self.producer = producer self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue( config.processing_sample_rate @@ -105,10 +85,6 @@ class CSIApplication: 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: @@ -148,12 +124,30 @@ class CSIApplication: self.processing_callback = func return func + def listen(self) -> None: + """ + Listen for incoming data from the selected ingestor, and call the appropriate + callback. + """ + data_gen = self.producer() + while True: + try: + sample = next(data_gen) + if self.csi_callback is not None: + self.csi_callback(sample.matrix) + if self.pre_merge_callback is not None: + self.pre_merge_callback(sample.frames) + except StopIteration: + break + except Exception as e: + self.logger.error(f"Error processing sample: {e}", exc_info=True) + break + def start(self) -> None: - for receiver in self.receivers: - receiver.thread.start() + # 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), + target=self.listen, ) self.buffer_thread.start() self.webapp_thread.start() @@ -170,16 +164,12 @@ class CSIApplication: 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 + self.producer.stop() if hasattr(self, "scheduler"): self.scheduler.stop() self.scheduler.thread.join() - self.buffer.active = False self.buffer_thread.join() self.webapp.active = False diff --git a/where_fi/cli/__init__.py b/where_fi/cli/__init__.py index efebbe6..9173d8b 100644 --- a/where_fi/cli/__init__.py +++ b/where_fi/cli/__init__.py @@ -1,8 +1,8 @@ import logging import multiprocessing as mp -from queue import Queue import threading import time +from queue import Queue from typing import Any, cast import numpy as np @@ -33,7 +33,7 @@ def antennas() -> None: """ from ..utils import antenna_order - if not globals.is_live: + if not globals.csi_producer.is_live: raise NotImplementedError("This command only works with live data") antenna_order.main() @@ -41,37 +41,28 @@ def antennas() -> None: @cli.command() def heatmap() -> None: + app = CSIApplication(globals.csi_producer) preprocessor = Preprocessor() aoa = AoA() - # Start webapp in background process - webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue( - config.processing_sample_rate - ) - webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) - webapp.start() - - 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: + @app.on_sample + def _(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.visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI) + processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data) + app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI) logger.info(f"Processed CSI data with shape {processed.shape}") processed_tensor = torch.tensor(processed, device=device) aoa.update(processed_tensor) - aoa.heatmap(visualiser=visualise_data) + aoa.heatmap(visualiser=app.visualise_data) - globals.csi_producer(csi_callback=callback) + app.start() logger.info("Finished processing CSI data") @cli.command() def phase_analysis( - subcarrier: int = 0, rx_antenna: int = 0, tx_antenna: int = 0 + subcarriers: list[int] = [0], rx_antenna: int = 0, tx_antenna: int = 0 ) -> None: """ Visualise the phase information in the CSI data received from the antennas. @@ -81,26 +72,34 @@ def phase_analysis( selected subcarrier and antenna. """ - app = CSIApplication() + app = CSIApplication(globals.csi_producer) preprocessor = Preprocessor() - subcarrier_phase: Queue[float] = Queue(config.collection_sample_rate) + if isinstance(subcarriers, int): + subcarriers = [subcarriers] + print("Starting phase analysis on subcarriers: ", subcarriers) + + subcarrier_phase: dict[int, Queue[float]] = { + x: Queue(config.collection_sample_rate) for x in subcarriers + } @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) + for subcarrier in subcarriers: + phase = cast(float, np.angle(processed[subcarrier, rx_antenna, tx_antenna])) + if subcarrier_phase[subcarrier].full(): + subcarrier_phase[subcarrier].get() + subcarrier_phase[subcarrier].put(phase) @app.on_process def _() -> None: - logger.info(f"Updating phase visualisation") + logger.info("Updating phase visualisation") app.visualise_data( - np.array(subcarrier_phase.queue), visualise.figures.Figure.PHASE_ANALYSIS + np.array([x.queue for x in subcarrier_phase.values()]), + visualise.figures.Figure.PHASE_ANALYSIS, ) # globals.csi_producer(csi_callback=callback) diff --git a/where_fi/cli/globals.py b/where_fi/cli/globals.py index cf5f0b9..d487bba 100644 --- a/where_fi/cli/globals.py +++ b/where_fi/cli/globals.py @@ -1,18 +1,15 @@ -from functools import partial from pathlib import Path from .. import collection from ..collection import file, ingest -csi_producer: collection.CSIProducer = collection.noop +csi_producer: collection.CSIProducer = collection.NoopCSIProducer() is_live = True def main(from_file: Path | None = None) -> None: global csi_producer, is_live if from_file: - csi_producer = partial(file.start_processing, file_path=from_file) - is_live = False + csi_producer = file.FileCSIPRoducer(path=from_file) else: - csi_producer = ingest.start_processing - is_live = True + csi_producer = ingest.RealtimeCSIProducer() diff --git a/where_fi/collection/__init__.py b/where_fi/collection/__init__.py index 7dc383c..8b54e48 100644 --- a/where_fi/collection/__init__.py +++ b/where_fi/collection/__init__.py @@ -1,14 +1,24 @@ +from typing import Iterator + import numpy as np import numpy.typing as npt from . import file, ingest -from .protocols import CSICallback, CSIProducer +from .protocols import CSIProducer, MergedCSI -def noop(csi_callback: CSICallback | None = None) -> None: - del csi_callback +class NoopCSIProducer: + """A no-op CSI producer that does nothing""" + + is_live = False + + def __call__(self) -> Iterator[MergedCSI]: + return iter([]) + + def stop(self) -> None: + pass CSIMatrix = npt.NDArray[np.complex64] -__all__ = ["file", "ingest", "noop", "CSIProducer", "CSICallback"] +__all__ = ["file", "ingest", "NoopCSIProducer", "CSIProducer"] diff --git a/where_fi/collection/file.py b/where_fi/collection/file.py index f83b9ca..3ef93cb 100644 --- a/where_fi/collection/file.py +++ b/where_fi/collection/file.py @@ -2,6 +2,7 @@ import logging import time from datetime import datetime, timedelta from pathlib import Path +from typing import Iterator import h5py @@ -10,30 +11,50 @@ from . import protocols logger = logging.getLogger(__name__) -def start_processing( - file_path: Path, - csi_callback: protocols.CSICallback | None = None, -) -> None: - logger.info(f"Replaying CSI data from {file_path}") - with h5py.File(file_path, "r") as file: - try: - for key in file: - datetime.fromisoformat(key) - except ValueError as e: - logger.exception( - "The file provided was not generated using this software", e - ) +class FileCSIPRoducer: + is_live = False - start_time = datetime.fromisoformat(list(file.keys())[0]) - target_offset = datetime.now() - start_time - for key in file: - logger.debug(f"Sending data from {key} at {datetime.now().isoformat()}") - csi_callback(file[key][:]) - curr_time_virtual = datetime.fromisoformat(key) - new_offset = datetime.now() - curr_time_virtual - logger.debug(f"New offset {new_offset}, target is {target_offset}") - if new_offset > target_offset + timedelta(seconds=1): - logger.warning( - f"Data is {new_offset - target_offset} behind, lagging behind..." + def __init__(self, path: Path) -> None: + self.path = path + if not self.path.exists(): + raise FileNotFoundError(f"File {self.path} does not exist") + + def __call__(self) -> Iterator[protocols.MergedCSI]: + """ + Replay the CSI data from the file. + """ + logger.info(f"Replaying CSI data from {self.path}") + with h5py.File(self.path, "r") as file: + try: + for key in file: + datetime.fromisoformat(key) + except ValueError as e: + logger.exception( + "The file provided was not generated using this software", e ) - time.sleep(max(0, (target_offset - new_offset).total_seconds())) + + start_time = datetime.fromisoformat(list(file.keys())[0]) + target_offset = datetime.now() - start_time + for key in file: + logger.debug(f"Sending data from {key} at {datetime.now().isoformat()}") + + # TODO: add raw frames + yield protocols.MergedCSI( + frames={}, + matrix=file[key][:], + ) + + curr_time_virtual = datetime.fromisoformat(key) + new_offset = datetime.now() - curr_time_virtual + logger.debug(f"New offset {new_offset}, target is {target_offset}") + if new_offset > target_offset + timedelta(seconds=1): + logger.warning( + f"Data is {new_offset - target_offset} behind, lagging behind..." + ) + time.sleep(max(0, (target_offset - new_offset).total_seconds())) + + def stop(self) -> None: + """ + Stop the producer. + """ + logger.info("Stopping file CSI producer") diff --git a/where_fi/collection/ingest.py b/where_fi/collection/ingest.py index d6c5400..70842c7 100644 --- a/where_fi/collection/ingest.py +++ b/where_fi/collection/ingest.py @@ -1,18 +1,18 @@ import logging import multiprocessing as mp +import selectors import socket -import struct import subprocess import threading import time from datetime import datetime -from typing import Callable, NamedTuple +from typing import Iterator import numpy as np from ..config import config from .csi_frame import CSI -from .protocols import CSICallback +from .protocols import MergedCSI Host = tuple[str, int] @@ -27,6 +27,7 @@ class FeitHost: self.active = True self.checker = threading.Thread(target=self.check_continuous) self.checker.start() + self.selectors: list[selectors.BaseSelector] = [] def check_connection(self) -> bool: try: @@ -46,6 +47,10 @@ class FeitHost: self.server.connect(self.host) self.server.send(b"stop\n") self.server.send(self.command.encode()) + + for selector in self.selectors: + selector.register(self.server, selectors.EVENT_READ, data=self) + self.logger.info(f"Connected to {self.host}") def check_continuous(self) -> None: @@ -83,40 +88,49 @@ class FeitTransmitter(FeitHost): class FeitReceiver(FeitHost): - def __init__(self, host: Host, queue: "mp.Queue[CSI]") -> None: + """ + A class used to connect to the host running FeitCSI and receive CSI data over a UDP + socket. + """ + + def __init__(self, host: Host) -> 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) -> None: - prev_time = datetime.now() - self.logger.info("Listening for CSI data") - 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 - # packet will depend on the frame format and channel width, which - # changes the number of subcarriers - data = self.server.recv(65535) - try: - csidata = CSI(data) - self.logger.debug( - f"Received CSI data after {datetime.now() - prev_time}" - ) - prev_time = datetime.now() - self.queue.put(csidata) - except struct.error: - self.logger.error("Failed to parse CSI data") - self.logger.info("Stopping CSI receiver") + def recv(self) -> CSI: + """ + Block until a CSI frame is received and return it. + """ + + # Receive the data from the socket + data = self.server.recv(65535) + # Decode the data using the CSI frame format + csi = CSI(data) + return csi + + def register(self, selector: selectors.BaseSelector) -> None: + """ + Register the receiver as a selector. This will allow us to multiplex the I/O + operations on a single thread. + + This allows us to block until any of the receivers receive data + """ + self.selectors.append(selector) + + if hasattr(self, "server"): + selector.register(self.server, selectors.EVENT_READ, data=self) -class CSIAntennaJoin: +class CSIAntennaArray: """ + Represents a set of receiver hosts that are used to receive CSI data, assumed to be + part of a linear antenna array. + 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. @@ -127,15 +141,12 @@ class CSIAntennaJoin: - `sample_callback`: Called with the merged data """ - def __init__( - self, - receiver_connections: dict[Host, "mp.Queue[CSI]"], - ) -> None: + def __init__(self, receivers: list[FeitReceiver]) -> None: self.pending_data: dict[Host, tuple[datetime, CSI]] = {} self.pending_data_lock = mp.Lock() self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") self.last_processed = datetime.now() - self.connections = receiver_connections + self.receivers = receivers self.active = True def add_data(self, host: Host, data: CSI) -> None: @@ -159,16 +170,13 @@ class CSIAntennaJoin: return False return True - def process_data( - self, - sample_callback: CSICallback | None = None, - pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, - ) -> None: + def process_data(self) -> None | MergedCSI: + if not self.is_ready(): + self.logger.debug("Not all data is ready") + return None self.last_processed = datetime.now() - if pre_merge_callback is not None: - pre_merge_callback( - {host: data[1] for host, data in self.pending_data.items()} - ) + + frames = {host: data[1] for host, data in self.pending_data.items()} antenna_data = [ np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2) for ip, antenna in config.antennas.order @@ -176,63 +184,45 @@ class CSIAntennaJoin: # We have data from all servers all_data = np.concat(antenna_data, axis=1) + return MergedCSI(frames=frames, matrix=all_data) - if sample_callback: - sample_callback(all_data) + def process_forever(self) -> Iterator[MergedCSI]: + # Register the receivers with the selector + self.logger.info("Starting to process data from receivers") + selector = selectors.DefaultSelector() + + for receiver in self.receivers: + receiver.register(selector) - def process_forever( - self, - callback: CSICallback | None = None, - pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, - ) -> None: 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) + events = selector.select(timeout=0.1) + if not events: + self.logger.warning("No events received in the last 0.1 seconds") + continue + + for key, _ in events: + receiver = key.data + self.add_data(receiver.host, receiver.recv()) + + sample = self.process_data() + if sample is not None: + yield sample -class Receiver(NamedTuple): - ip: Host - receiver: FeitReceiver - queue: "mp.Queue[CSI]" +class RealtimeCSIProducer: + def __init__(self) -> None: + self.receivers = [FeitReceiver(ip) for ip in config.receive_hosts] + self.transmitter = FeitTransmitter() + # Start injecting CSI frames + self.buffer = CSIAntennaArray(self.receivers) + self.is_live = True -def start_processing( - csi_callback: CSICallback | None = None, - pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None, -) -> None: - receivers = [ - Receiver(ip, FeitReceiver(ip), mp.Queue(config.collection_sample_rate)) - for ip in config.receive_hosts - ] + def __call__(self) -> Iterator[MergedCSI]: + return self.buffer.process_forever() - transmitter = FeitTransmitter() - - # Start injecting CSI frames - buffer = CSIAntennaJoin({r.ip: r.queue for r in receivers}) - - receiver_threads = [ - threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers - ] - - for proc in receiver_threads: - proc.start() - - buffering_thread = threading.Thread( - target=buffer.process_forever, args=(csi_callback, pre_merge_callback) - ) - buffering_thread.start() - try: - while True: - time.sleep(100) - except KeyboardInterrupt: - for r in receivers: - r.receiver.active = False - transmitter.active = False - buffer.active = False - return + def stop(self) -> None: + for r in self.receivers: + r.active = False + self.transmitter.active = False + self.buffer.active = False diff --git a/where_fi/collection/protocols.py b/where_fi/collection/protocols.py index 825f7a3..2ee591d 100644 --- a/where_fi/collection/protocols.py +++ b/where_fi/collection/protocols.py @@ -1,10 +1,22 @@ -from typing import Callable, Protocol +from typing import Iterator, NamedTuple, Protocol import numpy as np import numpy.typing as npt -CSICallback = Callable[[npt.NDArray[np.complex64]], None] +from .csi_frame import CSI + +CSIHost = tuple[str, int] + + +class MergedCSI(NamedTuple): + """Merged CSI data from all antennas""" + + frames: dict[CSIHost, CSI] + matrix: npt.NDArray[np.complex64] class CSIProducer(Protocol): - def __call__(self, csi_callback: CSICallback | None = None) -> None: ... + is_live: bool + + def __call__(self) -> Iterator[MergedCSI]: ... + def stop(self) -> None: ...