Start fixing tests

Author: jonathan-eq

src/ert/ensemble_evaluator/evaluator.py

@@ -11,13 +11,16 @@ class EventSentinel:
     pass
 
 
+class UserCancelled(Exception):
+    pass
+
+
 import zmq.asyncio
 
 from _ert.events import (
     EEEvent,
     EESnapshot,
     EESnapshotUpdate,
-    EETerminated,
     EEUserCancel,
     EEUserDone,
     EnsembleCancelled,
@@ -57,6 +60,7 @@ def __init__(
         self,
         ensemble: Ensemble,
         config: EvaluatorServerConfig,
+        send_to_brm: Callable[[Event], None],
     ) -> None:
         self._config: EvaluatorServerConfig = config
         self._ensemble: Ensemble = ensemble
@@ -66,7 +70,7 @@ def __init__(
 
         self._ee_tasks: list[asyncio.Task[None]] = []
         self._server_done: asyncio.Event = asyncio.Event()
-
+        self._running_loop = asyncio.get_running_loop()
         # batching section
         self._batch_processing_queue: asyncio.Queue[
             list[tuple[EVENT_HANDLER, Event]]
@@ -78,19 +82,37 @@ def __init__(
         self._dispatchers_connected: set[bytes] = set()
         self._dispatchers_empty: asyncio.Event = asyncio.Event()
         self._dispatchers_empty.set()
-        self._monitor_queue: asyncio.Queue[Event | EventSentinel] = asyncio.Queue()
         current_snapshot_dict = self._ensemble.snapshot.to_dict()
+        self._send_to_brm = send_to_brm
         event: Event = EESnapshot(
             snapshot=current_snapshot_dict,
             ensemble=self.ensemble.id_,
         )
-        self._monitor_queue.put_nowait(event)
+        self._send_to_brm(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_
         )
-        await self._monitor_queue.put(event)
+        self._send_to_brm(event)
+        if event.snapshot.get(ids.STATUS) in {
+            ENSEMBLE_STATE_STOPPED,
+            ENSEMBLE_STATE_FAILED,
+        }:
+            print("Ensemble was stopped")
+            logger.debug("observed evaluation stopped event, signal done")
+            logger.debug("monitor informing server monitor is done...")
+
+            done_event = EEUserDone()
+            await self.handle_client_event(done_event)
+            logger.debug("monitor informed server monitor is done")
+
+            if event.snapshot.get(ids.STATUS) == ENSEMBLE_STATE_CANCELLED:
+                logger.debug("observed evaluation cancelled event, exit drainer")
+                self._monitoring_result.set_exception(
+                    UserCancelled("Experiment cancelled by user during evaluation")
+                )
 
     async def _process_event_buffer(self) -> None:
         while True:
@@ -193,16 +215,24 @@ async def _failed_handler(self, events: Sequence[EnsembleFailed]) -> None:
     def ensemble(self) -> Ensemble:
         return self._ensemble
 
+    def cancel_gracefully(self) -> None:
+        cancel_event = EEUserCancel()
+        self._ee_tasks.append(
+            self._running_loop.create_task(self.handle_client_event(cancel_event))
+        )
+
     async def handle_client_event(self, event: EEEvent) -> None:
         if type(event) is EEUserCancel:
             print("EE GOT CANCEL EVENT")
             logger.debug("Client asked to cancel.")
             await self._signal_cancel()
+            self._monitoring_result.set_result(False)
             # self._clients_empty.set()
         elif type(event) is EEUserDone:
             print("EE GOT USER DONE EVENT")
             logger.debug("Client signalled done.")
             self.stop()
+            self._monitoring_result.set_result(True)
             # self._clients_empty.set()
 
     async def handle_dispatch(self, dealer: bytes, frame: bytes) -> None:
@@ -249,7 +279,7 @@ async def listen_for_messages(self) -> None:
 
     async def forward_checksum(self, event: Event) -> None:
         # clients still need to receive events via ws
-        await self._monitor_queue.put(event)
+        self._send_to_brm(event)
         await self._manifest_queue.put(event)
 
     async def _server(self) -> None:
@@ -288,9 +318,12 @@ async def _server(self) -> None:
             await self._events.join()
             await self._complete_batch.wait()
             await self._batch_processing_queue.join()
-            event = EETerminated(ensemble=self._ensemble.id_)
-            await self._monitor_queue.put(event)
-            print("PUT EETerminated")
+            try:
+                await asyncio.wait_for(self._monitoring_result, timeout=5)
+            except TimeoutError:
+                logger.warning(
+                    "Not all clients were disconnected when closing zmq server!"
+                )
             logger.debug("Async server exiting.")
         finally:
             try:

src/ert/run_models/base_run_model.py

@@ -16,16 +16,13 @@
 from contextlib import contextmanager
 from pathlib import Path
 from queue import SimpleQueue
-from typing import TYPE_CHECKING, Any, Protocol, cast
+from typing import TYPE_CHECKING, Any, Protocol
 
 import numpy as np
 
 from _ert.events import (
     EESnapshot,
     EESnapshotUpdate,
-    EETerminated,
-    EEUserCancel,
-    EEUserDone,
     Event,
 )
 from ert.analysis import ErtAnalysisError, smoother_update
@@ -51,13 +48,9 @@
     EvaluatorServerConfig,
     Realization,
 )
-from ert.ensemble_evaluator.evaluator import EventSentinel
-from ert.ensemble_evaluator.identifiers import STATUS
+from ert.ensemble_evaluator.evaluator import UserCancelled
 from ert.ensemble_evaluator.snapshot import EnsembleSnapshot
 from ert.ensemble_evaluator.state import (
-    ENSEMBLE_STATE_CANCELLED,
-    ENSEMBLE_STATE_FAILED,
-    ENSEMBLE_STATE_STOPPED,
     REALIZATION_STATE_FAILED,
     REALIZATION_STATE_FINISHED,
 )
@@ -105,10 +98,6 @@ def delete_runpath(run_path: str) -> None:
         shutil.rmtree(run_path)
 
 
-class UserCancelled(Exception):
-    pass
-
-
 class _LogAggregration(logging.Handler):
     def __init__(self, messages: MutableSequence[str]) -> None:
         self.messages = messages
@@ -214,12 +203,13 @@ def __init__(
         )
         self._iter_snapshot: dict[int, EnsembleSnapshot] = {}
         self._status_queue = status_queue
-        self._end_queue: SimpleQueue[str] = SimpleQueue()
+
         # This holds state about the run model
         self.minimum_required_realizations = minimum_required_realizations
         self.active_realizations = copy.copy(active_realizations)
         self.start_iteration = start_iteration
         self.restart = False
+        self._cancelled = False
 
     @property
     def api(self) -> BaseRunModelAPI:
@@ -238,7 +228,7 @@ def reports_dir(self, experiment_name: str) -> str:
 
     def log_at_startup(self) -> None:
         keys_to_drop = [
-            "_end_queue",
+            "_cancelled",
             "_queue_config",
             "_status_queue",
             "_storage",
@@ -324,7 +314,8 @@ def ensemble_size(self) -> int:
         return len(self._initial_realizations_mask)
 
     def cancel(self) -> None:
-        self._end_queue.put("END")
+        self._cancelled = True
+        self._ensemble_evaluator.cancel_gracefully()
 
     def has_failed_realizations(self) -> bool:
         return any(self._create_mask_from_failed_realizations())
@@ -552,113 +543,43 @@ def send_snapshot_event(self, event: Event, iteration: int) -> None:
                 )
             )
 
-    async def run_monitor(
-        self,
-        iteration: int,
-        evaluator: EnsembleEvaluator,
-    ) -> bool:
-        try:
-            heartbeat_interval_: float | None = 0.1
-            receiver_timeout: float = 60.0
-            closetracker_received: bool = False
-
-            while True:
-                try:
-                    event = await asyncio.wait_for(
-                        evaluator._monitor_queue.get(), timeout=heartbeat_interval_
-                    )
-                    evaluator._monitor_queue.task_done()
-                except TimeoutError:
-                    if closetracker_received:
-                        logger.error("Evaluator did not send the TERMINATED event!")
-                        break
-                    event = None
-                if isinstance(event, EventSentinel):
-                    closetracker_received = True
-                    heartbeat_interval_ = receiver_timeout
-                    print("JONAK - received sentinel")
-                    continue
-                if type(event) in {
-                    EESnapshot,
-                    EESnapshotUpdate,
-                }:
-                    event = cast(EESnapshot | EESnapshotUpdate, event)
-
-                    self.send_snapshot_event(event, iteration)
-
-                    if event.snapshot.get(STATUS) in {
-                        ENSEMBLE_STATE_STOPPED,
-                        ENSEMBLE_STATE_FAILED,
-                    }:
-                        print("Ensemble was stopped")
-                        logger.debug("observed evaluation stopped event, signal done")
-                        logger.debug("monitor informing server monitor is done...")
-
-                        done_event = EEUserDone()
-                        await evaluator.handle_client_event(done_event)
-                        logger.debug("monitor informed server monitor is done")
-                        await evaluator._monitor_queue.put(EventSentinel())
-
-                        if event.snapshot.get(STATUS) == ENSEMBLE_STATE_CANCELLED:
-                            logger.debug(
-                                "observed evaluation cancelled event, exit drainer"
-                            )
-                            raise UserCancelled(
-                                "Experiment cancelled by user during evaluation"
-                            )
-                elif type(event) is EETerminated:
-                    logger.debug("got terminated event")
-                    break
-
-                if not self._end_queue.empty():
-                    print("RUN MODEL WAS CANCELLED")
-                    logger.debug("Run model canceled - during evaluation")
-                    self._end_queue.get()
-                    logger.debug("monitor asking server to cancel...")
-                    cancel_event = EEUserCancel()
-                    await evaluator.handle_client_event(cancel_event)
-                    await evaluator._monitor_queue.put(EventSentinel())
-                    logger.debug("monitor asked server to cancel")
-                    logger.debug("Run model canceled - during evaluation - cancel sent")
-        except UserCancelled:
-            raise
-        except Exception as e:
-            logger.exception(f"unexpected error: {e}")
-            # We really don't know what happened...  shut down
-            # the thread and get out of here. The monitor has
-            # been stopped by the ctx-mgr
-            return False
-
-        return True
-
     async def run_ensemble_evaluator_async(
         self,
         run_args: list[RunArg],
         ensemble: Ensemble,
         ee_config: EvaluatorServerConfig,
     ) -> list[int]:
-        if not self._end_queue.empty():
+        if self._cancelled:
             logger.debug("Run model canceled - pre evaluation")
-            self._end_queue.get()
+            self._cancelled = False
             raise UserCancelled("Experiment cancelled by user in pre evaluation")
 
         ee_ensemble = self._build_ensemble(run_args, ensemble.experiment_id)
-        evaluator = EnsembleEvaluator(ee_ensemble, ee_config)
+        evaluator = EnsembleEvaluator(
+            ee_ensemble,
+            ee_config,
+            send_to_brm=functools.partial(
+                self.send_snapshot_event, iteration=ensemble.iteration
+            ),
+        )
+        self._ensemble_evaluator = evaluator
         evaluator_task = asyncio.create_task(
             evaluator.run_and_get_successful_realizations()
         )
         await evaluator._server_started
-        if not (await self.run_monitor(ensemble.iteration, evaluator)):
+
+        if not (await self._wait_for_evaluator_result(evaluator._monitoring_result)):
             await evaluator_task
             return []
 
         logger.debug("observed that model was finished, waiting tasks completion...")
         # The model has finished, we indicate this by sending a DONE
         logger.debug("tasks complete")
 
-        if not self._end_queue.empty():
+        if self._cancelled:
+            print("YEGA!")
             logger.debug("Run model canceled - post evaluation")
-            self._end_queue.get()
+            self._cancelled = False
             try:
                 await evaluator_task
             except Exception as e:
@@ -668,7 +589,8 @@ async def run_ensemble_evaluator_async(
                 ) from e
             print("RAISING USER_CANCELLED")
             raise UserCancelled("Experiment cancelled by user in post evaluation")
-
+        else:
+            print("HELL NAH!")
         await evaluator_task
         ensemble.refresh_ensemble_state()
 
@@ -849,6 +771,12 @@ def _evaluate_and_postprocess(
 
         return num_successful_realizations
 
+    async def _wait_for_evaluator_result(
+        self, monitoring_future: asyncio.Future[bool]
+    ) -> bool:
+        """This helper function is here for the sake of mocking in tests."""
+        return await monitoring_future
+
 
 class UpdateRunModel(BaseRunModel):
     def __init__(