py-neuromodulation 0.0.7__py3-none-any.whl → 0.1.0__py3-none-any.whl

py_neuromodulation/gui/backend/app_manager.py

@@ -0,0 +1,345 @@
+import multiprocessing as mp
+import threading
+import os
+import signal
+import time
+import platform
+
+import logging
+
+from .app_utils import force_terminate_process, create_logger, ansi_color, ansi_reset
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from multiprocessing.synchronize import Event
+
+
+# Shared memory configuration
+ARRAY_SIZE = 1000  # Adjust based on your needs
+
+SERVER_PORT = 50001
+DEV_SERVER_PORT = 54321
+
+
+def run_vite(
+    shutdown_event: "Event",
+    debug: bool = False,
+    dev_port: int = DEV_SERVER_PORT,
+    backend_port: int = SERVER_PORT,
+    dev_env_vars: dict = {},
+) -> None:
+    """Run Vite in a separate shell"""
+    import subprocess
+
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    logger = create_logger(
+        "Vite",
+        "magenta",
+        logging.DEBUG if debug else logging.INFO,
+    )
+
+    os.environ["VITE_BACKEND_PORT"] = str(backend_port)
+
+    for key, value in dev_env_vars.items():
+        os.environ["VITE_" + key] = value
+
+    def output_reader(shutdown_event: "Event", process: subprocess.Popen):
+        logger.debug("Initialized output stream")
+        color = ansi_color(color="magenta", bright=True, styles=["BOLD"])
+
+        def read_stream(stream, stream_name):
+            for line in iter(stream.readline, ""):
+                if shutdown_event.is_set():
+                    break
+                logger.info(f"{color}[{stream_name}]{ansi_reset} {line.strip()}")
+
+        stdout_thread = threading.Thread(
+            target=read_stream, args=(process.stdout, "stdout")
+        )
+        stderr_thread = threading.Thread(
+            target=read_stream, args=(process.stderr, "stderr")
+        )
+
+        stdout_thread.start()
+        stderr_thread.start()
+
+        shutdown_event.wait()
+
+        stdout_thread.join(timeout=2)
+        stderr_thread.join(timeout=2)
+
+        logger.debug("Output stream closed")
+
+    # Handle different operating systems
+    shutdown_signal = signal.CTRL_BREAK_EVENT if os.name == "nt" else signal.SIGINT
+    subprocess_flags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
+
+    process = subprocess.Popen(
+        "bun run dev --port " + str(dev_port),
+        cwd="gui_dev",
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        creationflags=subprocess_flags,
+        shell=True,
+    )
+
+    logging_thread = threading.Thread(
+        target=output_reader,
+        args=(shutdown_event, process),
+    )
+    logging_thread.start()
+
+    shutdown_event.wait()  # Wait for shutdown
+
+    logger.debug("Terminating Vite server...")
+    process.send_signal(shutdown_signal)
+
+    try:
+        process.wait(timeout=3)
+    except subprocess.TimeoutExpired:
+        logger.debug("Timeout expired, forcing termination...")
+        process.kill()
+
+    logging_thread.join(timeout=3)
+    if logging_thread.is_alive():
+        logger.debug("Logging thread did not finish in time")
+
+    logger.info("Development server stopped")
+
+
+def run_uvicorn(
+    debug: bool = False, reload=False, server_port: int = SERVER_PORT
+) -> None:
+    from uvicorn.server import Server
+    from uvicorn.config import LOGGING_CONFIG, Config
+
+    # Configure logging
+    color = ansi_color(color="green", bright=True, styles=["BOLD"])
+    log_level = "DEBUG" if debug else "INFO"
+    log_config = LOGGING_CONFIG.copy()
+    log_config["loggers"]["uvicorn"]["level"] = log_level
+    log_config["loggers"]["uvicorn.error"]["level"] = log_level
+    log_config["loggers"]["uvicorn.access"]["level"] = log_level
+    log_config["formatters"]["default"]["fmt"] = (
+        f"{color}[FastAPI %(levelname)s (%(asctime)s)]:{ansi_reset} %(message)s"
+    )
+    log_config["formatters"]["default"]["datefmt"] = "%H:%M:%S"
+    log_config["formatters"]["access"]["fmt"] = (
+        f"{color}[FastAPI access (%(asctime)s)]:{ansi_reset} %(message)s"
+    )
+    log_config["formatters"]["access"]["datefmt"] = "%H:%M:%S"
+
+    config = Config(
+        app="py_neuromodulation.gui.backend.app_backend:PyNMBackend",
+        host="localhost",
+        reload=reload,
+        factory=True,
+        port=server_port,
+        log_level="debug" if debug else "info",
+        log_config=log_config,
+        http="httptools",
+        ws_ping_interval=None,
+        ws_ping_timeout=None,
+        ws_max_size=1024 * 1024 * 1024,  # 1GB
+        loop="asyncio" if platform.system() == "Windows" else "uvloop",
+        lifespan="off",
+    )
+
+    server = Server(config=config)
+
+    if reload:
+        from uvicorn.supervisors import ChangeReload
+        from uvicorn._subprocess import get_subprocess
+
+        # Overload the restart method of uvicorn so that is does not kill all of our processes
+        # IMPORTANT: This is a hack and prevents shutdown events from triggering when the reloader is used
+        class CustomReloader(ChangeReload):
+            def restart(self) -> None:
+                self.process.terminate()  # Use terminate instead of os.kill
+                self.process.join()
+                self.process = get_subprocess(
+                    config=self.config, target=self.target, sockets=self.sockets
+                )
+                self.process.start()
+
+        sock = config.bind_socket()
+        server = CustomReloader(config, target=server.run, sockets=[sock])
+
+    server.run()
+
+
+def run_backend(
+    shutdown_event: "Event",
+    dev: bool = True,
+    debug: bool = False,
+    reload: bool = True,
+    server_port: int = SERVER_PORT,
+    dev_port: int = DEV_SERVER_PORT,
+) -> None:
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    # Pass create_backend parameters through environment variables
+    os.environ["PYNM_DEBUG"] = str(debug)
+    os.environ["PYNM_DEV"] = str(dev)
+    os.environ["PYNM_DEV_PORT"] = str(dev_port)
+
+    server_process = mp.Process(
+        target=run_uvicorn,
+        kwargs={"debug": debug, "reload": reload, "server_port": server_port},
+        name="Server",
+    )
+    server_process.start()
+    shutdown_event.wait()
+    server_process.join()
+
+
+class AppManager:
+    LAUNCH_FLAG = "PYNM_RUNNING"
+
+    def __init__(
+        self,
+        debug: bool = False,
+        dev: bool = True,
+        run_in_webview=False,
+        server_port=SERVER_PORT,
+        dev_port=DEV_SERVER_PORT,
+        dev_env_vars: dict = {},
+    ) -> None:
+        """_summary_
+
+        Args:
+            debug (bool, optional): If True, run the app in debug mode, which sets logging level to debug,
+                and starts uvicorn, FastAPI, and Vite in debug mode. Defaults to False.
+            dev (bool, optional): If True, run the app in development mode, which enables hot
+                reloading and runs the frontend in Vite server. If False, run the app in production mode,
+                which runs the frontend from the static files in the `frontend` directory. Defaults to True.
+            run_in_webview (bool, optional): If True, open a PyWebView window to display the app. Defaults to False.
+        """
+        self.debug = debug
+        self.dev = dev
+        self.run_in_webview = run_in_webview
+        self.server_port = server_port
+        self.dev_port = dev_port
+        self.dev_env_vars = dev_env_vars
+
+        self._reset()
+        # Prevent launching multiple instances of the app due to multiprocessing
+        # This allows the absence of a main guard in the main script
+        self.is_child_process = os.environ.get(self.LAUNCH_FLAG) == "TRUE"
+        os.environ[self.LAUNCH_FLAG] = "TRUE"
+
+        self.logger = create_logger(
+            "PyNM",
+            "yellow",
+            logging.DEBUG if self.debug else logging.INFO,
+        )
+
+    def _reset(self) -> None:
+        """Reset the AppManager to its initial state."""
+        # Flags to track the state of the application
+        self.shutdown_complete = False
+        self.shutdown_started = False
+
+        # Store background tasks
+        self.tasks: dict[str, mp.Process] = {}
+
+        # Events for multiprocessing synchronization
+        self.shutdown_event: Event = mp.Event()
+
+    def _terminate_app(self) -> None:
+        if self.shutdown_started:
+            self.logger.info("Termination already in progress. Skipping.")
+            return
+
+        self.shutdown_started = True
+
+        timeout = 10
+        deadline = time.time() + timeout
+
+        self.logger.info("App closed, cleaning up background tasks...")
+        self.shutdown_event.set()
+
+        for process_name, process in self.tasks.items():
+            remaining_time = max(deadline - time.time(), 0)
+            process.join(timeout=remaining_time)
+
+            if process.is_alive():
+                self.logger.info(
+                    f"{process_name} did not terminate in time. Forcing termination..."
+                )
+                force_terminate_process(process, process_name, logger=self.logger)
+
+            self.logger.info(f"Process {process.name} terminated.")
+
+        self.shutdown_complete = True
+        self.shutdown_event.clear()
+        self.logger.info("All background tasks succesfully terminated.")
+
+    def _sigint_handler(self, signum, frame):
+        if not self.shutdown_started:
+            self.logger.info("Received SIGINT. Initiating graceful shutdown...")
+            self._terminate_app()
+        else:
+            self.logger.info("SIGINT received again. Ignoring...")
+
+    def launch(self) -> None:
+        if self.is_child_process:
+            return
+
+        # Handle keyboard interrupt signals
+        signal.signal(signal.SIGINT, self._sigint_handler)
+        # signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+        # Create and start the subprocesses
+        if self.dev:
+            self.logger.info("Starting Vite server...")
+            self.tasks["vite"] = mp.Process(
+                target=run_vite,
+                kwargs={
+                    "shutdown_event": self.shutdown_event,
+                    "debug": self.debug,
+                    "dev_port": self.dev_port,
+                    "backend_port": self.server_port,
+                    "dev_env_vars": self.dev_env_vars,
+                },
+                name="Vite",
+            )
+
+        self.logger.info("Starting backend server...")
+        self.tasks["backend"] = mp.Process(
+            target=run_backend,
+            kwargs={
+                "shutdown_event": self.shutdown_event,
+                "debug": self.debug,
+                "reload": self.dev,
+                "dev": self.dev,
+                "server_port": self.server_port,
+                "dev_port": self.dev_port,
+            },
+            name="Backend",
+        )
+
+        for process in self.tasks.values():
+            process.start()
+
+        if self.run_in_webview:
+            from .app_window import WebViewWindow
+
+            self.logger.info("Starting PyWebView window...")
+            window = WebViewWindow(debug=self.debug)  # Must be called from main thread
+            window.register_event_handler("closed", self._terminate_app)
+            window.start()  # Start the window, this will block until the window is closed
+        else:
+            try:
+                while not self.shutdown_complete:
+                    time.sleep(0.1)
+            except KeyboardInterrupt:
+                pass  # The SIGINT handler will take care of termination
+
+        if not self.shutdown_complete:
+            self._terminate_app()
+
+        self.logger.info("All processes cleaned up. Exiting...")

py_neuromodulation/gui/backend/app_pynm.py

@@ -0,0 +1,244 @@
+import os
+import numpy as np
+from threading import Thread
+import time
+import asyncio
+import multiprocessing as mp
+from queue import Empty
+from pathlib import Path
+from py_neuromodulation.stream import Stream, NMSettings
+from py_neuromodulation.analysis.decode import RealTimeDecoder
+from py_neuromodulation.utils import set_channels
+from py_neuromodulation.utils.io import read_mne_data
+from py_neuromodulation.utils.types import _PathLike
+from py_neuromodulation import logger
+from py_neuromodulation.gui.backend.app_socket import WebsocketManager
+from py_neuromodulation.stream.backend_interface import StreamBackendInterface
+from py_neuromodulation import logger
+
+
+class PyNMState:
+    def __init__(
+        self,
+        log_queue_size: bool = False,
+    ) -> None:
+        self.log_queue_size = log_queue_size
+        self.lsl_stream_name: str = ""
+        self.is_stream_lsl: bool = False
+        self.experiment_name: str = "PyNM_Experiment"  # set by set_stream_params
+        self.out_dir: _PathLike = str(
+            Path.home() / "PyNM" / self.experiment_name
+        )  # set by set_stream_params
+        self.decoding_model_path: _PathLike | None = None
+        self.decoder: RealTimeDecoder | None = None
+
+        self.backend_interface: StreamBackendInterface | None = None
+        self.websocket_manager: WebsocketManager | None = None
+
+        # Note: sfreq and data are required for stream init
+        self.stream: Stream = Stream(sfreq=1500, data=np.random.random([1, 1]))
+
+        self.feature_queue = mp.Queue()
+        self.rawdata_queue = mp.Queue()
+        self.control_queue = mp.Queue()
+        self.stop_event = asyncio.Event()
+
+        self.messages_sent = 0
+
+    def start_run_function(
+        self,
+        websocket_manager: WebsocketManager | None = None,
+    ) -> None:
+
+        self.websocket_manager = websocket_manager
+
+        # Create decoder
+        if self.decoding_model_path is not None and self.decoding_model_path != "None":
+            if os.path.exists(self.decoding_model_path):
+                self.decoder = RealTimeDecoder(self.decoding_model_path)
+            else:
+                logger.debug("Passed decoding model path does't exist")
+
+        # Initialize the backend interface if not already done
+        if not self.backend_interface:
+            self.backend_interface = StreamBackendInterface(
+                self.feature_queue, self.rawdata_queue, self.control_queue
+            )
+
+        # The run_func_thread is terminated through the stream_handling_queue
+        # which initiates to break the data generator and save the features
+        stream_process = mp.Process(
+            target=self.stream.run,
+            kwargs={
+                "out_dir": "" if self.out_dir == "default" else self.out_dir,
+                "experiment_name": self.experiment_name,
+                "is_stream_lsl": self.is_stream_lsl,
+                "stream_lsl_name": self.lsl_stream_name,
+                "simulate_real_time": True,
+                "decoder": self.decoder,
+                "backend_interface": self.backend_interface,
+            },
+        )
+
+        stream_process.start()
+
+        # Start websocket sender process
+
+        if self.websocket_manager:
+            # TONI: Instead of having this function be not async and send the
+            # _process_queue function to the Uvicorn async loop, we could
+            # have this entire "start_run_function" function be async as well
+
+            # Get the current event loop and run the queue processor
+            loop = asyncio.get_running_loop()
+            queue_task = loop.create_task(self._process_queue())
+
+            # Store task reference for cleanup
+            self._queue_task = queue_task
+
+        # Store processes for cleanup
+        self.stream_process = stream_process
+
+    def stop_run_function(self) -> None:
+        """Stop the stream processing"""
+        if self.backend_interface:
+            self.backend_interface.send_command("stop")
+            self.stop_event.set()
+
+    def setup_lsl_stream(
+        self,
+        lsl_stream_name: str = "",
+        line_noise: float | None = None,
+    ):
+        from mne_lsl.lsl import resolve_streams
+
+        logger.info("resolving streams")
+        lsl_streams = resolve_streams()
+
+        for stream in lsl_streams:
+            if stream.name == lsl_stream_name:
+                logger.info(f"found stream {lsl_stream_name}")
+
+                ch_names = stream.get_channel_names()
+                if ch_names is None:
+                    ch_names = ["ch" + str(i) for i in range(stream.n_channels)]
+                logger.info(f"channel names: {ch_names}")
+
+                ch_types = stream.get_channel_types()
+                if ch_types is None:
+                    ch_types = ["eeg" for i in range(stream.n_channels)]
+
+                logger.info(f"channel types: {ch_types}")
+
+                info_ = stream.get_channel_info()
+                logger.info(f"channel info: {info_}")
+
+                channels = set_channels(
+                    ch_names=ch_names,
+                    ch_types=ch_types,
+                    used_types=["eeg", "ecog", "dbs", "seeg"],
+                )
+
+                # set all used column to 0
+                #channels.loc[:, "used"] = 0
+
+                logger.info(channels)
+                sfreq = stream.sfreq
+
+                self.stream: Stream = Stream(
+                    sfreq=sfreq,
+                    line_noise=line_noise,
+                    channels=channels,
+                )
+                logger.info("stream setup")
+                logger.info("settings setup")
+
+                self.lsl_stream_name = lsl_stream_name
+                self.is_stream_lsl = True
+                break
+        else:
+            logger.error(f"Stream {lsl_stream_name} not found")
+            self.is_stream_lsl = False
+            self.is_stream_lsl = ""
+            raise ValueError(f"Stream {lsl_stream_name} not found")
+
+    def setup_offline_stream(
+        self,
+        file_path: str,
+        line_noise: float,
+    ):
+        data, sfreq, ch_names, ch_types, bads = read_mne_data(file_path)
+
+        channels = set_channels(
+            ch_names=ch_names,
+            ch_types=ch_types,
+            bads=bads,
+            reference=None,
+            used_types=["eeg", "ecog", "dbs", "seeg"],
+            target_keywords=None,
+        )
+
+        self.stream: Stream = Stream(
+            sfreq=sfreq,
+            data=data,
+            channels=channels,
+            line_noise=line_noise,
+        )
+        self.is_stream_lsl = False
+        self.lsl_stream_name = ""
+
+    # Async function that will continuously run in the Uvicorn async loop
+    # and handle sending data through the websocket manager
+    async def _process_queue(self):
+        last_queue_check = time.time()
+
+        while not self.stop_event.is_set():
+            # Use asyncio.gather to process both queues concurrently
+            tasks = []
+            current_time = time.time()
+
+            # Check feature queue
+            while not self.feature_queue.empty():
+                try:
+                    data = self.feature_queue.get_nowait()
+                    tasks.append(self.websocket_manager.send_cbor(data))  # type: ignore
+                    self.messages_sent += 1
+                except Empty:
+                    break
+
+            # Check raw data queue
+            while not self.rawdata_queue.empty():
+                try:
+                    data = self.rawdata_queue.get_nowait()
+                    self.messages_sent += 1
+                    tasks.append(self.websocket_manager.send_cbor(data))  # type: ignore
+                except Empty:
+                    break
+
+            if tasks:
+                # Wait for all send operations to complete
+                await asyncio.gather(*tasks, return_exceptions=True)
+            else:
+                # Only sleep if we didn't process any messages
+                await asyncio.sleep(0.001)
+
+            # Log queue diagnostics every 5 seconds
+            if self.log_queue_size:
+                if current_time - last_queue_check > 5:
+                    logger.info(
+                        "\nQueue diagnostics:\n"
+                        f"\tMessages send to websocket: {self.messages_sent}.\n"
+                    )
+                    try:
+                        logger.info(
+                            f"\tFeature queue size: ~{self.feature_queue.qsize()}\n"
+                            f"\tRaw data queue size: ~{self.rawdata_queue.qsize()}"
+                        )
+                    except NotImplementedError:
+                        continue
+
+                last_queue_check = current_time
+
+            # Check if stream process is still alive
+            if not self.stream_process.is_alive():
+                break

py_neuromodulation/gui/backend/app_socket.py

@@ -0,0 +1,95 @@
+from fastapi import WebSocket
+import logging
+import cbor2
+import time
+
+
+class WebsocketManager:
+    """
+    Manages WebSocket connections and messages.
+    Perhaps in the future it will handle multiple connections.
+    """
+
+    def __init__(self):
+        self.active_connections: list[WebSocket] = []
+        self.logger = logging.getLogger("PyNM")
+        self.disconnected = []
+        self._queue_task = None
+        self._stop_event = None
+        self.loop = None
+        self.messages_sent = 0
+        self._last_diagnostic_time = time.time()
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+        client_address = websocket.client
+        if client_address:
+            self.logger.info(
+                f"Client connected with client ID: {client_address.port}:{client_address.port}"
+            )
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+        client_address = websocket.client
+        if client_address:
+            self.logger.info(
+                f"Client {client_address.port}:{client_address.port} disconnected."
+            )
+
+    # Combine IP and port to create a unique client ID
+    async def send_cbor(self, object: dict):
+        if not self.active_connections:
+            self.logger.warning("No active connection to send message.")
+            return
+
+        start_time = time.time()
+        cbor_data = cbor2.dumps(object)
+        serialize_time = time.time() - start_time
+
+        if serialize_time > 0.1:  # Log slow serializations
+            self.logger.warning(f"CBOR serialization took {serialize_time:.3f}s")
+
+        send_start = time.time()
+        for connection in self.active_connections:
+            try:
+                await connection.send_bytes(cbor_data)
+            except RuntimeError as e:
+                self.logger.error(f"Error sending CBOR message: {e}")
+                self.disconnected.append(connection)
+
+        send_time = time.time() - send_start
+        if send_time > 0.1:  # Log slow sends
+            self.logger.warning(f"WebSocket send took {send_time:.3f}s")
+
+        self.messages_sent += 1
+
+        # Log diagnostics every 5 seconds
+        current_time = time.time()
+        if current_time - self._last_diagnostic_time > 5:
+            self.logger.info(f"Messages sent: {self.messages_sent}")
+            self._last_diagnostic_time = current_time
+
+    async def send_message(self, message: str | dict):
+        if not self.active_connections:
+            self.logger.warning("No active connection to send message.")
+            return
+
+        self.logger.info(
+            f"Sending message within app_socket: {message.keys() if type(message) is dict else message}"
+        )
+        for connection in self.active_connections:
+            try:
+                if type(message) is dict:
+                    await connection.send_json(message)
+                elif type(message) is str:
+                    await connection.send_text(message)
+                self.logger.info(f"Message sent to {connection.client}")
+            except RuntimeError as e:
+                self.logger.error(f"Error sending message: {e}.")
+                self.active_connections.remove(connection)
+                await connection.close()
+
+    @property
+    def is_connected(self):
+        return self.active_connections is not None