Optimization should stop after n iteration (#34)

Hevagog · dJaniga · web-flow · commit 9b5d6db0ce45 · 2025-04-03T20:00:19.000+02:00
Co-authored-by: Damian Janiga &lt;janiga@agh.edu.pl&gt;
diff --git a/src/orchestration/risk_management_service/core/mappers/control_vector_mapper.py b/src/orchestration/risk_management_service/core/mappers/control_vector_mapper.py
@@ -1,5 +1,9 @@
-from services.problem_dispatcher_service import ControlVector as PDControlVector
-from services.solution_updater_service import ControlVector as SUControlVector
+from services.problem_dispatcher_service import (
+    ControlVector as PDControlVector,
+)
+from services.solution_updater_service import (
+    ControlVector as SUControlVector,
+)
 
 
 class ControlVectorMapper:
diff --git a/src/orchestration/risk_management_service/core/service/risk_management_service.py b/src/orchestration/risk_management_service/core/service/risk_management_service.py
@@ -3,19 +3,18 @@
 from logger.u_logger import get_logger
 from orchestration.risk_management_service.core.mappers import ControlVectorMapper
 from services.problem_dispatcher_service import (
-    ControlVector,
     ProblemDispatcherService,
     ServiceType,
 )
 from services.simulation_service import (
     SimulationService,
-    simulation_cluster_contex_manager,
+    simulation_cluster_context_manager,
 )
 from services.solution_updater_service import (
     OptimizationEngine,
     SolutionUpdaterService,
-    ensure_not_none,
 )
+from services.solution_updater_service.core.utils import ensure_not_none
 from services.well_management_service import WellManagementService
 
 logger = get_logger(__name__)
@@ -25,11 +24,13 @@ def run_risk_management(
     problem_definition: dict[str, Any],
     simulation_model_archive: bytes | str,
     n_size: int = 10,
+    max_generations: int = 10,
 ):
     """
     Main entry point for running risk management.
 
     Args:
+        max_generations: Maximum number of generations for the optimization process.
         problem_definition (dict[str, Any]): The problem definition used by the dispatcher.
         simulation_model_archive (bytes | str): The simulation model archive to transfer.
         n_size (int, optional): Number of samples for the dispatcher. Defaults to 10.
@@ -42,7 +43,7 @@ def run_risk_management(
         n_size,
     )
 
-    with simulation_cluster_contex_manager():
+    with simulation_cluster_context_manager():
         try:
             logger.info("Transferring simulation model archive to the cluster.")
             SimulationService.transfer_simulation_model(
@@ -53,7 +54,8 @@ def run_risk_management(
                 "Initializing SolutionUpdaterService and ProblemDispatcherService."
             )
             solution_updater = SolutionUpdaterService(
-                optimization_engine=OptimizationEngine.PSO
+                optimization_engine=OptimizationEngine.PSO,
+                max_generations=max_generations,
             )
             dispatcher = ProblemDispatcherService(
                 problem_definition=problem_definition, n_size=n_size
@@ -64,56 +66,66 @@ def run_risk_management(
             logger.debug("Boundaries retrieved: %s", boundaries)
 
             # Initialize solutions
-            next_solutions: list[ControlVector] | None = None
-
-            for iteration in range(5):
-                logger.info("Starting iteration %d for risk management.", iteration + 1)
+            next_solutions = None
+
+            loop_controller = solution_updater.loop_controller
+            try:
+                while loop_controller.running():
+                    logger.info(
+                        "Starting generation %d for risk management.",
+                        loop_controller.current_generation,
+                    )
 
-                # Generate or update solutions
-                solutions = dispatcher.process_iteration(next_solutions)
-                logger.debug("Generated solutions: %s", solutions)
+                    # Generate or update solutions
+                    solutions = dispatcher.process_iteration(next_solutions)
+                    logger.debug("Generated solutions: %s", solutions)
 
-                # Prepare simulation cases
-                sim_cases = _prepare_simulation_cases(solutions)
-                logger.debug("Prepared simulation cases: %s", sim_cases)
+                    # Prepare simulation cases
+                    sim_cases = _prepare_simulation_cases(solutions)
+                    logger.debug("Prepared simulation cases: %s", sim_cases)
 
-                # Process simulation with the simulation service
-                logger.info("Submitting simulation cases to SimulationService.")
-                completed_cases = SimulationService.process_request(
-                    {"simulation_cases": sim_cases}
-                )
-                logger.debug("Completed simulation cases: %s", completed_cases)
-
-                # Update solutions based on simulation results
-                updated_solutions = [
-                    {
-                        "control_vector": {"items": simulation_case.control_vector},
-                        "cost_function_results": {
-                            "values": ensure_not_none(
-                                simulation_case.results
-                            ).model_dump()
-                        },
-                    }
-                    for simulation_case in completed_cases.simulation_cases
-                ]
-                logger.debug(
-                    "Updated solutions for next iteration: %s", updated_solutions
-                )
+                    # Process simulation with the simulation service
+                    logger.info("Submitting simulation cases to SimulationService.")
+                    completed_cases = SimulationService.process_request(
+                        {"simulation_cases": sim_cases}
+                    )
+                    logger.debug("Completed simulation cases: %s", completed_cases)
 
-                # Map simulation service solutions to the ProblemDispatcherService format
-                next_solutions = ControlVectorMapper.convert_su_to_pd(
-                    solution_updater.process_request(
+                    # Update solutions based on simulation results
+                    updated_solutions = [
                         {
-                            "solution_candidates": updated_solutions,
-                            "optimization_constraints": {"boundaries": boundaries},
+                            "control_vector": {"items": simulation_case.control_vector},
+                            "cost_function_results": {
+                                "values": ensure_not_none(
+                                    simulation_case.results
+                                ).model_dump()
+                            },
                         }
-                    ).next_iter_solutions
-                )
+                        for simulation_case in completed_cases.simulation_cases
+                    ]
+                    logger.debug(
+                        "Updated solutions for next iteration: %s", updated_solutions
+                    )
+
+                    # Map simulation service solutions to the ProblemDispatcherService format
+                    next_solutions = ControlVectorMapper.convert_su_to_pd(
+                        solution_updater.process_request(
+                            {
+                                "solution_candidates": updated_solutions,
+                                "optimization_constraints": {"boundaries": boundaries},
+                            }
+                        ).next_iter_solutions
+                    )
+                    logger.info(
+                        "Generation %d successfully completed for risk management.",
+                        loop_controller.current_generation,
+                    )
+            except StopIteration as e:
                 logger.info(
-                    "Iteration %d successfully completed for risk management.",
-                    iteration + 1,
+                    "Loop controller stopped at generation %d: %s",
+                    loop_controller.current_generation,
+                    str(e),
                 )
-
         except Exception as e:
             logger.error("Error in risk management process: %s", str(e), exc_info=True)
             raise
diff --git a/src/services/simulation_service/__init__.py b/src/services/simulation_service/__init__.py
@@ -15,7 +15,7 @@
 
 from services.simulation_service.core.api import (  # noqa: F401, E402
     SimulationService,
-    simulation_cluster_contex_manager,
+    simulation_cluster_context_manager,
 )
 
-__all__ = ["SimulationService", "simulation_cluster_contex_manager"]
+__all__ = ["SimulationService", "simulation_cluster_context_manager"]
diff --git a/src/services/simulation_service/core/api.py b/src/services/simulation_service/core/api.py
@@ -1,6 +1,6 @@
 from services.simulation_service.core.service import (
     SimulationService,
-    simulation_cluster_contex_manager,
+    simulation_cluster_context_manager,
 )
 
-__all__ = ["SimulationService", "simulation_cluster_contex_manager"]
+__all__ = ["SimulationService", "simulation_cluster_context_manager"]
diff --git a/src/services/simulation_service/core/service/__init__.py b/src/services/simulation_service/core/service/__init__.py
@@ -5,7 +5,7 @@
 from logger.u_logger import configure_logger, get_logger
 from services.simulation_service.core.service.simulation_service import (
     SimulationService,
-    simulation_cluster_contex_manager,
+    simulation_cluster_context_manager,
 )
 
 configure_logger()
@@ -68,4 +68,4 @@ def _initialize_images() -> None:
 
 _initialize_images()
 
-__all__ = ["SimulationService", "simulation_cluster_contex_manager"]
+__all__ = ["SimulationService", "simulation_cluster_context_manager"]
diff --git a/src/services/simulation_service/core/service/simulation_service.py b/src/services/simulation_service/core/service/simulation_service.py
@@ -283,7 +283,7 @@ def _from_grpc(simulation: sm.Simulation) -> SimulationCase:
 
 
 @contextmanager
-def simulation_cluster_contex_manager():
+def simulation_cluster_context_manager():
     """
     Context manager for managing the simulation cluster lifecycle.
     """
diff --git a/src/services/solution_updater_service/core/service/solution_updater_service.py b/src/services/solution_updater_service/core/service/solution_updater_service.py
@@ -321,13 +321,39 @@ def _initiate_mapper_on_first_call(
         return _MapperState(control_vector_mapping, results_mapping, population_size)
 
 
+class _SolutionUpdaterServiceLoopController:
+    def __init__(self, max_generations: int) -> None:
+        """
+        Helper class to control the loop of the solution updater service.
+        This class manages the number of iterations for the optimization process, and raises StopIteration exception when convergence fails.
+        """
+        self._max_generations = max_generations
+        self.current_generation = 0
+        self._is_running = True
+
+    def running(self) -> bool:
+        """
+        Checks if the loop controller should run
+
+        Returns:
+            bool: True if the loop controller is running, False otherwise.
+        """
+        if self.current_generation >= self._max_generations:
+            self._is_running = False
+        self.current_generation += 1
+        return self._is_running
+
+
 class SolutionUpdaterService:
-    def __init__(self, optimization_engine: OptimizationEngine) -> None:
+    def __init__(
+        self, optimization_engine: OptimizationEngine, max_generations: int
+    ) -> None:
         self._mapper: _Mapper = _Mapper()
         self._engine: OptimizationEngineInterface = (
             OptimizationEngineFactory.get_engine(optimization_engine)
         )
         self._logger = get_logger(__name__)
+        self.loop_controller = _SolutionUpdaterServiceLoopController(max_generations)
 
     def process_request(
         self, request_dict: dict[str, Any]
@@ -385,6 +411,8 @@ def process_request(
         if not config.solution_candidates:
             raise RuntimeError("Nothing to optimize")
 
+        self._check_convergence(config.solution_candidates)
+
         control_vector, cost_function_values = self._mapper.to_numpy(
             config.solution_candidates
         )
@@ -399,4 +427,19 @@ def process_request(
 
         next_iter_solutions = self._mapper.to_control_vectors(updated_params)
         self._logger.info("Control vectors update request processed successfully.")
+
         return SolutionUpdaterServiceResponse(next_iter_solutions=next_iter_solutions)
+
+    def _check_convergence(
+        self, solution: list[SolutionCandidate], tol: float = 1e-4
+    ) -> None:
+        """
+        Should raise StopIteration exception when convergence reach desired value.
+        Args:
+            tol: function convergence tolerance
+            solution: list of SolutionCandidate
+
+        Returns:
+
+        """
+        pass
diff --git a/tests/solution_updater_service_test/test_solution_updater_service.py b/tests/solution_updater_service_test/test_solution_updater_service.py
@@ -84,7 +84,7 @@ def test_update_solution_for_next_iteration_single_call(  # type: ignore
     config_json, expected_result_parameters, mocked_engine, monkeypatch
 ):
     # Arrange
-    service = SolutionUpdaterService(optimization_engine=engine)
+    service = SolutionUpdaterService(optimization_engine=engine, max_generations=100)
 
     # Monkeypatch engine
     monkeypatch.setattr(service, "_engine", mocked_engine)
@@ -140,7 +140,7 @@ def test_update_solution_for_next_iteration_multiple_calls(  # type: ignore
     monkeypatch,
 ):
     # Arrange
-    service = SolutionUpdaterService(optimization_engine=engine)
+    service = SolutionUpdaterService(optimization_engine=engine, max_generations=100)
 
     # Monkeypatch engine
     monkeypatch.setattr(service, "_engine", mocked_engine)
@@ -210,7 +210,7 @@ def test_update_solution_with_boundaries_np(  # type: ignore
     config_json, expected_result_parameters, mocked_engine_with_bnb, monkeypatch
 ):
     # Arrange
-    service = SolutionUpdaterService(optimization_engine=engine)
+    service = SolutionUpdaterService(optimization_engine=engine, max_generations=100)
 
     # Monkeypatch engine to use mocked behavior
     monkeypatch.setattr(service, "_engine", mocked_engine_with_bnb)
@@ -312,9 +312,12 @@ def test_optimization_service_full_round(test_case):
         "optimization_constraints": {"boundaries": {k: [lb, ub] for k in param_names}},
     }
 
-    service = SolutionUpdaterService(optimization_engine=engine)
+    service = SolutionUpdaterService(
+        optimization_engine=engine, max_generations=iterations
+    )
+    loop_controller = service.loop_controller
 
-    for _ in range(iterations):
+    while loop_controller.running():
         result = service.process_request(config)
         positions = np.array(
             [get_numpy_values(vec.items) for vec in result.next_iter_solutions]

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@`
`15`	`15`
`16`	`16`	`from services.simulation_service.core.api import ( # noqa: F401, E402`
`17`	`17`	`SimulationService,`
`18`		`- simulation_cluster_contex_manager,`
	`18`	`+ simulation_cluster_context_manager,`
`19`	`19`	`)`
`20`	`20`
`21`		`-__all__ = ["SimulationService", "simulation_cluster_contex_manager"]`
	`21`	`+__all__ = ["SimulationService", "simulation_cluster_context_manager"]`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`from services.simulation_service.core.service import (`
`2`	`2`	`SimulationService,`
`3`		`- simulation_cluster_contex_manager,`
	`3`	`+ simulation_cluster_context_manager,`
`4`	`4`	`)`
`5`	`5`
`6`		`-__all__ = ["SimulationService", "simulation_cluster_contex_manager"]`
	`6`	`+__all__ = ["SimulationService", "simulation_cluster_context_manager"]`