Remove analysis config from base_run_model

sondreso · Nov 17, 2023 · 5ad54fd · 5ad54fd
1 parent f83a99e
commit 5ad54fd
Show file tree

Hide file tree

Showing 15 changed files with 56 additions and 49 deletions.
diff --git a/src/ert/cli/model_factory.py b/src/ert/cli/model_factory.py
@@ -104,6 +104,7 @@ def _setup_single_test_run(
             current_case=args.current_case,
             minimum_required_realizations=1,
             ensemble_size=1,
+            stop_long_running=config.analysis_config.stop_long_running,
         ),
         config,
         storage,
@@ -133,6 +134,7 @@ def _setup_ensemble_experiment(
             iter_num=int(args.iter_num),
             minimum_required_realizations=config.analysis_config.minimum_required_realizations,
             ensemble_size=config.model_config.num_realizations,
+            stop_long_running=config.analysis_config.stop_long_running,
         ),
         config,
         storage,
@@ -158,6 +160,7 @@ def _setup_ensemble_smoother(
             target_case=args.target_case,
             minimum_required_realizations=config.analysis_config.minimum_required_realizations,
             ensemble_size=config.model_config.num_realizations,
+            stop_long_running=config.analysis_config.stop_long_running,
         ),
         config,
         storage,
@@ -195,6 +198,7 @@ def _setup_multiple_data_assimilation(
             prior_ensemble=prior_ensemble,
             minimum_required_realizations=config.analysis_config.minimum_required_realizations,
             ensemble_size=config.model_config.num_realizations,
+            stop_long_running=config.analysis_config.stop_long_running,
         ),
         config,
         storage,
@@ -225,6 +229,7 @@ def _setup_iterative_ensemble_smoother(
             minimum_required_realizations=config.analysis_config.minimum_required_realizations,
             ensemble_size=config.model_config.num_realizations,
             num_retries_per_iter=config.analysis_config.num_retries_per_iter,
+            stop_long_running=config.analysis_config.stop_long_running,
         ),
         config,
         storage,

diff --git a/src/ert/ensemble_evaluator/_builder/_ensemble_builder.py b/src/ert/ensemble_evaluator/_builder/_ensemble_builder.py
@@ -2,15 +2,15 @@
 
 import copy
 import logging
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional
 
 from typing_extensions import Self
 
 from ._legacy import LegacyEnsemble
 from ._realization import RealizationBuilder
 
 if TYPE_CHECKING:
-    from ert.config import AnalysisConfig, QueueConfig
+    from ert.config import QueueConfig
 
     from ._ensemble import Ensemble
 
@@ -22,10 +22,9 @@ def __init__(self) -> None:
         self._reals: List[RealizationBuilder] = []
         self._forward_model: Optional[RealizationBuilder] = None
         self._size: int = 0
-        self._legacy_dependencies: Optional[
-            Tuple["QueueConfig", "AnalysisConfig"]
-        ] = None
-
+        self._legacy_dependencies: Optional["QueueConfig"] = None
+        self.stop_long_running = False
+        self.num_required_realizations = 0
         self._custom_port_range: Optional[range] = None
         self._max_running = 10000
         self._id: Optional[str] = None
@@ -52,9 +51,14 @@ def set_ensemble_size(self, size: int) -> Self:
         return self
 
     def set_legacy_dependencies(
-        self, queue_config: QueueConfig, analysis_config: AnalysisConfig
+        self,
+        queue_config: QueueConfig,
+        stop_long_running: bool,
+        num_required_realizations: int,
     ) -> Self:
-        self._legacy_dependencies = (queue_config, analysis_config)
+        self._legacy_dependencies = queue_config
+        self.stop_long_running = stop_long_running
+        self.num_required_realizations = num_required_realizations
         return self
 
     def set_custom_port_range(self, custom_port_range: range) -> Self:
@@ -93,4 +97,11 @@ def build(self) -> Ensemble:
 
         reals = [builder.build() for builder in real_builders]
 
-        return LegacyEnsemble(reals, {}, *self._legacy_dependencies, id_=self._id)
+        return LegacyEnsemble(
+            reals,
+            {},
+            self._legacy_dependencies,
+            self.stop_long_running,
+            self.num_required_realizations,
+            id_=self._id,
+        )
diff --git a/src/ert/ensemble_evaluator/_builder/_legacy.py b/src/ert/ensemble_evaluator/_builder/_legacy.py
@@ -17,7 +17,7 @@
 from ._ensemble import Ensemble
 
 if TYPE_CHECKING:
-    from ert.config import AnalysisConfig, QueueConfig
+    from ert.config import QueueConfig
 
     from ..config import EvaluatorServerConfig
     from ._realization import Realization
@@ -34,16 +34,16 @@ def __init__(
         reals: List[Realization],
         metadata: Dict[str, Any],
         queue_config: QueueConfig,
-        analysis_config: AnalysisConfig,
+        stop_long_running: bool,
+        min_required_realizations: int,
         id_: str,
     ) -> None:
         super().__init__(reals, metadata, id_)
         if not queue_config:
             raise ValueError(f"{self} needs queue_config")
-        if not analysis_config:
-            raise ValueError(f"{self} needs analysis_config")
         self._job_queue = JobQueue(queue_config)
-        self._analysis_config = analysis_config
+        self.stop_long_running = stop_long_running
+        self.min_required_realizations = min_required_realizations
         self._config: Optional[EvaluatorServerConfig] = None
 
     def generate_event_creator(
@@ -183,14 +183,11 @@ async def _evaluate_inner(  # pylint: disable=too-many-branches
             # commands to the task in order to have it killed/retried.
             # See https://github.com/equinor/ert/issues/1229
             queue_evaluators = None
-            if (
-                self._analysis_config.stop_long_running
-                and self._analysis_config.minimum_required_realizations > 0
-            ):
+            if self.stop_long_running and self.min_required_realizations > 0:
                 queue_evaluators = [
                     partial(
                         self._job_queue.stop_long_running_jobs,
-                        self._analysis_config.minimum_required_realizations,
+                        self.min_required_realizations,
                     )
                 ]
 

diff --git a/src/ert/run_models/base_run_model.py b/src/ert/run_models/base_run_model.py
@@ -359,7 +359,8 @@ def _build_ensemble(
     ) -> "Ensemble":
         builder = EnsembleBuilder().set_legacy_dependencies(
             self._queue_config,
-            self.ert_config.analysis_config,
+            self._simulation_arguments.stop_long_running,
+            self._simulation_arguments.minimum_required_realizations,
         )
 
         for iens, run_arg in enumerate(run_context):

diff --git a/src/ert/run_models/ensemble_smoother.py b/src/ert/run_models/ensemble_smoother.py
@@ -66,7 +66,7 @@ def run_experiment(
     ) -> RunContext:
         self.checkHaveSufficientRealizations(
             self._simulation_arguments.active_realizations.count(True),
-            self.update_settings.min_required_realizations
+            self._simulation_arguments.minimum_required_realizations,
         )
 
         log_msg = "Running ES"

diff --git a/src/ert/run_models/iterated_ensemble_smoother.py b/src/ert/run_models/iterated_ensemble_smoother.py
@@ -103,7 +103,7 @@ def run_experiment(
     ) -> RunContext:
         self.checkHaveSufficientRealizations(
             self._simulation_arguments.active_realizations.count(True),
-            self.update_settings.min_required_realizations
+            self._simulation_arguments.minimum_required_realizations,
         )
         iteration_count = self.facade.get_number_of_iterations()
         phase_count = iteration_count + 1

diff --git a/src/ert/run_models/multiple_data_assimilation.py b/src/ert/run_models/multiple_data_assimilation.py
@@ -68,7 +68,7 @@ def run_experiment(
     ) -> RunContext:
         self.checkHaveSufficientRealizations(
             self._simulation_arguments.active_realizations.count(True),
-            self.update_settings.min_required_realizations
+            self._simulation_arguments.minimum_required_realizations,
         )
         weights = self.parseWeights(self._simulation_arguments.weights)
 

diff --git a/src/ert/run_models/run_arguments.py b/src/ert/run_models/run_arguments.py
@@ -7,6 +7,7 @@ class SimulationArguments:
     random_seed: Optional[int]
     minimum_required_realizations: int
     ensemble_size: int
+    stop_long_running: bool
 
 
 @dataclass

diff --git a/tests/unit_tests/cli/test_model_factory.py b/tests/unit_tests/cli/test_model_factory.py
@@ -1,5 +1,6 @@
 import dataclasses
 from argparse import Namespace
+from unittest.mock import MagicMock
 from uuid import UUID
 
 import pytest
@@ -65,7 +66,7 @@ def test_setup_single_test_run(poly_case, storage):
     )
     assert isinstance(model, SingleTestRun)
     sim_args_as_dict = dataclasses.asdict(model._simulation_arguments)
-    assert len(sim_args_as_dict) == 5
+    assert len(sim_args_as_dict) == 6
 
 
 def test_setup_ensemble_experiment(poly_case, storage):
@@ -85,7 +86,7 @@ def test_setup_ensemble_experiment(poly_case, storage):
     assert isinstance(model, EnsembleExperiment)
 
     sim_args_as_dict = dataclasses.asdict(model._simulation_arguments)
-    assert len(sim_args_as_dict) == 8
+    assert len(sim_args_as_dict) == 9
     assert "active_realizations" in sim_args_as_dict
 
 
@@ -98,10 +99,7 @@ def test_setup_ensemble_smoother(poly_case, storage):
     )
 
     model = model_factory._setup_ensemble_smoother(
-        poly_case,
-        storage,
-        args,
-        UUID(int=0),
+        poly_case, storage, args, UUID(int=0), MagicMock()
     )
     assert isinstance(model, EnsembleSmoother)
     sim_args_as_dict = dataclasses.asdict(model._simulation_arguments)
@@ -121,10 +119,7 @@ def test_setup_multiple_data_assimilation(poly_case, storage):
     )
 
     model = model_factory._setup_multiple_data_assimilation(
-        poly_case,
-        storage,
-        args,
-        UUID(int=0),
+        poly_case, storage, args, UUID(int=0), MagicMock()
     )
     assert isinstance(model, MultipleDataAssimilation)
     sim_args_as_dict = dataclasses.asdict(model._simulation_arguments)
@@ -144,10 +139,7 @@ def test_setup_iterative_ensemble_smoother(poly_case, storage):
     )
 
     model = model_factory._setup_iterative_ensemble_smoother(
-        poly_case,
-        storage,
-        args,
-        UUID(int=0),
+        poly_case, storage, args, UUID(int=0), MagicMock()
     )
     assert isinstance(model, IteratedEnsembleSmoother)
     sim_args_as_dict = dataclasses.asdict(model._simulation_arguments)

diff --git a/tests/unit_tests/cli/test_model_hook_order.py b/tests/unit_tests/cli/test_model_hook_order.py
@@ -56,6 +56,7 @@ def test_hook_call_order_ensemble_smoother(monkeypatch):
         target_case="smooth",
         minimum_required_realizations=0,
         ensemble_size=1,
+        stop_long_running=False,
     )
     test_class = EnsembleSmoother(
         minimum_args,
@@ -92,6 +93,7 @@ def test_hook_call_order_es_mda(monkeypatch):
         prior_ensemble="",
         minimum_required_realizations=0,
         ensemble_size=1,
+        stop_long_running=False,
     )
     monkeypatch.setattr(multiple_data_assimilation, "sample_prior", MagicMock())
     monkeypatch.setattr(multiple_data_assimilation, "smoother_update", MagicMock())
@@ -151,6 +153,7 @@ def test_hook_call_order_iterative_ensemble_smoother(monkeypatch):
         num_retries_per_iter=1,
         minimum_required_realizations=0,
         ensemble_size=1,
+        stop_long_running=True,
     )
     test_class = IteratedEnsembleSmoother(
         minimum_args,

diff --git a/tests/unit_tests/cli/test_run_context.py b/tests/unit_tests/cli/test_run_context.py
@@ -25,6 +25,7 @@ def test_that_all_iterations_gets_correct_name_and_iteration_number(
         prior_ensemble="",
         minimum_required_realizations=1,
         ensemble_size=1,
+        stop_long_running=True,
     )
     ens_mock = MagicMock()
     ens_mock.iteration = 0

diff --git a/tests/unit_tests/ensemble_evaluator/conftest.py b/tests/unit_tests/ensemble_evaluator/conftest.py
@@ -139,17 +139,10 @@ def _make_ensemble_builder(monkeypatch, tmpdir, num_reals, num_jobs, job_sleep=0
                     )
                 )
 
-        analysis_config = Mock()
-        analysis_config.get_stop_long_running = Mock(return_value=False)
-        analysis_config.minimum_required_realizations = 0
-
         ecl_config = Mock()
         ecl_config.assert_restart = Mock()
 
-        builder.set_legacy_dependencies(
-            queue_config,
-            analysis_config,
-        )
+        builder.set_legacy_dependencies(queue_config, False, 0)
         builder.set_id("0")
         return builder
 

diff --git a/tests/unit_tests/ensemble_evaluator/test_ensemble_builder.py b/tests/unit_tests/ensemble_evaluator/test_ensemble_builder.py
@@ -13,9 +13,7 @@
 def test_build_ensemble(active_real):
     ensemble = (
         EnsembleBuilder()
-        .set_legacy_dependencies(
-            QueueConfig(queue_system=QueueSystem.LOCAL), MagicMock()
-        )
+        .set_legacy_dependencies(QueueConfig(queue_system=QueueSystem.LOCAL), False, 0)
         .add_realization(
             RealizationBuilder()
             .set_iens(2)

diff --git a/tests/unit_tests/run_models/test_base_run_model.py b/tests/unit_tests/run_models/test_base_run_model.py
@@ -14,7 +14,10 @@
 @pytest.fixture
 def base_arguments():
     return SimulationArguments(
-        random_seed=1234, minimum_required_realizations=0, ensemble_size=1
+        random_seed=1234,
+        minimum_required_realizations=0,
+        ensemble_size=1,
+        stop_long_running=False,
     )
 
 
@@ -103,6 +106,7 @@ def test_check_if_runpath_exists(
         iter_num=0,
         minimum_required_realizations=0,
         ensemble_size=1,
+        stop_long_running=False,
     )
 
     def get_run_path_mock(realizations, iteration=None):

diff --git a/tests/unit_tests/run_models/test_ensemble_experiment.py b/tests/unit_tests/run_models/test_ensemble_experiment.py
@@ -36,6 +36,7 @@ def test_check_if_runpath_exists(
         iter_num=iter_num,
         minimum_required_realizations=0,
         ensemble_size=1,
+        stop_long_running=False,
     )
 
     def get_run_path_mock(realizations, iteration=None):