Source code for GUIBRUSHR.Retrieval.ExofastMCMC.process_safety

"""
Process-safety helpers for the retrieval worker fan-out.

The retrieval main process forks worker processes (and a Manager server) with
the ``fork`` start method. They are non-daemon and the tree has no process
group / session isolation, so two failure modes need explicit handling:

1. **Parent dies -> children must die too.** If the main process crashes
   (segfault inside petitRADTRANS Fortran, OOM kill, ...) the forked children
   are reparented to ``init`` (PID 1) and keep burning CPU/RAM as orphans.
   ``install_parent_death_signal`` uses the Linux ``prctl(PR_SET_PDEATHSIG)``
   so the kernel kills the child the instant its parent dies.

2. **A child dies -> the parent must survive and report it.**
   ``describe_exitcode`` turns a ``multiprocessing.Process.exitcode`` into a
   human-readable cause (negative exitcode == killed by signal; SIGSEGV=-11,
   SIGKILL/OOM=-9, ...), so the parent can log *why* a worker died instead of
   failing silently.
"""

from __future__ import annotations

import multiprocessing
import os
import signal
import sys

# PR_SET_PDEATHSIG from <linux/prctl.h>. Linux-only; the call is a no-op on
# every other platform so the previous (orphaning) behaviour is unchanged
# there - acceptable because the production server is Linux.
_PR_SET_PDEATHSIG = 1




def mp_context():
    """Return a multiprocessing context safe for forking ``model_obj``.

    On Linux the default ``fork`` is used (CoW makes spawn cost ~0).
    On macOS Python 3.8+ defaults to ``spawn`` which would re-pickle the
    full pRT atmosphere (~hundreds of MB) for every Process, costing
    minutes-to-hours per run (see B8 audit). We force ``fork`` on macOS and
    set ``OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES`` so Apple's fork-safety
    check doesn't abort the workers - safe here because pRT workers do pure
    Python + Fortran and never touch Cocoa/CoreFoundation.

    Shared by map_optimizer and ModelData so the Manager server and the
    worker Processes always use the same (fork) start method.
    """
    if sys.platform == "darwin":
        os.environ.setdefault("OBJC_DISABLE_INITIALIZE_FORK_SAFETY", "YES")
        try:
            return multiprocessing.get_context("fork")
        except ValueError:
            return multiprocessing.get_context()
    return multiprocessing.get_context()





def install_parent_death_signal(sig: int = signal.SIGKILL) -> None:
    """Ask the kernel to send ``sig`` to this process when its parent dies.

    Linux only (uses ``prctl``). Intended to be called as the very first
    statement of a forked worker target (or as a Manager ``start`` initializer)
    so the child cannot outlive a crashed parent. On macOS / other platforms
    this is a no-op.

    Also closes the fork/prctl race: if the parent already died between
    ``fork()`` and the ``prctl`` call, this process was reparented to PID 1 and
    the death signal will never be delivered - detect that and exit now.
    """
    if not sys.platform.startswith("linux"):
        return
    try:
        import ctypes

        libc = ctypes.CDLL("libc.so.6", use_errno=True)
        libc.prctl(_PR_SET_PDEATHSIG, sig)
    except Exception:
        # Best-effort: if prctl is unavailable we keep the old behaviour
        # rather than crashing the worker over a safety nicety.
        return
    # Race guard: parent may have died before prctl was installed.
    if os.getppid() == 1:
        os._exit(1)



# Human-readable hints for the most common abnormal worker deaths. Keyed by
# raw signal number so the lookup works regardless of enum/int provenance.
_SIGNAL_HINTS = {
    int(signal.SIGSEGV): "SIGSEGV (segmentation fault - crash in native code, e.g. petitRADTRANS Fortran)",
    int(signal.SIGKILL): "SIGKILL (forced termination - typically out of RAM / OOM killer)",
    int(signal.SIGABRT): "SIGABRT (abort in native code)",
    int(signal.SIGBUS): "SIGBUS (bus error)",
    int(signal.SIGFPE): "SIGFPE (arithmetic error in native code)",
    int(signal.SIGTERM): "SIGTERM (terminated)",
    int(signal.SIGINT): "SIGINT (interrupted)",
}




def describe_exitcode(exitcode) -> str:
    """Return a human-readable explanation of a ``Process.exitcode``.

    ``None`` -> still running / not joined; ``0`` -> normal exit; negative ->
    killed by signal ``-exitcode``; positive -> non-zero exit status.
    """
    if exitcode is None:
        return "still alive / join not completed"
    if exitcode == 0:
        return "normal exit (code 0)"
    if exitcode < 0:
        sig = -exitcode
        return _SIGNAL_HINTS.get(sig, f"killed by signal {sig}")
    return f"exited with error code {exitcode}"