Cancel running workflow

docs/ert/reference/workflows/workflows.rst

@@ -80,3 +80,31 @@ state of the experiment that is running:
             print(f"Provided user arguments: {workflow_args}")
 
 For how to load internal workflow jobs into ERT, see: :ref:`installing workflows <legacy_ert_workflow_jobs>`
+
+Cancelling internal workflow jobs
+==================================
+
+A running workflow can be cancelled, either from the *Run workflow* tool in
+the GUI, or by terminating an experiment that is executing a hooked workflow.
+
+For an external workflow job this is straightforward: the invoked process is
+terminated.
+
+For an internal workflow job, cancellation is cooperative. Calling
+:code:`cancel()` only records that cancellation was requested; nothing
+forcibly interrupts a running :code:`run()` method. A job that does not
+check for this will run to completion regardless of the user's request.
+Jobs that loop over realizations or ensembles should check
+:code:`self.isCancelled()` between iterations and return early:
+
+.. code-block:: python
+
+    class MyJob(ErtScript):
+        def run(self, ensemble: Ensemble):
+            for realization in ensemble.get_realization_list():
+                if self.isCancelled():
+                    return
+                process(realization)
+
+A job that performs a single, indivisible piece of work has nothing
+meaningful to check between, and does not need to poll :code:`isCancelled()`.

src/ert/gui/ertwidgets/closabledialog.py

@@ -47,10 +47,16 @@ def keyPressEvent(self, a0: QKeyEvent | None) -> None:
         if self.close_button.isEnabled() or a0 is None or a0.key() != Qt.Key.Key_Escape:
             QDialog.keyPressEvent(self, a0)
 
-    def addButton(self, caption: str, listener: Callable[..., None]) -> QPushButton:
+    def addButton(
+        self,
+        caption: str,
+        listener: Callable[..., None],
+        *,
+        align_left: bool = False,
+    ) -> QPushButton:
         button = QPushButton(caption)
         button.setObjectName(str(caption).capitalize())
-        self.__button_layout.insertWidget(1, button)
+        self.__button_layout.insertWidget(0 if align_left else 1, button)
         button.clicked.connect(listener)
         return button
 

src/ert/gui/tools/workflows/run_workflow_widget.py

@@ -14,13 +14,12 @@
     QHBoxLayout,
     QLabel,
     QMessageBox,
-    QToolButton,
+    QPushButton,
     QWidget,
 )
 
 from _ert.threading import ErtThread
 from ert.gui.ertwidgets import EnsembleSelector
-from ert.gui.icon_utils import load_icon
 from ert.gui.tools.workflows.workflow_dialog import WorkflowDialog
 from ert.runpaths import Runpaths
 from ert.workflow_runner import WorkflowRunner
@@ -53,17 +52,10 @@ def __init__(self, config: ErtConfig, notifier: ErtNotifier) -> None:
         self.source_ensemble_selector = EnsembleSelector(notifier)
         layout.addRow("Ensemble", self.source_ensemble_selector)
 
-        self.run_button = QToolButton()
-        self.run_button.setIconSize(QSize(32, 32))
-        self.run_button.setText("Start workflow")
-        self.run_button.setIcon(load_icon("play_circle.svg"))
-        self.run_button.clicked.connect(self.startWorkflow)
-        self.run_button.setToolButtonStyle(Qt.ToolButtonStyle.ToolButtonTextBesideIcon)
-
-        layout.addRow(self.run_button)
-
         self.setLayout(layout)
 
+        self.run_button: QPushButton | None = None
+
         self._running_workflow_dialog: WorkflowDialog | None = None
 
         self.workflowSucceeded.connect(self.workflowFinished)
@@ -207,12 +199,6 @@ def workflowFinishedWithFail(self) -> None:
             self._running_workflow_dialog = None
 
     def workflowStoppedByUser(self) -> None:
-        workflow_name = self.getCurrentWorkflowName()
-        QMessageBox.information(
-            self,
-            "Workflow killed!",
-            f"The workflow '{workflow_name}' was killed successfully!",
-        )
         if self._running_workflow_dialog is not None:
             self._running_workflow_dialog.reject()
             self._running_workflow_dialog = None

src/ert/gui/tools/workflows/workflow_dialog.py

@@ -33,7 +33,7 @@ def __init__(
         layout.addWidget(widget)
 
         button_layout = QHBoxLayout()
-        self.close_button = QPushButton("Close")
+        self.close_button = QPushButton("Cancel")
         self.close_button.clicked.connect(self.closeButtonPressed.emit)
         button_layout.addStretch()
         button_layout.addWidget(self.close_button)

src/ert/gui/tools/workflows/workflows_tool.py

@@ -27,5 +27,8 @@ def __init__(self, config: ErtConfig, notifier: ErtNotifier) -> None:
     def trigger(self) -> None:
         run_workflow_widget = RunWorkflowWidget(self.config, self.notifier)
         dialog = ClosableDialog("Run workflow", run_workflow_widget, self.parent())  # type: ignore
+        run_workflow_widget.run_button = dialog.addButton(
+            "Start workflow", run_workflow_widget.startWorkflow, align_left=True
+        )
         dialog.exec()
         self.notifier.emitErtChange()  # workflow may have added new cases.

src/ert/plugins/hook_implementations/workflows/csv_export.py

@@ -84,6 +84,9 @@ def run(
         data = pd.DataFrame()
 
         for ensemble in ensembles:
+            if self.isCancelled():
+                break
+
             if not ensemble.has_data():
                 raise UserWarning(
                     f"The ensemble '{ensemble.name}' does not have any data!"

src/ert/plugins/hook_implementations/workflows/export_rft.py

@@ -42,6 +42,9 @@ def run(
         paths = run_paths.get_paths(realizations, iteration)
 
         for realization, runpath in zip(realizations, paths, strict=True):
+            if self.isCancelled():
+                return
+
             realization_data = observations_and_responses.filter(
                 observations_and_responses["realization"] == realization
             ).drop("realization")

src/ert/workflow_runner.py

@@ -2,6 +2,7 @@
 
 import datetime
 import logging
+import threading
 import types
 from concurrent import futures
 from concurrent.futures import Future
@@ -44,6 +45,8 @@ def __init__(self, workflow_job: WorkflowJob) -> None:
         self.job = workflow_job
         self.__running = False
         self.__script: ErtScript | None = None
+        self.__cancel_requested = False
+        self._lock = threading.Lock()
         self.stop_on_fail = False
 
     def run(
@@ -68,17 +71,26 @@ def run(
                     f"{self.job.max_args} arguments, {len(arguments)} given."
                 )
 
-            if isinstance(self.job, BaseErtScriptWorkflow):
-                ert_script_class = self.job.load_ert_script_class()
-                self.__script = ert_script_class()
-                # We let stop on fail either from class or config take precedence
-                self.stop_on_fail = self.job.stop_on_fail or self.__script.stop_on_fail
-
-            else:
-                self.__script = ExternalErtScript(
-                    self.job.executable,  # type: ignore
-                )
-                self.stop_on_fail = self.job.stop_on_fail
+            with self._lock:
+                if isinstance(self.job, BaseErtScriptWorkflow):
+                    ert_script_class = self.job.load_ert_script_class()
+                    self.__script = ert_script_class()
+                    # We let stop on fail either from class or config take
+                    # precedence
+                    self.stop_on_fail = (
+                        self.job.stop_on_fail or self.__script.stop_on_fail
+                    )
+
+                else:
+                    self.__script = ExternalErtScript(
+                        self.job.executable,  # type: ignore
+                    )
+                    self.stop_on_fail = self.job.stop_on_fail
+
+                # A cancellation requested before the script existed would
+                # be lost; apply it once there is a script to cancel.
+                if self.__cancel_requested:
+                    self.__script.cancel()
 
             return self.__script.initializeAndRun(
                 self.job.argument_types(), arguments, fixtures
@@ -97,8 +109,10 @@ def execution_type(self) -> str:
         return "external"
 
     def cancel(self) -> None:
-        if self.__script is not None:
-            self.__script.cancel()
+        with self._lock:
+            self.__cancel_requested = True
+            if self.__script is not None:
+                self.__script.cancel()
 
     def isRunning(self) -> bool:
         return self.__running
@@ -144,6 +158,7 @@ def __init__(
         self.__current_job: WorkflowJobRunner | None = None
         self.__status: dict[str, dict[str, Any]] = {}
         self.__job_results: list[WorkflowJobResult] = []
+        self._current_job_lock = threading.Lock()
 
     def __enter__(self) -> Self:
         self.run()
@@ -173,22 +188,24 @@ def run_blocking(self) -> None:
         self.__running = True
 
         for index, (job, args) in enumerate(self.__workflow):
-            if self.__cancelled:
-                # The workflow was cancelled before this job started
-                result = WorkflowJobResult(
-                    name=job.name,
-                    index=index,
-                    arguments=[str(arg) for arg in args],
-                    stdout="",
-                    stderr="",
-                    status=WorkflowJobStatus.CANCELLED,
-                )
-                self.__job_results.append(result)
-                logger.info(self._log_entry(result), extra=self._log_extra(result))
-                continue
+            with self._current_job_lock:
+                if self.__cancelled:
+                    # The workflow was cancelled before this job started
+                    result = WorkflowJobResult(
+                        name=job.name,
+                        index=index,
+                        arguments=[str(arg) for arg in args],
+                        stdout="",
+                        stderr="",
+                        status=WorkflowJobStatus.CANCELLED,
+                    )
+                    self.__job_results.append(result)
+                    logger.info(self._log_entry(result), extra=self._log_extra(result))
+                    continue
+
+                jobrunner = WorkflowJobRunner(job)
+                self.__current_job = jobrunner
 
-            jobrunner = WorkflowJobRunner(job)
-            self.__current_job = jobrunner
             logger.info(
                 f"Workflow job starting; {self._job_description(jobrunner.name, index)}"
             )
@@ -278,11 +295,11 @@ def isCancelled(self) -> bool:
         return self.__cancelled
 
     def cancel(self) -> None:
-        self.__cancelled = True
-        if self.__current_job is not None:
-            self.__current_job.cancel()
-        if self.isRunning() or self._workflow_job is not None:
-            self.wait()
+        with self._current_job_lock:
+            self.__cancelled = True
+            current_job = self.__current_job
+        if current_job is not None:
+            current_job.cancel()
 
     def exception(self) -> BaseException | None:
         if self._workflow_job is not None:

tests/ert/unit_tests/plugins/test_export_rft.py

@@ -167,3 +167,40 @@ def test_that_export_rft_uses_custom_filename():
         assert output_file.exists()
 
         assert pl.read_csv(output_file)["pressure"][0] == responses_real0["values"][0]
+
+
+@pytest.mark.usefixtures("use_tmpdir")
+def test_that_export_rft_stops_writing_files_once_cancelled(monkeypatch):
+    runpath_values = [
+        (Path("real0"), _create_rft_response_df()),
+        (Path("real1"), _create_rft_response_df(value=152.0)),
+    ]
+    for rp, _ in runpath_values:
+        rp.mkdir()
+
+    job = ExportRFTJob()
+    written_files = []
+    original_write_csv = pl.DataFrame.write_csv
+
+    def cancel_after_first_write(self, path, *args, **kwargs):
+        written_files.append(path)
+        job.cancel()
+        return original_write_csv(self, path, *args, **kwargs)
+
+    monkeypatch.setattr(pl.DataFrame, "write_csv", cancel_after_first_write)
+
+    with _create_rft_ensemble(ensemble_size=2) as ensemble:
+        for i, (_, response) in enumerate(runpath_values):
+            ensemble.save_response("rft", response, i)
+            ensemble.save_observation_location_metadata(
+                _create_rft_location_metadata_df(), i
+            )
+
+        job.run(_mock_runpath([str(rp) for rp, _ in runpath_values]), ensemble, [])
+
+    # Cancellation is noticed at the top of the next iteration, so the
+    # realization being written when cancel() was called still completes,
+    # but the one after it is skipped entirely.
+    assert len(written_files) == 1
+    assert (runpath_values[0][0] / "share/results/tables/rft_ert.csv").exists()
+    assert not (runpath_values[1][0] / "share/results/tables/rft_ert.csv").exists()

tests/ert/unit_tests/resources/test_csv_export.py

@@ -4,6 +4,7 @@
 
 from ert.plugins.hook_implementations.workflows.csv_export import CSVExportJob
 from ert.storage import open_storage
+from ert.storage.local_ensemble import LocalEnsemble
 from tests.ert.performance_tests.test_obs_and_responses_performance import (
     create_experiment_args,
 )
@@ -50,3 +51,61 @@ def test_that_csv_export_matches_snapshot(monkeypatch, tmp_path, snapshot):
         snapshot.assert_match(
             df.write_csv(include_header=True), "csv_export_result.csv"
         )
+
+
+def test_that_csv_export_stops_processing_ensembles_once_cancelled(
+    monkeypatch, tmp_path
+):
+    monkeypatch.chdir(tmp_path)
+    num_realizations = 2
+    info = create_experiment_args(
+        num_parameters=1,
+        num_gen_data_keys=1,
+        num_gen_data_report_steps=1,
+        num_gen_data_index=1,
+        num_gen_data_obs=1,
+        num_summary_keys=1,
+        num_summary_timesteps=1,
+        num_summary_obs=1,
+        num_realizations=num_realizations,
+    )
+
+    job = CSVExportJob()
+    processed_ensembles = []
+    original_has_data = LocalEnsemble.has_data
+
+    def tracking_has_data(self):
+        processed_ensembles.append(self.name)
+        job.cancel()
+        return original_has_data(self)
+
+    monkeypatch.setattr(LocalEnsemble, "has_data", tracking_has_data)
+
+    with open_storage(tmp_path / "storage", mode="w") as storage:
+        experiment = storage.create_experiment(
+            experiment_config={
+                "response_configuration": [info.gen_data_config, info.summary_config],
+                "parameter_configuration": info.gen_kw_configs,
+                "observations": info.gen_data_observations + info.summary_observations,
+            }
+        )
+        ensembles = [
+            experiment.create_ensemble(ensemble_size=num_realizations, name=name)
+            for name in ("first", "second")
+        ]
+        for ens in ensembles:
+            for real in range(num_realizations):
+                ens.save_response("summary", info.summary_responses.clone(), real)
+                ens.save_response("gen_data", info.gen_data_responses.clone(), real)
+            ens.save_parameters(dataset=info.genkw_data)
+
+        ensemble_list_json = json.dumps([str(ens.id) for ens in ensembles])
+        job.run(storage, ["the_export.csv", ensemble_list_json])
+
+    # Cancellation is noticed at the top of the next iteration, so the
+    # ensemble being processed when cancel() was called still completes,
+    # but the one after it is skipped entirely.
+    assert processed_ensembles == ["first"]
+
+    df = pl.read_csv("the_export.csv")
+    assert set(df["Ensemble"]) == {"first"}