[drivers] master-only raise hangs all other ranks instead of failing the job #24

Activity

1. the-hampel commented on Sep 15, 2026

   

   the-hampel

   on Sep 15, 2026

   Member

   More actions

   We hit this in solid_dmft years ago and it keeps coming back, so let me argue for fixing it
   one level down, in triqs core, rather than in dftkit — and for not writing the abort
   machinery ourselves at all, because mpi4py already ships it.

   Why not in dftkit

   The pattern isn't ours alone: is_master_node appears ~58× in dftkit, ~29× in dft_tools and
   ~90× in solid_dmft. The ten explicit raises listed above are only the visible tip — any
   exception in any of those blocks (a KeyError, an h5 read, a numpy error) hangs the job in
   exactly the same way, and those can't be enumerated. dft_tools/sumk_dft_transport.py already
   hand-rolls an MPI.COMM_WORLD.Abort(1), solid_dmft has four more, and dftkit carries a third
   variant in vasp/plovasp/sc_dmft.py. Four copies of the same idea, none shared.

   triqs.utility.mpi is also the only layer that knows whether we're running under mpi4py or
   under the serial mpi_nompi stub, so it's the only place the no-MPI fallback gets handled once.

   mpi4py already solved this

   mpi4py.run.set_abort_status(1) (what python -m mpi4py script.py uses internally) marks the
   job so that mpi4py calls MPI_Abort instead of MPI_Finalize when the interpreter exits.

   That deferral is the whole point, and it is strictly better than calling Abort() from inside
   an excepthook the way all our current copies do: the traceback gets printed, finally blocks
   and atexit handlers run, stdio is flushed — then the job is torn down. Verified: an
   atexit handler prints before OpenMPI's MPI_ABORT was invoked... banner. So the
   time.sleep(2) in solid_dmft/bin/solid_dmft.py, and the comment above it ("this sometimes
   weirdly suppresses error output completely"), both become unnecessary.

   It also means the VASP-driver cleanup can move into a finally / atexit and is then
   guaranteed to run before the abort — instead of today's "remember to call Driver.kill()
   before aborting", which is how you end up with an orphaned mpirun vasp still burning nodes.

   Concrete proposal for triqs.utility.mpi

   Two functions in mpi_mpi4py.py, with no-op/identity twins in mpi_nompi.py:

   def abort_on_exception(all_ranks_report=False):
       """Make an unhandled exception on *any* rank terminate the whole MPI job.
   
       Full traceback on master, one line on the other ranks, so an error that hits
       all ranks doesn't produce `size` interleaved tracebacks. Opt-in: applications
       call this from their entry point, libraries never do."""
       from mpi4py.run import set_abort_status
       previous = sys.excepthook
   
       def _hook(typ, value, tb):
           if is_master_node() or all_ranks_report:
               previous(typ, value, tb)
           else:
               print(f'[rank {rank}] {typ.__name__}: {str(value).splitlines()[0]}', file=sys.stderr)
           set_abort_status(1)          # abort at interpreter exit, after cleanup
   
       sys.excepthook = _hook
   
   
   @contextmanager
   def raise_on_all_ranks(comm=world):
       """Turn a failure inside a rank-local region into a raise on every rank."""
       error = None
       try:
           yield
       except Exception:
           error = f'rank {comm.Get_rank()}:\n{traceback.format_exc()}'
       errors = [e for e in comm.allgather(error) if e is not None]
       if errors:
           raise MPIError('failed on a subset of ranks --\n' + '\n'.join(errors)) from None
   
   
   def on_master(fn, *args, **kwargs):
       """Function-shaped variant; broadcasts fn's return value."""
       with raise_on_all_ranks():
           result = fn(*args, **kwargs) if is_master_node() else None
       return bcast(result)

   Note allgather, not bcast from root 0: a failure on rank 2 is caught just as well as one on
   the master, and the master's traceback then names the rank that actually failed.

   Tested

   Prototyped against triqs unstable, 4 ranks, OpenMPI 5.0.10 under mpirun, reproducing the
   _run_plo_converter shape (master-only failure, other ranks in mpi.barrier(poll_msec=100)):

   case	before	with the proposal
   master-only raise, unwrapped	hangs (killed at 25 s)	exit 1 in <1 s, 1 traceback
   master-only raise, wrapped in raise_on_all_ranks	—	exit 1 in <1 s, master shows the failing rank's stack
   failure on rank 2 only, wrapped	hangs	exit 1 in ~1 s, rank 2's traceback shown on master
   raise on all ranks	dies, 4 interleaved tracebacks	exit 1, 1 traceback + 3 one-liners
   no failure at all	—	exit 0, no spurious abort
   serial run, no launcher (mpi_nompi)	—	unchanged plain-Python traceback, exit 1

   I also confirmed the current solid_dmft/bin/solid_dmft.py hook does not help here: it
   aborts only on the non-master ranks, and in this scenario they never raise — they sit in the
   barrier. The master prints its traceback and then blocks in mpi4py's atexit MPI_Finalize,
   which is collective. Still a hang, measured.

   What dftkit then does

   # entry point, once
   mpi.abort_on_exception()

   and the ten sites keep their shape:

   with mpi.raise_on_all_ranks():
       if mpi.is_master_node():
           ...

   Three notes on the proposal above

   • with master_only(): cannot work as written. A context manager cannot skip its own body
     in Python — __enter__ has no way to say "don't run the block" short of sys.settrace
     hacks. Hence the shape above: the if mpi.is_master_node(): stays inside the with, and the
     context manager only handles the failure propagation. Same minimal diff at the call sites.
   • Don't stringify the exception. f"{type(exc).__name__}: {exc}" discards the traceback,
     which is the one thing you actually want from a remote rank. traceback.format_exc().
   • MPIHandler is the wrong home even after dedup: it's a dataclass about launching
     subprocesses (mpi_exec, env vars), and wien2k uses neither copy. Dedupe it on its own
     merits; MPI error policy belongs in triqs core.

   And a helper alone won't close this: it only covers the sites someone remembered to wrap. The
   excepthook is the backstop that catches the other ~170 master-only blocks; raise_on_all_ranks
   is for the failures that should stay catchable by a caller (a notebook, dft_tools driving the
   converter) — those never see an application's excepthook.

   Happy to open the triqs core issue/PR with the above if this direction looks right.
   (mpi4py.run.set_abort_status has been there since mpi4py 3.0; trivial to guard with a
   fallback to a plain Abort(1) if we care about older.)

   🤖 Analysis, prototype and measurements above generated with Claude Code