kernbench2/src/kernbench/runtime_api/multiprocessing.py

"""``torch.multiprocessing.spawn``-compatible namespace (ADR-0027 D1).

Real-PyTorch API *signature* parity only — execution model is a cooperative
greenlet scheduler in a single Python process (D1.0). Non-goals: process
isolation, independent address space, failure isolation, OS-level scheduler
fairness, mp.Queue/Lock.

Attached to ``RuntimeContext`` as ``ctx.multiprocessing`` in
``__post_init__`` (D1.3).
"""
from __future__ import annotations

from typing import Any, Callable


class SpawnException(RuntimeError):
    """Raised from ``_MultiprocessingNamespace.spawn`` on worker failure.

    ``errors`` contains only root-cause ranks — the rank(s) whose body
    raised. Sibling greenlets terminated via ``throw(SystemExit)`` during
    cleanup are NOT recorded (SystemExit does not satisfy ``except
    Exception`` in the entry wrapper).
    """

    def __init__(self, errors: dict[int, Exception]):
        self.errors = errors
        first = next(iter(errors.items()), None)
        msg = (
            f"spawn failed on ranks {sorted(errors.keys())}"
            + (
                f": rank {first[0]} raised {first[1]!r}"
                if first is not None
                else ""
            )
        )
        super().__init__(msg)


def _drain_pending(ctx: Any) -> None:
    """Drain worker-wait + collective-pending queues in main context (D0.4/D0.5).

    Loop-until-empty: runs until both queues are simultaneously empty. Safe
    under the current model where main-context ``ctx.wait`` never re-enqueues
    (D0.5 main-context non-reentrance invariant); also safe under future
    extensions where drain can add sub-handles (SimPy causality gives finite
    depth).
    """
    distributed = getattr(ctx, "distributed", None)
    backend = getattr(distributed, "_backend", None) if distributed else None

    def _collective_nonempty() -> bool:
        if backend is None:
            return False
        pending = getattr(backend, "_pending_collective_handles", None)
        return bool(pending)

    while ctx._pending_worker_waits or _collective_nonempty():
        # (a) Worker-driven waits (D0.1). FIFO.
        while ctx._pending_worker_waits:
            h = ctx._pending_worker_waits.pop(0)
            if h not in ctx._completed:
                wait_fn = getattr(ctx.engine, "wait", None)
                if wait_fn is not None:
                    wait_fn(h)
                # Populate _completed so fast-path in ctx.wait short-circuits
                # on the return leg.
                ctx._completed.add(h)
        # (b) Collective backend queue (ADR-0027 D0.4-(2)).
        if backend is not None:
            pending_list = getattr(backend, "_pending_collective_handles", None)
            if pending_list is not None:
                while pending_list:
                    h, _sip_id, meta = pending_list.pop(0)
                    # Main context: ctx.wait drives engine directly and does
                    # NOT re-enqueue (D0.5 invariant).
                    ctx.wait(h, _meta=meta)


class _MultiprocessingNamespace:
    """torch.multiprocessing-compat facade bound to a RuntimeContext."""

    def __init__(self, ctx: Any) -> None:
        self._ctx = ctx

    def spawn(
        self,
        fn: Callable,
        args: tuple = (),
        nprocs: int = 1,
        join: bool = True,
    ) -> None:
        """Spawn ``nprocs`` worker greenlets, each calling ``fn(rank, *args)``.

        Mirrors ``torch.multiprocessing.spawn`` signature (minus ``daemon``).
        Runs the D0.4 round-robin scheduler loop until all workers finish,
        draining pending queues between rounds.
        """
        from greenlet import greenlet

        ctx = self._ctx
        dist = ctx.distributed
        gs: list = []
        errors: dict[int, Exception] = {}

        for rank in range(nprocs):
            def _entry(r: int = rank) -> None:
                try:
                    fn(r, *args)
                except Exception as e:
                    errors[r] = e
                    raise

            g = greenlet(_entry)
            if dist is not None and hasattr(dist, "_bind_rank"):
                dist._bind_rank(g, rank)
            gs.append(g)

        try:
            while True:
                alive = [g for g in gs if not g.dead]
                if not alive:
                    break
                for g in alive:
                    if not g.dead:
                        g.switch()
                _drain_pending(ctx)
        except Exception as outer:
            # D0.4-(4) sibling cleanup. Abort live greenlets, clear state.
            for other in gs:
                if not other.dead:
                    try:
                        other.throw(SystemExit)
                    except BaseException:
                        # SystemExit inherits BaseException; greenlet.throw
                        # re-raises in caller if target doesn't catch it.
                        # Silent — we're already in cleanup.
                        pass
            backend = getattr(dist, "_backend", None)
            if backend is not None:
                if hasattr(backend, "_barrier") and hasattr(backend._barrier, "reset"):
                    try:
                        backend._barrier.reset()
                    except Exception:
                        pass
                pending_collective = getattr(
                    backend, "_pending_collective_handles", None,
                )
                if pending_collective is not None:
                    pending_collective.clear()
            ctx._pending_worker_waits.clear()
            raise SpawnException(errors) from outer
        # join=True: we already waited for all workers above.