Function _distributed_worker

src/fused_mm_sampling/tp_info.py:83–107

(rank: int, world_size: int, port: int, fn: Callable, args: tuple)

        fn(*args)
    finally:
        dist.destroy_process_group()


def _distributed_worker(rank: int, world_size: int, port: int, fn: Callable, args: tuple) -> None:
    if rank == 0:
        _print_gpu_topology()
    print(
        f"Rank {rank}: available CPUs: {sorted(os.sched_getaffinity(0))}",
        flush=True,
    )
    device_id = rank % torch.cuda.device_count()
    torch.cuda.set_device(device_id)
    backend = "nccl" if torch.cuda.device_count() >= world_size else "gloo"
    if rank == 0:
        print(f"Using distributed backend: '{backend}'")

    dist.init_process_group(
        backend=backend,
        init_method=f"tcp://localhost:{port}",
        rank=rank,
        world_size=world_size,
        device_id=device_id,
    )
    try:
        fn(*args)
    finally:
        dist.destroy_process_group()