def extract_max_avg_bandwidth_cluster()

in optimum/amd/topology_utils.py [0:0]

• 15 lines of code
• 7 McCabe index (conditional complexity)

def extract_max_avg_bandwidth_cluster(bandwidth_matrix, cluster_num_devices):
    """
    Returns the cluster of a given number of devices that has the maximum average bandwidth between devices.
    """
    if len(bandwidth_matrix) < cluster_num_devices:
        raise ValueError("Number of devices in the cluster cannot be greater than the number of devices in the system")

    if cluster_num_devices == 1:
        return [0], float("inf")

    num_devices = range(len(bandwidth_matrix))

    max_avg_bandwidth = 0
    max_bandwidth_cluster = None

    for cluster in combinations(num_devices, cluster_num_devices):
        curr_bandwidth_matrix = [[bandwidth_matrix[i][j] for i in cluster] for j in cluster]
        curr_avg_bandwidth = np.mean(curr_bandwidth_matrix, where=~np.eye(len(curr_bandwidth_matrix), dtype=bool))
        if curr_avg_bandwidth > max_avg_bandwidth:
            max_avg_bandwidth = curr_avg_bandwidth
            max_bandwidth_cluster = list(cluster)

    return max_bandwidth_cluster, max_avg_bandwidth