Coverage for pySDC/implementations/convergence_controller_classes/basic_restarting.py: 99%
98 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1from pySDC.core.convergence_controller import ConvergenceController, Pars
2from pySDC.implementations.convergence_controller_classes.spread_step_sizes import (
3 SpreadStepSizesBlockwise,
4)
5from pySDC.core.errors import ConvergenceError
6import numpy as np
9class BasicRestarting(ConvergenceController):
10 """
11 Restart every step after one that requests a restart, and limit how often a step may be restarted in a row.
12 The specific functions are:
14 - Telling each step after one that requested a restart to get restarted as well
15 - Allowing each step to be restarted a limited number of times in a row, before raising a ``ConvergenceError``
16 or, with ``crash_after_max_restarts=False``, moving on
18 Default control order is 95.
19 """
21 @classmethod
22 def get_implementation(cls, useMPI):
23 """
24 Retrieve the implementation for a specific flavor of this class.
26 Args:
27 useMPI (bool): Whether or not the controller uses MPI
29 Returns:
30 cls: The child class that implements the desired flavor
31 """
32 if useMPI:
33 return BasicRestartingMPI
34 else:
35 return BasicRestartingNonMPI
37 def __init__(self, controller, params, description, **kwargs):
38 """
39 Initialization routine
41 Args:
42 controller (pySDC.Controller): The controller
43 params (dict): Parameters for the convergence controller
44 description (dict): The description object used to instantiate the controller
45 """
46 super().__init__(controller, params, description)
47 self.buffers = Pars({"restart": False, "max_restart_reached": False})
49 def setup(self, controller, params, description, **kwargs):
50 """
51 Define parameters here.
53 Default parameters are:
54 - control_order (int): The order relative to other convergence controllers
55 - max_restarts (int): Maximum number of restarts we allow each step before we just move on with whatever we
56 have
57 - step_size_spreader (pySDC.ConvergenceController): A convergence controller that takes care of distributing
58 the steps sizes between blocks
60 Args:
61 controller (pySDC.Controller): The controller
62 params (dict): The params passed for this specific convergence controller
63 description (dict): The description object used to instantiate the controller
65 Returns:
66 (dict): The updated params dictionary
67 """
68 defaults = {
69 "control_order": 95,
70 "max_restarts": 10,
71 "crash_after_max_restarts": True,
72 "restart_from_first_step": False,
73 "step_size_spreader": SpreadStepSizesBlockwise.get_implementation(useMPI=params['useMPI']),
74 }
76 from pySDC.implementations.hooks.log_restarts import LogRestarts
78 controller.add_hook(LogRestarts)
80 return {**defaults, **super().setup(controller, params, description, **kwargs)}
82 def setup_status_variables(self, *args, **kwargs):
83 """
84 Add status variables for whether to restart now and how many times the step has been restarted in a row to the
85 Steps
87 Returns:
88 None
89 """
90 self.add_status_variable_to_step('restart', False)
91 self.add_status_variable_to_step('restarts_in_a_row', 0)
93 def reset_status_variables(self, *args, **kwargs):
94 """
95 Add status variables for whether to restart now and how many times the step has been restarted in a row to the
96 Steps
98 Returns:
99 None
100 """
101 self.set_step_status_variable('restart', False)
103 def dependencies(self, controller, description, **kwargs):
104 """
105 Load a convergence controller that spreads the step sizes between steps.
107 Args:
108 controller (pySDC.Controller): The controller
109 description (dict): The description object used to instantiate the controller
111 Returns:
112 None
113 """
114 spread_step_sizes_params = {
115 'spread_from_first_restarted': not self.params.restart_from_first_step,
116 }
117 controller.add_convergence_controller(
118 self.params.step_size_spreader, description=description, params=spread_step_sizes_params
119 )
120 return None
122 def determine_restart(self, controller, S, **kwargs):
123 """
124 Restart all steps after the first one which wants to be restarted as well, but also check if we lost patience
125 with the restarts and want to move on anyways.
127 Args:
128 controller (pySDC.Controller): The controller
129 S (pySDC.Step): The current step
131 Returns:
132 None
133 """
134 raise NotImplementedError("Please implement a function to determine if we need a restart here!")
137class BasicRestartingNonMPI(BasicRestarting):
138 """
139 Basic restarting for the non-MPI controller, which passes restart requests between steps through shared buffers.
140 """
142 def reset_buffers_nonMPI(self, controller, **kwargs):
143 """
144 Reset all variables with are used to simulate communication here
146 Args:
147 controller (pySDC.Controller): The controller
149 Returns:
150 None
151 """
152 self.buffers.restart = False
153 self.buffers.max_restart_reached = False
155 return None
157 def determine_restart(self, controller, S, MS, **kwargs):
158 """
159 Restart all steps after the first one which wants to be restarted as well, but also check if we lost patience
160 with the restarts and want to move on anyways.
162 Args:
163 controller (pySDC.Controller): The controller
164 S (pySDC.Step): The current step
165 MS (list): List of active steps
167 Returns:
168 None
169 """
170 # check if we performed too many restarts
171 if S.status.first:
172 self.buffers.max_restart_reached = S.status.restarts_in_a_row >= self.params.max_restarts
174 if self.buffers.max_restart_reached and S.status.restart:
175 if self.params.crash_after_max_restarts:
176 raise ConvergenceError(f"Restarted {S.status.restarts_in_a_row} time(s) already, surrendering now.")
177 self.log(
178 f"Step(s) restarted {S.status.restarts_in_a_row} time(s) already, maximum reached, moving \
179on...",
180 S,
181 )
183 self.buffers.restart = S.status.restart or self.buffers.restart
184 S.status.restart = (S.status.restart or self.buffers.restart) and not self.buffers.max_restart_reached
186 if S.status.last and self.params.restart_from_first_step and not self.buffers.max_restart_reached:
187 for step in MS:
188 step.status.restart = self.buffers.restart
190 return None
192 def prepare_next_block(self, controller, S, size, time, Tend, MS, **kwargs):
193 """
194 Update restarts in a row for all steps.
196 Args:
197 controller (pySDC.Controller): The controller
198 S (pySDC.Step): The current step
199 size (int): The number of ranks
200 time (list): List containing the time of all the steps
201 Tend (float): Final time of the simulation
202 MS (list): List of active steps
204 Returns:
205 None
206 """
207 # update all steps at once, since the new counts are read from the old ones
208 if S is not MS[0]:
209 return None
211 restart_from = min([me.status.slot for me in MS if me.status.restart] + [size - 1])
213 new = [me.status.restarts_in_a_row + 1 if me.status.restart else 0 for me in MS[restart_from:]]
214 for i, step in enumerate(MS):
215 step.status.restarts_in_a_row = new[i] if i < len(new) else 0
217 return None
220class BasicRestartingMPI(BasicRestarting):
221 """
222 Basic restarting for the MPI controller, which passes restart requests on to the following ranks with MPI.
223 """
225 def __init__(self, controller, params, description, **kwargs):
226 """
227 Initialization routine. Adds a buffer.
229 Args:
230 controller (pySDC.Controller): The controller
231 params (dict): Parameters for the convergence controller
232 description (dict): The description object used to instantiate the controller
233 """
234 from mpi4py import MPI
236 self.OR = MPI.LOR
238 super().__init__(controller, params, description)
239 self.buffers = Pars({"restart": False, "max_restart_reached": False, 'restart_earlier': False})
241 def determine_restart(self, controller, S, comm, **kwargs):
242 """
243 Restart all steps after the first one which wants to be restarted as well, but also check if we lost patience
244 with the restarts and want to move on anyways.
246 Args:
247 controller (pySDC.Controller): The controller
248 S (pySDC.Step): The current step
249 comm (mpi4py.MPI.Intracomm): Communicator
251 Returns:
252 None
253 """
254 crash_now = False
256 if S.status.first:
257 # check if we performed too many restarts
258 self.buffers.max_restart_reached = S.status.restarts_in_a_row >= self.params.max_restarts
259 self.buffers.restart_earlier = False # there is no earlier step
261 if self.buffers.max_restart_reached and S.status.restart:
262 if self.params.crash_after_max_restarts:
263 crash_now = True
264 self.log(
265 f"Step(s) restarted {S.status.restarts_in_a_row} time(s) already, maximum reached, moving \
266on...",
267 S,
268 )
269 elif not S.status.prev_done and not self.params.restart_from_first_step:
270 # receive information about restarts from earlier ranks
271 buff = np.empty(3, dtype=bool)
272 self.Recv(comm=comm, source=S.status.slot - 1, buffer=[buff, self.MPI_BOOL])
273 self.buffers.restart_earlier = buff[0]
274 self.buffers.max_restart_reached = buff[1]
275 crash_now = buff[2]
277 # decide whether to restart
278 S.status.restart = (S.status.restart or self.buffers.restart_earlier) and not self.buffers.max_restart_reached
280 # send information about restarts forward
281 if not S.status.last and not self.params.restart_from_first_step:
282 buff = np.empty(3, dtype=bool)
283 buff[0] = S.status.restart
284 buff[1] = self.buffers.max_restart_reached
285 buff[2] = crash_now
286 self.Send(comm, dest=S.status.slot + 1, buffer=[buff, self.MPI_BOOL])
288 if self.params.restart_from_first_step:
289 # only the first step knows whether we lost patience, and whether to crash
290 max_restart_reached, crash_now = comm.bcast((self.buffers.max_restart_reached, crash_now), root=0)
291 S.status.restart = comm.allreduce(S.status.restart, op=self.OR) and not max_restart_reached
293 if crash_now:
294 raise ConvergenceError("Surrendering because of too many restarts...")
296 return None
298 def prepare_next_block(self, controller, S, size, time, Tend, comm, **kwargs):
299 """
300 Update restarts in a row for all steps.
302 Args:
303 controller (pySDC.Controller): The controller
304 S (pySDC.Step): The current step
305 size (int): The number of ranks
306 time (list): List containing the time of all the steps
307 Tend (float): Final time of the simulation
308 comm (mpi4py.MPI.Intracomm): Communicator
310 Returns:
311 None
312 """
314 restart_from = min(comm.allgather(S.status.slot if S.status.restart else S.status.time_size - 1))
316 # send "backward" the number of restarts in a row
317 if S.status.slot >= restart_from:
318 buff = np.empty(1, dtype=int)
319 buff[0] = int(S.status.restarts_in_a_row + 1 if S.status.restart else 0)
320 self.Send(
321 comm,
322 dest=S.status.slot - restart_from,
323 buffer=[buff, self.MPI_INT],
324 blocking=False,
325 )
327 # receive new number of restarts in a row
328 if S.status.slot + restart_from < size:
329 buff = np.empty(1, dtype=int)
330 self.Recv(comm, source=(S.status.slot + restart_from), buffer=[buff, self.MPI_INT])
331 S.status.restarts_in_a_row = buff[0]
332 else:
333 S.status.restarts_in_a_row = 0
335 return None