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