Coverage for pySDC/implementations/controller_classes/controller_MPI.py: 84%
299 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-19 17:32 +0000
1import numpy as np
2from pySDC.core.controller import Controller
3from pySDC.core.errors import ControllerError
4from pySDC.core.step import Step
5from pySDC.implementations.convergence_controller_classes.basic_restarting import BasicRestarting
8class controller_MPI(Controller):
9 """
11 PFASST controller, running parallel version of PFASST in blocks (MG-style)
13 """
15 def __init__(self, controller_params, description, comm):
16 """
17 Initialization routine for PFASST controller
19 Args:
20 controller_params: parameter set for the controller and the step class
21 description: all the parameters to set up the rest (levels, problems, transfer, ...)
22 comm: MPI communicator
23 """
25 # call parent's initialization routine
26 super().__init__(controller_params, description, useMPI=True)
28 # create single step per processor
29 self.S: Step = Step(description)
31 # pass communicator for future use
32 self.comm = comm
34 num_procs = self.comm.Get_size()
35 rank = self.comm.Get_rank()
37 # insert data on time communicator to the steps (helpful here and there)
38 self.S.status.time_size = num_procs
40 self.base_convergence_controllers += [BasicRestarting.get_implementation(useMPI=True)]
41 for convergence_controller in self.base_convergence_controllers:
42 self.add_convergence_controller(convergence_controller, description)
44 if self.params.dump_setup and rank == 0:
45 self.dump_setup(step=self.S, controller_params=controller_params, description=description)
47 num_levels = len(self.S.levels)
49 # add request handler for status send
50 self.req_status = None
51 # add request handle container for isend
52 self.req_send = [None] * num_levels
54 if num_procs > 1 and num_levels > 1:
55 for L in self.S.levels:
56 if not L.sweep.coll.right_is_node or L.sweep.params.do_coll_update:
57 raise ControllerError("For PFASST to work, we assume uend^k = u_M^k")
59 # `it_coarse` sweeps the coarsest level exactly once. Single-level Gauss-like MSSDC routes
60 # through it too. Check here rather than asserting mid-sweep: by then every rank has posted
61 # receives, so a failure hangs the job instead of raising. `mssdc_jac` only decides the
62 # routing when there is more than one step: a single step is plain SDC and always goes
63 # through `it_fine`, which honours nsweeps.
64 if self.S.levels[-1].params.nsweeps > 1 and (num_levels > 1 or (num_procs > 1 and not self.params.mssdc_jac)):
65 raise ControllerError('this controller cannot do multiple sweeps on coarsest level')
67 self.check_variable_coefficients(num_procs)
69 if num_levels == 1 and self.params.predict_type is not None:
70 self.logger.warning(
71 'you have specified a predictor type but only a single level.. predictor will be ignored'
72 )
74 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
75 C.setup_status_variables(self, comm=comm)
77 self.stages = self.get_stages()
79 def run(self, u0, t0, Tend):
80 """
81 Main driver for running the parallel version of SDC, MSSDC, MLSDC and PFASST
83 Args:
84 u0: initial values
85 t0: starting time
86 Tend: ending time
88 Returns:
89 end values on the finest level
90 stats object containing statistics for each step, each level and each iteration
91 """
93 # reset stats to prevent double entries from old runs
94 for hook in self.hooks:
95 hook.reset_stats()
97 # setup time initially
98 all_dt = self.comm.allgather(self.S.dt)
99 time = t0 + sum(all_dt[: self.comm.rank])
101 active = self.step_is_active(time, t0, Tend)
102 comm_active = self.comm.Split(active)
103 self.S.status.slot = comm_active.rank
105 if self.comm.rank == 0 and not active:
106 raise ControllerError('Nothing to do, check t0, dt and Tend!')
108 # initialize block of steps with u0
109 self.restart_block(comm_active.size, time, u0, comm=comm_active)
110 uend = u0
112 # call post-setup hook
113 for hook in self.hooks:
114 hook.post_setup(step=None, level_number=None)
116 # call pre-run hook
117 for hook in self.hooks:
118 hook.pre_run(step=self.S, level_number=0)
120 comm_active.Barrier()
122 # while any process still active...
123 while active:
124 while not self.S.status.done:
125 self.pfasst(comm_active, comm_active.size)
127 # determine where to restart
128 restarts = comm_active.allgather(self.S.status.restart)
130 # communicate time and solution to be used as next initial conditions
131 if True in restarts:
132 restart_at = np.where(restarts)[0][0]
133 uend = self.S.levels[0].u[0].bcast(root=restart_at, comm=comm_active)
134 tend = comm_active.bcast(self.S.time, root=restart_at)
135 self.logger.info(f'Starting next block with initial conditions from step {restart_at}')
137 else:
138 uend = self.S.levels[0].uend.bcast(root=comm_active.size - 1, comm=comm_active)
139 tend = comm_active.bcast(self.S.time + self.S.dt, root=comm_active.size - 1)
141 # do convergence controller stuff
142 if not self.S.status.restart:
143 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
144 C.post_step_processing(self, self.S, comm=comm_active)
146 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
147 C.prepare_next_block(self, self.S, self.S.status.time_size, tend, Tend, comm=comm_active)
149 # set new time
150 all_dt = comm_active.allgather(self.S.dt)
151 time = tend + sum(all_dt[: self.S.status.slot])
153 active = self.step_is_active(time, tend, Tend)
155 # check if we need to split the communicator
156 if tend + sum(all_dt[: comm_active.size - 1]) >= Tend - 10 * np.finfo(float).eps:
157 comm_active_new = comm_active.Split(active)
158 comm_active.Free()
159 comm_active = comm_active_new
161 self.S.status.slot = comm_active.rank
163 # initialize block of steps with u0
164 if active:
165 self.restart_block(comm_active.size, time, uend, comm=comm_active)
167 # call post-run hook
168 for hook in self.hooks:
169 hook.post_run(step=self.S, level_number=0)
171 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
172 C.post_run_processing(self, self.S, comm=self.comm)
174 comm_active.Free()
176 return uend, self.return_stats()
178 def restart_block(self, size, time, u0, comm):
179 """
180 Helper routine to reset/restart block of (active) steps
182 Args:
183 size: number of active time steps
184 time: current time
185 u0: initial value to distribute across the steps
186 comm: the communicator
188 Returns:
189 block of (all) steps
190 """
192 # store link to previous step
193 self.S.prev = (self.S.status.slot - 1) % size
194 self.S.next = (self.S.status.slot + 1) % size
196 # resets step
197 self.S.reset_step()
198 # determine whether I am the first and/or last in line
199 self.S.status.first = self.S.prev == size - 1
200 self.S.status.last = self.S.next == 0
201 # initialize step with u0
202 self.S.init_step(u0)
203 # reset some values
204 self.S.status.done = False
205 self.S.status.iter = 0
206 self.S.status.stage = 'SPREAD'
207 for l in self.S.levels:
208 l.tag = None
209 self.req_status = None
210 self.req_send = [None] * len(self.S.levels)
211 self.S.status.prev_done = False
212 self.S.status.force_done = False
214 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
215 C.reset_status_variables(self, comm=comm)
217 self.S.status.time_size = size
219 for lvl in self.S.levels:
220 lvl.status.time = time
221 lvl.status.sweep = 1
223 def recv(self, target, source, tag=None, comm=None):
224 """
225 Receive function
227 Args:
228 target: level which will receive the values
229 source: level which initiated the send
230 tag: identifier to check if this message is really for me
231 comm: communicator
232 """
233 req = target.u[0].irecv(source=source, tag=tag, comm=comm)
234 self.wait_for_request(request=req)
235 if self.S.status.force_done:
236 return None
237 # re-evaluate f on left interval boundary
238 target.f[0] = target.prob.eval_f(target.u[0], target.time)
240 def send_full(self, comm=None, blocking=False, level=None, add_to_stats=False):
241 """
242 Function to perform the send, including bookkeeping and logging
244 Args:
245 comm: the communicator
246 blocking: flag to indicate that we need blocking communication
247 level: the level number
248 add_to_stats: a flag to end recording data in the hooks (defaults to False)
250 Note:
251 Computing the end point is this function's job, not the caller's. Callers must not do it
252 themselves.
253 """
254 for hook in self.hooks:
255 hook.pre_comm(step=self.S, level_number=level)
257 if not blocking:
258 self.wait_for_request(request=self.req_send[level])
259 if self.S.status.force_done:
260 return None
262 self.S.levels[level].sweep.compute_end_point()
264 if not self.S.status.last:
265 self.logger.debug(
266 'isend data: process %s, stage %s, time %s, target %s, tag %s, iter %s'
267 % (
268 self.S.status.slot,
269 self.S.status.stage,
270 self.S.time,
271 self.S.next,
272 level * 100 + self.S.status.iter,
273 self.S.status.iter,
274 )
275 )
276 self.req_send[level] = self.S.levels[level].uend.isend(
277 dest=self.S.next, tag=level * 100 + self.S.status.iter, comm=comm
278 )
279 if blocking:
280 self.wait_for_request(request=self.req_send[level])
281 if self.S.status.force_done:
282 return None
284 for hook in self.hooks:
285 hook.post_comm(step=self.S, level_number=level, add_to_stats=add_to_stats)
287 def recv_full(self, comm, level=None, add_to_stats=False):
288 """
289 Function to perform the recv, including bookkeeping and logging
291 Args:
292 comm: the communicator
293 level: the level number
294 add_to_stats: a flag to end recording data in the hooks (defaults to False)
295 """
297 for hook in self.hooks:
298 hook.pre_comm(step=self.S, level_number=level)
299 if not self.S.status.first and not self.S.status.prev_done:
300 self.logger.debug(
301 'recv data: process %s, stage %s, time %s, source %s, tag %s, iter %s'
302 % (
303 self.S.status.slot,
304 self.S.status.stage,
305 self.S.time,
306 self.S.prev,
307 level * 100 + self.S.status.iter,
308 self.S.status.iter,
309 )
310 )
311 self.recv(target=self.S.levels[level], source=self.S.prev, tag=level * 100 + self.S.status.iter, comm=comm)
313 for hook in self.hooks:
314 hook.post_comm(step=self.S, level_number=level, add_to_stats=add_to_stats)
316 def wait_for_request(self, request):
317 """
318 Wait for a non-blocking communication to complete.
320 This used to poll for an interrupt while waiting, so that a rank could be told mid-wait that
321 the iteration estimator had decided everyone was done. That estimator has been removed, and
322 with it the only thing that could ever have interrupted a wait, so this is now a plain wait.
323 `force_done` is still honoured by the callers -- convergence controllers and hooks set it
324 between stages -- but nothing sets it *during* a wait any more.
326 Args:
327 request: request to wait for
328 """
329 if request is not None:
330 request.Wait()
332 def pfasst(self, comm, num_procs):
333 """
334 Main function including the stages of SDC, MLSDC and PFASST (the "controller")
336 For the workflow of this controller, check out one of our PFASST talks or the pySDC paper
338 Args:
339 comm: communicator
340 num_procs (int): number of parallel processes
341 """
343 stage = self.S.status.stage
345 self.logger.debug(stage + ' - process ' + str(self.S.status.slot))
347 self.stages.get(stage, self.default)(comm, num_procs)
349 def get_stages(self):
350 """
351 The stages this controller can be in, and what to run in each.
353 A subclass that iterates differently replaces the iteration stages here and says which one
354 to enter in `next_iteration_stage`; everything around the iteration is the same for any
355 algorithm this controller runs.
357 Returns:
358 dict: stage name -> the method that runs it
359 """
360 return {
361 'SPREAD': self.spread,
362 'PREDICT': self.predict,
363 'IT_CHECK': self.it_check,
364 'IT_FINE': self.it_fine,
365 'IT_DOWN': self.it_down,
366 'IT_COARSE': self.it_coarse,
367 'IT_UP': self.it_up,
368 }
370 def next_iteration_stage(self, S):
371 """
372 The stage that starts one iteration of the algorithm.
374 Args:
375 S (pySDC.Step.step): The current step
377 Returns:
378 str: name of the stage to enter
379 """
380 if len(S.levels) > 1: # MLSDC or PFASST
381 return 'IT_DOWN'
382 elif S.status.time_size == 1 or self.params.mssdc_jac: # SDC or parallel MSSDC (Jacobi-like)
383 return 'IT_FINE'
384 else:
385 return 'IT_COARSE' # serial MSSDC (Gauss-like)
387 def spread(self, comm, num_procs):
388 """
389 Spreading phase
390 """
392 # first stage: spread values
393 for hook in self.hooks:
394 hook.pre_step(step=self.S, level_number=0)
396 # call predictor from sweeper
397 self.S.levels[0].sweep.predict()
399 self.compute_residual_after_spread(self.S)
401 # update stage
402 if len(self.S.levels) > 1: # MLSDC or PFASST with predict
403 self.S.status.stage = 'PREDICT'
404 else:
405 self.S.status.stage = 'IT_CHECK'
407 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
408 C.post_spread_processing(self, self.S, comm=comm)
410 def predict(self, comm, num_procs):
411 """
412 Predictor phase
413 """
415 for hook in self.hooks:
416 hook.pre_predict(step=self.S, level_number=0)
418 if self.params.predict_type is None:
419 pass
421 elif self.params.predict_type == 'fine_only':
422 # do a fine sweep only
423 self.S.levels[0].sweep.update_nodes()
425 elif self.params.predict_type == 'pfasst_burnin':
426 # restrict to coarsest level
427 for l in range(1, len(self.S.levels)):
428 self.S.transfer(source=self.S.levels[l - 1], target=self.S.levels[l])
430 for p in range(self.S.status.slot + 1):
431 if not p == 0:
432 self.recv_full(comm=comm, level=len(self.S.levels) - 1)
433 if self.S.status.force_done:
434 return None
436 # do the sweep with new values
437 self.S.levels[-1].sweep.update_nodes()
438 self.S.levels[-1].sweep.compute_end_point()
440 self.send_full(
441 comm=comm, blocking=True, level=len(self.S.levels) - 1, add_to_stats=(p == self.S.status.slot)
442 )
443 if self.S.status.force_done:
444 return None
446 # interpolate back to finest level
447 for l in range(len(self.S.levels) - 1, 0, -1):
448 self.S.transfer(source=self.S.levels[l], target=self.S.levels[l - 1])
450 self.send_full(comm=comm, level=0)
451 if self.S.status.force_done:
452 return None
454 self.recv_full(comm=comm, level=0)
455 if self.S.status.force_done:
456 return None
458 # end this with a fine sweep
459 self.S.levels[0].sweep.update_nodes()
461 elif self.params.predict_type == 'fmg':
462 # TODO: implement FMG predictor
463 raise NotImplementedError('FMG predictor is not yet implemented')
465 else:
466 raise ControllerError('Wrong predictor type, got %s' % self.params.predict_type)
468 for hook in self.hooks:
469 hook.post_predict(step=self.S, level_number=0)
471 # update stage
472 self.S.status.stage = 'IT_CHECK'
474 def it_check(self, comm, num_procs):
475 """
476 Key routine to check for convergence/termination
477 """
479 self.prepare_convergence_check(comm)
481 if self.S.status.force_done:
482 return None
484 if self.S.status.iter > 0:
485 for hook in self.hooks:
486 hook.post_iteration(step=self.S, level_number=0)
488 # decide if the step is done, needs to be restarted and other things convergence related
489 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
490 C.post_iteration_processing(self, self.S, comm=comm)
491 C.convergence_control(self, self.S, comm=comm)
493 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
494 C.post_iteration_processing_block(self, comm=comm)
496 # if not ready, keep doing stuff
497 if not self.S.status.done:
498 # increment iteration count here (and only here)
499 self.S.status.iter += 1
501 for hook in self.hooks:
502 hook.pre_iteration(step=self.S, level_number=0)
503 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
504 C.pre_iteration_processing(self, self.S, comm=comm)
506 self.S.status.stage = self.next_iteration_stage(self.S)
508 else:
509 # Need to finish all pending isend requests. These will occur for the first active process, since
510 # in the last iteration the wait statement will not be called ("send and forget")
511 for req in self.req_send:
512 if req is not None:
513 req.Wait()
514 if self.req_status is not None:
515 self.req_status.Wait()
517 for hook in self.hooks:
518 hook.post_step(step=self.S, level_number=0)
519 self.S.status.stage = 'DONE'
521 def compute_residual_after_spread(self, S):
522 """
523 Make the residual current right after the initial guess, for algorithms that need it there.
525 This controller computes the residual at the top of `it_check` instead, so there is nothing
526 to do; an algorithm whose residual is not available at that point overrides this. Called
527 before `post_spread_processing`, because convergence controllers there may still change the
528 initial iterate.
530 Args:
531 S (pySDC.Step.step): The current step
532 """
533 pass
535 def prepare_convergence_check(self, comm):
536 """
537 Make current what `it_check` is about to read: the end point and the residual.
539 For a sweep-based algorithm both come out of the exchange with the neighbours, so this is
540 the send and receive plus the residual it enables. An algorithm that gets its residual from
541 the iteration instead overrides this, and still owes the end point.
543 Args:
544 comm: the communicator
545 """
546 self.send_full(comm=comm, level=0)
547 if self.S.status.force_done:
548 return None
550 self.recv_full(comm=comm, level=0)
551 if self.S.status.force_done:
552 return None
554 self.S.levels[0].sweep.compute_residual(stage='IT_CHECK')
556 def it_fine(self, comm, num_procs):
557 """
558 Fine sweeps
559 """
561 nsweeps = self.S.levels[0].params.nsweeps
563 self.S.levels[0].status.sweep = 0
565 # do fine sweep
566 for k in range(nsweeps):
567 self.S.levels[0].status.sweep += 1
569 # send values forward
570 self.send_full(comm=comm, level=0)
571 if self.S.status.force_done:
572 return None
574 # recv values from previous
575 self.recv_full(comm=comm, level=0, add_to_stats=(k == nsweeps - 1))
576 if self.S.status.force_done:
577 return None
579 for hook in self.hooks:
580 hook.pre_sweep(step=self.S, level_number=0)
582 self.S.levels[0].sweep.updateVariableCoeffs(k + 1) # update QDelta coefficients if variable preconditioner
583 self.S.levels[0].sweep.update_nodes()
584 self.S.levels[0].sweep.compute_residual(stage='IT_FINE')
586 for hook in self.hooks:
587 hook.post_sweep(step=self.S, level_number=0)
589 # update stage
590 self.S.status.stage = 'IT_CHECK'
592 def it_down(self, comm, num_procs):
593 """
594 Go down the hierarchy from finest to coarsest level
595 """
597 self.S.transfer(source=self.S.levels[0], target=self.S.levels[1])
599 # sweep and send on middle levels (not on finest, not on coarsest, though)
600 for l in range(1, len(self.S.levels) - 1):
601 nsweeps = self.S.levels[l].params.nsweeps
603 for _ in range(nsweeps):
604 self.send_full(comm=comm, level=l)
605 if self.S.status.force_done:
606 return None
608 self.recv_full(comm=comm, level=l)
609 if self.S.status.force_done:
610 return None
612 for hook in self.hooks:
613 hook.pre_sweep(step=self.S, level_number=l)
615 self.S.levels[l].sweep.update_nodes()
616 self.S.levels[l].sweep.compute_residual(stage='IT_DOWN')
617 for hook in self.hooks:
618 hook.post_sweep(step=self.S, level_number=l)
620 # transfer further down the hierarchy
621 self.S.transfer(source=self.S.levels[l], target=self.S.levels[l + 1])
623 # update stage
624 self.S.status.stage = 'IT_COARSE'
626 def it_coarse(self, comm, num_procs):
627 """
628 Coarse sweep
629 """
631 # receive from previous step (if not first)
632 self.recv_full(comm=comm, level=len(self.S.levels) - 1)
633 if self.S.status.force_done:
634 return None
636 # do the sweep
637 for hook in self.hooks:
638 hook.pre_sweep(step=self.S, level_number=len(self.S.levels) - 1)
639 self.S.levels[-1].sweep.update_nodes()
640 self.S.levels[-1].sweep.compute_residual(stage='IT_COARSE')
641 for hook in self.hooks:
642 hook.post_sweep(step=self.S, level_number=len(self.S.levels) - 1)
644 # send to next step (`send_full` computes the end point itself)
645 self.send_full(comm=comm, blocking=True, level=len(self.S.levels) - 1, add_to_stats=True)
646 if self.S.status.force_done:
647 return None
649 # update stage
650 if len(self.S.levels) > 1: # MLSDC or PFASST
651 self.S.status.stage = 'IT_UP'
652 else:
653 self.S.status.stage = 'IT_CHECK' # MSSDC
655 def it_up(self, comm, num_procs):
656 """
657 Prolong corrections up to finest level (parallel)
658 """
660 # receive and sweep on middle levels (except for coarsest level)
661 for l in range(len(self.S.levels) - 1, 0, -1):
662 # prolong values
663 self.S.transfer(source=self.S.levels[l], target=self.S.levels[l - 1])
665 # on middle levels: do sweep as usual
666 if l - 1 > 0:
667 nsweeps = self.S.levels[l - 1].params.nsweeps
669 for k in range(nsweeps):
670 self.send_full(comm, level=l - 1)
671 if self.S.status.force_done:
672 return None
674 self.recv_full(comm=comm, level=l - 1, add_to_stats=(k == nsweeps - 1))
675 if self.S.status.force_done:
676 return None
678 for hook in self.hooks:
679 hook.pre_sweep(step=self.S, level_number=l - 1)
680 self.S.levels[l - 1].sweep.update_nodes()
681 self.S.levels[l - 1].sweep.compute_residual(stage='IT_UP')
682 for hook in self.hooks:
683 hook.post_sweep(step=self.S, level_number=l - 1)
685 # update stage
686 self.S.status.stage = 'IT_FINE'
688 def default(self, num_procs):
689 """
690 Default routine to catch wrong status
691 """
692 raise ControllerError('Weird stage, got %s' % self.S.status.stage)