checkpoint

Author: jonathan-eq

src/_ert/events.py

@@ -217,7 +217,7 @@ class EEUserDone(BaseEvent):
 EEEvent = EESnapshot | EESnapshotUpdate | EETerminated | EEUserCancel | EEUserDone
 
 Event = FMEvent | ForwardModelStepChecksum | RealizationEvent | EEEvent | EnsembleEvent
-
+EventForBrm = EESnapshot | EESnapshotUpdate | ForwardModelStepChecksum
 DispatchEvent = FMEvent | ForwardModelStepChecksum | RealizationEvent | EnsembleEvent
 
 _DISPATCH_EVENTS_ANNOTATION = Annotated[

src/ert/ensemble_evaluator/evaluator.py

@@ -28,6 +28,7 @@ class UserCancelled(Exception):
     EnsembleStarted,
     EnsembleSucceeded,
     Event,
+    EventForBrm,
     FMEvent,
     ForwardModelStepChecksum,
     RealizationEvent,
@@ -60,7 +61,7 @@ def __init__(
         self,
         ensemble: Ensemble,
         config: EvaluatorServerConfig,
-        send_to_brm: Callable[[Event], None],
+        event_handler: Callable[[EventForBrm], None] = lambda _: _,
     ) -> None:
         self._config: EvaluatorServerConfig = config
         self._ensemble: Ensemble = ensemble
@@ -83,19 +84,19 @@ def __init__(
         self._dispatchers_empty: asyncio.Event = asyncio.Event()
         self._dispatchers_empty.set()
         current_snapshot_dict = self._ensemble.snapshot.to_dict()
-        self._send_to_brm = send_to_brm
-        event: Event = EESnapshot(
+        self._handle_event = event_handler
+        event = EESnapshot(
             snapshot=current_snapshot_dict,
             ensemble=self.ensemble.id_,
         )
-        self._send_to_brm(event)
+        self._handle_event(event)
         self._monitoring_result: asyncio.Future[bool] = asyncio.Future()
 
     async def _append_message(self, snapshot_update_event: EnsembleSnapshot) -> None:
         event = EESnapshotUpdate(
             snapshot=snapshot_update_event.to_dict(), ensemble=self._ensemble.id_
         )
-        self._send_to_brm(event)
+        self._handle_event(event)
         if event.snapshot.get(ids.STATUS) in {
             ENSEMBLE_STATE_STOPPED,
             ENSEMBLE_STATE_FAILED,
@@ -215,11 +216,12 @@ async def _failed_handler(self, events: Sequence[EnsembleFailed]) -> None:
     def ensemble(self) -> Ensemble:
         return self._ensemble
 
-    def cancel_gracefully(self) -> None:
+    def cancel_gracefully_synced(self) -> None:
+        self._ee_tasks.append(self._running_loop.create_task(self.cancel_gracefully()))
+
+    async def cancel_gracefully(self):
         cancel_event = EEUserCancel()
-        self._ee_tasks.append(
-            self._running_loop.create_task(self.handle_client_event(cancel_event))
-        )
+        await self.handle_client_event(cancel_event)
 
     async def handle_client_event(self, event: EEEvent) -> None:
         if type(event) is EEUserCancel:
@@ -277,9 +279,9 @@ async def listen_for_messages(self) -> None:
             except asyncio.CancelledError:
                 return
 
-    async def forward_checksum(self, event: Event) -> None:
+    async def forward_checksum(self, event: ForwardModelStepChecksum) -> None:
         # clients still need to receive events via ws
-        self._send_to_brm(event)
+        self._handle_event(event)
         await self._manifest_queue.put(event)
 
     async def _server(self) -> None:

src/ert/run_models/base_run_model.py

@@ -315,7 +315,7 @@ def ensemble_size(self) -> int:
 
     def cancel(self) -> None:
         self._cancelled = True
-        self._ensemble_evaluator.cancel_gracefully()
+        self._ensemble_evaluator.cancel_gracefully_synced()
 
     def has_failed_realizations(self) -> bool:
         return any(self._create_mask_from_failed_realizations())