diff --git a/gui/src/aaregui/threads/automation_worker.py b/gui/src/aaregui/threads/automation_worker.py new file mode 100644 index 00000000..c75fe68a --- /dev/null +++ b/gui/src/aaregui/threads/automation_worker.py @@ -0,0 +1,205 @@ +import sys +import asyncio +from enum import Enum, auto +from dataclasses import dataclass, field +from typing import List, Dict, Any +from PySide6.QtCore import QObject, Signal, Slot, QThread +from pydantic import BaseModel + +class State(Enum): + MOUNT = "Mount" + LOOP_CENTRE = "Loop Centre" + RASTER = "Raster" + DATA_COLLECTION = "Data Collection" + +class AllStates(BaseModel): + ALL_STATES: list[State] = [State.MOUNT, State.LOOP_CENTRE, State.RASTER, State.DATA_COLLECTION] + +class FSMWorker(QObject): + """ + Runs the FSM in a separate thread using an asyncio event loop. + The GUI interacts via signals. + """ + finished = Signal() + progress = Signal(int, str) # (sample_index, message) + state_changed = Signal(int, str) # (sample_index, state_name) + error = Signal(int, str) # (sample_index, error_message) + sample_complete = Signal(int) # (sample_index) + + # Control signals from GUI + pause_requested = False + stop_requested = False + next_sample_requested = False + skip_state_requested = False + + def __init__(self, queue: List[Sample], server_client): + """ + server_client: any object that exposes async methods: + async def mount(sample) + async def loop_centre(sample) + async def raster(sample) + async def collect_data(sample) + Replace with wrappers if your functions are synchronous. + """ + super().__init__() + self.queue = queue + self.server = server_client + self.all_states = AllStates() + # internal control + self._pause_event = asyncio.Event() + self._pause_event.set() # initially not paused + + # --- External control API (called from GUI thread) --- + @Slot() + def request_pause(self): + self.pause_requested = True + + @Slot() + def request_resume(self): + self.pause_requested = False + # allow worker to continue + try: + # safe to call from GUI thread - schedule on loop via asyncio.run_coroutine_threadsafe if needed. + self._pause_event.set() + except Exception: + pass + + @Slot() + def request_stop(self): + self.stop_requested = True + # un-pause so it can see the stop quickly + try: + self._pause_event.set() + except Exception: + pass + + @Slot() + def request_next_sample(self): + self.next_sample_requested = True + # also un-pause + try: + self._pause_event.set() + except Exception: + pass + + @Slot() + def request_skip_state(self): + self.skip_state_requested = True + try: + self._pause_event.set() + except Exception: + pass + + # ------------------------- + # FSM main loop + # ------------------------- + def run(self): + """Entry point for QThread: create asyncio loop and run main coroutine.""" + asyncio.run(self._main()) + + async def _main(self): + try: + for idx, sample in enumerate(self.queue): + if self.stop_requested: + self.progress.emit(idx, "Stopped by user before starting sample.") + break + + self.progress.emit(idx, f"Starting sample {sample.id}") + # iterate states in order + for state in self.all_states: + # check global control flags: + if self.stop_requested: + self.progress.emit(idx, "Stopped by user.") + break + + # Respect per-sample control (user can toggle which states to run) + if not sample.run_states.get(state, True): + self.progress.emit(idx, f"State {state.value} skipped (user disabled).") + continue + + # Wait if paused + await self._maybe_pause_loop(idx) + + # Check if user asked to skip to next sample: + if self.next_sample_requested: + self.progress.emit(idx, "Next-sample requested: skipping remaining states.") + break + + if self.skip_state_requested: + self.progress.emit(idx, f"Skip-state requested: skipping {state.value}.") + self.skip_state_requested = False + continue + + # Run state + self.state_changed.emit(idx, state.value) + self.progress.emit(idx, f"Running {state.value}...") + + try: + # Call the server method corresponding to the state. + # These server_* calls are async. Replace them with your implementations. + if state == State.MOUNT: + await self._safe_call(self.server.mount, sample, idx, state) + elif state == State.LOOP_CENTRE: + await self._safe_call(self.server.loop_centre, sample, idx, state) + elif state == State.RASTER: + await self._safe_call(self.server.raster, sample, idx, state) + elif state == State.DATA_COLLECTION: + await self._safe_call(self.server.collect_data, sample, idx, state) + else: + raise RuntimeError("Unknown state") + except asyncio.CancelledError: + self.progress.emit(idx, f"{state.value} cancelled.") + raise + except Exception as e: + # emit error but continue or stop depending on preferences + self.error.emit(idx, f"Error during {state.value}: {e}") + # by default, stop the queue on error (you can change this) + self.progress.emit(idx, "Stopping queue due to error.") + self.stop_requested = True + break + + # after data collection, if not last state, continue + # loop continues + + # finished sample iteration + if self.stop_requested: + self.progress.emit(idx, "Stopping after sample due to request or error.") + break + + if self.next_sample_requested: + self.progress.emit(idx, "Moving to next sample (user requested).") + self.next_sample_requested = False + + self.progress.emit(idx, f"Sample {sample.id} complete.") + self.sample_complete.emit(idx) + + self.progress.emit(-1, "Queue processing finished.") + finally: + self.finished.emit() + + async def _maybe_pause_loop(self, idx: int): + """ + Pause behavior: whenever pause_requested is True, we clear the event and + await it being set again. This allows a resume to set the event. + """ + if self.pause_requested: + self.progress.emit(idx, "Paused. Waiting for resume...") + self._pause_event.clear() + await self._pause_event.wait() + + async def _safe_call(self, coroutine_func, sample: Sample, idx: int, state: State, timeout: float = 60.0): + """ + Call server coroutine_func(sample) with a timeout and friendly error messages. + + coroutine_func is expected to be an async callable accepting sample. + If your server functions are synchronous, wrap them with asyncio.to_thread or run_in_executor. + """ + try: + # set a timeout to avoid hanging forever + await asyncio.wait_for(coroutine_func(sample), timeout=timeout) + self.progress.emit(idx, f"{state.value} completed successfully.") + except asyncio.TimeoutError: + raise RuntimeError(f"{state.value} timed out after {timeout} s") + except Exception as e: + raise +