communication.py

import torch
import torch.distributed as dist
from utils.logs_utils import log
# torch.manual_seed(0)

from threading import BrokenBarrierError  # multiprocessing.Barrier raises this

def synchronize_all_clients(
    cuda_event,
    cuda_stream,
    process_group,          # leaders-only group (created on local-rank-0 of each GPU)
    group_barrier,          # mp.Barrier shared by ALL local procs on THIS GPU
    i_rank,                 # local rank on this GPU: 0..(n_clients_per_gpu-1)
    barrier_timeout_s: float = 600.0,
):
    """
    1) Fence THIS process's GPU work up to `cuda_event` on `cuda_stream`.
    2) Inter-GPU barrier across leaders (i_rank == 0) using torch.distributed.
    3) Intra-GPU barrier across all local processes on THIS GPU using mp.Barrier.

    Notes:
      - `group_barrier` must be a multiprocessing.Barrier created with parties=n_clients_per_gpu
        and passed to the local processes of the same GPU.
      - Do NOT call group_barrier.reset(); it's cyclic already.
      - Every local process must reach this function the same number of times.
    """
    # --- (1) GPU fence for this process ---
    if cuda_stream is not None:
        cuda_event.record(cuda_stream)
    else:
        cuda_event.record()
    cuda_event.synchronize()  # CPU waits until work before event completes on that stream

    # --- (2) Inter-GPU (leaders only) ---
    if i_rank == 0:
        # Only local-rank-0 participates in the leaders group sync
        dist.barrier(group=process_group)
        

    # --- (3) Intra-GPU (all local processes on this GPU) ---
    try:
        group_barrier.wait(timeout=barrier_timeout_s)
    except BrokenBarrierError as e:
        # Optional: add logging here to help identify which local rank didn't arrive.
        raise e

    # Final device drain for this process (safety)
    torch.cuda.synchronize()

def synchronize_with_rank0_first(cuda_event, cuda_stream, process_group, group_barrier, i_rank, rank, is_first_call=True):
    """
    Two-phase synchronization that allows rank 0 to proceed first, then others.
    
    Args:
        cuda_event: CUDA event for GPU synchronization
        cuda_stream: CUDA stream for GPU operations
        process_group: Process group for distributed communication
        group_barrier: Multiprocessing barrier for process synchronization
        i_rank: Local rank within the GPU (0 to n_clients_per_gpu-1)
        rank: Global rank across all GPUs and processes
        is_first_call: If True, only rank 0 proceeds. If False, remaining ranks proceed.
    """
    # First synchronize GPU operations
    cuda_event.record(cuda_stream)
    cuda_event.wait(cuda_stream)
    cuda_event.synchronize()
    
    if is_first_call:
        # First phase: Everyone syncs, only rank 0 proceeds
        if i_rank == 0:
            # First ensure all processes reach this point
            dist.barrier(group=process_group)
        # Then use the group barrier to coordinate within the group
        group_barrier.wait()
        if i_rank == 0:
            group_barrier.reset()
            # Signal completion of first phase
            dist.barrier(group=process_group)
            
    else:
        # Second phase: Let all remaining ranks proceed
        if i_rank == 0:
            # Wait for first phase completion signal
            dist.barrier(group=process_group)
        # Synchronize all processes before proceeding
        group_barrier.wait()
        if i_rank == 0:
            group_barrier.reset()
            # Final sync to ensure all processes have completed
            dist.barrier(group=process_group)
    
    # Final GPU sync
    torch.cuda.synchronize()

def wait_event(com_event, cuda_stream):
    com_event.record(cuda_stream)
    com_event.wait(cuda_stream)
    com_event.synchronize()
    
@torch.no_grad()
def decentralized_communication_step(
    i_rank,
    graph,
    params,
    communication_buffer,
    group_ranks,
    group_communication_buffer_send,
    group_communication_buffer_rcv,
    process_group,
    com_events,
    cuda_stream,
):
    # re-gather the rank
    rank = group_ranks[i_rank]
    # gather the communication schedule for this client
    schedule = graph.nodes[rank]["com_schedule"]
    # if we are client 0 of this GPU, we organize the outside communications for the group in addition to our own communications
    if i_rank == 0:
        # gather the list of communications for the group
        group_schedule = graph.nodes[rank]["group_schedule"]
        for (i,j) in group_schedule:
            # if the communication is outside
            if i not in group_ranks or j not in group_ranks:
                # sort the inside and outside ranks
                rank_out = i if i not in group_ranks else j
                rank_in = i if i in group_ranks else j
                i_rank_in = rank_in%len(group_ranks)
                # if it is not ourselves that need to commmunicate outside
                if rank_in != rank:
                    # gather the mp.Queue objects
                    queue_get = graph.edges[(rank, rank_in)]["com_queue"][(rank_in, rank)]
                    # get the other client's params
                    params_other = queue_get.get()
                    # copies it in the group's buffer
                    group_communication_buffer_send.copy_(params_other)
                    # wait for the buffer to be effectively filled before sending anything
                    wait_event(com_events[i_rank_in], cuda_stream)
                    # release memory, as advised in https://pytorch.org/docs/stable/multiprocessing.html
                    del params_other
                else:
                    # copies our params in the group's buffer
                    group_communication_buffer_send.copy_(params)
                    # wait for the buffer to be effectively filled before sending anything
                    wait_event(com_events[i_rank_in], cuda_stream)
                # retrieve the gpu rank hosting the other client
                rank_gpu_other = rank_out//len(group_ranks)
                # makes the p2p communication
                if rank_in > rank_out:
                    # sends the params
                    dist.send(group_communication_buffer_send, rank_gpu_other, process_group)
                    # receives the params from the other in the group's buffer
                    dist.recv(group_communication_buffer_rcv, rank_gpu_other, process_group)
                else:
                    # receives the params from the other in the group's buffer
                    dist.recv(group_communication_buffer_rcv, rank_gpu_other, process_group)
                    # sends our params
                    dist.send(group_communication_buffer_send, rank_gpu_other, process_group)
                # if it was not for ourselves
                if rank_in != rank:
                    # gather the mp.Queue objects to communicate back to the group client
                    queue_put = graph.edges[(rank, rank_in)]["com_queue"][(rank, rank_in)]
                    queue_put.put(group_communication_buffer_rcv.clone())
                # if it was for ourself
                else:
                    communication_buffer.add_(group_communication_buffer_rcv, alpha=graph.edges[(rank, rank_out)]["metropolis_weight"])
            # else, if the communication is inside the group and concerns us
            # we perform the communication normally
            elif rank in (i,j):
                # retrieve the other's rank
                rank_other = i if i!=rank else j
                i_rank_other = rank_other%len(group_ranks)
                # gather the mp.Queues
                queue_put = graph.edges[(i,j)]["com_queue"][(rank, rank_other)]
                queue_get = graph.edges[(i,j)]["com_queue"][(rank_other, rank)]
                # perform the p2p communication
                queue_put.put(params.clone())
                params_other = queue_get.get()
                communication_buffer.add_(params_other, alpha=graph.edges[(rank, rank_other)]["metropolis_weight"])
                # wait for the buffer to be effectively filled before going to next com
                wait_event(com_events[i_rank_other], cuda_stream)
                del params_other
    # if we are not the client 0 of this GPU
    else:
        # gather the rank of the 0'th client
        rank_0 = group_ranks[0]
        for rank_other in schedule:
            # if the other client is not in our group, client 0 is in charge of the out communications
            if rank_other not in group_ranks:
                # gather the mp.Queues
                queue_put = graph.edges[(rank, rank_0)]["com_queue"][(rank, rank_0)]
                queue_get = graph.edges[(rank, rank_0)]["com_queue"][(rank_0, rank)]
            # else, directly communicate with our peer using queues
            else:
                # gather the mp.Queues
                queue_put = graph.edges[(rank, rank_other)]["com_queue"][(rank, rank_other)]
                queue_get = graph.edges[(rank, rank_other)]["com_queue"][(rank_other, rank)]
            # perform the p2p communication
            queue_put.put(params.clone())
            params_other = queue_get.get()
            # adds the communication in our buffer
            communication_buffer.add_(params_other, alpha=graph.edges[(rank, rank_other)]["metropolis_weight"])
            # wait that the communication finished before beginning a new one
            wait_event(com_events[0], cuda_stream)
            # release memory
            del params_other


@torch.no_grad()
def complete_communication_step(i_rank, model_mask, graph, params_or_grads, communication_buffer, group_ranks, process_group, com_events, cuda_stream, num_clients_in_group):
    # re-gather the rank
    rank = group_ranks[i_rank]
    # if we are the client 0 of this GPU
    if i_rank == 0:
        ##print("Now Processing Process RANK 0")
        # we receive the params of every other client hosted here
        for idx_j, j in enumerate(group_ranks[1:num_clients_in_group]):
            # gather the mp.Queue objects
            queue_get_j = graph.edges[(rank, j)]["com_queue"][(j,rank)]
            # get the other client's params
            params_or_grads_other, mask_other = queue_get_j.get()
            ##print(f"Other Params: {params_other[:15]}")
            # copies it in our buffer
            # print(f'Grads after multiplication with mask: {params_or_grads_other * mask_other} | Any NaN: {torch.isnan(params_or_grads_other * mask_other).any()}')
            communication_buffer.add_(params_or_grads_other * mask_other)
            wait_event(com_events[idx_j], cuda_stream)
            # release memory, as advised in https://pytorch.org/docs/stable/multiprocessing.html
            del params_or_grads_other
            del mask_other
        # add our params to the buffer
        communication_buffer.add_(params_or_grads * model_mask)
        wait_event(com_events[0], cuda_stream)
        # all reduce the buffers of each GPU
        dist.all_reduce(communication_buffer, group=process_group, op=dist.ReduceOp.SUM)
        wait_event(com_events[-1], cuda_stream)
        # send the final version to all other clients hosted in this GPU
        for j in group_ranks[1:num_clients_in_group]:
            # gather the mp.Queue objects
            queue_put = graph.edges[(rank,j)]["com_queue"][(rank,j)]
            queue_put.put(communication_buffer.clone())
    # if we are not the client 0
    elif i_rank < num_clients_in_group:
        # gather the client 0 rank
        rank_0 = group_ranks[0]
        # sends our parameters to client 0 using the queue
        # Undirected graph
        queue_put = graph.edges[(rank, rank_0)]["com_queue"][(rank, rank_0)]
        queue_put.put((params_or_grads.clone(), model_mask.clone()))
        # get the final version of the params from client 0
        queue_get = graph.edges[(rank, rank_0)]["com_queue"][(rank_0, rank)]
        all_reduced_params = queue_get.get()
        # copies it in our buffer
        communication_buffer.copy_(all_reduced_params)
        # wait that the communication finished before deleting
        wait_event(com_events[0], cuda_stream)
        # release memory
        del all_reduced_params

@torch.no_grad()
def communicate_ref_worker_with_other_clients(graph, group_ranks, client_stream, process_group, i_rank, params_ref_worker, num_clients_in_group, com_events, ref_worker_communication):
    """
    Communicate the parameters of the reference worker with other clients in the group.
    Args:
    """
    
    rank_0 = group_ranks[0]
    rank = group_ranks[i_rank]
    if i_rank == 0: # First worker on the gpu
        ref_worker_communication.add_(params_ref_worker)
        wait_event(com_events[0], client_stream)
        dist.all_reduce(ref_worker_communication, group=process_group, op=dist.ReduceOp.SUM)
        wait_event(com_events[-1], client_stream)
        for idx_j, j in enumerate(group_ranks[1:num_clients_in_group]):
            queue_put = graph.edges[(rank, j)]["com_queue"][(rank, j)]
            queue_put.put(ref_worker_communication.clone())
        
    # Now obtain this params_ref to other processes of the GPU
    else:
        queue_get = graph.edges[(rank_0, rank)]["com_queue"][(rank_0, rank)]
        all_reduced_ref_worker = queue_get.get()
        ref_worker_communication.copy_(all_reduced_ref_worker)
        wait_event(com_events[0], client_stream)
        del all_reduced_ref_worker
        
@torch.no_grad()
def communicate_shuffled_mask_with_other_clients(graph, group_ranks, client_stream, process_group, i_rank, new_mask, num_clients_in_group, com_events, masks_shuffle_communication):
    """
    Communicate the shuffled mask with other clients in the group.
    
    """
    rank_0 = group_ranks[0]
    rank = group_ranks[i_rank]    
    # synchronize with rank 0 to get the shuffled mask
    if i_rank == 0: 
        masks_shuffle_communication.add_(new_mask)
        wait_event(com_events[0], client_stream)
        dist.all_reduce(masks_shuffle_communication, group=process_group, op=dist.ReduceOp.SUM)
        wait_event(com_events[-1], client_stream)
        for idx_j, j in enumerate(group_ranks[1:num_clients_in_group]):
            queue_put = graph.edges[(rank, j)]["com_queue"][(rank, j)]
            queue_put.put(masks_shuffle_communication.clone())

    else:
        queue_get = graph.edges[(rank_0, rank)]["com_queue"][(rank_0, rank)]
        all_reduced_new_masks = queue_get.get()
        masks_shuffle_communication.copy_(all_reduced_new_masks)
        wait_event(com_events[0], client_stream)
        del all_reduced_new_masks

@torch.no_grad()
def communicate_summed_mask_with_other_clients(graph, group_ranks, client_stream, process_group, i_rank, summed_mask, num_clients_in_group, com_events, summed_mask_communication):
    """
    Communicate the summed mask with other clients in the group.
    
    """
    rank_0 = group_ranks[0]
    rank = group_ranks[i_rank]    
    # synchronize with rank 0 to get the shuffled mask
    if i_rank == 0: 
        for idx_j, j in enumerate(group_ranks[1:num_clients_in_group]):
            queue_get = graph.edges[(rank, j)]["com_queue"][(j, rank)]
            summed_mask_from_other_clients = queue_get.get()
            summed_mask_communication.add_(summed_mask_from_other_clients.to(summed_mask_communication.dtype))
            wait_event(com_events[idx_j], client_stream)
            del summed_mask_from_other_clients

        summed_mask_communication.add_(summed_mask.to(summed_mask_communication.dtype))
        wait_event(com_events[0], client_stream)
        dist.all_reduce(summed_mask_communication, group=process_group, op=dist.ReduceOp.SUM)
        wait_event(com_events[-1], client_stream)
        
        for idx_j, j in enumerate(group_ranks[1:num_clients_in_group]):
            queue_put = graph.edges[(rank, j)]["com_queue"][(rank, j)]
            queue_put.put(summed_mask_communication.clone())

    else:
        rank_0 = group_ranks[0]
        queue_put = graph.edges[(rank, rank_0)]["com_queue"][(rank, rank_0)]
        queue_put.put(summed_mask.clone())
        queue_get = graph.edges[(rank, rank_0)]["com_queue"][(rank_0, rank)]
        
        all_reduced_summed_masks = queue_get.get()
        summed_mask_communication.copy_(all_reduced_summed_masks)
        wait_event(com_events[0], client_stream)
        del all_reduced_summed_masks


@torch.no_grad()
def communicate_optimizer_states(
    model,
    i_rank,
    graph,
    optimizer,
    communication_buffer_exp_avg,
    communication_buffer_exp_avg_sq,
    group_ranks,
    process_group,
    com_events,
    cuda_stream,
    num_clients_in_group,
):
    def extract_buffers():
        exp_avg_list = []
        exp_avg_sq_list = []
        for name, param in model.named_parameters():
            state = optimizer.state.get(param, None)
            if state is not None and 'exp_avg' in state and 'exp_avg_sq' in state:
                exp_avg_list.append(state['exp_avg'].clone())
                exp_avg_sq_list.append(state['exp_avg_sq'].clone())
        return exp_avg_list, exp_avg_sq_list

    def set_communication_buffers(new_exp_avg, new_exp_avg_sq):
        for i in range(len(local_exp_avg)):
            communication_buffer_exp_avg[i].copy_(new_exp_avg[i])
            communication_buffer_exp_avg_sq[i].copy_(new_exp_avg_sq[i])
    
    # re-gather the rank
    rank = group_ranks[i_rank]
    local_exp_avg, local_exp_avg_sq = extract_buffers()    # if we are client 0 of this GPU, we organize the outside communications for the group in addition to our own communications
    
    if i_rank == 0:
        # gather the list of communications for the group
        for idx_j, j in enumerate(group_ranks[1:num_clients_in_group]):
            queue_get_j = graph.edges[(rank, j)]["com_queue"][(j,rank)]
            exp_avg, exp_avg_sq = queue_get_j.get()
            
            for i in range(len(local_exp_avg)):
                communication_buffer_exp_avg[i].add(exp_avg[i])
                communication_buffer_exp_avg_sq[i].add(exp_avg_sq[i])
            
            wait_event(com_events[idx_j], cuda_stream)
            del exp_avg
            del exp_avg_sq
        
        # print_optimizer_state_by_model_order(model, optimizer)

        for i in range(len(local_exp_avg)):
            communication_buffer_exp_avg[i].add(local_exp_avg[i])
            communication_buffer_exp_avg_sq[i].add(local_exp_avg_sq[i])
        
        wait_event(com_events[0], cuda_stream)
        
        for i in range(len(local_exp_avg)):
            dist.all_reduce(communication_buffer_exp_avg[i], group=process_group, op=dist.ReduceOp.SUM)
            dist.all_reduce(communication_buffer_exp_avg_sq[i], group=process_group, op=dist.ReduceOp.SUM)

        wait_event(com_events[-1], cuda_stream)

        for j in group_ranks[1:num_clients_in_group]:
            # gather the mp.Queue objects
            queue_put = graph.edges[(rank,j)]["com_queue"][(rank,j)]
            queue_put.put((
                            [t.clone() for t in communication_buffer_exp_avg],
                            [t.clone() for t in communication_buffer_exp_avg_sq]
                        ))
            # queue_put.put((communication_buffer_exp_avg, communication_buffer_exp_avg_sq))
    
    elif i_rank < num_clients_in_group:
        rank_0 = group_ranks[0]
        queue_put = graph.edges[(rank, rank_0)]["com_queue"][(rank, rank_0)]
        queue_put.put((local_exp_avg, local_exp_avg_sq))
        queue_get = graph.edges[(rank, rank_0)]["com_queue"][(rank_0, rank)]
        
        all_reduced_exp_avg, all_reduced_exp_sq = queue_get.get()
        
        set_communication_buffers(all_reduced_exp_avg, all_reduced_exp_sq)
        wait_event(com_events[0], cuda_stream)
        
        del all_reduced_exp_avg
        del all_reduced_exp_sq

def print_optimizer_state_by_model_order(model, optimizer):
    print("\n========== Optimizer State Sanity Check (Ordered by model.named_parameters()) ==========\n")
    for name, param in model.named_parameters():
        print(f"Param '{name}' (shape: {param.shape})")

        if param.grad is None:
            print("  ➤ grad: None")
        else:
            print(f"  ➤ grad: mean={param.grad.mean():.6f}, std={param.grad.std():.6f}")

        state = optimizer.state.get(param, {})
        step = state.get("step", "N/A")
        exp_avg = state.get("exp_avg", None)
        exp_avg_sq = state.get("exp_avg_sq", None)

        print(f"  ➤ step: {step}")
        if exp_avg is not None:
            print(f"  ➤ exp_avg    : mean={exp_avg.mean():.6f}, std={exp_avg.std():.6f}")
        else:
            print("  ➤ exp_avg    : Not initialized")

        if exp_avg_sq is not None:
            print(f"  ➤ exp_avg_sq : mean={exp_avg_sq.mean():.6f}, std={exp_avg_sq.std():.6f}")
        else:
            print("  ➤ exp_avg_sq : Not initialized")
        print()