#!/usr/bin/env python
# coding: utf-8

# In[5]:


from functools import partial
import ipyparallel as ipp


def crash():
    import ctypes

    ctypes.CDLL(None).time(-1)


def notice_mpiexec_stopped(stop_info, launcher):
    output = launcher.get_output()
    print("mpiexec stopped: pid={pid}, exit_code={exit_code}".format(**stop_info))
    # extract mpiexec output, which starts with =
    # could also dump the full thing, which is a lot,
    # or tail a number of lines
    for line in output.splitlines(True):
        if line.startswith("="):
            sys.stdout.write(line)


with ipp.Cluster(engine_launcher_class="mpi", n=4) as rc:
    cluster = rc.cluster
    for launcher in cluster.engines.values():
        launcher.on_stop(partial(notice_mpiexec_stopped, launcher=launcher))
    ar = rc[-1].apply_async(crash)
    try:
        ar.get()
    except Exception as e:
        print(f"error in task (expected): {e}")