diff --git a/docs/ert/reference/workflows/workflows.rst b/docs/ert/reference/workflows/workflows.rst index 99c7a0a9e37..113d15b4cca 100644 --- a/docs/ert/reference/workflows/workflows.rst +++ b/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 ` + +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()`. diff --git a/src/ert/gui/ertwidgets/closabledialog.py b/src/ert/gui/ertwidgets/closabledialog.py index 82050dda2a3..f5872d04354 100644 --- a/src/ert/gui/ertwidgets/closabledialog.py +++ b/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 diff --git a/src/ert/gui/tools/workflows/run_workflow_widget.py b/src/ert/gui/tools/workflows/run_workflow_widget.py index c64700e0641..8791b504bdc 100644 --- a/src/ert/gui/tools/workflows/run_workflow_widget.py +++ b/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 diff --git a/src/ert/gui/tools/workflows/workflow_dialog.py b/src/ert/gui/tools/workflows/workflow_dialog.py index 8e1ded1835d..9ccad28ce01 100644 --- a/src/ert/gui/tools/workflows/workflow_dialog.py +++ b/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) diff --git a/src/ert/gui/tools/workflows/workflows_tool.py b/src/ert/gui/tools/workflows/workflows_tool.py index 8d3e3195dcc..b22c4195616 100644 --- a/src/ert/gui/tools/workflows/workflows_tool.py +++ b/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. diff --git a/src/ert/plugins/hook_implementations/workflows/csv_export.py b/src/ert/plugins/hook_implementations/workflows/csv_export.py index f14ccb54f56..aecb06386d9 100644 --- a/src/ert/plugins/hook_implementations/workflows/csv_export.py +++ b/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!" diff --git a/src/ert/plugins/hook_implementations/workflows/export_rft.py b/src/ert/plugins/hook_implementations/workflows/export_rft.py index 82a6ef23567..89cdf649b7a 100644 --- a/src/ert/plugins/hook_implementations/workflows/export_rft.py +++ b/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") diff --git a/src/ert/workflow_runner.py b/src/ert/workflow_runner.py index d8e6961cf13..56d0610cce5 100644 --- a/src/ert/workflow_runner.py +++ b/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: diff --git a/tests/ert/unit_tests/plugins/test_export_rft.py b/tests/ert/unit_tests/plugins/test_export_rft.py index 2b572234e33..a3c04f4dd45 100644 --- a/tests/ert/unit_tests/plugins/test_export_rft.py +++ b/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() diff --git a/tests/ert/unit_tests/resources/test_csv_export.py b/tests/ert/unit_tests/resources/test_csv_export.py index 1c3f7c54e11..3a3b4130885 100644 --- a/tests/ert/unit_tests/resources/test_csv_export.py +++ b/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"} diff --git a/tests/ert/unit_tests/workflow_runner/test_workflow_runner.py b/tests/ert/unit_tests/workflow_runner/test_workflow_runner.py index d18c9a04e40..21ef1033800 100644 --- a/tests/ert/unit_tests/workflow_runner/test_workflow_runner.py +++ b/tests/ert/unit_tests/workflow_runner/test_workflow_runner.py @@ -1,4 +1,6 @@ import logging +import threading +import time from pathlib import Path from textwrap import dedent from unittest.mock import patch @@ -517,3 +519,79 @@ def test_that_job_runner_stops_reporting_it_is_running_when_arguments_are_reject runner.run([1]) assert not runner.isRunning() + + +@pytest.mark.slow +@pytest.mark.usefixtures("use_tmpdir") +@pytest.mark.filterwarnings("ignore:.*Deprecated keywords, SCRIPT and INTERNAL") +def test_that_cancel_does_not_block_on_an_uncooperative_internal_job(): + WorkflowCommon.createUncancellableWaitJob() + + wait_job = workflow_job_from_file( + "uncancellable_wait_job", name="UNCANCELLABLE_WAIT", origin="user" + ) + workflow = Workflow.from_file( + "uncancellable_wait_workflow", {}, {"UNCANCELLABLE_WAIT": wait_job} + ) + + workflow_runner = WorkflowRunner(workflow, fixtures={}) + + workflow_runner.run() + wait_until(lambda: Path("uncancellable_wait_started_0").exists()) + + start = time.time() + workflow_runner.cancel() + elapsed = time.time() - start + + # cancel() cannot forcibly interrupt an internal job that never checks + # isCancelled(); it must return immediately regardless, rather than + # blocking until the uncooperative job finishes on its own (5 seconds). + assert elapsed < 1 + assert workflow_runner.isCancelled() + + wait_until(lambda: Path("uncancellable_wait_finished_0").exists(), timeout=10) + workflow_runner.wait() + + +@pytest.mark.usefixtures("use_tmpdir") +@pytest.mark.filterwarnings("ignore:.*Deprecated keywords, SCRIPT and INTERNAL") +def test_that_a_cancel_call_racing_with_job_startup_still_cancels_the_job( + monkeypatch, +): + """Regression test for a race between cancel() and a job publishing + itself as the current job but not yet having created its ErtScript. + + cancel() reaching WorkflowJobRunner in that window must not be a + silent no-op: the cancellation has to be applied as soon as the + script is created. + """ + WorkflowCommon.createWaitJob() + wait_job = workflow_job_from_file("wait_job", name="WAIT", origin="user") + workflow = Workflow.from_file("wait_workflow", {}, {"WAIT": wait_job}) + + job_runner_published = threading.Event() + may_create_script = threading.Event() + original_run = WorkflowJobRunner.run + + def delayed_run(self, *args, **kwargs): + # At this point the WorkflowRunner has already published this + # WorkflowJobRunner as its current job, but self.__script does + # not exist yet. + job_runner_published.set() + assert may_create_script.wait(timeout=10) + return original_run(self, *args, **kwargs) + + monkeypatch.setattr(WorkflowJobRunner, "run", delayed_run) + + workflow_runner = WorkflowRunner(workflow, fixtures={}) + workflow_runner.run() + + assert job_runner_published.wait(timeout=10) + workflow_runner.cancel() + may_create_script.set() + + wait_until(lambda: Path("wait_cancelled_0").exists(), timeout=10) + workflow_runner.wait() + + assert workflow_runner.isCancelled() + assert not Path("wait_started_1").exists() diff --git a/tests/ert/unit_tests/workflow_runner/workflow_common.py b/tests/ert/unit_tests/workflow_runner/workflow_common.py index 427599fdd84..a77bcd96f1f 100644 --- a/tests/ert/unit_tests/workflow_runner/workflow_common.py +++ b/tests/ert/unit_tests/workflow_runner/workflow_common.py @@ -126,3 +126,40 @@ def createWaitJob(): Path("fast_wait_workflow").write_text( "WAIT 0 1\nEXTERNAL_WAIT 1 1\n", encoding="utf-8" ) + + @staticmethod + def createUncancellableWaitJob(): + Path("uncancellable_wait_job.py").write_text( + "import time\n" + "from pathlib import Path\n" + "\n" + "from ert import ErtScript\n" + "\n" + "class UncancellableWaitScript(ErtScript):\n" + " def dump(self, filename, content):\n" + " Path(filename).write_text(content, encoding='utf-8')\n" + "\n" + " def run(self, *argv):\n" + " number, wait_time = argv\n" + " self.dump('uncancellable_wait_started_%d' % number, 'text')\n" + # Deliberately does not poll isCancelled(): this simulates an + # internal job that has not adopted the cancellation contract. + " time.sleep(wait_time)\n" + " self.dump('uncancellable_wait_finished_%d' % number, 'text')\n" + " return None\n", + encoding="utf-8", + ) + + Path("uncancellable_wait_job").write_text( + "INTERNAL True\n" + "SCRIPT uncancellable_wait_job.py\n" + "MIN_ARG 2\n" + "MAX_ARG 2\n" + "ARG_TYPE 0 INT\n" + "ARG_TYPE 1 INT\n", + encoding="utf-8", + ) + + Path("uncancellable_wait_workflow").write_text( + "UNCANCELLABLE_WAIT 0 5\n", encoding="utf-8" + )