Supports cancelling mxcube flows

mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py

@@ -1,10 +1,8 @@
-import asyncio
 import logging
 import os
 import pprint
 import time
 from os import path
-from time import perf_counter
 from typing import Literal
 
 import gevent
@@ -316,27 +314,34 @@ def abort(self) -> None:
         None
         """
         logging.getLogger("HWR").info("Aborting current workflow")
-        loop = asyncio.get_event_loop()
-        if loop.is_running():
-            # TODO: MXcube takes some significant time to cancel workflows (investigate why)
-            timeout = 30
-            logging.getLogger("HWR").warning(
-                f"Loop is still running, waiting for {timeout} s to complete"
+
+        if (
+            self.current_flow.prefect_client is not None
+            and self.current_flow.prefect_client.flow_run_id is not None
+        ):
+            try:
+                self.current_flow.prefect_client.cancel_flow_run(wait=False)
+                logging.getLogger("user_level_log").info(
+                    "Prefect workflow aborted successfully."
+                )
+            except Exception as ex:
+                logging.getLogger("HWR").error(
+                    f"Error aborting prefect flow: {str(ex)}"
+                )
+                return
+        else:
+            logging.getLogger("user_level_log").warning(
+                "No prefect client available to abort the workflow"
             )
-            t = perf_counter()
-            while loop.is_running():
-                gevent.sleep(1)
-                logging.getLogger("HWR").warning(f"Loop is still running")
-                if perf_counter() > t + timeout:
-                    raise QueueExecutionException("Asyncio Loop is still running", self)
+            return
 
         # If necessary unblock dialog
         if not self.gevent_event.is_set():
             self.gevent_event.set()
 
         if self.workflow_name == PrefectFlows.grid_scan:
-            self.raster_flow.prefect_flow_aborted = True
-            self.raster_flow.mxcubecore_workflow_aborted = True
+            self.current_flow.prefect_flow_aborted = True
+            self.current_flow.mxcubecore_workflow_aborted = True
 
         self.state.value = "ON"
 
@@ -428,86 +433,84 @@ def start_prefect_workflow(self, sample_id: int | None) -> None:
 
         if self.workflow_name == PrefectFlows.screen_sample:
             logging.getLogger("HWR").info(f"Starting workflow: {self.workflow_name}")
-            self.screening_flow = ScreeningFlow(
+            self.current_flow = ScreeningFlow(
                 state=self._state, resolution=self.resolution, sample_id=sample_id
             )
-            dialog_box_parameters = self.open_dialog(self.screening_flow.dialog_box())
+            dialog_box_parameters = self.open_dialog(self.current_flow.dialog_box())
             if dialog_box_parameters:
                 logging.getLogger("HWR").info(
                     f"Dialog box parameters: {dialog_box_parameters}"
                 )
                 self.state.value = "ON"
-                self.screening_flow.run(dialog_box_parameters=dialog_box_parameters)
+                self.current_flow.run(dialog_box_parameters=dialog_box_parameters)
             else:
                 self.state.value = "ON"
                 raise QueueExecutionException("dialog_box_parameters is empty", self)
 
         elif self.workflow_name == PrefectFlows.collect_dataset:
             logging.getLogger("HWR").info(f"Starting workflow: {self.workflow_name}")
-            self.full_dataset_flow = FullDatasetFlow(
+            self.current_flow = FullDatasetFlow(
                 state=self._state, resolution=self.resolution, sample_id=sample_id
             )
-            dialog_box_parameters = self.open_dialog(
-                self.full_dataset_flow.dialog_box()
-            )
+            dialog_box_parameters = self.open_dialog(self.current_flow.dialog_box())
             if dialog_box_parameters:
                 logging.getLogger("HWR").info(
                     f"Dialog box parameters: {dialog_box_parameters}"
                 )
-                self.full_dataset_flow.run(dialog_box_parameters=dialog_box_parameters)
+                self.current_flow.run(dialog_box_parameters=dialog_box_parameters)
             else:
                 self.state.value = "ON"
                 raise QueueExecutionException("dialog_box_parameters is empty", self)
 
         elif self.workflow_name == PrefectFlows.grid_scan:
             logging.getLogger("HWR").info(f"Starting workflow: {self.workflow_name}")
-            self.raster_flow = GridScanFlow(
+            self.current_flow = GridScanFlow(
                 sample_view=self.sample_view,
                 state=self._state,
                 redis_connection=self.redis_connection,
                 resolution=self.resolution,
                 sample_id=sample_id,
             )
-            dialog_box_parameters = self.open_dialog(self.raster_flow.dialog_box())
+            dialog_box_parameters = self.open_dialog(self.current_flow.dialog_box())
             if dialog_box_parameters:
                 logging.getLogger("HWR").info(
                     f"Dialog box parameters: {dialog_box_parameters}"
                 )
 
-                self.raster_flow.run(dialog_box_parameters=dialog_box_parameters)
+                self.current_flow.run(dialog_box_parameters=dialog_box_parameters)
             else:
                 self.state.value = "ON"
                 raise QueueExecutionException("dialog_box_parameters is empty", self)
 
         elif self.workflow_name == PrefectFlows.one_shot:
             logging.getLogger("HWR").info(f"Starting workflow: {self.workflow_name}")
-            self.one_shot_flow = OneShotFlow(
+            self.current_flow = OneShotFlow(
                 state=self._state, resolution=self.resolution, sample_id=sample_id
             )
-            dialog_box_parameters = self.open_dialog(self.one_shot_flow.dialog_box())
+            dialog_box_parameters = self.open_dialog(self.current_flow.dialog_box())
             if dialog_box_parameters:
                 logging.getLogger("HWR").info(
                     f"Dialog box parameters: {dialog_box_parameters}"
                 )
-                self.one_shot_flow.run(dialog_box_parameters=dialog_box_parameters)
+                self.current_flow.run(dialog_box_parameters=dialog_box_parameters)
             else:
                 self.state.value = "ON"
                 raise QueueExecutionException("dialog_box_parameters is empty", self)
 
         elif self.workflow_name == PrefectFlows.partial_udc:
             logging.getLogger("HWR").info(f"Starting workflow: {self.workflow_name}")
 
-            self.partial_udc_flow = PartialUDCFlow(
+            self.current_flow = PartialUDCFlow(
                 state=self._state,
                 resolution=self.resolution,
                 sample_id=sample_id,
             )
-            dialog_box_parameters = self.open_dialog(self.partial_udc_flow.dialog_box())
+            dialog_box_parameters = self.open_dialog(self.current_flow.dialog_box())
             if dialog_box_parameters:
                 logging.getLogger("HWR").info(
                     f"Dialog box parameters: {dialog_box_parameters}"
                 )
-                self.partial_udc_flow.run(dialog_box_parameters=dialog_box_parameters)
+                self.current_flow.run(dialog_box_parameters=dialog_box_parameters)
             else:
                 self.state.value = "ON"
                 raise QueueExecutionException("dialog_box_parameters is empty", self)

mxcubecore/HardwareObjects/ANSTO/prefect_flows/abstract_flow.py

@@ -28,8 +28,8 @@
 from .schemas.full_dataset import FullDatasetDialogBox
 from .schemas.grid_scan import GridScanDialogBox
 from .schemas.one_shot import OneShotDialogBox
-from .schemas.prefect_workflow import PrefectFlows
 from .schemas.screening import ScreeningDialogBox
+from .sync_prefect_client import MX3SyncPrefectClient
 
 
 class AbstractPrefectWorkflow(ABC):
@@ -69,6 +69,7 @@ def __init__(self, state, resolution: Resolution, sample_id: int | None) -> None
         self.prefect_flow_aborted = False
         self.mxcubecore_workflow_aborted = False
         self.project_id_lab_name_map = None
+        self.prefect_client: None | MX3SyncPrefectClient = None
 
         if settings.BL_ACTIVE:
             self.robot_client = Client(host=settings.ROBOT_HOST, readonly=False)

mxcubecore/HardwareObjects/ANSTO/prefect_flows/full_dataset_collection_flow.py

@@ -98,15 +98,15 @@ def run(self, dialog_box_parameters: dict) -> None:
             f"parameters sent to prefect flow {prefect_parameters}"
         )
 
-        full_dataset_flow = MX3SyncPrefectClient(
+        self.prefect_client = MX3SyncPrefectClient(
             name=settings.FULL_DATASET_DEPLOYMENT_NAME, parameters=prefect_parameters
         )
 
         # Remember the collection params for the next collection
         self._save_dialog_box_params_to_redis(dialog_box_model)
 
         try:
-            full_dataset_flow.trigger_data_collection(sample_id)
+            self.prefect_client.trigger_data_collection(sample_id)
             logging.getLogger("user_level_log").info(
                 "Dataset collection completed successfully."
             )

mxcubecore/HardwareObjects/ANSTO/prefect_flows/grid_scan_flow.py

@@ -182,12 +182,12 @@ def start_prefect_flow_and_get_results_from_redis(
             with the exception message.
         """
 
-        grid_scan_flow = MX3SyncPrefectClient(
+        self.prefect_client = MX3SyncPrefectClient(
             name=settings.GRID_SCAN_DEPLOYMENT_NAME,
             parameters=prefect_parameters.model_dump(exclude_none=True),
         )
         try:
-            flow_run_uuid = grid_scan_flow.trigger_grid_scan()
+            flow_run_uuid = self.prefect_client.trigger_grid_scan()
 
             logging.getLogger("HWR").info("Getting spotfinder results from redis...")
             logging.getLogger("HWR").info(
@@ -238,6 +238,10 @@ def start_prefect_flow_and_get_results_from_redis(
                 self.sample_view.set_grid_data(
                     grid_id, heatmap_dict, data_file_path=None
                 )
+            else:
+                self._state.value = "ON"
+                self.mxcubecore_workflow_aborted = False
+                raise RuntimeError("Workflow aborted by the user.")
 
             self._state.value = "ON"
             self.mxcubecore_workflow_aborted = False
@@ -246,7 +250,7 @@ def start_prefect_flow_and_get_results_from_redis(
             self._state.value = "ON"
             self.mxcubecore_workflow_aborted = False
 
-            grid_scan_flow.check_flow_state()
+            self.prefect_client.check_flow_state()
 
             logging.getLogger("user_level_log").error(
                 f"Grid scan flow was not successful: {str(ex)}"

mxcubecore/HardwareObjects/ANSTO/prefect_flows/one_shot_flow.py

@@ -94,15 +94,15 @@ def run(self, dialog_box_parameters: dict) -> None:
             f"parameters sent to prefect flow {prefect_parameters}"
         )
 
-        one_shot_flow = MX3SyncPrefectClient(
+        self.prefect_client = MX3SyncPrefectClient(
             name=settings.ONE_SHOT_DEPLOYMENT_NAME, parameters=prefect_parameters
         )
 
         # Remember the collection params for the next collection
         self._save_dialog_box_params_to_redis(dialog_box_model)
 
         try:
-            one_shot_flow.trigger_data_collection(sample_id)
+            self.prefect_client.trigger_data_collection(sample_id)
             logging.getLogger("user_level_log").info("One-shot completed successfully.")
         except Exception as ex:
             raise QueueExecutionException(str(ex), self) from ex

mxcubecore/HardwareObjects/ANSTO/prefect_flows/partial_udc.py

@@ -176,7 +176,7 @@ def run(self, dialog_box_parameters: dict) -> None:
                 exposure_time=screening_model.exposure_time,
                 number_of_passes=1,
                 count_time=None,
-                number_of_frames=screening_model.number_of_frames,
+                degrees_per_frame=screening_model.degrees_per_frame,
                 detector_distance=self._resolution_to_distance(
                     screening_model.resolution,
                     energy=photon_energy,
@@ -194,7 +194,7 @@ def run(self, dialog_box_parameters: dict) -> None:
                 exposure_time=full_dataset_model.exposure_time,
                 number_of_passes=1,
                 count_time=None,
-                number_of_frames=full_dataset_model.number_of_frames,
+                degrees_per_frame=full_dataset_model.degrees_per_frame,
                 detector_distance=self._resolution_to_distance(
                     full_dataset_model.resolution,
                     energy=photon_energy,
@@ -266,11 +266,11 @@ def run(self, dialog_box_parameters: dict) -> None:
                 1 if dialog_box_model.run_full_dataset else 0,
             )
 
-        partial_udc_flow = MX3SyncPrefectClient(
+        self.prefect_client = MX3SyncPrefectClient(
             name=settings.PARTIAL_UDC_DEPLOYMENT_NAME, parameters=prefect_parameters
         )
         try:
-            partial_udc_flow.trigger_data_collection(sample_id, mode="partial_udc")
+            self.prefect_client.trigger_data_collection(sample_id, mode="partial_udc")
             logging.getLogger("user_level_log").info(
                 "Partial UDC completed successfully."
             )
@@ -347,9 +347,9 @@ def dialog_box(self) -> dict:
                                 "required": [
                                     "exposure_time",
                                     "omega_range",
-                                    "number_of_frames",
                                     "resolution",
                                     "transmission",
+                                    "degrees_per_frame",
                                 ],
                             }
                         }
@@ -366,10 +366,10 @@ def dialog_box(self) -> dict:
                                 "required": [
                                     "exposure_time",
                                     "omega_range",
-                                    "number_of_frames",
                                     "resolution",
                                     "transmission",
                                     "processing_pipeline",
+                                    "degrees_per_frame",
                                 ],
                             }
                         }