Function run_maybe_distributed

src/fused_mm_sampling/tp_info.py:35–57

Run fn(*args) in a single process or spawn n_procs distributed workers. fn and args must be picklable (top-level functions, dataclasses, etc.) because mp.spawn serializes them across processes. Do not pass lambdas. If the process was launched by torchrun (RANK env var is set), skips

(fn: Callable, n_procs: int, *args)

def run_maybe_distributed(fn: Callable, n_procs: int, *args) -> None:
    """Run fn(*args) in a single process or spawn n_procs distributed workers.

    fn and args must be picklable (top-level functions, dataclasses, etc.)
    because mp.spawn serializes them across processes. Do not pass lambdas.

    If the process was launched by torchrun (RANK env var is set), skips
    mp.spawn and uses the torchrun-provided environment directly.
    """
    if _is_torchrun():
        _torchrun_worker(fn, args)
    elif n_procs > 1:
        print(f"Running {fn.__name__} with {n_procs} processes")

        mp.spawn(
            _distributed_worker,
            args=(n_procs, _find_free_port(), fn, args),
            nprocs=n_procs,
            join=True,
        )
    else:
        print(f"Running {fn.__name__} single-process")
        fn(*args)


def _is_torchrun() -> bool: