Source code for bmtk.builder.builder_utils

import numpy as np
import hashlib
import logging
import uuid

try:
    from mpi4py import MPI

    comm = MPI.COMM_WORLD
    mpi_rank = comm.Get_rank()
    mpi_size = comm.Get_size()
    barrier = comm.barrier

except ImportError:
    comm = None
    mpi_rank = 0
    mpi_size = 1
    barrier = lambda: None


MPI_fail_params_nonuniform = True  # throw exception unless params across MPI ranks are the same.
logger = logging.getLogger(__name__)


def add_hdf5_attrs(hdf5_handle):
    # TODO: move this as a utility function
    hdf5_handle['/'].attrs['magic'] = np.uint32(0x0A7A)
    hdf5_handle['/'].attrs['version'] = [np.uint32(0), np.uint32(1)]


def list_to_hash(str_list):
    str_list = str_list.copy()
    str_list.sort()
    combined_keys = ':'.join(str_list).encode('utf-8')
    return hashlib.md5(combined_keys).hexdigest()


def check_properties_across_ranks(properties, graph_type='node'):
    """Checks that a properties table is consistant across all MPI ranks. Mainly used by add_nodes() and add_edges()
    method due to bug where using random generator without rng_seed was causing issues building the network properties
    consistantly using multiple cores.

    Will throw an Exception or a warning message (if MPI_fail_params_nonuniform is false)

    :param properties: A dictionary
    :param graph_type: 'node' or 'edge', used in error message. default 'node'.
    """
    if mpi_size < 2:
        return

    # Check that model_properties have the same number of items and the keys match
    n_args = len(properties)
    ranked_args = comm.allgather(n_args)
    if len(set(ranked_args)) > 1:
        err_msg = '{} properties are not the same across all ranks.'.format(graph_type)
        if not MPI_fail_params_nonuniform:
            logger.warning(err_msg)
        else:
            raise IndexError(err_msg)

    if n_args == 0:
        return

    # create a string/id that will be uniform across all ranks, even if the dict on one rank returns keys out-of-order.
    prop_keys = list(properties.keys())
    prop_keys.sort()
    combined_keys = ':'.join(prop_keys).encode('utf-8')
    hash_id = hashlib.md5(combined_keys).hexdigest()
    ranked_keys = comm.allgather(hash_id)
    if len(set(ranked_keys)) > 1:
        err_msg = '{} properties are not the same across all ranks.'.format(graph_type)
        if not MPI_fail_params_nonuniform:
            logger.warning(err_msg)
        else:
            raise IndexError(err_msg)

    # For each item in model_properties dictionary try to check that values are the same
    for pkey in prop_keys:
        # Don't use Dict.items() method since it is possible the ret order is different across ranks.
        pval = properties[pkey]

        try:
            if isinstance(pval, bytes):
                phash = hashlib.md5(pval).hexdigest()

            elif isinstance(pval, str):
                phash = hashlib.md5(pval.encode('utf-8')).hexdigest()

            elif isinstance(pval, (int, float, bool)):
                phash = pval

            elif isinstance(pval, (list, tuple)):
                joined_keys = ':'.join([str(p) for p in pval]).encode('utf-8')
                phash = hashlib.md5(joined_keys).hexdigest()

            elif isinstance(pval, np.ndarray):
                phash = hashlib.md5(pval.data.tobytes()).hexdigest()

            else:
                continue

        except TypeError as te:
            # If the hashing fails assume there is no MPI data issue and continue with the next property.
            continue

        ranked_vals = comm.allgather(phash)
        if len(set(ranked_vals)) > 1:
            err_msg = '{} property "{}" varies across ranks, please make sure parameter value is uniform across all' \
                      'ranks or set bmtk.builder.MPI_fail_params_nonuniform to False'.format(graph_type, pkey)
            if not MPI_fail_params_nonuniform:
                logger.warning(err_msg)
            else:
                raise TypeError(err_msg)


__build_time_uuid = None


def build_time_uuid():
    """Produces a unique identifier that will be unique for each time a network building script is executed.
    """
    global __build_time_uuid
    if __build_time_uuid is not None:
        return __build_time_uuid

    elif comm is not None and mpi_size > 1:
        barrier()
        
        try:
            if mpi_rank == 0:
                bcast_data = str(uuid.uuid4().hex)
            else:
                bcast_data = None

            bcast_data = comm.bcast(bcast_data, root=0)
            __build_time_uuid = bcast_data
        except Exception as e:
            __build_time_uuid = 'tmp'
    
    else:
        __build_time_uuid = str(uuid.uuid4().hex)

    # __build_time_uuid = str(uuid.uuid4().hex)
    return __build_time_uuid