Coverage for pySDC/implementations/controller_classes/controller_MPI.py: 85%
302 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +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 Controller running SDC, MLSDC and PFASST with each time step of a block on its own MPI rank.
13 PFASST controller, running parallel version of PFASST in blocks (MG-style)
15 """
17 def __init__(self, controller_params, description, comm):
18 """
19 Initialization routine for PFASST controller
21 Args:
22 controller_params: parameter set for the controller and the step class
23 description: all the parameters to set up the rest (levels, problems, transfer, ...)
24 comm: MPI communicator
25 """
27 # call parent's initialization routine
28 super().__init__(controller_params, description, useMPI=True)
30 # create single step per processor
31 self.S: Step = Step(description)
33 # pass communicator for future use
34 self.comm = comm
36 num_procs = self.comm.Get_size()
37 rank = self.comm.Get_rank()
39 # insert data on time communicator to the steps (helpful here and there)
40 self.S.status.time_size = num_procs
42 self.base_convergence_controllers += [BasicRestarting.get_implementation(useMPI=True)]
43 for convergence_controller in self.base_convergence_controllers:
44 self.add_convergence_controller(convergence_controller, description)
46 if self.params.dump_setup and rank == 0:
47 self.dump_setup(step=self.S, controller_params=controller_params, description=description)
49 num_levels = len(self.S.levels)
51 # add request handler for status send
52 self.req_status = None
53 # add request handle container for isend
54 self.req_send = [None] * num_levels
56 if num_procs > 1 and num_levels > 1:
57 for L in self.S.levels:
58 if not L.sweep.coll.right_is_node or L.sweep.params.do_coll_update:
59 raise ControllerError("For PFASST to work, we assume uend^k = u_M^k")
61 # `it_coarse` sweeps the coarsest level exactly once. Single-level Gauss-like MSSDC routes
62 # through it too. Check here rather than asserting mid-sweep: by then every rank has posted
63 # receives, so a failure hangs the job instead of raising. `mssdc_jac` only decides the
64 # routing when there is more than one step: a single step is plain SDC and always goes
65 # through `it_fine`, which honours nsweeps.
66 if self.S.levels[-1].params.nsweeps > 1 and (num_levels > 1 or (num_procs > 1 and not self.params.mssdc_jac)):
67 raise ControllerError('this controller cannot do multiple sweeps on coarsest level')
69 self.check_variable_coefficients(num_procs)
71 if num_levels == 1 and self.params.predict_type is not None:
72 self.logger.warning(
73 'you have specified a predictor type but only a single level.. predictor will be ignored'
74 )
76 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
77 C.setup_status_variables(self, comm=comm)
79 self.stages = self.get_stages()
81 def run(self, u0, t0, Tend):
82 """
83 Main driver for running the parallel version of SDC, MSSDC, MLSDC and PFASST
85 Args:
86 u0: initial values
87 t0: starting time
88 Tend: ending time
90 Returns:
91 end values on the finest level
92 stats object containing statistics for each step, each level and each iteration
93 """
95 # reset stats to prevent double entries from old runs
96 for hook in self.hooks:
97 hook.reset_stats()
99 # start from the step size in the description, not the one adaptivity left from a previous run
100 for L in self.S.levels:
101 L.params.dt = L.params.dt_initial
103 # setup time initially
104 all_dt = self.comm.allgather(self.S.dt)
105 time = t0 + sum(all_dt[: self.comm.rank])
107 active = self.step_is_active(time, block_start=t0, t0=t0, Tend=Tend, dt=self.S.dt)
108 comm_active = self.comm.Split(active)
109 self.S.status.slot = comm_active.rank
111 if self.comm.rank == 0 and not active:
112 raise ControllerError('Nothing to do, check t0, dt and Tend!')
114 # initialize block of steps with u0
115 self.restart_block(comm_active.size, time, u0, comm=comm_active)
116 uend = u0
118 # call post-setup hook
119 for hook in self.hooks:
120 hook.post_setup(step=None, level_number=None)
122 # call pre-run hook
123 for hook in self.hooks:
124 hook.pre_run(step=self.S, level_number=0)
126 comm_active.Barrier()
128 # while any process still active...
129 while active:
130 while not self.S.status.done:
131 self.pfasst(comm_active, comm_active.size)
133 # determine where to restart
134 restarts = comm_active.allgather(self.S.status.restart)
136 # communicate time and solution to be used as next initial conditions
137 if True in restarts:
138 restart_at = np.where(restarts)[0][0]
139 uend = self.S.levels[0].u[0].bcast(root=restart_at, comm=comm_active)
140 tend = comm_active.bcast(self.S.time, root=restart_at)
141 self.logger.info(f'Starting next block with initial conditions from step {restart_at}')
143 else:
144 uend = self.S.levels[0].uend.bcast(root=comm_active.size - 1, comm=comm_active)
145 tend = comm_active.bcast(self.S.time + self.S.dt, root=comm_active.size - 1)
147 # do convergence controller stuff
148 if not self.S.status.restart:
149 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
150 C.post_step_processing(self, self.S, comm=comm_active)
152 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
153 C.prepare_next_block(self, self.S, self.S.status.time_size, tend, Tend, comm=comm_active)
155 # set new time
156 all_dt = comm_active.allgather(self.S.dt)
157 time = tend + sum(all_dt[: self.S.status.slot])
159 active = self.step_is_active(time, block_start=tend, t0=t0, Tend=Tend, dt=self.S.dt)
161 # check if we need to split the communicator
162 if not all(comm_active.allgather(active)):
163 comm_active_new = comm_active.Split(active)
164 comm_active.Free()
165 comm_active = comm_active_new
167 self.S.status.slot = comm_active.rank
169 # initialize block of steps with u0
170 if active:
171 self.restart_block(comm_active.size, time, uend, comm=comm_active)
173 # ranks that sat out the last blocks hold an older value, but rank 0 takes part in every block
174 if self.comm.rank > 0:
175 uend = self.S.levels[0].prob.dtype_u(uend)
176 uend = uend.bcast(root=0, comm=self.comm)
178 # call post-run hook
179 for hook in self.hooks:
180 hook.post_run(step=self.S, level_number=0)
182 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
183 C.post_run_processing(self, self.S, comm=self.comm)
185 comm_active.Free()
187 return uend, self.return_stats()
189 def restart_block(self, size, time, u0, comm):
190 """
191 Helper routine to reset/restart block of (active) steps
193 Args:
194 size: number of active time steps
195 time: current time
196 u0: initial value to distribute across the steps
197 comm: the communicator
199 Returns:
200 block of (all) steps
201 """
203 # store link to previous step
204 self.S.prev = (self.S.status.slot - 1) % size
205 self.S.next = (self.S.status.slot + 1) % size
207 # resets step
208 self.S.reset_step()
209 # determine whether I am the first and/or last in line
210 self.S.status.first = self.S.prev == size - 1
211 self.S.status.last = self.S.next == 0
212 # initialize step with u0
213 self.S.init_step(u0)
214 # reset some values
215 self.S.status.done = False
216 self.S.status.iter = 0
217 self.S.status.stage = 'SPREAD'
218 for l in self.S.levels:
219 l.tag = None
220 self.req_status = None
221 self.req_send = [None] * len(self.S.levels)
222 self.S.status.prev_done = False
223 self.S.status.force_done = False
225 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
226 C.reset_status_variables(self, comm=comm)
228 self.S.status.time_size = size
230 for lvl in self.S.levels:
231 lvl.status.time = time
232 lvl.status.sweep = 1
234 def recv(self, target, source, tag=None, comm=None):
235 """
236 Receive function
238 Args:
239 target: level which will receive the values
240 source: level which initiated the send
241 tag: identifier to check if this message is really for me
242 comm: communicator
243 """
244 req = target.u[0].irecv(source=source, tag=tag, comm=comm)
245 self.wait_for_request(request=req)
246 if self.S.status.force_done:
247 return None
248 # re-evaluate f on left interval boundary
249 target.f[0] = target.prob.eval_f(target.u[0], target.time)
251 def send_full(self, comm=None, blocking=False, level=None, add_to_stats=False):
252 """
253 Function to perform the send, including bookkeeping and logging
255 Args:
256 comm: the communicator
257 blocking: flag to indicate that we need blocking communication
258 level: the level number
259 add_to_stats: a flag to end recording data in the hooks (defaults to False)
261 Note:
262 Computing the end point is this function's job, not the caller's. Callers must not do it
263 themselves.
264 """
265 for hook in self.hooks:
266 hook.pre_comm(step=self.S, level_number=level)
268 if not blocking:
269 self.wait_for_request(request=self.req_send[level])
270 if self.S.status.force_done:
271 return None
273 self.S.levels[level].sweep.compute_end_point()
275 if not self.S.status.last:
276 self.logger.debug(
277 'isend data: process %s, stage %s, time %s, target %s, tag %s, iter %s'
278 % (
279 self.S.status.slot,
280 self.S.status.stage,
281 self.S.time,
282 self.S.next,
283 level * 100 + self.S.status.iter,
284 self.S.status.iter,
285 )
286 )
287 self.req_send[level] = self.S.levels[level].uend.isend(
288 dest=self.S.next, tag=level * 100 + self.S.status.iter, comm=comm
289 )
290 if blocking:
291 self.wait_for_request(request=self.req_send[level])
292 if self.S.status.force_done:
293 return None
295 for hook in self.hooks:
296 hook.post_comm(step=self.S, level_number=level, add_to_stats=add_to_stats)
298 def recv_full(self, comm, level=None, add_to_stats=False):
299 """
300 Function to perform the recv, including bookkeeping and logging
302 Args:
303 comm: the communicator
304 level: the level number
305 add_to_stats: a flag to end recording data in the hooks (defaults to False)
306 """
308 for hook in self.hooks:
309 hook.pre_comm(step=self.S, level_number=level)
310 if not self.S.status.first and not self.S.status.prev_done:
311 self.logger.debug(
312 'recv data: process %s, stage %s, time %s, source %s, tag %s, iter %s'
313 % (
314 self.S.status.slot,
315 self.S.status.stage,
316 self.S.time,
317 self.S.prev,
318 level * 100 + self.S.status.iter,
319 self.S.status.iter,
320 )
321 )
322 self.recv(target=self.S.levels[level], source=self.S.prev, tag=level * 100 + self.S.status.iter, comm=comm)
324 for hook in self.hooks:
325 hook.post_comm(step=self.S, level_number=level, add_to_stats=add_to_stats)
327 def wait_for_request(self, request):
328 """
329 Wait for a non-blocking communication to complete.
331 This used to poll for an interrupt while waiting, so that a rank could be told mid-wait that
332 the iteration estimator had decided everyone was done. That estimator has been removed, and
333 with it the only thing that could ever have interrupted a wait, so this is now a plain wait.
334 `force_done` is still honoured by the callers -- convergence controllers and hooks set it
335 between stages -- but nothing sets it *during* a wait any more.
337 Args:
338 request: request to wait for
339 """
340 if request is not None:
341 request.Wait()
343 def pfasst(self, comm, num_procs):
344 """
345 Main function including the stages of SDC, MLSDC and PFASST (the "controller")
347 For the workflow of this controller, check out one of our PFASST talks or the pySDC paper
349 Args:
350 comm: communicator
351 num_procs (int): number of parallel processes
352 """
354 stage = self.S.status.stage
356 self.logger.debug(stage + ' - process ' + str(self.S.status.slot))
358 self.stages.get(stage, self.default)(comm, num_procs)
360 def get_stages(self):
361 """
362 The stages this controller can be in, and what to run in each.
364 A subclass that iterates differently replaces the iteration stages here and says which one
365 to enter in `next_iteration_stage`; everything around the iteration is the same for any
366 algorithm this controller runs.
368 Returns:
369 dict: stage name -> the method that runs it
370 """
371 return {
372 'SPREAD': self.spread,
373 'PREDICT': self.predict,
374 'IT_CHECK': self.it_check,
375 'IT_FINE': self.it_fine,
376 'IT_DOWN': self.it_down,
377 'IT_COARSE': self.it_coarse,
378 'IT_UP': self.it_up,
379 }
381 def next_iteration_stage(self, S):
382 """
383 The stage that starts one iteration of the algorithm.
385 Args:
386 S (pySDC.Step.step): The current step
388 Returns:
389 str: name of the stage to enter
390 """
391 if len(S.levels) > 1: # MLSDC or PFASST
392 return 'IT_DOWN'
393 elif S.status.time_size == 1 or self.params.mssdc_jac: # SDC or parallel MSSDC (Jacobi-like)
394 return 'IT_FINE'
395 else:
396 return 'IT_COARSE' # serial MSSDC (Gauss-like)
398 def spread(self, comm, num_procs):
399 """
400 Spreading phase
401 """
403 # first stage: spread values
404 for hook in self.hooks:
405 hook.pre_step(step=self.S, level_number=0)
407 # call predictor from sweeper
408 self.S.levels[0].sweep.predict()
410 self.compute_residual_after_spread(self.S)
412 # update stage
413 if len(self.S.levels) > 1: # MLSDC or PFASST with predict
414 self.S.status.stage = 'PREDICT'
415 else:
416 self.S.status.stage = 'IT_CHECK'
418 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
419 C.post_spread_processing(self, self.S, comm=comm)
421 def predict(self, comm, num_procs):
422 """
423 Predictor phase
424 """
426 for hook in self.hooks:
427 hook.pre_predict(step=self.S, level_number=0)
429 if self.params.predict_type is None:
430 pass
432 elif self.params.predict_type == 'fine_only':
433 # do a fine sweep only
434 self.S.levels[0].sweep.update_nodes()
436 elif self.params.predict_type == 'pfasst_burnin':
437 # restrict to coarsest level
438 for l in range(1, len(self.S.levels)):
439 self.S.transfer(source=self.S.levels[l - 1], target=self.S.levels[l])
441 for p in range(self.S.status.slot + 1):
442 if not p == 0:
443 self.recv_full(comm=comm, level=len(self.S.levels) - 1)
444 if self.S.status.force_done:
445 return None
447 # do the sweep with new values
448 self.S.levels[-1].sweep.update_nodes()
449 self.S.levels[-1].sweep.compute_end_point()
451 self.send_full(
452 comm=comm, blocking=True, level=len(self.S.levels) - 1, add_to_stats=(p == self.S.status.slot)
453 )
454 if self.S.status.force_done:
455 return None
457 # interpolate back to finest level
458 for l in range(len(self.S.levels) - 1, 0, -1):
459 self.S.transfer(source=self.S.levels[l], target=self.S.levels[l - 1])
461 self.send_full(comm=comm, level=0)
462 if self.S.status.force_done:
463 return None
465 self.recv_full(comm=comm, level=0)
466 if self.S.status.force_done:
467 return None
469 # end this with a fine sweep
470 self.S.levels[0].sweep.update_nodes()
472 elif self.params.predict_type == 'fmg':
473 # TODO: implement FMG predictor
474 raise NotImplementedError('FMG predictor is not yet implemented')
476 else:
477 raise ControllerError('Wrong predictor type, got %s' % self.params.predict_type)
479 for hook in self.hooks:
480 hook.post_predict(step=self.S, level_number=0)
482 # update stage
483 self.S.status.stage = 'IT_CHECK'
485 def it_check(self, comm, num_procs):
486 """
487 Key routine to check for convergence/termination
488 """
490 self.prepare_convergence_check(comm)
492 if self.S.status.force_done:
493 return None
495 if self.S.status.iter > 0:
496 for hook in self.hooks:
497 hook.post_iteration(step=self.S, level_number=0)
499 # decide if the step is done, needs to be restarted and other things convergence related
500 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
501 C.post_iteration_processing(self, self.S, comm=comm)
502 C.convergence_control(self, self.S, comm=comm)
504 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
505 C.post_iteration_processing_block(self, comm=comm)
507 # share convergence only once every controller has decided, since some reopen `done`
508 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
509 C.communicate_convergence(self, self.S, comm=comm)
511 # if not ready, keep doing stuff
512 if not self.S.status.done:
513 # increment iteration count here (and only here)
514 self.S.status.iter += 1
516 for hook in self.hooks:
517 hook.pre_iteration(step=self.S, level_number=0)
518 for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
519 C.pre_iteration_processing(self, self.S, comm=comm)
521 self.S.status.stage = self.next_iteration_stage(self.S)
523 else:
524 # Need to finish all pending isend requests. These will occur for the first active process, since
525 # in the last iteration the wait statement will not be called ("send and forget")
526 for req in self.req_send:
527 if req is not None:
528 req.Wait()
529 if self.req_status is not None:
530 self.req_status.Wait()
532 for hook in self.hooks:
533 hook.post_step(step=self.S, level_number=0)
534 self.S.status.stage = 'DONE'
536 def compute_residual_after_spread(self, S):
537 """
538 Make the residual current right after the initial guess, for algorithms that need it there.
540 This controller computes the residual at the top of `it_check` instead, so there is nothing
541 to do; an algorithm whose residual is not available at that point overrides this. Called
542 before `post_spread_processing`, because convergence controllers there may still change the
543 initial iterate.
545 Args:
546 S (pySDC.Step.step): The current step
547 """
548 pass
550 def prepare_convergence_check(self, comm):
551 """
552 Make current what `it_check` is about to read: the end point and the residual.
554 For a sweep-based algorithm both come out of the exchange with the neighbours, so this is
555 the send and receive plus the residual it enables. An algorithm that gets its residual from
556 the iteration instead overrides this, and still owes the end point.
558 Args:
559 comm: the communicator
560 """
561 self.send_full(comm=comm, level=0)
562 if self.S.status.force_done:
563 return None
565 self.recv_full(comm=comm, level=0)
566 if self.S.status.force_done:
567 return None
569 self.S.levels[0].sweep.compute_residual(stage='IT_CHECK')
571 def it_fine(self, comm, num_procs):
572 """
573 Fine sweeps
574 """
576 nsweeps = self.S.levels[0].params.nsweeps
578 self.S.levels[0].status.sweep = 0
580 # do fine sweep
581 for k in range(nsweeps):
582 self.S.levels[0].status.sweep += 1
584 # send values forward
585 self.send_full(comm=comm, level=0)
586 if self.S.status.force_done:
587 return None
589 # recv values from previous
590 self.recv_full(comm=comm, level=0, add_to_stats=(k == nsweeps - 1))
591 if self.S.status.force_done:
592 return None
594 for hook in self.hooks:
595 hook.pre_sweep(step=self.S, level_number=0)
597 # update QDelta coefficients if variable preconditioner, indexed by the sweep count across iterations
598 self.S.levels[0].sweep.updateVariableCoeffs((self.S.status.iter - 1) * nsweeps + k + 1)
599 self.S.levels[0].sweep.update_nodes()
601 for hook in self.hooks:
602 hook.post_sweep(step=self.S, level_number=0)
604 # update stage
605 self.S.status.stage = 'IT_CHECK'
607 def it_down(self, comm, num_procs):
608 """
609 Go down the hierarchy from finest to coarsest level
610 """
612 self.S.transfer(source=self.S.levels[0], target=self.S.levels[1])
614 # sweep and send on middle levels (not on finest, not on coarsest, though)
615 for l in range(1, len(self.S.levels) - 1):
616 nsweeps = self.S.levels[l].params.nsweeps
618 for _ in range(nsweeps):
619 self.send_full(comm=comm, level=l)
620 if self.S.status.force_done:
621 return None
623 self.recv_full(comm=comm, level=l)
624 if self.S.status.force_done:
625 return None
627 for hook in self.hooks:
628 hook.pre_sweep(step=self.S, level_number=l)
630 self.S.levels[l].sweep.update_nodes()
631 for hook in self.hooks:
632 hook.post_sweep(step=self.S, level_number=l)
634 # transfer further down the hierarchy
635 self.S.transfer(source=self.S.levels[l], target=self.S.levels[l + 1])
637 # update stage
638 self.S.status.stage = 'IT_COARSE'
640 def it_coarse(self, comm, num_procs):
641 """
642 Coarse sweep
643 """
645 # receive from previous step (if not first)
646 self.recv_full(comm=comm, level=len(self.S.levels) - 1)
647 if self.S.status.force_done:
648 return None
650 # do the sweep
651 for hook in self.hooks:
652 hook.pre_sweep(step=self.S, level_number=len(self.S.levels) - 1)
653 self.S.levels[-1].sweep.update_nodes()
654 for hook in self.hooks:
655 hook.post_sweep(step=self.S, level_number=len(self.S.levels) - 1)
657 # send to next step (`send_full` computes the end point itself)
658 self.send_full(comm=comm, blocking=True, level=len(self.S.levels) - 1, add_to_stats=True)
659 if self.S.status.force_done:
660 return None
662 # update stage
663 if len(self.S.levels) > 1: # MLSDC or PFASST
664 self.S.status.stage = 'IT_UP'
665 else:
666 self.S.status.stage = 'IT_CHECK' # MSSDC
668 def it_up(self, comm, num_procs):
669 """
670 Prolong corrections up to finest level (parallel)
671 """
673 # receive and sweep on middle levels (except for coarsest level)
674 for l in range(len(self.S.levels) - 1, 0, -1):
675 # prolong values
676 self.S.transfer(source=self.S.levels[l], target=self.S.levels[l - 1])
678 # on middle levels: do sweep as usual
679 if l - 1 > 0:
680 nsweeps = self.S.levels[l - 1].params.nsweeps
682 for k in range(nsweeps):
683 self.send_full(comm, level=l - 1)
684 if self.S.status.force_done:
685 return None
687 self.recv_full(comm=comm, level=l - 1, add_to_stats=(k == nsweeps - 1))
688 if self.S.status.force_done:
689 return None
691 for hook in self.hooks:
692 hook.pre_sweep(step=self.S, level_number=l - 1)
693 self.S.levels[l - 1].sweep.update_nodes()
694 for hook in self.hooks:
695 hook.post_sweep(step=self.S, level_number=l - 1)
697 # update stage
698 self.S.status.stage = 'IT_FINE'
700 def default(self, num_procs):
701 """
702 Default routine to catch wrong status
703 """
704 raise ControllerError('Weird stage, got %s' % self.S.status.stage)