From 786c00e32d287700d3bd47dc9a397816b099c353 Mon Sep 17 00:00:00 2001 From: Jonathan Karlsen Date: Thu, 5 Dec 2024 10:31:59 +0100 Subject: [PATCH] Fix progress bar not updating realization count for new iterations This commit fixes the bug introduced in 31e607b066ab79415671f83f2d57c7400c4d4e98, where the status reporting in GUI was done the same way when rerunning failed realizations, and running new iterations. This is an issue because when rerunning failed realizations, we want to show all realizations and add the finished/failed count from the previous run, while new iterations should drop the failed realizations altogether. --- src/ert/run_models/base_run_model.py | 12 +- src/ert/run_models/ensemble_experiment.py | 1 + src/ert/run_models/ensemble_smoother.py | 1 + src/ert/run_models/evaluate_ensemble.py | 1 + src/ert/run_models/everest_run_model.py | 1 + .../run_models/iterated_ensemble_smoother.py | 2 +- .../run_models/multiple_data_assimilation.py | 1 + .../run_models/test_base_run_model.py | 129 +++++++++++++++--- 8 files changed, 123 insertions(+), 25 deletions(-) diff --git a/src/ert/run_models/base_run_model.py b/src/ert/run_models/base_run_model.py index f587313cee5..6943812f6f3 100644 --- a/src/ert/run_models/base_run_model.py +++ b/src/ert/run_models/base_run_model.py @@ -188,6 +188,7 @@ def __init__( self.minimum_required_realizations = minimum_required_realizations self.active_realizations = copy.copy(active_realizations) self.start_iteration = start_iteration + self.restart = False self.validate_active_realizations_count() def log_at_startup(self) -> None: @@ -404,7 +405,10 @@ def get_current_status(self) -> dict[str, int]: for real in all_realizations.values(): status[str(real["status"])] += 1 - status["Finished"] += self._get_number_of_finished_realizations_from_reruns() + if self.restart: + status["Finished"] += ( + self._get_number_of_finished_realizations_from_reruns() + ) return status def _get_number_of_finished_realizations_from_reruns(self) -> int: @@ -644,7 +648,11 @@ def get_number_of_existing_runpaths(self) -> int: return [real_path.exists() for real_path in realization_set].count(True) def get_number_of_active_realizations(self) -> int: - return self._initial_realizations_mask.count(True) + return ( + self._initial_realizations_mask.count(True) + if self.restart + else self.active_realizations.count(True) + ) def get_number_of_successful_realizations(self) -> int: return self.active_realizations.count(True) diff --git a/src/ert/run_models/ensemble_experiment.py b/src/ert/run_models/ensemble_experiment.py index 13d8a3a8411..7fa54a454c7 100644 --- a/src/ert/run_models/ensemble_experiment.py +++ b/src/ert/run_models/ensemble_experiment.py @@ -63,6 +63,7 @@ def run_experiment( restart: bool = False, ) -> None: self.log_at_startup() + self.restart = restart if not restart: self.experiment = self._storage.create_experiment( name=self.experiment_name, diff --git a/src/ert/run_models/ensemble_smoother.py b/src/ert/run_models/ensemble_smoother.py index b745b2df0a5..c27c6de0481 100644 --- a/src/ert/run_models/ensemble_smoother.py +++ b/src/ert/run_models/ensemble_smoother.py @@ -62,6 +62,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() + self.restart = restart ensemble_format = self.target_ensemble_format experiment = self._storage.create_experiment( parameters=self.ert_config.ensemble_config.parameter_configuration, diff --git a/src/ert/run_models/evaluate_ensemble.py b/src/ert/run_models/evaluate_ensemble.py index aa67f42b07c..ffe6a8c55e8 100644 --- a/src/ert/run_models/evaluate_ensemble.py +++ b/src/ert/run_models/evaluate_ensemble.py @@ -64,6 +64,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() + self.restart = restart ensemble = self.ensemble experiment = ensemble.experiment self.set_env_key("_ERT_EXPERIMENT_ID", str(experiment.id)) diff --git a/src/ert/run_models/everest_run_model.py b/src/ert/run_models/everest_run_model.py index a92d8cf10a3..f2cf9a4c94f 100644 --- a/src/ert/run_models/everest_run_model.py +++ b/src/ert/run_models/everest_run_model.py @@ -391,6 +391,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() + self.restart = restart simulator = Simulator( self.everest_config, self.ert_config, diff --git a/src/ert/run_models/iterated_ensemble_smoother.py b/src/ert/run_models/iterated_ensemble_smoother.py index 14f57e97c45..6eda13e40df 100644 --- a/src/ert/run_models/iterated_ensemble_smoother.py +++ b/src/ert/run_models/iterated_ensemble_smoother.py @@ -122,7 +122,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() - + self.restart = restart target_ensemble_format = self.target_ensemble_format experiment = self._storage.create_experiment( parameters=self.ert_config.ensemble_config.parameter_configuration, diff --git a/src/ert/run_models/multiple_data_assimilation.py b/src/ert/run_models/multiple_data_assimilation.py index ac30649c6dc..800af2a1161 100644 --- a/src/ert/run_models/multiple_data_assimilation.py +++ b/src/ert/run_models/multiple_data_assimilation.py @@ -85,6 +85,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() + self.restart = restart if self.restart_run: id = self.prior_ensemble_id try: diff --git a/tests/ert/unit_tests/run_models/test_base_run_model.py b/tests/ert/unit_tests/run_models/test_base_run_model.py index 1653379d4d3..00b062f4979 100644 --- a/tests/ert/unit_tests/run_models/test_base_run_model.py +++ b/tests/ert/unit_tests/run_models/test_base_run_model.py @@ -202,6 +202,49 @@ def test_num_cpu_is_propagated_from_config_to_ensemble(run_args): assert ensemble.reals[1].num_cpu == 42 +@pytest.mark.parametrize( + "real_status_dict, expected_result", + [ + pytest.param( + {"0": "Finished", "1": "Finished", "2": "Finished"}, + {"Finished": 3}, + id="ran_all_realizations_and_all_succeeded", + ), + pytest.param( + {"0": "Finished", "1": "Finished", "2": "Failed"}, + {"Finished": 2, "Failed": 1}, + id="ran_all_realizations_and_some_failed", + ), + pytest.param( + {"0": "Finished", "1": "Running", "2": "Failed"}, + {"Finished": 1, "Failed": 1, "Running": 1}, + id="ran_all_realizations_and_result_was_mixed", + ), + ], +) +def test_get_current_status( + real_status_dict, + expected_result, +): + config = ErtConfig.from_file_contents("NUM_REALIZATIONS 3") + initial_active_realizations = [True] * 3 + new_active_realizations = [True] * 3 + brm = BaseRunModel( + config=config, + storage=MagicMock(spec=Storage), + queue_config=config.queue_config, + status_queue=MagicMock(spec=SimpleQueue), + active_realizations=initial_active_realizations, + ) + snapshot_dict_reals = {} + for index, realization_status in real_status_dict.items(): + snapshot_dict_reals[index] = {"status": realization_status} + iter_snapshot = EnsembleSnapshot.from_nested_dict({"reals": snapshot_dict_reals}) + brm._iter_snapshot[0] = iter_snapshot + brm.active_realizations = new_active_realizations + assert dict(brm.get_current_status()) == expected_result + + @pytest.mark.parametrize( "initial_active_realizations, new_active_realizations, real_status_dict, expected_result", [ @@ -240,27 +283,6 @@ def test_num_cpu_is_propagated_from_config_to_ensemble(run_args): {"Finished": 0}, id="did_not_run_all_realizations_and_none_succeeded", ), - pytest.param( - [True, True, True], - [True, True, True], - {"0": "Finished", "1": "Finished", "2": "Finished"}, - {"Finished": 3}, - id="ran_all_realizations_and_all_succeeded", - ), - pytest.param( - [True, True, True], - [True, True, True], - {"0": "Finished", "1": "Finished", "2": "Failed"}, - {"Finished": 2, "Failed": 1}, - id="ran_all_realizations_and_some_failed", - ), - pytest.param( - [True, True, True], - [True, True, True], - {"0": "Finished", "1": "Running", "2": "Failed"}, - {"Finished": 1, "Failed": 1, "Running": 1}, - id="ran_all_realizations_and_result_was_mixed", - ), pytest.param( [True, True, True], [True, True, False], @@ -277,7 +299,7 @@ def test_num_cpu_is_propagated_from_config_to_ensemble(run_args): ), ], ) -def test_get_current_status( +def test_get_current_status_when_rerun( initial_active_realizations, new_active_realizations, real_status_dict, @@ -292,6 +314,7 @@ def test_get_current_status( status_queue=MagicMock(spec=SimpleQueue), active_realizations=initial_active_realizations, ) + brm.restart = True snapshot_dict_reals = {} for index, realization_status in real_status_dict.items(): snapshot_dict_reals[index] = {"status": realization_status} @@ -299,3 +322,65 @@ def test_get_current_status( brm._iter_snapshot[0] = iter_snapshot brm.active_realizations = new_active_realizations assert dict(brm.get_current_status()) == expected_result + + +def test_get_current_status_for_new_iteration_when_realization_failed_in_previous_run(): + """Active realizations gets changed when we run next iteration, and the failed realizations from + the previous run should not be present in the current_status.""" + initial_active_realizations = [True] * 5 + # Realization 0,1, and 3 failed in the previous iteration + new_active_realizations = [False, False, True, False, True] + config = ErtConfig.from_file_contents("NUM_REALIZATIONS 5") + brm = BaseRunModel( + config=config, + storage=MagicMock(spec=Storage), + queue_config=config.queue_config, + status_queue=MagicMock(spec=SimpleQueue), + active_realizations=initial_active_realizations, + ) + snapshot_dict_reals = { + "2": {"status": "Running"}, + "4": {"status": "Finished"}, + } + iter_snapshot = EnsembleSnapshot.from_nested_dict({"reals": snapshot_dict_reals}) + brm._iter_snapshot[0] = iter_snapshot + brm.active_realizations = new_active_realizations + + assert brm.restart is False + assert dict(brm.get_current_status()) == {"Running": 1, "Finished": 1} + + +@pytest.mark.parametrize( + "new_active_realizations, was_rerun, expected_result", + [ + pytest.param( + [False, False, False, True, False], + True, + 5, + id="rerun_so_total_realization_count_is_not_affected_by_previous_failed_realizations", + ), + pytest.param( + [True, True, False, False, False], + False, + 2, + id="new_iteration_so_total_realization_count_is_only_previously_successful_realizations", + ), + ], +) +def test_get_number_of_active_realizations_varies_when_rerun_or_new_iteration( + new_active_realizations, was_rerun, expected_result +): + """When rerunning, we include all realizations in the total amount of active realization. + When running a new iteration based on the result of the previous iteration, we only include the successful realizations.""" + initial_active_realizations = [True] * 5 + config = ErtConfig.from_file_contents("NUM_REALIZATIONS 5") + brm = BaseRunModel( + config=config, + storage=MagicMock(spec=Storage), + queue_config=config.queue_config, + status_queue=MagicMock(spec=SimpleQueue), + active_realizations=initial_active_realizations, + ) + brm.active_realizations = new_active_realizations + brm.restart = was_rerun + assert brm.get_number_of_active_realizations() == expected_result