def call_once_and_share(func, shape, dtype, rank=0):
    """Invoke the function in a single process of the PyTorch distributed process group,
    and share the result with other processes.

    Parameters
    ----------
    func : callable
        Any callable that accepts no arguments and returns an arbitrary object.
    shape : tuple[int]
        The shape of the shared tensor.  Must match the output of :attr:`func`.
    dtype : torch.dtype
        The data type of the shared tensor.  Must match the output of :attr:`func`.
    rank : int, optional
        The process ID to actually execute the function.
    """
    current_rank = torch.distributed.get_rank()
    dist_buf = torch.LongTensor([1])

    if torch.distributed.get_backend() == "nccl":
        # Use .cuda() to transfer it to the correct device.  Should be OK since
        # PyTorch recommends the users to call set_device() after getting inside
        # torch.multiprocessing.spawn()
        dist_buf = dist_buf.cuda()

    # Process with the given rank creates and populates the shared memory array.
    if current_rank == rank:
        # PyTorch Lightning 1.6+ seems to set the random seed during process spawning
        # to the same seed value.
        random_ = random.Random()
        id_ = random_.getrandbits(32)
        name = _get_shared_mem_name(id_)
        result = create_shared_mem_array(name, shape, dtype)
        result[:] = func()
        dist_buf[0] = id_

    # Broadcasts the name of the shared array to other processes.
    torch.distributed.broadcast(dist_buf, rank)
    # If no exceptions, other processes open the same shared memory object.
    if current_rank != rank:
        id_ = dist_buf.item()
        name = _get_shared_mem_name(id_)
        result = get_shared_mem_array(name, shape, dtype)

    return result
Invoke the function in a single process of the PyTorch distributed process group,
    and share the result with other processes.

    Parameters
    ----------
    func : callable
        Any callable that accepts no arguments and returns an arbitrary object.
    shape : tuple[int]
        The shape of the shared tensor.  Must match the output of :attr:`func`.
    dtype : torch.dtype
        The data type of the shared tensor.  Must match the output of :attr:`func`.
    rank : int, optional
        The process ID to actually execute the function.
    
python/dgl/multiprocessing/pytorch.py